Zygote and SDE failure: cannot differentiate through Lux NN

evolbio · July 13, 2022, 1:58pm

I am optimizing the parameters of a differential equation with a Lux.jl neural network (NN) component, using Optimization.jl. The code for the ODE part is below.

This works well when there is no stochasticity, ie, an ODE with Tsit5() solver, with either Zygote or ForwardDiff for AD.

With stochasticity, as SDE with ISSEM() solver, this works well but slowly with ForwardDiff. It fails with Zygote.

Failure for SDE and Zygote at line:
f_val = tf(u_p,re(p_nn),state)[1]

where (w/Lux)

tf = Chain(Dense(S.n => 10*S.n, tanh), Dense(10*S.n => S.n, sigmoid))
ps, state = Lux.setup(Random.default_rng(), tf)
p_node, re = destructure(ps)

In essence, p_node = p_nn is vector for optimization, and ps = re(p_nn) is structured parameter form for feeding into Lux NN.

I have tried a variety of SDE solvers. They give a variety error messages but none work.

ODE code:

function node!(du, u, p, t, S, tf, re, state, G)
	n = S.n
	u_m = @view u[1:n]			# mRNA
	u_p = @view u[n+1:2n]		# protein
	p_rates = @view p[1:4n]
	p_nn = @view p[4n+1:end]
	
	# first 4n parameters are rates
	pp = linear_sigmoid.(p_rates, S.d, S.k1, S.k2)	# normalizes on [0,1] w/linear_sigmoid 
	# set min on rates m_a, m_d, p_a, p_d, causes top to be p_max + p_min
	ppp = (pp .* S.p_mult) .+ S.p_min

	m_a = @view ppp[1:n]
	m_d = @view ppp[n+1:2n]
	p_a = @view ppp[2n+1:3n]
	p_d = @view ppp[3n+1:4n]
	
	# remainder of p is for NN
	f_val = tf(u_p,re(p_nn),state)[1]
	du[1:n] = m_a .* f_val .- m_d .* u_m		# mRNA level
	du[n+1:2n] = p_a .* u_m .- p_d .* u_p		# protein level
	light = G.circadian_val(G,t)			# noisy circadian input
	# fast extra production rate by post-translation modification or allostery
	du[n+2] += S.light_prod_rate * intensity(light)
end

Here is part of the error message. I vaguely understand the issue with dual numbers and the type issue, but not enough to be able to resolve the problem or even to decide if that would be sufficient.

ERROR: MethodError: ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, Nothing}(::ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}) is ambiguous. Candidates:
  (T::Type{<:Real})(x::ForwardDiff.Dual) in Tracker at /opt/julia/packages/Tracker/9xWLl/src/lib/real.jl:110
  ReverseDiff.TrackedReal{V, D, O}(value) where {V, D, O} in ReverseDiff at /opt/julia/packages/ReverseDiff/5MMPp/src/tracked.jl:56
Possible fix, define
  ReverseDiff.TrackedReal{V, D, O}(::ForwardDiff.Dual) where {V, D, O}
Stacktrace:
  [1] convert(#unused#::Type{ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, Nothing}}, x::Int64)
    @ ReverseDiff /opt/julia/packages/ReverseDiff/5MMPp/src/tracked.jl:264
  [2] oftype(x::ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, Nothing}, y::Int64)
    @ Base ./essentials.jl:391
  [3] zero(x::ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, Nothing})
    @ Base ./number.jl:296
  [4] generic_matvecmul!(C::Vector{ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, Nothing}}, tA::Char, A::Matrix{ReverseDiff.TrackedReal{Float64, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ReverseDiff.TrackedArray{Float64, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, 1, Vector{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}}}}, B::SubArray{ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ReverseDiff.TrackedArray{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, 1, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}}}, 1, ReverseDiff.TrackedArray{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, 1, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}}, Tuple{UnitRange{Int64}}, true}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /Library/Local/Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:805
  [5] mul!
    @ /Library/Local/Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:81 [inlined]
  [6] mul!
    @ /Library/Local/Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:276 [inlined]
  [7] *
    @ /Library/Local/Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:56 [inlined]
  [8] Dense
    @ /opt/julia/packages/Lux/lEqCI/src/layers/basic.jl:627 [inlined]
  [9] macro expansion
    @ /opt/julia/packages/Lux/lEqCI/src/layers/basic.jl:0 [inlined]
 [10] applychain(layers::NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(NNlib.tanh_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}, x::SubArray{ReverseDiff.TrackedReal{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ReverseDiff.TrackedArray{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, 1, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}}}, 1, ReverseDiff.TrackedArray{ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}, 
MORE FOLLOWS ...

avikpal · July 14, 2022, 5:43am

Any particular reason to use restructure instead of ComponentArrays? Also, you seem to be using ReverseDiff, not Zygote, which is not very well tested with Lux.

ChrisRackauckas · July 14, 2022, 6:56am

It’s doing ReverseDiff for the in-place mutation. @evolbio change this to the out of place form and you should be fine.

evolbio · July 14, 2022, 2:54pm

Thanks, @ChrisRackauckas . I changed to out of place for both drift and diffusion functions and avoided all mutation. As before, ODE works fine. For SDE, the code has switched to Zygote (thanks @avikpal), but throws error
UndefVarError: J not defined
triggered by this line
solve(prob, S.solver, u0=u_init, p=p[S.ddim+1:end])
for prob as SDEProblem with no special args, and solver as ISSEM(). Seems like something in StochasticDiffEq but I was not able to find the issue. Any ideas?

ERROR: UndefVarError: J not defined
Stacktrace:
  [1] perform_step!(integrator::StochasticDiffEq.SDEIntegrator{StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}, false, Vector{Float64}, Float64, Float64, Float64, Vector{Float64}, Float64, Float64, Float64, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, Vector{Float64}, SciMLBase.RODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:maxiters, :reltol, :abstol, :callback), Tuple{Float64, Float64, Float64, SciMLBase.CallbackSet{Tuple{}, Tuple{}}}}}, Nothing}, StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}, StochasticDiffEq.LinearInterpolationData{Vector{Vector{Float64}}, Vector{Float64}}, DiffEqBase.DEStats}, StochasticDiffEq.ISSEMConstantCache{OrdinaryDiffEq.NLSolver{DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, false, Vector{Float64}, Float64, Nothing, Float64, OrdinaryDiffEq.NLNewtonConstantCache{Float64, Float64, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, SciMLBase.UDerivativeWrapper{SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Float64, Vector{Float64}}}}}, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Nothing, StochasticDiffEq.SDEOptions{Float64, Float64, OrdinaryDiffEq.PIController{Float64}, typeof(DiffEqBase.ODE_DEFAULT_NORM), Bool, SciMLBase.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, Nothing, Float64, Float64, Float64, Float64, Tuple{}, Tuple{}, Tuple{}}, Nothing, Float64, Nothing, Nothing}, cache::StochasticDiffEq.ISSEMConstantCache{OrdinaryDiffEq.NLSolver{DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, false, Vector{Float64}, Float64, Nothing, Float64, OrdinaryDiffEq.NLNewtonConstantCache{Float64, Float64, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, SciMLBase.UDerivativeWrapper{SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Float64, Vector{Float64}}}}})
    @ StochasticDiffEq /opt/julia/packages/StochasticDiffEq/Qloc0/src/perform_step/implicit_split_step.jl:60
  [2] solve!(integrator::StochasticDiffEq.SDEIntegrator{StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}, false, Vector{Float64}, Float64, Float64, Float64, Vector{Float64}, Float64, Float64, Float64, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, Vector{Float64}, SciMLBase.RODESolution{Float64, 2, Vector{Vector{Float64}}, Nothing, Nothing, Vector{Float64}, DiffEqNoiseProcess.NoiseProcess{Float64, 2, Float64, Vector{Float64}, Nothing, Nothing, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Vector{Float64}, Nothing}, false}, DiffEqNoiseProcess.RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:maxiters, :reltol, :abstol, :callback), Tuple{Float64, Float64, Float64, SciMLBase.CallbackSet{Tuple{}, Tuple{}}}}}, Nothing}, StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}, StochasticDiffEq.LinearInterpolationData{Vector{Vector{Float64}}, Vector{Float64}}, DiffEqBase.DEStats}, StochasticDiffEq.ISSEMConstantCache{OrdinaryDiffEq.NLSolver{DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, false, Vector{Float64}, Float64, Nothing, Float64, OrdinaryDiffEq.NLNewtonConstantCache{Float64, Float64, Matrix{Float64}, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, SciMLBase.UDerivativeWrapper{SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Float64, Vector{Float64}}}}}, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Nothing, StochasticDiffEq.SDEOptions{Float64, Float64, OrdinaryDiffEq.PIController{Float64}, typeof(DiffEqBase.ODE_DEFAULT_NORM), Bool, SciMLBase.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, Nothing, Float64, Float64, Float64, Float64, Tuple{}, Tuple{}, Tuple{}}, Nothing, Float64, Nothing, Nothing})
    @ StochasticDiffEq /opt/julia/packages/StochasticDiffEq/Qloc0/src/solve.jl:611
  [3] #__solve#110
    @ /opt/julia/packages/StochasticDiffEq/Qloc0/src/solve.jl:7 [inlined]
  [4] solve_call(_prob::SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:maxiters, :reltol, :abstol, :callback), Tuple{Float64, Float64, Float64, SciMLBase.CallbackSet{Tuple{}, Tuple{}}}}}, Nothing}, args::StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}; merge_callbacks::Bool, kwargshandle::DiffEqBase.KeywordArgError, kwargs::Base.Pairs{Symbol, Real, NTuple{7, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :verbose, :reltol, :abstol, :maxiters), Tuple{Bool, Bool, Bool, Bool, Float64, Float64, Float64}}})
    @ DiffEqBase /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:433
  [5] #solve_up#34
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:772 [inlined]
  [6] #solve#33
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:756 [inlined]
  [7] _concrete_solve_adjoint(::SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{5, Symbol}, NamedTuple{(:reltol, :abstol, :callback, :saveat, :maxiters), Tuple{Float64, Float64, Nothing, Vector{Float64}, Float64}}}, Nothing}, ::StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive}, ::DiffEqSensitivity.InterpolatingAdjoint{0, true, Val{:central}, DiffEqSensitivity.ZygoteVJP}, ::Vector{Float64}, ::Vector{Float64}, ::SciMLBase.ChainRulesOriginator; save_start::Bool, save_end::Bool, saveat::Vector{Float64}, save_idxs::Nothing, kwargs::Base.Pairs{Symbol, Union{Nothing, Real}, NTuple{5, Symbol}, NamedTuple{(:verbose, :reltol, :abstol, :callback, :maxiters), Tuple{Bool, Float64, Float64, Nothing, Float64}}})
    @ DiffEqSensitivity /opt/julia/packages/DiffEqSensitivity/Pn9H4/src/concrete_solve.jl:164
  [8] #_concrete_solve_adjoint#224
    @ /opt/julia/packages/DiffEqSensitivity/Pn9H4/src/concrete_solve.jl:102 [inlined]
  [9] #_solve_adjoint#57
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:1142 [inlined]
 [10] _solve_adjoint
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:1111 [inlined]
 [11] #rrule#55
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:1095 [inlined]
 [12] rrule
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:1091 [inlined]
 [13] rrule
    @ /opt/julia/packages/ChainRulesCore/16PWJ/src/rules.jl:134 [inlined]
 [14] chain_rrule
    @ /opt/julia/packages/Zygote/IoW2g/src/compiler/chainrules.jl:217 [inlined]
 [15] macro expansion
    @ /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0 [inlined]
 [16] _pullback
    @ /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:9 [inlined]
 [17] _apply
    @ ./boot.jl:816 [inlined]
 [18] adjoint
    @ /opt/julia/packages/Zygote/IoW2g/src/lib/lib.jl:204 [inlined]
 [19] _pullback
    @ /opt/julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [20] _pullback
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:756 [inlined]
 [21] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve#33", ::Nothing, ::Vector{Float64}, ::Vector{Float64}, ::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(CommonSolve.solve), ::SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{5, Symbol}, NamedTuple{(:reltol, :abstol, :callback, :saveat, :maxiters), Tuple{Float64, Float64, Nothing, Vector{Float64}, Float64}}}, Nothing}, ::StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [22] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [23] adjoint
    @ /opt/julia/packages/Zygote/IoW2g/src/lib/lib.jl:204 [inlined]
 [24] _pullback
    @ /opt/julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [25] _pullback
    @ /opt/julia/packages/DiffEqBase/aAyno/src/solve.jl:749 [inlined]
 [26] _pullback(::Zygote.Context, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:u0, :p), Tuple{Vector{Float64}, Vector{Float64}}}, ::typeof(CommonSolve.solve), ::SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{5, Symbol}, NamedTuple{(:reltol, :abstol, :callback, :saveat, :maxiters), Tuple{Float64, Float64, Nothing, Vector{Float64}, Float64}}}, Nothing}, ::StochasticDiffEq.ISSEM{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:central}, true, nothing, DiffEqBase.NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [27] _pullback
    @ ~/sim/zzOtherLang/julia/projects/OptTF/src/OptTF.jl:241 [inlined]
 [28] _pullback(::Zygote.Context, ::OptTF.var"#predict_ode_dummy#13"{Settings, Float64, Float64, Float64}, ::Vector{Float64}, ::SciMLBase.SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, Vector{Float64}, Nothing, SciMLBase.SDEFunction{false, OptTF.var"#24#26"{Settings, loss_args, OptTF.OptTF_settings.OptTF_data.Circadian}, typeof(OptTF.ode_noise), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, typeof(OptTF.ode_noise), Base.Pairs{Symbol, Any, NTuple{5, Symbol}, NamedTuple{(:reltol, :abstol, :callback, :saveat, :maxiters), Tuple{Float64, Float64, Nothing, Vector{Float64}, Float64}}}, Nothing})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [29] _pullback
    @ ~/sim/zzOtherLang/julia/projects/OptTF/src/OptTF.jl:320 [inlined]
 [30] _pullback(::Zygote.Context, ::typeof(loss), ::Vector{Float64}, ::Settings, ::loss_args)
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [31] _pullback
    @ ~/sim/zzOtherLang/julia/projects/OptTF/src/OptTF.jl:351 [inlined]
 [32] _pullback(::Zygote.Context, ::OptTF.var"#35#36"{Settings, loss_args}, ::Vector{Float64}, ::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [33] _apply
    @ ./boot.jl:816 [inlined]
 [34] adjoint
    @ /opt/julia/packages/Zygote/IoW2g/src/lib/lib.jl:204 [inlined]
 [35] _pullback
    @ /opt/julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [36] _pullback
    @ /opt/julia/packages/SciMLBase/TqBga/src/scimlfunctions.jl:2891 [inlined]
 [37] _pullback(::Zygote.Context, ::SciMLBase.OptimizationFunction{true, Optimization.AutoZygote, OptTF.var"#35#36"{Settings, loss_args}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ::Vector{Float64}, ::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [38] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [39] adjoint
    @ /opt/julia/packages/Zygote/IoW2g/src/lib/lib.jl:204 [inlined]
 [40] _pullback
    @ /opt/julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [41] _pullback
    @ /opt/julia/packages/Optimization/i9NGR/src/function/zygote.jl:30 [inlined]
 [42] _pullback(ctx::Zygote.Context, f::Optimization.var"#128#138"{SciMLBase.OptimizationFunction{true, Optimization.AutoZygote, OptTF.var"#35#36"{Settings, loss_args}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Vector{Float64}}, args::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [43] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [44] adjoint
    @ /opt/julia/packages/Zygote/IoW2g/src/lib/lib.jl:204 [inlined]
 [45] _pullback
    @ /opt/julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [46] _pullback
    @ /opt/julia/packages/Optimization/i9NGR/src/function/zygote.jl:32 [inlined]
 [47] _pullback(ctx::Zygote.Context, f::Optimization.var"#131#141"{Tuple{}, Optimization.var"#128#138"{SciMLBase.OptimizationFunction{true, Optimization.AutoZygote, OptTF.var"#35#36"{Settings, loss_args}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Vector{Float64}}}, args::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0
 [48] _pullback(f::Function, args::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface.jl:34
 [49] pullback(f::Function, args::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface.jl:40
 [50] gradient(f::Function, args::Vector{Float64})
    @ Zygote /opt/julia/packages/Zygote/IoW2g/src/compiler/interface.jl:75
 [51] (::Optimization.var"#129#139"{Optimization.var"#128#138"{SciMLBase.OptimizationFunction{true, Optimization.AutoZygote, OptTF.var"#35#36"{Settings, loss_args}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Vector{Float64}}})(::Vector{Float64}, ::Vector{Float64})
    @ Optimization /opt/julia/packages/Optimization/i9NGR/src/function/zygote.jl:32
 [52] macro expansion
    @ /opt/julia/packages/OptimizationOptimisers/XLPqT/src/OptimizationOptimisers.jl:35 [inlined]
 [53] macro expansion
    @ /opt/julia/packages/Optimization/i9NGR/src/utils.jl:35 [inlined]
 [54] __solve(prob::SciMLBase.OptimizationProblem{true, SciMLBase.OptimizationFunction{true, Optimization.AutoZygote, OptTF.var"#35#36"{Settings, loss_args}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, Vector{Float64}, Nothing, Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, opt::Optimisers.Adam{Float64}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers /opt/julia/packages/OptimizationOptimisers/XLPqT/src/OptimizationOptimisers.jl:33
 [55] #solve#516
    @ /opt/julia/packages/SciMLBase/TqBga/src/solve.jl:71 [inlined]
 [56] fit_diffeq(S::Settings; noise::Float64, new_rseed::Bool, init_on::Bool, offset::Bool, noise_wait::Float64, hill_k_init::Float64)
    @ OptTF ~/sim/zzOtherLang/julia/projects/OptTF/src/OptTF.jl:425
 [57] top-level scope
    @ REPL[15]:1

ChrisRackauckas · July 14, 2022, 6:04pm

Odd, I need to take a look at that. Open an issue?

evolbio · July 14, 2022, 10:16pm

@ChrisRackauckas: opened issue, thanks

https://github.com/SciML/StochasticDiffEq.jl/issues/481#issue-1305341306