Skip to content

Error on differentiation with a terminate! callback #720

@axsk

Description

@axsk

Trying to autodiff (Zygote.jl) through a neural SDE (StochasticDiffEq.jl, Lux.jl) with a ContinuousCallback triggering a terminate! leads to the following error

julia> test()
ERROR: Event was triggered but no corresponding event in ContinuousCallback was found. Please report this error.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] get_indx(cb::ContinuousCallback{var"#9#10"{Int64, Int64}, SciMLSensitivity.TrackedAffect{Float64, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, typeof(terminate!), SciMLSensitivity.ImplicitCorrection{Vector{Float64}, Vector{Float64}, SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, ForwardDiff.DerivativeConfig{ForwardDiff.Tag{SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 1}}}, ForwardDiff.GradientConfig{ForwardDiff.Tag{SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 2}}}, var"#9#10"{Int64, Int64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Base.RefValue{Int64}}, Int64}, SciMLSensitivity.TrackedAffect{Float64, Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, typeof(terminate!), SciMLSensitivity.ImplicitCorrection{Vector{Float64}, Vector{Float64}, SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, ForwardDiff.DerivativeConfig{ForwardDiff.Tag{SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SciMLSensitivity.ConditionTimeWrapper{var"#9#10"{Int64, Int64}, Vector{Float64}, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 1}}}, ForwardDiff.GradientConfig{ForwardDiff.Tag{SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{SciMLSensitivity.ConditionUWrapper{var"#9#10"{Int64, Int64}, Float64, SciMLSensitivity.FakeIntegrator{Vector{Float64}, ComponentArrays.ComponentVector{Float32, Vector{Float32}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:20, Axis(weight = ViewAxis(1:10, ShapedAxis((10, 1), NamedTuple())), bias = ViewAxis(11:20, ShapedAxis((10, 1), NamedTuple())))), layer_2 = ViewAxis(21:31, Axis(weight = ViewAxis(1:10, ShapedAxis((1, 10), NamedTuple())), bias = ViewAxis(11:11, ShapedAxis((1, 1), NamedTuple())))))}}}, Float64, Float64}}, Float64}, Float64, 2}}}, var"#9#10"{Int64, Int64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Base.RefValue{Int64}}, Int64}, typeof(SciMLBase.INITIALIZE_DEFAULT), typeof(SciMLBase.FINALIZE_DEFAULT), Float64, Int64, Rational{Int64}, Nothing, Int64}, t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/0ZY5u/src/callback_tracking.jl:365

The full stacktrace is here. The full code is here (call test())

The callback is

termination(ub, lb) = ContinuousCallback((u,t,int)->(u[1]-lb) * (ub-u[1]), terminate!)
prob = SDEProblem(_drift, _noise, xy0, T, p, noise_rate_prototype = n0, callback=termination(1,-1))

and I solve with

function msolve(prob; ps=prob.p, dt=0.01, salg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(), noisemixing=true))
    s = solve(prob, EM(), sensealg=salg, dt=dt, p=ps)
end

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions