Add forward enzyme rules for init and solve #416

sharanry · 2023-11-05T18:09:16Z

Followup to Add Enzyme extension #377

TODO:

Test against other algorithms
Error on unspported algorithms.

codecov · 2023-11-05T18:12:44Z

Codecov Report

Merging #416 (0935919) into main (b9da6ac) will decrease coverage by 1.11%.
The diff coverage is 0.00%.

@@            Coverage Diff             @@
##             main     #416      +/-   ##
==========================================
- Coverage   64.94%   63.83%   -1.11%     
==========================================
  Files          26       26              
  Lines        2068     2099      +31     
==========================================
- Hits         1343     1340       -3     
- Misses        725      759      +34

Files	Coverage Δ
ext/LinearSolveEnzymeExt.jl	`0.00% <0.00%> (ø)`

... and 1 file with indirect coverage changes

📣 Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in Chrome or Firefox today!

sharanry · 2023-11-05T18:20:48Z

test/enzyme.jl

These tests fail due to numerical imprecision. Not entirely sure why.

df/db:

manual_jac = [-5.061916901336408, 1.5770609499033128, 5.446411839233853, 0.612648432464526] fd_jac = [-5.061916947364807, 1.5770610570907593, 5.446411848068237, 0.6126484870910645] en_jac = [-5.061916901336408, 1.5770609499033128, 5.446411839233853, 0.612648432464526]

That's an expected difference due to summation order.

addressed in 31745cc

ChrisRackauckas · 2023-11-05T19:33:44Z

ext/LinearSolveEnzymeExt.jl

+    end
+
+    dres = deepcopy(res)
+    invA = inv(A)


never use inv

addressed in d010911

Not fully, see #416 (comment). Your version has an extra factorization and is using the wrong linear solver.

addressed in 31745cc

ChrisRackauckas · 2023-11-05T19:36:43Z

ext/LinearSolveEnzymeExt.jl

+    invA = inv(A)
+    db = linsolve.dval.b
+    dA = linsolve.dval.A
+    dres.u .= invA * (db - dA * res.u)


There's no need to refactorize here. This is just A * u = db - dA * u, or u = A \ (db - dA * u). But this is just a linsolve call. Not only that, but it's the same operator as the one used in the forward pass, so you don't need to refactorize A. Therefore, this should simply use the same linsolve and do linsolve.b = db - dA * u and then solve!.

But I don't think the formula is correct. Isn't it just linsolve.b = db and then solve!(linsolve) then du = linsolve.u?

But linsolves.A seem to mutate after the solve. That is why I had to add all this other stuff.

MWE:

using LinearSolve A = rand(5,5) b = rand(5) prob = LinearProblem(A,b) linsolve = init(prob) @show linsolve.A sol = solve!(linsolve) @show linsolve.A

5×5 Matrix{Float64}: 0.207897 0.0737386 0.220551 0.935437 0.883482 0.362979 0.687049 0.220086 0.216771 0.145246 0.538937 0.577133 0.965879 0.438704 0.689019 0.348145 0.174776 0.0491639 0.145817 0.155529 0.92258 0.637729 0.508441 0.166917 0.218296 5×5 Matrix{Float64}: 0.92258 0.637729 0.508441 0.166917 0.218296 0.393439 0.436142 0.0200455 0.1511 0.0593595 0.584162 0.469106 0.659464 0.270316 0.533653 0.225343 -0.160427 0.16558 0.877305 0.755451 0.37736 -0.151045 -0.211799 0.185687 0.0548682

Yes, the mutation that it's doing is lu!(A), i.e. its using the same memory in the representation of the LU-factorization. You want to let it use exactly the same mutated A via cache.cacheval in order to do the next solve, which is precisely what's done in the caching interface. See https://docs.sciml.ai/LinearSolve/stable/tutorials/caching_interface/

Sounds good. And the simple derivation of the formula is mentioned in https://scicomp.stackexchange.com/a/29421.

Also, the correctness of the formula is checked with FiniteDiff in the tests - https://github.com/SciML/LinearSolve.jl/pull/416/files#diff-7c5a302dc0e55153407f2354959d201e28d2b9a61c4ef6bc86f05697a52b4cacR180-R196

addressed in 31745cc 19a54fe

ChrisRackauckas · 2023-11-07T00:20:23Z

This looks like a real test failure.

wsmoses · 2023-11-07T00:23:53Z

ext/LinearSolveEnzymeExt.jl

+    if RT <: DuplicatedNoNeed
+        return dres
+    elseif RT <: Duplicated
+        return Duplicated(res, dres)


This doesn't handle batching atm

Throw a good error for now? It would be good to get something merged even if it doesn't handle every case, as long as the errors are clear.

tried to address this in 585cbb8. We can make a followup PR for batch support.

sharanry · 2023-11-07T13:44:43Z

@ChrisRackauckas The test failures seem unrelated.

… due to summation ordering

…rted

…d random failures

sharanry force-pushed the enzyme_forward branch from 76d30f8 to e89be05 Compare November 5, 2023 18:15

sharanry commented Nov 5, 2023

View reviewed changes

sharanry force-pushed the enzyme_forward branch from 4bcdc63 to 3f88629 Compare November 5, 2023 18:23

ChrisRackauckas reviewed Nov 5, 2023

View reviewed changes

ChrisRackauckas approved these changes Nov 5, 2023

View reviewed changes

sharanry marked this pull request as ready for review November 6, 2023 17:30

wsmoses reviewed Nov 7, 2023

View reviewed changes

sharanry added 10 commits November 8, 2023 00:24

Add forward enzyme rules for init and solve

d5e64bd

Remove EnzymeTestUtils as it is not used

6cfe69a

Show computed jacobians in tests

f7b7c14

Avoid inv

8b8f8c4

Avoid refactorization and cleanup tests to allow for numerical errors…

447fb61

… due to summation ordering

Remove unneeded deepcopy

59edb4e

Switch to reltol from abstol for tests

60afb5d

Add tests for other algs and handle cases of algs currently not suppo…

c14aeb2

…rted

Error on unsupported return type and relax tolerence on tests to avoi…

5abd3f2

…d random failures

Fix bug in handling Keylov solvers

4265053

ChrisRackauckas force-pushed the enzyme_forward branch from 2803bda to 4265053 Compare November 7, 2023 23:25

Update Project.toml

0935919

ChrisRackauckas merged commit 9b540ba into SciML:main Nov 8, 2023
10 of 13 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add forward enzyme rules for init and solve #416

Add forward enzyme rules for init and solve #416

sharanry commented Nov 5, 2023 •

edited

Loading

codecov bot commented Nov 5, 2023 •

edited

Loading

sharanry Nov 5, 2023

sharanry Nov 5, 2023 •

edited

Loading

ChrisRackauckas Nov 5, 2023

sharanry Nov 5, 2023

ChrisRackauckas Nov 5, 2023

sharanry Nov 5, 2023

ChrisRackauckas Nov 5, 2023

sharanry Nov 5, 2023

ChrisRackauckas Nov 5, 2023

sharanry Nov 5, 2023

ChrisRackauckas Nov 5, 2023

sharanry Nov 5, 2023

sharanry Nov 5, 2023

sharanry Nov 5, 2023 •

edited

Loading

ChrisRackauckas commented Nov 7, 2023

wsmoses Nov 7, 2023

ChrisRackauckas Nov 7, 2023

sharanry Nov 7, 2023

sharanry commented Nov 7, 2023

Add forward enzyme rules for init and solve #416

Add forward enzyme rules for init and solve #416

Conversation

sharanry commented Nov 5, 2023 • edited Loading

codecov bot commented Nov 5, 2023 • edited Loading

Codecov Report

Choose a reason for hiding this comment

sharanry Nov 5, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sharanry Nov 5, 2023 • edited Loading

Choose a reason for hiding this comment

ChrisRackauckas commented Nov 7, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sharanry commented Nov 7, 2023

sharanry commented Nov 5, 2023 •

edited

Loading

codecov bot commented Nov 5, 2023 •

edited

Loading

sharanry Nov 5, 2023 •

edited

Loading

sharanry Nov 5, 2023 •

edited

Loading