Gradient for Cholesky decomposition with CuArrays with Zygote.jl

General Usage
, , ,
1

Here’s a minimal example:

using Zygote, CUDA, LinearAlgebra
CUDA.allowscalar(false)
C = cu(rand(5,5))
f(x) = sum(cholesky(x’ * x).L)
Zygote.gradient(f, C)

While the CPU version works fine, this will throw the scalar indexing error. It seems like this is a recurring issue:

and there’s a purposed change with a customized adjoint:

But that doesn’t work for the above case.

Does anyone know how to get around this?