Add Enzyme extension #3327
Add Enzyme extension #3327
Conversation
|
On Julia < 1.9 you need requires [which is standard practice by extension packages], if Oceananigans enforces 1.9+, it's not necessary. Wasn't sure so added both here, and you can remove the requires dependency if desired. |
|
Oceananigans requires Julia v1.9 |
|
Please add a small description of this PR in the top comment. Question: is Oceananigans the right place for the extension to be or an extension for Oceananigans in Enzyme? (I'm not arguing for one or the other; I'm just wondering.) |
Done
So I think here is the right place, since code has a dependency on Oceananigans itself and EnzymeCore (which is dependency free). It also is how other packages (like NNlib, LinearSolve, SciMLSensitivity, etc) work with Enzyme and other AD tools. |
|
OK. How are these methods going to be tested and or maintained? |
|
So I added a unit test showing that checks that should test its functionality. Once we get more of Oceananigans differentiated we can add some bigger integration tests. |
wsmoses
force-pushed
the
enzymeext
branch
2 times, most recently
from
c4ea879
to
8ba01ee
Compare
|
@wsmoses should we add a new buildkite run dedicated solely to the enzyme extension? It might help for debugging / interpreting failures. (If that makes sense, I can shuffle things around to make that work.) |
|
What's happening presently is KA is forcing an older version of Enzyme to be used in the tests, that doesn't have the inactive_type used by the testing infra. This PR JuliaGPU/KernelAbstractions.jl#426 added relevant parts to KA, which currently is waiting for the julia registry to release a new patch release: JuliaRegistries/General#93272 Once that lands, this should succeed. |
Ok great! Also let me know if its ok with you if I set up a new buildkite step for the enzyme extension. |
|
@glwagner the prerequisites have landed with releases cut. Hopefully this is happy (and finds the correct versions of KA/Enzyme), but it was just tagged a few hours ago, so we'll see |
|
Might need to update the Manifest |
This PR adds an extension package to Oceananigans to describing the custom derivative behavior of Oceananigans types and functions.