Sensitivity analysis on Non-diagonal noise SDEs

Specific Domains Modelling & Simulations
, ,
1

I am trying to do a parameter optimization of an SDE with non-diagonal constant noise. A simplified version of the implementation is (SDE reduced to the simplest possible version which still have the problem)

using StochasticDiffEq, DiffEqSensitivity
using Flux

a(du,u,_p,t) = (du .= -u) 

function b(du,u,_p,t)
    KR, KI = _p[1:2]

    du[1,1] = KR
    du[2,1] = KI
end

p = [1.,0.]
ps = params(p)
prob = SDEProblem(a,b,[0.,0.],(0.0,1.0),p,noise_rate_prototype=eltype(p).(zeros(2,2)))

function loss()
    sol = solve(prob, EM(),dt=1e-2)
    return sum(Array(sol))
end

Flux.train!(loss, ps, Iterators.repeated((), 2), ADAM(0.01))

where I get the error

ERROR: LoadError: TypeError: in typeassert, expected Float64, got a value of type ForwardDiff.Dual{Nothing, Float64, 2}
Stacktrace:
  [1] setindex!
    @ ./array.jl:841 [inlined]
  [2] b(du::Matrix{Float64}, u::Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, _p::Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, t::Float64)
    @ Main ~/Developer/Stochastic/MLKernel/testing.jl:9
  [3] perform_step!(integrator::StochasticDiffEq.SDEIntegrator{EM{true}, true, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Float64, Float64, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Float64, Float64, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE), true, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, RODESolution{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, 2, Vector{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}}, Nothing, Nothing, Vector{Float64}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE), true, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEProblem{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Tuple{Float64, Float64}, true, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Nothing, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Matrix{Float64}}, EM{true}, StochasticDiffEq.LinearInterpolationData{Vector{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}}, Vector{Float64}}, DiffEqBase.DEStats}, StochasticDiffEq.EMCache{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Matrix{Float64}}, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Nothing, StochasticDiffEq.SDEOptions{Float64, Float64, OrdinaryDiffEq.PIController{Float64}, typeof(DiffEqBase.ODE_DEFAULT_NORM), Nothing, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, Nothing, UnitRange{Int64}, Int64, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Tuple{}, Vector{Float64}, Tuple{}}, Nothing, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Nothing, Nothing}, cache::StochasticDiffEq.EMCache{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Matrix{Float64}})
    @ StochasticDiffEq ~/.julia/packages/StochasticDiffEq/PtS8q/src/perform_step/low_order.jl:40
  [4] solve!(integrator::StochasticDiffEq.SDEIntegrator{EM{true}, true, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Float64, Float64, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Float64, Float64, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE), true, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, RODESolution{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, 2, Vector{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}}, Nothing, Nothing, Vector{Float64}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE), true, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, true}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEProblem{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Tuple{Float64, Float64}, true, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Nothing, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Matrix{Float64}}, EM{true}, StochasticDiffEq.LinearInterpolationData{Vector{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}}, Vector{Float64}}, DiffEqBase.DEStats}, StochasticDiffEq.EMCache{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Matrix{Float64}}, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Nothing, StochasticDiffEq.SDEOptions{Float64, Float64, OrdinaryDiffEq.PIController{Float64}, typeof(DiffEqBase.ODE_DEFAULT_NORM), Nothing, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, Nothing, UnitRange{Int64}, Int64, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Tuple{}, Vector{Float64}, Tuple{}}, Nothing, ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}, Nothing, Nothing})
    @ StochasticDiffEq ~/.julia/packages/StochasticDiffEq/PtS8q/src/solve.jl:607
  [5] #__solve#100
    @ ~/.julia/packages/StochasticDiffEq/PtS8q/src/solve.jl:7 [inlined]
  [6] #solve_call#42
    @ ~/.julia/packages/DiffEqBase/WKucm/src/solve.jl:61 [inlined]
  [7] solve_up(prob::SDEProblem{Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Tuple{Float64, Float64}, true, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, Nothing, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Matrix{Float64}}, sensealg::Nothing, u0::Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, p::Vector{ForwardDiff.Dual{ForwardDiff.Tag{SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}}, args::EM{true}; kwargs::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :save_idxs, :dt), Tuple{Vector{Float64}, UnitRange{Int64}, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/WKucm/src/solve.jl:87
  [8] #solve#43
    @ ~/.julia/packages/DiffEqBase/WKucm/src/solve.jl:73 [inlined]
  [9] (::DiffEqSensitivity.var"#225#231"{0, FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Vector{Float64}, Base.Iterators.Pairs{Symbol, Float64, Tuple{Symbol}, NamedTuple{(:dt,), Tuple{Float64}}}, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, Nothing, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Matrix{Float64}}, EM{true}, ForwardDiffSensitivity{0, nothing}, Vector{Float64}, Vector{Float64}, Tuple{}, UnitRange{Int64}})()
    @ DiffEqSensitivity ~/.julia/packages/DiffEqSensitivity/uLSNO/src/concrete_solve.jl:467
 [10] unthunk
    @ ~/.julia/packages/ChainRulesCore/BYuIz/src/differentials/thunks.jl:192 [inlined]
 [11] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/0da6K/src/compiler/chainrules.jl:55 [inlined]
 [12] map
    @ ./tuple.jl:215 [inlined]
 [13] map (repeats 3 times)
    @ ./tuple.jl:216 [inlined]
 [14] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/0da6K/src/compiler/chainrules.jl:56 [inlined]
 [15] ZBack
    @ ~/.julia/packages/Zygote/0da6K/src/compiler/chainrules.jl:91 [inlined]
 [16] kw_zpullback
    @ ~/.julia/packages/Zygote/0da6K/src/compiler/chainrules.jl:117 [inlined]
 [17] #193
    @ ~/.julia/packages/Zygote/0da6K/src/lib/lib.jl:203 [inlined]
 [18] (::Zygote.var"#1746#back#195"{Zygote.var"#193#194"{Tuple{NTuple{6, Nothing}, Tuple{Nothing}}, Zygote.var"#kw_zpullback#40"{DiffEqSensitivity.var"#forward_sensitivity_backpass#228"{0, Vector{Float64}, Base.Iterators.Pairs{Symbol, Float64, Tuple{Symbol}, NamedTuple{(:dt,), Tuple{Float64}}}, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, Nothing, SDEFunction{true, typeof(a), typeof(b), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, typeof(b), Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Matrix{Float64}}, EM{true}, ForwardDiffSensitivity{0, nothing}, Vector{Float64}, Vector{Float64}, Tuple{}, UnitRange{Int64}}}}})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
 [19] Pullback
    @ ~/.julia/packages/DiffEqBase/WKucm/src/solve.jl:73 [inlined]
 [20] (::typeof(∂(#solve#43)))(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface2.jl:0
 [21] (::Zygote.var"#193#194"{Tuple{NTuple{6, Nothing}, Tuple{Nothing}}, typeof(∂(#solve#43))})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/lib/lib.jl:203
 [22] (::Zygote.var"#1746#back#195"{Zygote.var"#193#194"{Tuple{NTuple{6, Nothing}, Tuple{Nothing}}, typeof(∂(#solve#43))}})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
 [23] Pullback
    @ ~/.julia/packages/DiffEqBase/WKucm/src/solve.jl:68 [inlined]
 [24] (::typeof(∂(solve##kw)))(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface2.jl:0
 [25] Pullback
    @ ~/Developer/Stochastic/MLKernel/testing.jl:52 [inlined]
 [26] (::typeof(∂(loss)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface2.jl:0
 [27] #193
    @ ~/.julia/packages/Zygote/0da6K/src/lib/lib.jl:203 [inlined]
 [28] #1746#back
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
 [29] Pullback
    @ ~/.julia/packages/Flux/0c9kI/src/optimise/train.jl:102 [inlined]
 [30] (::Zygote.var"#74#75"{Zygote.Params, typeof(∂(#39)), Zygote.Context})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface.jl:255
 [31] gradient(f::Function, args::Zygote.Params)
    @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface.jl:59
 [32] macro expansion
    @ ~/.julia/packages/Flux/0c9kI/src/optimise/train.jl:101 [inlined]
 [33] macro expansion
    @ ~/.julia/packages/Juno/n6wyj/src/progress.jl:134 [inlined]
 [34] train!(loss::Function, ps::Zygote.Params, data::Base.Iterators.Take{Base.Iterators.Repeated{Tuple{}}}, opt::ADAM; cb::Flux.Optimise.var"#40#46")
    @ Flux.Optimise ~/.julia/packages/Flux/0c9kI/src/optimise/train.jl:99
 [35] train!(loss::Function, ps::Zygote.Params, data::Base.Iterators.Take{Base.Iterators.Repeated{Tuple{}}}, opt::ADAM)
    @ Flux.Optimise ~/.julia/packages/Flux/0c9kI/src/optimise/train.jl:97
 [36] top-level scope
    @ ~/Developer/Stochastic/MLKernel/testing.jl:56

It looks like it the pre-allocated noise coefficient matrix du in b(du,u,_p,t) is not compatible with Dual numbers for the back propagation, but I can’t seem to find a way to redefine it. I tried, as can be seen in the code, to use the noise_rate_prototype=eltype(p).(zeros(2,2)), but this does not seem to change anything. Is there a way to do this which I have overlooked?

2

This looks like an error that would happen when AutoForwardDiff is mixed with non-diagonal noise SDEs. Could you open an issue on DiffEqSensitivity.jl?