ChrisRackauckas
commented
1 month ago Note that NonlinearSolve is currently missing Enzyme overloads. Only LinearSolve and ODEProblems have gotten that treatment so far.
Open m-bossart opened 1 month ago
ChrisRackauckas
commented
1 month ago Note that NonlinearSolve is currently missing Enzyme overloads. Only LinearSolve and ODEProblems have gotten that treatment so far.
wsmoses
commented
1 month ago @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
ChrisRackauckas
commented
1 month ago 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.
wsmoses
commented
1 month ago Sorry you're going to have to give me some more concrete pointers to functions which should have a rule.
m-bossart
commented
1 month ago I believe this is the dispatch Chris is referring to: (https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/src/solve.jl#L1066C1-L1083C4)
For the ODE case, the enzyme rule is for the solve_up function, so I'm guessing the same should be true for the case of nonlinear problems:
https://github.com/SciML/DiffEqBase.jl/blob/08fbaf0561d797359027849bac7797ce10948846/ext/DiffEqBaseEnzymeExt.jl#L9C1-L37C4
I'm not sure though, @ChrisRackauckas can you confirm?
ChrisRackauckas
commented
1 month ago 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.
m-bossart
commented
1 month ago I tried making NonlinearProblem a mutable struct but am still getting an error:
Closest candidates are:
copyto!(::OrdinaryDiffEq.ArrayFuse, ::Base.Broadcast.Broadcasted)
@ OrdinaryDiffEq C:\Users\Matt Bossart\.julia\packages\OrdinaryDiffEq\tAI61\src\wrappers.jl:26
copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}})
@ Base broadcast.jl:959
copyto!(::AbstractArray, ::Base.Broadcast.Broadcasted)
@ Base broadcast.jl:956
...It is failing at this line: https://github.com/SciML/DiffEqBase.jl/blob/1658ed0c18e5e8e6776e7da67c52a56af5076fa7/ext/DiffEqBaseEnzymeExt.jl#L33
wsmoses
commented
1 month ago 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?
m-bossart
commented
1 month ago The type of the variable being zero'd is Float64. The error I hit above as the same error you get if you do (trying to broadcast a value to a float):
x = 1.0
x .= 0.0The part I'm unclear on is why this works for an ODEProblemand not a NonlinearProblem. @ChrisRackauckas is it expected that the return from _solve_adjoint would be different for those two?
wsmoses
commented
1 month ago Oh that's definitely it.
Here's a solution. Replace the deepcopy with an Enzyme.make_zero and get rid of the for loop
avik-pal
commented
1 month ago 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?
ChrisRackauckas
commented
1 month ago 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.
m-bossart
commented
1 month ago I can confirm that by using Enzyme.make_zero instead of the for loop a simple case runs. It seems to work as well even without making NonlinearProblem mutable. What is the underlying reason for needing it to be mutable?
ChrisRackauckas
commented
1 month ago Limitation of Enzyme mixed activities.
wsmoses
commented
1 month ago 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 https://github.com/SciML/NonlinearSolve.jl/issues/439#issuecomment-2150575969, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXCT4AV2UTFAW4BR5FLZF5CGZAVCNFSM6AAAAABIJ3YE3WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJQGU3TKOJWHE . You are receiving this because you commented.Message ID: @.***>
m-bossart
commented
1 month ago 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?
m-bossart
commented
1 month ago I ask because I've tried it and still get Mixed activity errors when passing the reference instead
m-bossart
commented
4 weeks ago 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?
m-bossart
commented
3 weeks ago @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.
wsmoses
commented
3 weeks ago @m-bossart see here: https://github.com/EnzymeAD/Enzyme.jl/pull/1518 maybe test it out and see if it works [and in interim getting things ready for to release]
wsmoses
commented
3 weeks ago 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 https://github.com/SciML/NonlinearSolve.jl/issues/439#issuecomment-2156083306, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXC5ZVA4KEWQI4MHP33ZGMR2DAVCNFSM6AAAAABIJ3YE3WVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJWGA4DGMZQGY . You are receiving this because you were mentioned.Message ID: @.***>
m-bossart
commented
3 weeks ago 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?
wsmoses
commented
3 weeks ago 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()