Adjoints for LinearSolve

[avik-pal]

Simple non-working example.

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

We need SciML/LinearSolve.jl#322 before we can work on the adjoint code.

[avik-pal]

This is done.