Use SparseDiffTools v2 for steadystateadjoint

[vpuri3]

part of: SciML/SciMLOperators.jl#142

[vpuri3]

TODO- update OrdinaryDiffEq to use sparsedifftools v2

[avik-pal]

This is not really equivalent. IIRC VecJac recomputes the pullback everytime a call to mul! is made. In this case, we have a fixed input, only the seeding changes so we compute the pullback once and just reevaluate it multiple times.

[ChrisRackauckas]

Is this actually fixed?

[avik-pal]

Yes, see JuliaDiff/SparseDiffTools.jl#245

[vpuri3]

retriggering CI

[vpuri3]

We should add a test case to ensure that Zygote is able to propagate gradients past solve(::LinearProblem)

[avik-pal]

We should add a test case to ensure that Zygote is able to propagate gradients past solve(::LinearProblem)

Does this work yet? I was using LinearSolve.jl, and I think the adjoints are not implemented yet. I need it (quite desperately), so I will probably implement something by next week.

using Zygote, ForwardDiff, SciMLSensitivity, SciMLBase, LinearSolve, ComponentArrays,
      FiniteDiff

function loss_function(θ)
    (; A, b) = θ
    prob = LinearProblem(A, b)
    sol = solve(prob, nothing)
    return sum(sol.u)
end

function loss_function_chainrules(θ)
    (; A, b) = θ
    x = A \ b
    return sum(x)
end

A = Float32[1 0; 1 -2]; b = Float32[32; -4];
θ = ComponentArray(; A, b)

loss_function(θ) ≈ loss_function_chainrules(θ)  # true

Zygote.gradient(loss_function, θ)  # fails
Zygote.gradient(loss_function_chainrules, θ)  # works

ForwardDiff.gradient(loss_function, θ)  # fails
ForwardDiff.gradient(loss_function_chainrules, θ)  # works

FiniteDiff.finite_difference_gradient(loss_function, θ)  # works

[vpuri3]

with this branch, your example is erroring on some try/catch block which is incompatible with zygote.

ERROR: LoadError: Compiling Tuple{LinearSolve.var"##solve#32", Base.Pairs{Symbol, Union
{}, Tuple{}, NamedTuple{(), Tuple{}}}, typeof(solve), LinearSolve.LinearCache{Base.Resh
apedArray{Float32, 2, SubArray{Float32, 1, Vector{Float32}, Tuple{UnitRange{Int64}}, tr
ue}, Tuple{}}, SubArray{Float32, 1, Vector{Float32}, Tuple{UnitRange{Int64}}, true}, Ve
ctor{Float32}, SciMLBase.NullParameters, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}
, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Krylov.GmresSolver{Fl
oat32, Float32, Vector{Float32}}, SciMLOperators.IdentityOperator, SciMLOperators.Ident
ityOperator, Float32, true, LinearSolve.OperatorCondition.IllConditioned}, KrylovJL{typ
eof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(),
 Tuple{}}}}}: try/catch is not supported.                                              
Refer to the Zygote documentation for fixes.                                           
https://fluxml.ai/Zygote.jl/latest/limitations                                         ```

[avik-pal]

Yes, ik I gave that as a testcase that doesn't work. The linear problem doesn't have tests because the adjoints aren't implemented.

Also, taking a look at how linearsolve works, it might a bit refactoring before we can include adjoints since it dispatches on solve instead of __solve

[ChrisRackauckas]

We can make a higher level in SciMLBase and dispatch it on __solve.

[avik-pal]

We can get this PR merged, right? The linear solve issue is entirely tangential and needs to be handled downstream first.

[ChrisRackauckas]

someone needs to resolve the merge.

[ChrisRackauckas]

No Enzyme or ReverseDiff though?