Sensitivity analysis on Non-diagonal noise SDEs

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

``````

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}}})
[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}}})
[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
[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
@ 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
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?
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?