Skip to content

Support Enzyme through EnzymeRules#186

Draft
vchuravy wants to merge 4 commits into
JuliaDecisionFocusedLearning:mainfrom
vchuravy:vc/enzyme_rule
Draft

Support Enzyme through EnzymeRules#186
vchuravy wants to merge 4 commits into
JuliaDecisionFocusedLearning:mainfrom
vchuravy:vc/enzyme_rule

Conversation

@vchuravy

Copy link
Copy Markdown

Took a stab at defining a rule for Enzyme, still need to port the reverse rule

@gdalle

gdalle commented Nov 14, 2025

Copy link
Copy Markdown
Member

Hi, thanks for the contribution! I didn't think anyone was actually using this package.
Are you sure you pushed all the commits?

@vchuravy

Copy link
Copy Markdown
Author

sigh missed the extension itself.

I was looking at ImplicitDifferentiation for AD through non linear problems and adjoint models for implicit time integration.

@gdalle

gdalle commented Nov 14, 2025

Copy link
Copy Markdown
Member

I was looking at ImplicitDifferentiation for AD through non linear problems and adjoint models for implicit time integration.

Feel free to suggest any improvements!
Right now the biggest limitations are

  • the out-of-place linear operator
  • the allocating linear solver of KrylovKit (but it was the one that came closest to idiomatic Julia without type restrictions, as opposed to Krylov which has lots of weird conversions)

Comment thread 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

@gdalle gdalle Nov 14, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@gdalle

gdalle commented Nov 15, 2025

Copy link
Copy Markdown
Member

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?

@tpapp

tpapp commented Feb 24, 2026

Copy link
Copy Markdown

@vchuravy, do you plan to continue this PR? (no pressure, just asking/friendly ping). I would find the functionality useful.

@gdalle

gdalle commented Jun 26, 2026

Copy link
Copy Markdown
Member

@copilot resolve the merge conflicts in this pull request

@vchuravy

Copy link
Copy Markdown
Author

Do you plan to continue this PR?

I don't have a current use myself so feel free to hijack this PR :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants