Using callback during adjoint method for NeuralODE

xk_y · October 20, 2021, 12:33am

I am training a neural ode with a callback function to add a certain value at every integration timestep. The forward pass works but an error occur during backpropagation. I was wondering if the adjoint method supports a user-defined callback function. My code is following:

using Flux, DiffEqFlux, OrdinaryDiffEq

training_data = Flux.glorot_uniform(3,5,6)
callback_data = ones(3,5,6)*1e-2
train_dataloader = Flux.Data.DataLoader((x = training_data, c=callback_data),batchsize = 2)
timelength=80 
timesteps = LinRange(0, timelength, 6)

model =  Flux.Chain(
    Flux.Dense(3, 8, gelu),
    Flux.Dense(8, 8, gelu),
    Flux.Dense(8, 3))

mutable struct Affect{T}
  callback_data::T
end

function callback_(callback_data)
    affect! = Affect(callback_data)
    condition(u,t,integrator) = integrator.t > 0
    DiscreteCallback(condition,affect!,save_positions=(false,false))
end

function (cb::Affect)(integrator)
    integrator.u = integrator.u + cb.callback_data[:,Int(integrator.t÷20+1),:] .* (integrator.t-integrator.tprev)
end

basic_tgrad(u,p,t) = zero(u)
struct _NeuralODE{M,P,RE,T,A,K} 
    model::M
    p::P
    re::RE
    tspan::T
    args::A
    kwargs::K

    function _NeuralODE(model,tspan,args...;p = nothing,kwargs...)
        _p,re = Flux.destructure(model)
        if p === nothing
            p = _p
        end
        new{typeof(model),typeof(p),typeof(re),
            typeof(tspan),typeof(args),typeof(kwargs)}(
            model,p,re,tspan,args,kwargs)
    end
end

function (n::_NeuralODE)(x,callback_data,p=n.p)
    dudt_(u,p,t) = n.re(p)(u) 
    ff = ODEFunction{false}(dudt_,tgrad=basic_tgrad)
    prob = ODEProblem{false}(ff,x,getfield(n,:tspan),p)
    sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP())
    solve(prob,n.args...;sense=sense,callback = callback_(callback_data),saveat=20)
end

Flux.trainable(m::_NeuralODE) = (m.p,)

prob_neuralode = _NeuralODE(model,(0.0,timelength),OrdinaryDiffEq.Tsit5())

function predict(x,callback_data)
    pred = Array(prob_neuralode(x[:,1,:],callback_data))
    pred = permutedims(pred, [1,3,2])
end

function loss(x,callback_data)
    Flux.mse(predict(x,callback_data), x)
end
loss(train_dataloader.data.x,train_dataloader.data.c) # Forward pass is successful

opt = ADAM(3e-4)
epochs=2

function train!(loss, ps, data, opt)
  ps = Params(ps)
  for (x,c) in data
      gs = gradient(ps) do
        loss(x,c)
      end
      Flux.update!(opt, ps, gs)
  end
end
@Flux.epochs epochs train!(loss, Flux.params(prob_neuralode), 
                                train_dataloader, opt) 
# Back propagation failed

MethodError: no method matching Vector{Float32}(::Matrix{Float32})
Closest candidates are:
  Array{T, N}(::AbstractArray{S, N}) where {T, N, S} at array.jl:540
  Vector{T}() where T at boot.jl:467
  Array{T, N}(::Core.Compiler.BitArray{N}) where {T, N} at bitarray.jl:494

When training on GPU with CUDA.jl, the forward pass also works, but BP is failed due to the following error:

MethodError: no method matching CuArray{Float32, 1}(::UndefInitializer, ::Tuple{Int64, Int64})
Closest candidates are:
  CuArray{T, N}(::UndefInitializer, ::Tuple{Vararg{Int64, N}}) where {T, N} at /projects/julia_1.6.1/packages/CUDA/Ozu5O/src/array.jl:19

ChrisRackauckas · October 20, 2021, 11:12am

Post the stacktraces. What was posted cuts off all of the debugging information.

xk_y · October 20, 2021, 1:03pm

MethodError: no method matching Vector{Float32}(::Matrix{Float32})
Closest candidates are:
  Array{T, N}(::AbstractArray{S, N}) where {T, N, S} at array.jl:540
  Vector{T}() where T at boot.jl:467
  Vector{T}(::Core.Compiler.AbstractRange{T}) where T at range.jl:1042
  ...

Stacktrace:
  [1] convert(#unused#::Type{Vector{Float32}}, a::Matrix{Float32})
    @ Base .\array.jl:532
  [2] push!(a::Vector{Vector{Float32}}, item::Matrix{Float32})
    @ Base .\array.jl:928
  [3] (::DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}})(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Matrix{Float32}, Nothing, Float64, Vector{Float32}, Float64, Float32, Float32, Float64, Vector{Matrix{Float32}}, ODESolution{Float32, 3, Vector{Matrix{Float32}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Symbol, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}}}}, SciMLBase.StandardODEProblem}, Tsit5, OrdinaryDiffEq.InterpolationData{ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Matrix{Float32}}, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}, OrdinaryDiffEq.DEOptions{Float32, Float32, Float32, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Matrix{Float32}, Float32, Nothing, OrdinaryDiffEq.DefaultInit})
    @ DiffEqSensitivity C:\Users\Administrator\.julia\packages\DiffEqSensitivity\mhKek\src\callback_tracking.jl:28
  [4] apply_discrete_callback!
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\callbacks.jl:830 [inlined]
  [5] handle_callbacks!(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Matrix{Float32}, Nothing, Float64, Vector{Float32}, Float64, Float32, Float32, Float64, Vector{Matrix{Float32}}, ODESolution{Float32, 3, Vector{Matrix{Float32}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Symbol, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}}}}, SciMLBase.StandardODEProblem}, Tsit5, OrdinaryDiffEq.InterpolationData{ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Matrix{Float32}}, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}, OrdinaryDiffEq.DEOptions{Float32, Float32, Float32, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Matrix{Float32}, Float32, Nothing, OrdinaryDiffEq.DefaultInit})
    @ OrdinaryDiffEq C:\Users\Administrator\.julia\packages\OrdinaryDiffEq\2AoGt\src\integrators\integrator_utils.jl:261
  [6] _loopfooter!(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Matrix{Float32}, Nothing, Float64, Vector{Float32}, Float64, Float32, Float32, Float64, Vector{Matrix{Float32}}, ODESolution{Float32, 3, Vector{Matrix{Float32}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Symbol, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}}}}, SciMLBase.StandardODEProblem}, Tsit5, OrdinaryDiffEq.InterpolationData{ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Matrix{Float32}}, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}, OrdinaryDiffEq.DEOptions{Float32, Float32, Float32, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Matrix{Float32}, Float32, Nothing, OrdinaryDiffEq.DefaultInit})
    @ OrdinaryDiffEq C:\Users\Administrator\.julia\packages\OrdinaryDiffEq\2AoGt\src\integrators\integrator_utils.jl:204
  [7] loopfooter!
    @ C:\Users\Administrator\.julia\packages\OrdinaryDiffEq\2AoGt\src\integrators\integrator_utils.jl:168 [inlined]
  [8] solve!(integrator::OrdinaryDiffEq.ODEIntegrator{Tsit5, false, Matrix{Float32}, Nothing, Float64, Vector{Float32}, Float64, Float32, Float32, Float64, Vector{Matrix{Float32}}, ODESolution{Float32, 3, Vector{Matrix{Float32}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Symbol, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}}}}, SciMLBase.StandardODEProblem}, Tsit5, OrdinaryDiffEq.InterpolationData{ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{Matrix{Float32}}, Vector{Float64}, Vector{Vector{Matrix{Float32}}}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}}, DiffEqBase.DEStats}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.Tsit5ConstantCache{Float32, Float64}, OrdinaryDiffEq.DEOptions{Float32, Float32, Float32, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Bool, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}, Matrix{Float32}, Float32, Nothing, OrdinaryDiffEq.DefaultInit})
    @ OrdinaryDiffEq C:\Users\Administrator\.julia\packages\OrdinaryDiffEq\2AoGt\src\solve.jl:479
  [9] #__solve#464
    @ C:\Users\Administrator\.julia\packages\OrdinaryDiffEq\2AoGt\src\solve.jl:5 [inlined]
 [10] #solve_call#56
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:61 [inlined]
 [11] solve_up(prob::ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Symbol, CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{var"#condition#1", DiffEqSensitivity.TrackedAffect{Float64, Vector{Float32}, Matrix{Float32}, Affect{Array{Float64, 3}}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}}}}, SciMLBase.StandardODEProblem}, sensealg::Nothing, u0::Matrix{Float32}, p::Vector{Float32}, args::Tsit5; kwargs::Base.Iterators.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :sense), Tuple{Bool, Bool, Bool, BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}}}})
    @ DiffEqBase C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:82
 [12] #solve#57
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:70 [inlined]
 [13] _concrete_solve_adjoint(::ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, ::Matrix{Float32}, ::Vector{Float32}; save_start::Bool, save_end::Bool, saveat::Int64, save_idxs::Nothing, kwargs::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:sense, :callback), Tuple{BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, DiscreteCallback{var"#condition#1", Affect{Array{Float64, 3}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}}}})
    @ DiffEqSensitivity C:\Users\Administrator\.julia\packages\DiffEqSensitivity\mhKek\src\concrete_solve.jl:120
 [14] #_concrete_solve_adjoint#195
    @ C:\Users\Administrator\.julia\packages\DiffEqSensitivity\mhKek\src\concrete_solve.jl:66 [inlined]
 [15] #_solve_adjoint#75
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:315 [inlined]
 [16] #rrule#73
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:301 [inlined]
 [17] chain_rrule_kw
    @ C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\chainrules.jl:115 [inlined]
 [18] macro expansion
    @ C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0 [inlined]
 [19] _pullback(::Zygote.Context, ::DiffEqBase.var"#solve_up##kw", ::NamedTuple{(:sense, :callback, :saveat), Tuple{BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, DiscreteCallback{var"#condition#1", Affect{Array{Float64, 3}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}, Int64}}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Nothing, ::Matrix{Float32}, ::Vector{Float32}, ::Tsit5)
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:9
 [20] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [21] adjoint
    @ C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [22] _pullback
    @ C:\Users\Administrator\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [23] _pullback
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:70 [inlined]
 [24] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve#57", ::Nothing, ::Nothing, ::Nothing, ::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:sense, :callback, :saveat), Tuple{BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, DiscreteCallback{var"#condition#1", Affect{Array{Float64, 3}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}, Int64}}}, ::typeof(solve), ::ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5)
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [25] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804

xk_y · October 20, 2021, 1:04pm

 [26] adjoint
    @ C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [27] _pullback
    @ C:\Users\Administrator\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [28] _pullback
    @ C:\Users\Administrator\.julia\packages\DiffEqBase\yFYIc\src\solve.jl:68 [inlined]
 [29] _pullback(::Zygote.Context, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:sense, :callback, :saveat), Tuple{BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, DiscreteCallback{var"#condition#1", Affect{Array{Float64, 3}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}, Int64}}, ::typeof(solve), ::ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::Tsit5)
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [30] _apply(::Function, ::Vararg{Any, N} where N)
    @ Core .\boot.jl:804
 [31] adjoint
    @ C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\lib\lib.jl:200 [inlined]
 [32] adjoint(::Zygote.Context, ::typeof(Core._apply_iterate), ::typeof(iterate), ::Function, ::Tuple{NamedTuple{(:sense, :callback, :saveat), Tuple{BacksolveAdjoint{0, true, Val{:central}, ZygoteVJP, Bool}, DiscreteCallback{var"#condition#1", Affect{Array{Float64, 3}}, typeof(DiffEqBase.INITIALIZE_DEFAULT), typeof(DiffEqBase.FINALIZE_DEFAULT)}, Int64}}, typeof(solve), ODEProblem{Matrix{Float32}, Tuple{Float64, Float64}, false, Vector{Float32}, ODEFunction{false, var"#dudt_#4"{_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ::Tuple{Tsit5})
    @ Zygote .\none:0
 [33] _pullback
    @ C:\Users\Administrator\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:57 [inlined]
 [34] _pullback
    @ .\In[1]:53 [inlined]
 [35] _pullback(::Zygote.Context, ::_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::Matrix{Float32}, ::Array{Float64, 3}, ::Vector{Float32})
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [36] _pullback
    @ .\In[1]:49 [inlined]
 [37] _pullback(::Zygote.Context, ::_NeuralODE{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#61#63"{Chain{Tuple{Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(gelu), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Float64, Int64}, Tuple{Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::Matrix{Float32}, ::Array{Float64, 3})
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [38] _pullback
    @ .\In[1]:61 [inlined]
 [39] _pullback(::Zygote.Context, ::typeof(predict), ::Array{Float32, 3}, ::Array{Float64, 3})
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [40] _pullback
    @ .\In[1]:66 [inlined]
 [41] _pullback(::Zygote.Context, ::typeof(loss), ::Array{Float32, 3}, ::Array{Float64, 3})
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [42] _pullback
    @ .\In[1]:77 [inlined]
 [43] _pullback(::Zygote.Context, ::var"#5#6"{typeof(loss), Array{Float64, 3}, Array{Float32, 3}})
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
 [44] pullback(f::Function, ps::Params)
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface.jl:343
 [45] gradient(f::Function, args::Params)
    @ Zygote C:\Users\Administrator\.julia\packages\Zygote\TaBlo\src\compiler\interface.jl:75
 [46] train!(loss::typeof(loss), ps::Params, data::Flux.Data.DataLoader{NamedTuple{(:x, :c), Tuple{Array{Float32, 3}, Array{Float64, 3}}}, Random._GLOBAL_RNG}, opt::ADAM)
    @ Main .\In[1]:76
 [47] macro expansion
    @ C:\Users\Administrator\.julia\packages\Flux\0c9kI\src\optimise\train.jl:136 [inlined]
 [48] top-level scope
    @ C:\Users\Administrator\.julia\packages\Juno\n6wyj\src\progress.jl:134
 [49] eval
    @ .\boot.jl:360 [inlined]
 [50] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base .\loading.jl:1094

ChrisRackauckas · October 20, 2021, 6:16pm

Open an issue to DiffEqSensitivity?

@frankschae