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:
- Need a GPU compatible `rrule` for Cholesky · Issue #629 · JuliaDiff/ChainRules.jl · GitHub
- Adjoint of `cholesky` is hard-coded for the CPU · Issue #1210 · FluxML/Zygote.jl · GitHub
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?