Support Enzyme through EnzymeRules#186
Support Enzyme through EnzymeRules#186vchuravy wants to merge 3 commits intoJuliaDecisionFocusedLearning:mainfrom
Conversation
|
Hi, thanks for the contribution! I didn't think anyone was actually using this package. |
|
sigh missed the extension itself. I was looking at ImplicitDifferentiation for AD through non linear problems and adjoint models for implicit time integration. |
Feel free to suggest any improvements!
|
ext/ImplicitDifferentiationEnzyme.jl
Outdated
| EnzymeRules.forward(config, implicit, RT, Const(prep), x, args...) | ||
| end | ||
|
|
||
| @inline function EnzymeRules.forward(config, implicit::Const{<:ImplicitFunction}, RT::Type, prep::Const{<:ImplicitFunctionPreparation{R}}, x, args...) where R |
There was a problem hiding this comment.
I'm not convinced we should define a rule for the version with a prep. The main reason that one exists in the ForwardDiff extension was to enable the use of sparse Jacobian matrices inside the linear system, where the sparsity pattern and coloring are contained in prep. But in theory, including prep in an Enzyme rule means we should define derivatives of the result with respect to prep, which seems complicated?
|
Thanks for the contribution! Before I review, can you add Enzyme to the systematic tests, both as an inner and as an outer backend, to see if they run? |
|
@vchuravy, do you plan to continue this PR? (no pressure, just asking/friendly ping). I would find the functionality useful. |
3 participants
Took a stab at defining a rule for Enzyme, still need to port the reverse rule