UDE Training does not work with AutoZygote
#20
Comments
|
This is SciML/SciMLSensitivity.jl#1010 |
|
after updates to SciMLSensitivity.jl I reran the Friction example with AutoZygote and I got the same error as before: Closest candidates are: Stacktrace: |
|
@DhairyaLGandhi can you work on updating the example here to use AutoZygote? |
|
yeah this is the UDEs piece. As a workaround it is possible to define something like function Base.length(p::ModelingToolkit.MTKParameters)
return length(p.tunable)
endThe basic issue is that |
|
@DhairyaLGandhi I tried your suggestion, but then it failed on a similar issue with similar(::ModelingToolkit.MTKParameters{…}) so I made another dispatch for similar: but now the run end with an error: Closest candidates are: Stacktrace: p.tunable is not compatible with similar |
|
Can you try against https://github.com/SciML/SciMLSensitivity.jl/tree/dg/fwd |
|
Tried your branch. Same error. It seems that all the SciMLSensitivity function expect a flat vector for p, but when I look at p.tunable then I get: |
|
I even tried to get around this by: but then it just fails at the next function `ERROR: MethodError: no method matching jacobian!(::Matrix{…}, ::SciMLBase.ParamJacobianWrapper{…}, ::ModelingToolkit.MTKParameters{…}, ::Vector{…}, ::GaussAdjoint{…}, ::FiniteDiff.JacobianCache{…}) Closest candidates are: Stacktrace: |
|
Looking at the stacktrace, the culprit is in the default sensealg selection with SDEProblem. https://github.com/SciML/SciMLSensitivity.jl/tree/dg/length should fix that. Running julia> res = solve(op, Adam(), maxiters = 5000)#, callback = plot_cb)
retcode: Default
u: 57-element Vector{Float64}:
0.17041097372357306
0.06721610894572924
1.5036621719846985
0.6223167933686936
2.406843634370448
-1.2392234160442255
1.0211360153466225
-0.26749421630122894
1.2072691961676691
-0.3280020513366658
-0.3199580480281966
0.15962631397576563
-0.12814730955995285
0.25890635293263664
-0.9340225939644783
-0.7582786846498972
⋮
0.33400507029273274
0.37703694871113425
-0.38265915595848343
0.2532281365300615
-1.0286697496982982
1.9856277316142081
-1.2304071436862438
2.7505848772204997
-1.663046900479545
2.3058247613525467
1.4008403661001196
-2.0467838051849423
-0.6461702206670253
2.6279530122972967
-0.5972921571545664
1.723921054838542 |
|
Why SDEProblem? it's an ODE? |
|
@DhairyaLGandhi I tried your branch "dg/length" on the test/lotka_volterra.jl using AutoZygote, and this is what I got: ERROR: ForwardDiffSensitivity assumes the Stacktrace: |
|
@DhairyaLGandhi I omitted these messages before the final fail in my previous attempt with my own simular function: ┌ Warning: Potential performance improvement omitted. EnzymeVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add ┌ Warning: Potential performance improvement omitted. ReverseDiffVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add ┌ Warning: Reverse-Mode AD VJP choices all failed. Falling back to numerical VJPs Closest candidates are: Stacktrace: |
It's the same dispatch for both. This is the automatic sensealg. Apologies for the confusion. Yes I think I understand what is happening there. I'll add a fix for this. |
|
I'm getting with Zygote and default auto sensealg. If I specify sensealg = @DhairyaLGandhi I think ReverseDiffAdjoint also needs to be updated to use tunables a la your PR? If I use sensealg = |
|
One Zygote issue is FluxML/Zygote.jl#1517 |
|
I have updated SciML/SciMLSensitivity.jl#1085 to actually get this example running which doesn't run into differentiating the integrator interface. (That's what causes the zygote failure in the previous comments) It needs adding an extra |
|
Issues:
ForwardDiffSensitivity is slow with many parameters though, so it'd be really nice if ZygoteAdjoint and/or ReverseDiffAdjoint worked |
|
Why not one of the adjoint methods, like GaussAdjoint? This is really an orthogonal discussion to the purpose of the thread: if ForwardDiffSensitivity and GaussAdjoint work then the standard AutoZygote will work in any default chosen case or recommended case, so ReverseDiffAdjoint and ZygoteAdjoint really don't have a place in the story. |
ReverseDiffAdjoint should do better now, since it now always canonicalizes the tunables to a vector. So it's at least in principle possible to handle that now. It might need a SciMLStructures part added to its concrete_solve dispatch though to grab that vector, which isn't the hardest PR. Basically the same thing as SciML/SciMLSensitivity.jl#1085 |
Describe the bug 🐞
UDE Training does not work with
AutoZygoteExpected behavior
It should work.
Minimal Reproducible Example 👇
The lotka volterra test in
tests/lotka_volterra.jlError & Stacktrace⚠️
The text was updated successfully, but these errors were encountered: