Open sathvikbhagavan opened 2 months ago
after updates to SciMLSensitivity.jl
pkg> st [336ed68f] CSV v0.10.14 [13f3f980] CairoMakie v0.12.4 [b0b7db55] ComponentArrays v0.15.14 [a93c6f00] DataFrames v1.6.1 [0c46a032] DifferentialEquations v7.13.0 [a98d9a8b] Interpolations v0.15.1 [033835bb] JLD2 v0.4.48 [b2108857] Lux v0.5.59 [961ee093] ModelingToolkit v9.23.0 [f162e290] ModelingToolkitNeuralNets v1.0.2 [16a59e39] ModelingToolkitStandardLibrary v2.7.2 [7f7a1694] Optimization v3.26.3 [36348300] OptimizationOptimJL v0.3.2 [42dfb2eb] OptimizationOptimisers v0.2.1 [1dea7af3] OrdinaryDiffEq v6.85.0 [18e31ff7] Peaks v0.5.2 [1ed8b502] SciMLSensitivity v7.62.0 [53ae85a6] SciMLStructures v1.4.1 [860ef19b] StableRNGs v1.0.2 [2efcf032] SymbolicIndexingInterface v0.3.22 [e88e6eb3] Zygote v0.6.70 [9e88b42a] SerializationI
I reran the Friction example with AutoZygote and I got the same error as before: `ERROR: MethodError: no method matching length(::ModelingToolkit.MTKParameters{…})
Closest candidates are: length(::Combinatorics.FixedPartitions) @ Combinatorics ~/.julia/packages/Combinatorics/Udg6X/src/partitions.jl:96 length(::Core.Compiler.InstructionStream) @ Base show.jl:2777 length(::SymbolicUtils.Code.AtIndex) @ SymbolicUtils ~/.julia/packages/SymbolicUtils/dtCid/src/utils.jl:227 ...
Stacktrace:
[1] _empty(x::ModelingToolkit.MTKParameters{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:171
[2] map
@ ./tuple.jl:291 [inlined]
[3] adjoint
@ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:204 [inlined]
[4] adjoint(context::Zygote.Context{…}, 463::typeof(Core._apply_iterate), 464::typeof(iterate), f::Function, args::ModelingToolkit.MTKParameters{…})
@ Zygote ./none:0
[5] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[6] replace
@ ~/.julia/packages/ModelingToolkit/nadwo/src/systems/parameter_buffer.jl:288 [inlined]
[7] _pullback(::Zygote.Context{…}, ::typeof(SciMLStructures.replace), ::Tunable, ::ModelingToolkit.MTKParameters{…}, ::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
[8] loss
@ ~/Documents/programming/NeurobloxSciML/friction.jl:73 [inlined]
[9] _pullback(::Zygote.Context{…}, ::typeof(loss), ::Vector{…}, ::Tuple{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
[10] _apply
@ ./boot.jl:838 [inlined]
[11] adjoint
@ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
[12] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[13] OptimizationFunction
@ ~/.julia/packages/SciMLBase/rR75x/src/scimlfunctions.jl:3763 [inlined]
[14] _pullback(::Zygote.Context{…}, ::OptimizationFunction{…}, ::Vector{…}, ::Tuple{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
[15] _apply(::Function, ::Vararg{Any})
@ Core ./boot.jl:838
[16] adjoint
@ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
[17] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[18] #37
@ ~/.julia/packages/OptimizationBase/32Mb0/ext/OptimizationZygoteExt.jl:90 [inlined]
[19] _pullback(ctx::Zygote.Context{…}, f::OptimizationZygoteExt.var"#37#55"{…}, args::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
[20] _apply(::Function, ::Vararg{Any})
@ Core ./boot.jl:838
[21] adjoint
@ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
[22] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[23] #39
@ ~/.julia/packages/OptimizationBase/32Mb0/ext/OptimizationZygoteExt.jl:93 [inlined]
[24] _pullback(ctx::Zygote.Context{…}, f::OptimizationZygoteExt.var"#39#57"{…}, args::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
[25] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:90
[26] pullback
@ ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:88 [inlined]
[27] gradient(f::Function, args::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:147
[28] (::OptimizationZygoteExt.var"#38#56"{OptimizationZygoteExt.var"#37#55"{…}})(::Vector{Float64}, ::Vector{Float64})
@ OptimizationZygoteExt ~/.julia/packages/OptimizationBase/32Mb0/ext/OptimizationZygoteExt.jl:93
[29] macro expansion
@ ~/.julia/packages/OptimizationOptimisers/AOkbT/src/OptimizationOptimisers.jl:68 [inlined]
[30] macro expansion
@ ~/.julia/packages/Optimization/EmxXu/src/utils.jl:32 [inlined]
[31] __solve(cache::OptimizationCache{…})
@ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/AOkbT/src/OptimizationOptimisers.jl:66
[32] solve!(cache::OptimizationCache{…})
@ SciMLBase ~/.julia/packages/SciMLBase/rR75x/src/solve.jl:188
[33] solve(::OptimizationProblem{…}, ::Adam; kwargs::@Kwargs{…})
@ SciMLBase ~/.julia/packages/SciMLBase/rR75x/src/solve.jl:96
[34] top-level scope
@ ~/Documents/programming/NeurobloxSciML/friction.jl:101
Some type information was truncated. Use show(err)
to see complete types.`
@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)
end
The basic issue is that MTKParameters
defines getindex
but not length
.
@DhairyaLGandhi I tried your suggestion, but then it failed on a similar issue with similar(::ModelingToolkit.MTKParameters{…})
so I made another dispatch for similar:
function Base.similar(p::ModelingToolkit.MTKParameters) return similar(p.tunable) end
but now the run end with an error: `ERROR: MethodError: no method matching similar(::Tuple{Vector{Vector{Float64}}})
Closest candidates are: similar(::VSCodeServer.JuliaInterpreter.Compiled, ::Any) @ VSCodeServer ~/.vscode/extensions/julialang.language-julia-1.79.2/scripts/packages/JuliaInterpreter/src/types.jl:7 similar(::Type{SA}, ::Type{T}, ::StaticArraysCore.Size{S}) where {SA<:StaticArraysCore.SizedArray, T, S} @ StaticArrays ~/.julia/packages/StaticArrays/MSJcA/src/abstractarray.jl:135 similar(::Type{A}, ::Type{T}, ::StaticArraysCore.Size{S}) where {A<:Array, T, S} @ StaticArrays ~/.julia/packages/StaticArrays/MSJcA/src/abstractarray.jl:136 ...
Stacktrace: [1] similar(p::ModelingToolkit.MTKParameters{…}) @ Main ~/Documents/programming/NeurobloxSciML/friction.jl:23 [2] build_param_jac_config(alg::GaussAdjoint{…}, pf::Function, u::Vector{…}, p::ModelingToolkit.MTKParameters{…}) @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/derivative_wrappers.jl:1069 [3] SciMLSensitivity.GaussIntegrand(sol::ODESolution{…}, sensealg::GaussAdjoint{…}, checkpoints::Vector{…}, dgdp::Nothing) @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/gauss_adjoint.jl:440`
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:
prob.p.tunable ([[-0.42492082715034485, -0.6530026793479919, 0.6973646879196167, 0.058748241513967514, 0.1943797767162323, -0.05171017348766327, -0.08400869369506836, 0.2643197178840637, -0.5101186037063599, -0.7106459736824036 … -0.6610037684440613, 0.1330592930316925, 0.03395846113562584, 0.6738154292106628, -0.41331109404563904, 0.14921322464942932, 0.6819839477539062, 0.3220810294151306, -0.3412023186683655, 0.21613739430904388]],)
which is a Tuple{Vector{Vector{Float64}}}
I even tried to get around this by:
function Base.similar(p::ModelingToolkit.MTKParameters) return similar(p.tunable[1][1]) end
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: jacobian!(::AbstractMatrix{<:Number}, ::Any, ::AbstractArray{<:Number}, ::Union{Nothing, AbstractArray{<:Number}}, ::SciMLBase.AbstractOverloadingSensitivityAlgorithm, ::Any) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/derivative_wrappers.jl:147
Stacktrace: [1] vec_pjac!(out::RecursiveArrayTools.ArrayPartition{…}, λ::Vector{…}, y::Vector{…}, t::Float64, S::SciMLSensitivity.GaussIntegrand{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:469 [2] GaussIntegrand @ ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:519 [inlined] [3] (::SciMLSensitivity.var"#262#263"{…})(out::RecursiveArrayTools.ArrayPartition{…}, u::Vector{…}, t::Float64, integrator::OrdinaryDiffEq.ODEIntegrator{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:560 [4] (::DiffEqCallbacks.SavingIntegrandSumAffect{…})(integrator::OrdinaryDiffEq.ODEIntegrator{…})`
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 tests/lotka_volterra.jl
as in the OP
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 AbstractArray
interface for p
. Thus while
DifferentialEquations.jl can support any parameter struct type, usage
with ForwardDiffSensitivity requires that p
could be a valid
type for being the initial condition u0
of an array. This means that
many simple types, such as Tuple
s and NamedTuple
s, will work as
parameters in normal contexts but will fail during ForwardDiffSensitivity
construction. To work around this issue for complicated cases like nested structs,
look into defining p
using AbstractArray
libraries such as RecursiveArrayTools.jl
or ComponentArrays.jl.
Stacktrace: [1] _concrete_solve_adjoint(::ODEProblem{…}, ::Rodas4{…}, ::ForwardDiffSensitivity{…}, ::Vector{…}, ::ModelingToolkit.MTKParameters{…}, ::SciMLBase.ChainRulesOriginator; saveat::Vector{…}, kwargs::@Kwargs{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:788 [2] _concrete_solve_adjoint(::ODEProblem{…}, ::Rodas4{…}, ::Nothing, ::Vector{…}, ::ModelingToolkit.MTKParameters{…}, ::SciMLBase.ChainRulesOriginator; verbose::Bool, kwargs::@Kwargs{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:278 [3] _concrete_solve_adjoint @ ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:245 [inlined] [4] #_solve_adjoint#75 @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:1537 [inlined] [5] _solve_adjoint @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:1510 [inlined]
@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 verbose = false
to the solve
call.
└ @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:24
┌ 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 verbose = false
to the solve
call.
└ @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:67
┌ Warning: Reverse-Mode AD VJP choices all failed. Falling back to numerical VJPs └ @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/concrete_solve.jl:207 ERROR: MethodError: no method matching jacobian!(::Matrix{…}, ::SciMLBase.ParamJacobianWrapper{…}, ::ModelingToolkit.MTKParameters{…}, ::Vector{…}, ::GaussAdjoint{…}, ::FiniteDiff.JacobianCache{…})
Closest candidates are: jacobian!(::AbstractMatrix{<:Number}, ::Any, ::AbstractArray{<:Number}, ::Union{Nothing, AbstractArray{<:Number}}, ::SciMLBase.AbstractOverloadingSensitivityAlgorithm, ::Any) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/derivative_wrappers.jl:147
Stacktrace: [1] vec_pjac!(out::RecursiveArrayTools.ArrayPartition{…}, λ::Vector{…}, y::Vector{…}, t::Float64, S::SciMLSensitivity.GaussIntegrand{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:469 [2] GaussIntegrand @ ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:519 [inlined] [3] (::SciMLSensitivity.var"#262#263"{…})(out::RecursiveArrayTools.ArrayPartition{…}, u::Vector{…}, t::Float64, integrator::OrdinaryDiffEq.ODEIntegrator{…}) @ SciMLSensitivity ~/Documents/programming/SciMLSensitivity.jl/src/gauss_adjoint.jl:560 [4] (::DiffEqCallbacks.SavingIntegrandSumAffect{…})(integrator::OrdinaryDiffEq.ODEIntegrator{…}) @ DiffEqCallbacks ~/.julia/packages/DiffEqCallbacks/9fKPq/src/integrating_sum.jl:50
Why SDEProblem? it's an ODE?
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.
Describe the bug 🐞
UDE Training does not work with
AutoZygote
Expected behavior
It should work.
Minimal Reproducible Example 👇
The lotka volterra test in
tests/lotka_volterra.jl
Error & Stacktrace ⚠️