UDE Training does not work with `AutoZygote`

[sathvikbhagavan]

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 ⚠️

julia> res = solve(op, Adam(), maxiters = 10)
ERROR: MethodError: no method matching length(::ModelingToolkit.MTKParameters{Tuple{Vector{Vector{Float64}}}, Tuple{}, Tuple{Vector{Float64}}, Tuple{}, Tuple{Vector{DataType}}, Nothing, Nothing})

Closest candidates are:
  length(::LaTeXStrings.LaTeXString)
   @ LaTeXStrings ~/.julia/packages/LaTeXStrings/ZtSdh/src/LaTeXStrings.jl:115
  length(::SymbolicUtils.Code.Assignment)
   @ SymbolicUtils ~/.julia/packages/SymbolicUtils/c0xQb/src/utils.jl:225
  length(::CSTParser.EXPR)
   @ CSTParser ~/.julia/packages/CSTParser/mVfZt/src/spec.jl:278
  ...

Stacktrace:
  [1] automatic_sensealg_choice(prob::ODEProblem{…}, u0::Vector{…}, p::ModelingToolkit.MTKParameters{…}, verbose::Bool)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/rXkM4/src/concrete_solve.jl:84
  [2] _concrete_solve_adjoint(::ODEProblem{…}, ::Rodas4{…}, ::Nothing, ::Vector{…}, ::ModelingToolkit.MTKParameters{…}, ::SciMLBase.ChainRulesOriginator; verbose::Bool, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/rXkM4/src/concrete_solve.jl:218
  [3] _solve_adjoint(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::ModelingToolkit.MTKParameters{…}, originator::SciMLBase.ChainRulesOriginator, args::Rodas4{…}; merge_callbacks::Bool, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/O8cUq/src/solve.jl:1537
  [4] rrule(::typeof(DiffEqBase.solve_up), prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::ModelingToolkit.MTKParameters{…}, args::Rodas4{…}; kwargs::@Kwargs{…})
    @ DiffEqBaseChainRulesCoreExt ~/.julia/packages/DiffEqBase/O8cUq/ext/DiffEqBaseChainRulesCoreExt.jl:26
  [5] kwcall(::@NamedTuple{…}, ::typeof(ChainRulesCore.rrule), ::Zygote.ZygoteRuleConfig{…}, ::Function, ::ODEProblem{…}, ::Nothing, ::Vector{…}, ::ModelingToolkit.MTKParameters{…}, ::Rodas4{…})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/zgT0R/src/rules.jl:140
  [6] chain_rrule_kw
    @ ~/.julia/packages/Zygote/jxHJc/src/compiler/chainrules.jl:235 [inlined]
  [7] macro expansion
    @ ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0 [inlined]
  [8] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::Nothing, ::Vector{…}, ::ModelingToolkit.MTKParameters{…}, ::Rodas4{…})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:81
  [9] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [10] adjoint
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:203 [inlined]
 [11] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [12] #solve#51
    @ ~/.julia/packages/DiffEqBase/O8cUq/src/solve.jl:1003 [inlined]
 [13] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve#51", ::Nothing, ::Nothing, ::Nothing, ::Val{…}, ::@Kwargs{…}, ::typeof(solve), ::ODEProblem{…}, ::Rodas4{…})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [14] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [15] adjoint
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:203 [inlined]
 [16] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [17] solve
    @ ~/.julia/packages/DiffEqBase/O8cUq/src/solve.jl:993 [inlined]
 [18] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(solve), ::ODEProblem{…}, ::Rodas4{…})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [19] loss
    @ ./REPL[30]:5 [inlined]
 [20] _pullback(::Zygote.Context{…}, ::typeof(loss), ::Vector{…}, ::Tuple{…})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [21] _apply
    @ ./boot.jl:838 [inlined]
 [22] adjoint
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:203 [inlined]
 [23] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [24] OptimizationFunction
    @ ~/.julia/packages/SciMLBase/NjslX/src/scimlfunctions.jl:3649 [inlined]
 [25] _pullback(::Zygote.Context{…}, ::OptimizationFunction{…}, ::Vector{…}, ::Tuple{…})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [26] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [27] adjoint
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:203 [inlined]
 [28] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [29] #37
    @ ~/.julia/packages/OptimizationBase/rRpJs/ext/OptimizationZygoteExt.jl:90 [inlined]
 [30] _pullback(ctx::Zygote.Context{false}, f::OptimizationZygoteExt.var"#37#55"{OptimizationFunction{…}, OptimizationBase.ReInitCache{…}}, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [31] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [32] adjoint
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:203 [inlined]
 [33] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [34] #39
    @ ~/.julia/packages/OptimizationBase/rRpJs/ext/OptimizationZygoteExt.jl:93 [inlined]
 [35] _pullback(ctx::Zygote.Context{false}, f::OptimizationZygoteExt.var"#39#57"{Tuple{}, OptimizationZygoteExt.var"#37#55"{…}}, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [36] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface.jl:90
 [37] pullback
    @ ~/.julia/packages/Zygote/jxHJc/src/compiler/interface.jl:88 [inlined]
 [38] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface.jl:147
 [39] (::OptimizationZygoteExt.var"#38#56"{OptimizationZygoteExt.var"#37#55"{OptimizationFunction{…}, OptimizationBase.ReInitCache{…}}})(::Vector{Float64}, ::Vector{Float64})
    @ OptimizationZygoteExt ~/.julia/packages/OptimizationBase/rRpJs/ext/OptimizationZygoteExt.jl:93
 [40] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/AOkbT/src/OptimizationOptimisers.jl:68 [inlined]
 [41] macro expansion
    @ ~/.julia/packages/Optimization/5DEdF/src/utils.jl:32 [inlined]
 [42] __solve(cache::OptimizationCache{…})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/AOkbT/src/OptimizationOptimisers.jl:66
 [43] solve!(cache::OptimizationCache{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/NjslX/src/solve.jl:180
 [44] solve(::OptimizationProblem{…}, ::Adam; kwargs::@Kwargs{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/NjslX/src/solve.jl:96
 [45] top-level scope
    @ REPL[35]:1
Some type information was truncated. Use `show(err)` to see complete types.

[ChrisRackauckas]

This is SciML/SciMLSensitivity.jl#1010

[hstrey]

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.`

[ChrisRackauckas]

@DhairyaLGandhi can you work on updating the example here to use AutoZygote?

[DhairyaLGandhi]

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.

[hstrey]

@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

[DhairyaLGandhi]

Can you try against https://github.com/SciML/SciMLSensitivity.jl/tree/dg/fwd

[hstrey]

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}}}

[hstrey]

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{…})`

[DhairyaLGandhi]

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

[ChrisRackauckas]

Why SDEProblem? it's an ODE?

[hstrey]

@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 Tuples and NamedTuples, 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]

[hstrey]

@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

[DhairyaLGandhi]

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.