Support more of Enzyme - Githubissues

JuliaSmoothOptimizers / ADNLPModels.jl

Other

29 stars 12 forks source link

Support more of Enzyme #228

Open kiranshila opened 1 month ago

kiranshila commented 1 month ago

Hey there!

I've been using your packages to work on an inverse design problem for a paper (and my thesis) with great success. I'm trying to squeeze some more performance out of the solver, though. So, I've read Enzyme can be more performant than the default ForwardDiff. Enzyme supports both forward and reverse mode, as well as jacobians and hessians (via forward over reverse). Currently, ADNLPModels only support Enzyme's reverse mode. It seems it would be useful to support all the features Enzyme has to offer has a higher-performance alternative to the currently defined AD backends.

tmigot commented 1 month ago

Hi @kiranshila ! Thanks for the suggestion. Tbh, I haven't completely followed last updates of Enzyme. Currently, we only have the reversed gradient, but for sure it would be great to add more.

We are a bit under-staffed at the moment, if you have some time to help on this that's would greatly speed-up answering your request? I can give guidance on how to integrate changes in the package.

kiranshila commented 1 month ago

Sure, I'll see what I can do!

Enzyme has made significant progress recently, and from what I've heard from the SciML folks, they really want to start moving people in the direction of Enzyme. It's still not without its rough edges, but seems to be the best performance you can get (on type stable code).

tmigot commented 1 month ago

In general, in ADNLPModels we are interested in computing derivatives of:

f(x): 1st and 2nd order derivative.
Directional derivatives of f, i.e. something like ForwardDiff.derivative(t -> ForwardDiff.gradient(f, x + t * v), 0)
derivatives (and directional derivatives) of in-place functions, ie x -> c!(cx, x).

If you manage to make Enzyme with the desired mode work on any of these, I can help with ADNLPModel-specific part.