Adjoints for Linear Solve

[avik-pal]

Fixes #198, Fixes #322

TODOs:

• Preconditioning. These need to be added to the Adjoint Sensitivity Struct. Can we use a left preconditioner for the forward problem as the transpose right preconditioner? Need someone with a better grasp of linear algebra to chime in here
• Support the linsolve !== nothing case. This is useful if we know that $A^T$ has a structure exploitable by a different solver
• Tests -- Take from Enzyme and ChainRules
• Fix literal_getproperty for LinearSolution SciMLBase.jl#566

Example:

using LinearSolve, Zygote

A = rand(4, 4)
b = rand(4)

test_func_1(A, b) = sum(abs2, A \ b)

test_func_1(A, b)

∂A_1, ∂b_1 = @btime Zygote.gradient(test_func_1, copy(A), copy(b))
display(∂A_1)
display(∂b_1)

function test_func_2(A, b)
    prob = LinearProblem(A, b)
    sol = solve(prob)
    return sum(abs2, sol.u)
end

test_func_2(A, b)

∂A_2, ∂b_2 = @btime Zygote.gradient(test_func_2, copy(A), copy(b))
display(∂A_2)
display(∂b_2)

In the following case the cache stores the correct gradients but they are not propagated to A and b. @ChrisRackauckas any idea how to fix this?

cache = init(LinearProblem(copy(A), copy(b)), nothing);
function test_func_3(cache, A, b)
    cache.A = A
    cache.b = b
    sol = solve!(cache)
    return sum(abs2, sol.u)
end

test_func_3(cache, copy(A), copy(b))

∂cache, ∂A_3, ∂b_3 = @btime Zygote.gradient(test_func_3, cache, copy(A), copy(b))
∂cache.A
∂cache.b
display(∂A_3)  # nothing
display(∂b_3)  # nothing

[codecov]

Codecov Report

Attention: Patch coverage is 6.66667% with 42 lines in your changes are missing coverage. Please review.

Project coverage is 22.96%. Comparing base (a206054) to head (06c09a3).

❗ Current head 06c09a3 differs from pull request most recent head 7671369. Consider uploading reports for the commit 7671369 to get more accurate results

Files	Patch %	Lines
src/adjoint.jl	2.32%	42 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #449       +/-   ##
===========================================
- Coverage   66.12%   22.96%   -43.17%     
===========================================
  Files          27       28        +1     
  Lines        2146     2147        +1     
===========================================
- Hits         1419      493      -926     
- Misses        727     1654      +927     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ChrisRackauckas]

Don't we technically have to deepcopy in here?

[avik-pal]

I guess so, it can be problematic if there are 2 subsequent solve calls on the cache.

[ChrisRackauckas]

In the following case the cache stores the correct gradients but they are not propagated to A and b. @ChrisRackauckas any idea how to fix this?

Is this not just an inherent limitation of Zygote with mutation? I would presume we just need to stay away from that and only support solve with CRC.