Zygote adjoints for cholesky factorization on sparse arrays

Specific Domains Machine Learning
1

Hi. Is there a quick and dirty way of getting AD to work when dealing with cholesky factorizations? The example below performs factorization (must be done inside) followed by solving a system of equations. The goal is to update a scalar parameter dt:

function f(A,b,dt) 
    # C = cholesky(Matrix(A)) # this works
    D = I + dt * A
    C = cholesky(D) # Sparse array
    x = C \ b
    sum(x)
end
function ldiv!(x::AbstractVector, F::SparseArrays.CHOLMOD.Factor{Float64}, b::Vector{Float64})
    x .= A \ b
end
function ldiv!(F::SparseArrays.CHOLMOD.Factor{Float64}, b::Vector{Float64})
    x = similar(b)
    ldiv!(x, F, b)
end

SGD:

n = 5
A = sprand(n,n,0.2)
A = A * A' + .01 * I
b = rand(n)
t = rand()

# println(f(A,b))
for i=1:300
    grad = Zygote.gradient(f,A,b,t)
    b -= 0.0001 * grad[3]
    if i % 1 == 0
        println(f(A,b))
    end
end
println(f(A,b))

I needed to define ldiv! for CHOLMOD factors because it apparently isn’t supported yet: missing `ldiv!` overload for sparse Cholesky · Issue #319 · JuliaSparse/SparseArrays.jl · GitHub. Now, Zygote still complains that Need an adjoint for constructor SparseArrays.CHOLMOD.Dense{Float64}. Gradient is of type FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}} which I suppose means I need to define a function with signature akin to Zygote.@adjoint function SparseArrays.CHOLMOD.Sparse{Float64}(x). I don’t know why I need to do this or how I would even attempt to do this. Any insight is appreciated.