Enzyme Implementation

[SCiarella]

This branch implements AD using Enzyme.jl.

In particular, we can now compute $$\dfrac{\partial f(u,t)}{\partial u}$$ where $$f(u,t)$$ is the right hand side of INS, projected in order to be compatible with SciML. Notice that at the moment there is no learnable closure in the loop, and the derivative is calculated only for $u$.

---

Main result

In src/sciml.jl there are now create_right_hand_side (out of place) and right_hand_side! (in place).
The in-place version is ~30% faster and uses only ~40% of the memory compared to the out-of-place, as one can see from test/enzyme_integration. This discrepancy becomes even more drastic when $$f$$ is used in SciML.solve(), however this has not been tested yet.

To do AD, Zygote requires out-of-place operations, so I have compared the pullback speed of Zygote+out vs Enzyme+in . test/enzyme_integration reports that Enzyme is ~30% faster and uses ~40% of the memory.

---

Observations

In test/chainrules_enzyme I compare and test the pullback rules of all the functions that are also tested in chainrules.
However Enzyme has a significant limitation which is the fact that it can not differentiate functions that return vectors. To overcome this problem I introduced enzyme_wrap() that executes the function in a local scope and returns nothing. In this way Enzyme works and produces the same gradients as Zygote with better performance.

But still, due to this limitation one can not really reuse the rules defined for the single functions because it is not a good idea to wrap all the pieces in the right hand side, so the pullback rule for the full rhs has been defined separately.

---

Todo before merge

• add docstrings
• remove the @info in the tests that compare the time
• check the consistency of the code

---

Bugs fixed

src/operators.jl Line 841 applybodyforce! : If I am not mistaken, in case of static force, it has to be applied only once outside the loop. Please check if this is correct.

[agdestein]

Great! It would also be good to have an example of computing gradients with Enzyme, something like this:

https://github.com/agdestein/IncompressibleNavierStokes.jl/blob/main/docs/src/manual/differentiability.md

[agdestein]

I also added a citation file, you can add yourself there.

[codecov]

Codecov Report

Attention: Patch coverage is 92.15686% with 4 lines in your changes missing coverage. Please review.

Project coverage is 59.39%. Comparing base (5ed9a60) to head (6a582de).
Report is 18 commits behind head on main.

Files with missing lines	Patch %	Lines
src/operators.jl	83.33%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #116      +/-   ##
==========================================
+ Coverage   57.39%   59.39%   +2.00%     
==========================================
  Files          29       30       +1     
  Lines        3666     3699      +33     
==========================================
+ Hits         2104     2197      +93     
+ Misses       1562     1502      -60     

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

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests