Implicit Function Adjoint Rules for Enzyme #439
Comments
|
Note that NonlinearSolve is currently missing Enzyme overloads. Only LinearSolve and ODEProblems have gotten that treatment so far. |
avik-pal
added
enhancement
avik-pal
changed the title
|
@ChrisRackauckas I can help with the EnzymeRules infra, is there an adjoint function lying around for this. Similarly what would be the right function to hook into with a rule |
|
It's just like the ODE infrastructure in that it just lowers to call SciMLSensitivity for all of the real work. The current dispatch is in DiffEqBase.jl right next to the ODE one. |
|
Sorry you're going to have to give me some more concrete pointers to functions which should have a rule. |
|
I believe this is the dispatch Chris is referring to: For the ODE case, the enzyme rule is for the I'm not sure though, @ChrisRackauckas can you confirm? |
|
Yup and then the adjoint just passes it on: https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/ext/DiffEqBaseChainRulesCoreExt.jl#L22C28-L29 . Looking at the ODE one, the NonlinearProblem one might just work if NonlinearProblem is made a mutable struct. |
|
I tried making NonlinearProblem a mutable struct but am still getting an error: It is failing at this line: |
|
What is the type of the variable being zero's there Enzyme semantics require it be zeroed but it looks like it may not be defined for that type? |
|
The type of the variable being zero'd is The part I'm unclear on is why this works for an |
|
Oh that's definitely it. Here's a solution. Replace the deepcopy with an Enzyme.make_zero and get rid of the for loop |
|
Will https://github.com/SciML/SimpleNonlinearSolve.jl/blob/6b682af1e2155298064d2a035939e6aa373edbed/test/gpu/cuda_tests.jl#L40-L72 these still work if we make the struct mutable? |
|
No, which makes this a major PITA. With DiffEqGPU we had to create a separate immutable problem that gets used internally to keep CUDA kernel compatibility. It converts under the hood so the user doesn't know, but it's definitely code smell. |
|
I can confirm that by using |
|
Limitation of Enzyme mixed activities. |
|
At one point it’s on the roadmap to remove the mixed activity limitation,
but some other GPU, type stability and blas stuff is currently higher
priority
…On Wed, Jun 5, 2024 at 7:20 PM Christopher Rackauckas < ***@***.***> wrote:
Limitation of Enzyme mixed activities.
—
Reply to this email directly, view it on GitHub
<#439 (comment)>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAJTUXCT4AV2UTFAW4BR5FLZF5CGZAVCNFSM6AAAAABIJ3YE3WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJQGU3TKOJWHE>
.
You are receiving this because you commented.Message ID:
***@***.***>
|
|
Interesting, I think I am hitting this same mixed activity limitation for another non-mutable struct in my application. @wsmoses would you expect passing a reference to the struct as in the docs here: https://enzyme.mit.edu/julia/stable/faq/#Mixed-activity would be a suitable workaround for the mixed activity limitation? |
|
I ask because I've tried it and still get Mixed activity errors when passing the reference instead |
|
I started two draft PRs above. @ChrisRackauckas I just added the code to SimpleNonlinearSolve because that is where the tests are. Should it be moved somewhere else? I'm hitting a method error when trying to convert which I don't understand; I mentioned it in the PR. Can you see what I'm doing wrong there? |
|
@wsmoses let me know when the new |
|
@m-bossart see here: EnzymeAD/Enzyme.jl#1518 maybe test it out and see if it works [and in interim getting things ready for to release] |
|
It is released
…On Sat, Jun 8, 2024 at 11:49 AM Matt Bossart ***@***.***> wrote:
@wsmoses <https://github.com/wsmoses> let me know when the new make_zero!
is available. Then I will fix the enzyme rule in DiffEqBase.jl and add a
test in SciMLSensitivity.jl and we should be all set.
—
Reply to this email directly, view it on GitHub
<#439 (comment)>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAJTUXC5ZVA4KEWQI4MHP33ZGMR2DAVCNFSM6AAAAABIJ3YE3WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJWGA4DGMZQGY>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
|
Ok I now have 4 PRs open that should address this. @ChrisRackauckas can you review these collectively? @wsmoses is the use of make_zero and make_zero correct? |
|
Reviewed the ones I can offer comments on |
I am trying to use Enzyme to differentiate a function which builds and solves a NonlinearProblem. I believ the issue is with the enzyme extension which is why I'm opening the issue here.
Expected behavior
I expect the SciMLSensitivity overload for NonlinearProblem to be called by Enzyme when differentiating.
Minimal Reproducible Example 👇
Error & Stacktrace⚠️
I believe this is the problematic line: https://github.com/SciML/DiffEqBase.jl/blob/1658ed0c18e5e8e6776e7da67c52a56af5076fa7/ext/DiffEqBaseEnzymeExt.jl#L33
Environment (please complete the following information):
using Pkg; Pkg.status()using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)versioninfo()The text was updated successfully, but these errors were encountered: