Runtime dispatch in train!

[Vilin97]

Package Version

v0.13.7

Julia Version

1.8.2

OS / Environment

Windows 11

Describe the bug

I followed the example of fitting a straight line from the Flux's docs, and used JET to analyze train!. It found 19 possible runtime dispatch errors.

Steps to Reproduce

using Flux, JET
actual(x) = 4x + 2
x_train, x_test = hcat(0:5...), hcat(6:10...)
y_train, y_test = actual.(x_train), actual.(x_test)
predict = Dense(1 => 1)
loss_(x, y) = Flux.Losses.mse(predict(x), y);
opt = Descent()
data = [(x_train, y_train)]
parameters = Flux.params(predict)
Flux.train!(loss_, parameters, data, opt)  # [edit: qualify train!]
@report_opt Flux.train!(loss_, parameters, data, opt) # ═════ 19 possible errors found ═════ 
# runtime dispatch detected: isequal(%1::Any, v::Task)::Bool

Expected Results

I expected to find no errors with JET

Observed Results

I found runtime dispatch errors.

Relevant log output

julia> @report_opt train!(loss_, parameters, data, opt)
═════ 19 possible errors found ═════
┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:125 Flux.Optimise.:(var"#train!#36")(#38, #self#, loss, ps, data, opt)      
│┌ @ logging.jl:376 logger = Base.CoreLogging.current_logger_for_env(std_level, group, _module)
││┌ @ logging.jl:499 Base.CoreLogging.env_override_minlevel(group, _module)
│││┌ @ logging.jl:565 Base.moduleroot(_module)
││││┌ @ reflection.jl:45 Base.is_root_module(m)
│││││┌ @ lock.jl:221 lock(temp)
││││││┌ @ lock.jl:103 slowlock(rl)
│││││││┌ @ lock.jl:112 wait(c)
││││││││┌ @ condition.jl:126 Base.list_deletefirst!(ct.queue, ct)
│││││││││┌ @ linked_list.jl:145 isequal(h.value, val)
││││││││││┌ @ gcutils.jl:4 isequal(%1, v)
│││││││││││ runtime dispatch detected: isequal(%1::Any, v::Task)::Bool    
││││││││││└────────────────
││││││││┌ @ condition.jl:126 Base.list_deletefirst!(%45, %39)
│││││││││ runtime dispatch detected: Base.list_deletefirst!(%45::Any, %39::Task)::Any
││││││││└────────────────────
│││││┌ @ lock.jl:225 unlock(temp)
││││││┌ @ lock.jl:133 _unlock(rl)
│││││││┌ @ lock.jl:139 notifywaiters(rl)
││││││││┌ @ lock.jl:143  = notify(cond_wait)
│││││││││┌ @ condition.jl:142 #self#(c, Base.nothing)
││││││││││┌ @ condition.jl:142 Base.:(var"#notify#586")(true, false, #self#, c, arg)
│││││││││││┌ @ condition.jl:142 notify(c, arg, all, error)
││││││││││││┌ @ condition.jl:148 Core.kwfunc(schedule)(NamedTuple{(:error,)}(tuple(error)), schedule, t, arg)
│││││││││││││┌ @ task.jl:789 Base.:(var"#schedule#613")(error, _3, t, arg)
││││││││││││││┌ @ task.jl:793 %10(%11, t)
│││││││││││││││ runtime dispatch detected: %10::typeof(Base.list_deletefirst!)(%11::Any, t::Task)::Any
││││││││││││││└───────────────
│┌ @ logging.jl:364 Base.CoreLogging.logging_error(logger, level, _module, group, id, file, line, err, true)
││┌ @ logging.jl:463 (%51)
│││ runtime dispatch detected: ::NamedTuple{(:exception,)}(%51::Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}})::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}
││└──────────────────
││┌ @ logging.jl:463 handle_message##kw(%52, Base.CoreLogging.handle_message, %37, Base.CoreLogging.Error, %44, %38, :logevent_error, %39, %40, %41)
│││ runtime dispatch detected: handle_message##kw(%52::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}, Base.CoreLogging.handle_message, %37::Base.CoreLogging.AbstractLogger, Base.CoreLogging.Error, %44::Union{LazyString, String}, %38::Module, :logevent_error, %39::Base.UUID, %40::String, %41::Int64)::Any     
││└──────────────────
│┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:131 Flux.Optimise.withgradient(#37, ps)
││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:132 pullback(tuple(f), args...)
│││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:384 Zygote._pullback(cx, f)
││││┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:132 Zygote._pullback(ctx, Core._apply_iterate, iterate, Zygote._pullback(ctx, Zygote.literal_getfield, f, Val{:loss}())[1], Zygote._pullback(ctx, Flux.Optimise.batchmemaybe, Zygote._pullback(ctx, Zygote.literal_getfield, f, Val{:d}())[1])[1])
│││││┌ @ C:\Users\Math User\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:65 ZygoteRules.adjoint(tuple(__context__, 550, 551, f), args...)       
││││││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\lib\lib.jl:203 Core._apply(tuple(Zygote._pullback, tuple(__context__, f)), args...)
│││││││┌ @ boot.jl:816 Core._apply_iterate(tuple(Main.Base.iterate), x...)
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 Zygote._pullback(ctx, Zygote.literal_getproperty, Flux.Losses, Val{:mse}())
│││││││││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\lib\literal_getproperty.jl:83 Zygote._pullback(cx, Zygote.getproperty, x, :mse)
││││││││││┌ @ Base.jl:31 Zygote._pullback(ctx, Base.getfield, Base.getfield(args, 1), Base.getfield(args, 2))
│││││││││││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\lib\lib.jl:244 Zygote.Val(field_name)
││││││││││││┌ @ essentials.jl:714 %1()
│││││││││││││ runtime dispatch detected: %1::Type{Val{_A}} where _A()::Val
││││││││││││└─────────────────────
│││││││││││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\lib\lib.jl:244 Zygote._pullback(cx, Zygote.literal_getfield, x, %1)
││││││││││││ runtime dispatch detected: Zygote._pullback(cx::Zygote.Context{true}, Zygote.literal_getfield, x::Module, %1::Val)::Tuple{Any, Zygote.var"#2077#back#218"}
│││││││││││└───────────────────────────────────────────────────────────   
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 Zygote._pullback(ctx, %12, %1)
│││││││││ runtime dispatch detected: Zygote._pullback(ctx::Zygote.Context{true}, %12::Any, %1::Matrix{Int64})::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 %21[1]
│││││││││ runtime dispatch detected: (%21::Any)[1]::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 %21[2]
│││││││││ runtime dispatch detected: (%21::Any)[2]::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 Zygote._pullback(ctx, %5, %22, %2)
│││││││││ runtime dispatch detected: Zygote._pullback(ctx::Zygote.Context{true}, %5::Any, %22::Any, %2::Matrix{Int64})::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 %24[1]
│││││││││ runtime dispatch detected: (%24::Any)[1]::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││││││││┌ @ c:\Users\Math User\Github\Math\src\fokker_planck_probability_flow.jl:170 %24[2]
│││││││││ runtime dispatch detected: (%24::Any)[2]::Any
││││││││└────────────────────────────────────────────────────────────────────────────
││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:133 grad = back(Zygote.sensitivity(y))
│││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:390 Zygote.Grads(getfield(#self#, :cx).cache, getfield(#self#, :ps))
││││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:281
 convert(, grads)
│││││ runtime dispatch detected: convert(::IdDict{Any, Any}, grads::Nothing)
││││└──────────────────────────────────────────────────────────────────────
││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:133 Zygote.sensitivity(%5)
│││ runtime dispatch detected: Zygote.sensitivity(%5::Any)::Any
││└────────────────────────────────────────────────────────────────────── 
││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:133 %10(%11)
│││ runtime dispatch detected: %10::Zygote.var"#99#100"{Zygote.Params{Zygote.Buffer{Any, Vector{Any}}}, _A, Zygote.Context{true}} where _A(%11::Any)::Zygote.Grads
││└────────────────────────────────────────────────────────────────────── 
││┌ @ C:\Users\Math User\.julia\dev\Zygote\src\compiler\interface.jl:135 (%13)
│││ runtime dispatch detected: ::NamedTuple{(:val, :grad)}(%13::Tuple{Any, Zygote.Grads})::NamedTuple{(:val, :grad), _A} where _A<:Tuple{Any, Zygote.Grads}
││└────────────────────────────────────────────────────────────────────── 
│┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:135 string("Loss is ", l, " on data item ", i, ", stopping training")      
││┌ @ strings/io.jl:185 Base.print_to_string(xs...)
│││┌ @ strings/io.jl:144 print(%83, %86)
││││ runtime dispatch detected: print(%83::IOBuffer, %86::Any)::Any       
│││└─────────────────────
│┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:137 update!(opt, ps, gs)
││┌ @ C:\Users\Math User\.julia\packages\Flux\nJ0IB\src\optimise\train.jl:24 update!(opt, %32, %65)
│││ runtime dispatch detected: update!(opt::Descent, %32::Any, %65::Any)::Any
││└──────────────────────────────────────────────────────────────────────────

[ToucheSir]

Flux.train! was written at a time before JET and type stability were big discussion points in the community. Thus, it really depends on the dynamic nature of the language and doesn't exactly play well with type inference. That said, any performance or latency hit you get from this should be so minimal that it doesn't even register in practice. You may be interested in the ongoing PRs (#2082, #2083) to revamp train! for the modern day. More pragmatically though, I'd recommend just using a custom training loop for now: it's not many more lines of code, more flexible and can be as type stable as you want.

[mcabbott]

For what it's worth, here's the result on master:

julia> loss3(m, x, y) = Flux.Losses.mse(m(x), y);

julia> state = Flux.setup(Descent(), predict)
(weight = Leaf(Descent{Float64}(0.1), nothing), bias = Leaf(Descent{Float64}(0.1), nothing), σ = ())

julia> Flux.train!(loss3, predict, data, state)

julia> @report_opt Flux.train!(loss3, predict, data, state)
═════ 47 possible errors found ═════

julia> @report_opt Flux.train!(loss3, predict, data, state)
═════ 47 possible errors found ═════
┌ @ /Users/me/.julia/dev/Flux/src/train.jl:97 Flux.Train.:(var"#train!#3")(Flux.Train.nothing, #self#, loss, model, data, opt)
│┌ @ logging.jl:376 logger = Base.CoreLogging.current_logger_for_env(std_level, group, _module)
││┌ @ logging.jl:502 Base.CoreLogging.env_override_minlevel(group, _module)
│││┌ @ logging.jl:568 Base.moduleroot(_module)
││││┌ @ reflection.jl:45 Base.is_root_module(m)
│││││┌ @ lock.jl:265 lock(temp)
││││││┌ @ lock.jl:147 slowlock(rl)
│││││││┌ @ lock.jl:156 wait(c)
││││││││┌ @ condition.jl:125 Base.:(var"#wait#624")(false, #self#, c)
│││││││││┌ @ condition.jl:132 Base.list_deletefirst!(ct.queue, ct)
││││││││││┌ @ linked_list.jl:145 isequal(h.value, val)
│││││││││││┌ @ gcutils.jl:4 isequal(%1, v)
││││││││││││ runtime dispatch detected: isequal(%1::Any, v::Task)::Bool
│││││││││││└────────────────
│││││││││┌ @ condition.jl:132 Base.list_deletefirst!(%43, %37)
││││││││││ runtime dispatch detected: Base.list_deletefirst!(%43::Any, %37::Task)::Any
│││││││││└────────────────────
│││││┌ @ lock.jl:269 unlock(temp)
││││││┌ @ lock.jl:177 _unlock(rl)
│││││││┌ @ lock.jl:183 notifywaiters(rl)
││││││││┌ @ lock.jl:187  = notify(cond_wait)
│││││││││┌ @ condition.jl:148 #self#(c, Base.nothing)
││││││││││┌ @ condition.jl:148 Base.:(var"#notify#625")(true, false, #self#, c, arg)
│││││││││││┌ @ condition.jl:148 notify(c, arg, all, error)
││││││││││││┌ @ condition.jl:154 Core.kwcall(NamedTuple{(:error,)}(tuple(error)), schedule, t, arg)
│││││││││││││┌ @ task.jl:827 Base.:(var"#schedule#652")(_7, _3, t, arg)
││││││││││││││┌ @ task.jl:831 Base.list_deletefirst!(%10, t)
│││││││││││││││ runtime dispatch detected: Base.list_deletefirst!(%10::Any, t::Task)::Any
││││││││││││││└───────────────
│┌ @ logging.jl:364 Base.CoreLogging.logging_error(logger, level, _module, group, id, file, line, err, true)
││┌ @ logging.jl:466 Core.kwcall(%54, Base.CoreLogging.handle_message, %37, Base.CoreLogging.Error, %44, %38, :logevent_error, %39, %40, %41)
│││ runtime dispatch detected: Core.kwcall(%54::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}, Base.CoreLogging.handle_message, %37::Base.CoreLogging.AbstractLogger, Base.CoreLogging.Error, %44::Union{LazyString, String}, %38::Module, :logevent_error, %39::Base.UUID, %40::String, %41::Int64)::Any
││└──────────────────
│┌ @ /Users/me/.julia/dev/Flux/src/train.jl:102 withgradient(#4, model)
││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:123 pullback(tuple(f), args...)
│││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:42 pullback(tuple(f, Zygote.Context()), args...)
││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:44 Zygote._pullback(tuple(cx, f), args...)
│││││┌ @ /Users/me/.julia/dev/Flux/src/train.jl:102 Zygote._pullback(ctx, Core._apply_iterate, iterate, Zygote._pullback(ctx, Zygote.literal_getfield, f, Val{:loss}())[1], Zygote._pullback(ctx, tuple, Base.getfield(args, 1))[1], Zygote._pullback(ctx, Zygote.literal_getfield, f, Val{:d_splat}())[1])
││││││┌ @ /Users/me/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 ZygoteRules.adjoint(tuple(__context__, 530, 531, f), args...)
│││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/lib/lib.jl:203 Core._apply(tuple(Zygote._pullback, tuple(__context__, f)), args...)
││││││││┌ @ boot.jl:836 Core._apply_iterate(tuple(Main.Base.iterate), x...)
│││││││││┌ @ REPL[23]:1 Zygote._pullback(ctx, Base.getfield(args, 1), Base.getfield(args, 2))
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/layers/basic.jl:172 Zygote._pullback(ctx, Flux.:*, Zygote._pullback(ctx, Zygote.literal_getproperty, f, Val{:weight}())[1], Base.getfield(args, 1))
│││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0 Zygote.chain_rrule(tuple(Zygote.ZygoteRuleConfig(ctx), f), args...)
││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:218 Zygote.rrule(tuple(config, f), args...)
│││││││││││││┌ @ /Users/me/.julia/packages/ChainRulesCore/C73ay/src/rules.jl:134 rrule(args...)
││││││││││││││┌ @ /Users/me/.julia/packages/ChainRules/hVHC4/src/rulesets/Base/arraymath.jl:40 A ChainRules.:* B
│││││││││││││││┌ @ /Users/me/.julia/dev/julia/usr/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:141 mul!(LinearAlgebra.similar(B, TS, tuple(LinearAlgebra.size(A, 1), LinearAlgebra.size(B, 2))), A, B)
││││││││││││││││┌ @ /Users/me/.julia/dev/julia/usr/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:276 mul!(C, A, B, true, false)
│││││││││││││││││┌ @ /Users/me/.julia/dev/julia/usr/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:303 LinearAlgebra.generic_matmatmul!(C, N, N, A, B, LinearAlgebra.MulAddMul(alpha, beta))
││││││││││││││││││┌ @ /Users/me/.julia/dev/julia/usr/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:839 LinearAlgebra.matmul2x2!(C, tA, tB, A, B, _add)
│││││││││││││││││││┌ @ /Users/me/.julia/dev/julia/usr/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:1032 LinearAlgebra.DimensionMismatch(LazyString("A has size ", LinearAlgebra.size(A), ", B has size ", LinearAlgebra.size(B), ", C has size ", LinearAlgebra.size(C)))
││││││││││││││││││││┌ @ array.jl:12 convert(fieldtype(DimensionMismatch, 1), msg)
│││││││││││││││││││││┌ @ strings/basic.jl:232 T(s)
││││││││││││││││││││││┌ @ strings/lazy.jl:83 str = sprint(#103)
│││││││││││││││││││││││┌ @ strings/io.jl:107 Base.:(var"#sprint#487")(tuple(Base.nothing, 0, #self#, f), args...)
││││││││││││││││││││││││┌ @ strings/io.jl:114 f(tuple(s), args...)
│││││││││││││││││││││││││┌ @ strings/lazy.jl:85 print(io, %15)
││││││││││││││││││││││││││ runtime dispatch detected: print(io::IOBuffer, %15::Any)::Any
│││││││││││││││││││││││││└──────────────────────
│││││││││┌ @ REPL[23]:1 Zygote._pullback(ctx, %6, %53, %3)
││││││││││ runtime dispatch detected: Zygote._pullback(ctx::Zygote.Context{false}, %6::Any, %53::Matrix{Float32}, %3::Matrix{Int64})::Any
│││││││││└──────────────
│││││││││┌ @ REPL[23]:1 %62[1]
││││││││││ runtime dispatch detected: (%62::Any)[1]::Any
│││││││││└──────────────
│││││││││┌ @ REPL[23]:1 %62[2]
││││││││││ runtime dispatch detected: (%62::Any)[2]::Any
│││││││││└──────────────
││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:125  = Zygote.map(Zygote._project, args, grad)
│││┌ @ tuple.jl:298 f(t[1], s[1])
││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:184 Zygote.zygote2differential(dx, x)
│││││┌ @ logging.jl:364 Base.CoreLogging.logging_error(logger, level, _module, group, id, file, line, err, true)
││││││┌ @ logging.jl:466 Core.kwcall(%54, Base.CoreLogging.handle_message, %37, Base.CoreLogging.Error, %44, %38, :logevent_error, %39, %40, %41)
│││││││ runtime dispatch detected: Core.kwcall(%54::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}, Base.CoreLogging.handle_message, %37::Base.CoreLogging.AbstractLogger, Base.CoreLogging.Error, %44::Union{LazyString, String}, %38::Module, :logevent_error, %39::Symbol, %40::String, %41::Int64)::Any
││││││└──────────────────
││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:124 Zygote.sensitivity(%5)
│││ runtime dispatch detected: Zygote.sensitivity(%5::Any)::Any
││└────────────────────────────────────────────────────────────────────────
││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:124 %9(%10)
│││ runtime dispatch detected: %9::Zygote.var"#60#61"(%10::Any)::Union{Nothing, Tuple}
││└────────────────────────────────────────────────────────────────────────
│┌ @ /Users/me/.julia/dev/Flux/src/train.jl:104 string("Loss is ", l, " on data item ", i, ", stopping training")
││┌ @ strings/io.jl:185 Base.print_to_string(xs...)
│││┌ @ strings/io.jl:144 print(%74, %77)
││││ runtime dispatch detected: print(%74::IOBuffer, %77::Any)::Any
│││└─────────────────────
│┌ @ /Users/me/.julia/dev/Flux/src/train.jl:106 Optimisers.update!(opt, model, gs[1])
││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:69 Optimisers._grads!(tuple(grads, tree, model, grad), higher...)
│││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:112 Optimisers.foreach(tuple(#15, tree, x′), x̄s′...)
││││┌ @ abstractarray.jl:3068 zip(itrs...)
│││││┌ @ iterators.jl:370 Base.Iterators.Zip(a)
││││││ runtime dispatch detected: Base.Iterators.Zip(a::Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Optimisers.Leaf{Optimisers.Descent{Float64}, Nothing}, Optimisers.Leaf{Optimisers.Descent{Float64}, Nothing}, Tuple{}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Matrix{Float32}, Vector{Float32}, typeof(identity)}}, NamedTuple{(:weight, :bias, :σ), _A} where _A<:Tuple{Any, Any, Any}})::Base.Iterators.Zip
│││││└────────────────────
││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:72 newmodel = Core.kwcall(NamedTuple{(:grads, :params)}(tuple(grads, Optimisers.IdDict())), Optimisers._update!, tree, model)
│││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:76 Optimisers.:(var"#_update!#8")(grads, params, _3, tree, x)
││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Optimisers.map(#9, tree, x′)
│││││┌ @ namedtuple.jl:219 map(tuple(f), map(Base.Tuple, tuple(tuple(nt), nts...))...)
││││││┌ @ tuple.jl:302 f(t[1], s[1])
│││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Core.kwcall(NamedTuple{(:grads, :params)}(tuple(getfield(#self#, :grads), getfield(#self#, :params))), Optimisers._update!, tᵢ, xᵢ)
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:87 Optimisers.:(var"#_update!#11")(grads, params, _3, ℓ, x)
│││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92  = Optimisers.subtract!(x, x̄′)
││││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:98 Base.broadcasted(Optimisers.:-, x, x̄)
│││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:339 Zygote.map(#91, gs)
││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:354 Zygote.map!(tuple(f, gsout, gs1), gss...)
│││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 Base.Generator(%41, gss)
││││││││││││││ runtime dispatch detected: Base.Generator(%41::Zygote.var"#96#98", gss::Tuple{Zygote.Grads})::Base.Generator{Tuple{Zygote.Grads}}
│││││││││││││└────────────────────────────────────────────────────────────────────────
│││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 f(%46)
││││││││││││││ runtime dispatch detected: f::Zygote.var"#91#92"{typeof(-), Matrix{Float32}}(%46::Any)::Any
│││││││││││││└────────────────────────────────────────────────────────────────────────
│││││││││││┌ @ broadcast.jl:1317 Base.Broadcast.broadcasted(tuple(Base.Broadcast.combine_styles(tuple(arg1′, arg2′), args′...), f, arg1′, arg2′), args′...)
││││││││││││┌ @ /Users/me/.julia/dev/Flux/src/deprecations.jl:22 Base.depwarn("Broadcasting is not safe to use with RNNs, as it does not guarantee an iteration order.\nRe-writing this as a comprehension would be better.", :broadcasted)
│││││││││││││┌ @ deprecated.jl:121 Base.:(var"#depwarn#1006")(false, #self#, msg, funcsym)
││││││││││││││┌ @ deprecated.jl:132 Base.firstcaller(bt, funcsym)
│││││││││││││││┌ @ deprecated.jl:152 Base.firstcaller(bt, tuple(funcsym))
││││││││││││││││┌ @ deprecated.jl:157 lkups = Base.StackTraces.lookup(ip)
│││││││││││││││││┌ @ stacktraces.jl:114 StackFrame(%63, %66, %69, %70, %73, %76, %4)
││││││││││││││││││ runtime dispatch detected: StackFrame(%63::Symbol, %66::Symbol, %69::Int64, %70::Any, %73::Bool, %76::Bool, %4::UInt64)::Base.StackTraces.StackFrame
│││││││││││││││││└──────────────────────
│││││││││││││││││┌ @ stacktraces.jl:146 Base.IRShow.method_name(lineinfo)
││││││││││││││││││┌ @ compiler/ssair/show.jl:178 Base.IRShow.normalize_method_name(%1)
│││││││││││││││││││ runtime dispatch detected: Base.IRShow.normalize_method_name(%1::Any)::Any
││││││││││││││││││└──────────────────────────────
│││││││││││││││││┌ @ stacktraces.jl:142 %56[%59]
││││││││││││││││││ runtime dispatch detected: (%56::Any)[%59::Int32]::Any
│││││││││││││││││└──────────────────────
│││││││││││││││││┌ @ stacktraces.jl:148 %107[%108]
││││││││││││││││││ runtime dispatch detected: (%107::Any)[%108::Int32]::Any
│││││││││││││││││└──────────────────────
││││││││││││││┌ @ logging.jl:364 Base.CoreLogging.logging_error(logger, level, _module, group, id, file, line, err, true)
│││││││││││││││┌ @ logging.jl:466 Core.kwcall(%54, Base.CoreLogging.handle_message, %37, Base.CoreLogging.Error, %44, %38, :logevent_error, %39, %40, %41)
││││││││││││││││ runtime dispatch detected: Core.kwcall(%54::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}, Base.CoreLogging.handle_message, %37::Base.CoreLogging.AbstractLogger, Base.CoreLogging.Error, %44::Union{LazyString, String}, %38::Module, :logevent_error, %39::Tuple{Union{Ptr{Nothing}, Base.InterpreterIP}, Symbol}, %40::String, %41::Int64)::Any
│││││││││││││││└──────────────────
││││││││││││┌ @ /Users/me/.julia/dev/Flux/src/deprecations.jl:24 Flux.map(tuple(f), args...)
│││││││││││││┌ @ abstractarray.jl:3377 Base.Generator(tuple(f), iters...)
││││││││││││││┌ @ generator.jl:36 zip(tuple(I1, I2), Is...)
│││││││││││││││┌ @ iterators.jl:370 Base.Iterators.Zip(a)
││││││││││││││││ runtime dispatch detected: Base.Iterators.Zip(a::Tuple{typeof(-), Matrix{Float32}, Any})::Base.Iterators.Zip
│││││││││││││││└────────────────────
│││││││││││││┌ @ abstractarray.jl:3377 collect(Base.Generator(tuple(f), iters...))
││││││││││││││┌ @ array.jl:759 et = Base.promote_typejoin_union(T)
│││││││││││││││┌ @ promotion.jl:182 Base.promote_typejoin(Base.promote_typejoin_union(T.a), Base.promote_typejoin_union(T.b))
││││││││││││││││┌ @ promotion.jl:171 c = typejoin(Base._promote_typesubtract(a), Base._promote_typesubtract(b))
│││││││││││││││││┌ @ promotion.jl:34 typejoin(a.body, b)
││││││││││││││││││┌ @ promotion.jl:36 typejoin(a, b.body)
│││││││││││││││││││┌ @ promotion.jl:121 %359.var
││││││││││││││││││││ runtime dispatch detected: (%359::Any).var::Any
│││││││││││││││││││└────────────────────
│││││││││││││││││││┌ @ promotion.jl:122 %359.body
││││││││││││││││││││ runtime dispatch detected: (%359::Any).body::Any
│││││││││││││││││││└────────────────────
│││││││││││││││││││┌ @ promotion.jl:126 Base.UnionAll(%426, %428)
││││││││││││││││││││ runtime dispatch detected: Base.UnionAll(%426::Any, %428::Any)::Any
│││││││││││││││││││└────────────────────
││││││││││││││││││┌ @ promotion.jl:121 %354.var
│││││││││││││││││││ runtime dispatch detected: (%354::Any).var::Any
││││││││││││││││││└────────────────────
││││││││││││││││││┌ @ promotion.jl:122 %354.body
│││││││││││││││││││ runtime dispatch detected: (%354::Any).body::Any
││││││││││││││││││└────────────────────
││││││││││││││││││┌ @ promotion.jl:126 Base.UnionAll(%421, %423)
│││││││││││││││││││ runtime dispatch detected: Base.UnionAll(%421::Any, %423::Any)::Any
││││││││││││││││││└────────────────────
│││││││││││││││││┌ @ promotion.jl:121 %349.var
││││││││││││││││││ runtime dispatch detected: (%349::Any).var::Any
│││││││││││││││││└────────────────────
│││││││││││││││││┌ @ promotion.jl:122 %349.body
││││││││││││││││││ runtime dispatch detected: (%349::Any).body::Any
│││││││││││││││││└────────────────────
│││││││││││││││││┌ @ promotion.jl:126 Base.UnionAll(%416, %418)
││││││││││││││││││ runtime dispatch detected: Base.UnionAll(%416::Any, %418::Any)::Any
│││││││││││││││││└────────────────────
│││││││││││││││┌ @ promotion.jl:210 Base.promote_typejoin_union(%45)
││││││││││││││││ runtime dispatch detected: Base.promote_typejoin_union(%45::Any)::Type
│││││││││││││││└────────────────────
││││││││││││││┌ @ array.jl:779 Base.grow_to!(Core.apply_type(Vector, et)(), itr)
│││││││││││││││┌ @ array.jl:859 Base.grow_to!(dest2, itr, y[2])
││││││││││││││││┌ @ array.jl:883 new = Base.push_widen(dest, el)
│││││││││││││││││┌ @ array.jl:864 Base.promote_typejoin(eltype(dest), Base.typeof(el))
││││││││││││││││││┌ @ promotion.jl:171 c = typejoin(Base._promote_typesubtract(a), Base._promote_typesubtract(b))
│││││││││││││││││││┌ @ promotion.jl:121 %336.var
││││││││││││││││││││ runtime dispatch detected: (%336::Any).var::Any
│││││││││││││││││││└────────────────────
│││││││││││││││││││┌ @ promotion.jl:122 %336.body
││││││││││││││││││││ runtime dispatch detected: (%336::Any).body::Any
│││││││││││││││││││└────────────────────
│││││││││││││││││││┌ @ promotion.jl:126 Base.UnionAll(%403, %405)
││││││││││││││││││││ runtime dispatch detected: Base.UnionAll(%403::Any, %405::Any)::Any
│││││││││││││││││││└────────────────────
│││││││││││││││││┌ @ array.jl:869 append!(new, dest)
││││││││││││││││││┌ @ array.jl:1122 copyto!(a, length(a) - n + 1, items, first(itemindices), n)
│││││││││││││││││││┌ @ array.jl:319 Base._copyto_impl!(dest, doffs, src, soffs, n)
││││││││││││││││││││┌ @ array.jl:327 unsafe_copyto!(dest, doffs, src, soffs, n)
│││││││││││││││││││││┌ @ array.jl:280 destp = pointer(dest, doffs)
││││││││││││││││││││││┌ @ abstractarray.jl:1243 Base._memory_offset(x, i)
│││││││││││││││││││││││┌ @ abstractarray.jl:1247 Base.elsize(x)
││││││││││││││││││││││││┌ @ abstractarray.jl:251 Base.elsize(Base.typeof(A))
│││││││││││││││││││││││││┌ @ array.jl:213 Base.aligned_sizeof(T)
││││││││││││││││││││││││││┌ @ reflection.jl:392 Base.datatype_alignment(T)
│││││││││││││││││││││││││││ runtime dispatch detected: Base.datatype_alignment(T::Type)::Int64
││││││││││││││││││││││││││└─────────────────────
│││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92 Optimisers.subtract!(x, %52)
││││││││││ runtime dispatch detected: Optimisers.subtract!(x::Matrix{Float32}, %52::Any)::Any
│││││││││└──────────────────────────────────────────────────────────────────
││││││┌ @ tuple.jl:302 map(f, Base.tail(t), Base.tail(s))
│││││││┌ @ tuple.jl:299 f(t[1], s[1])
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Core.kwcall(NamedTuple{(:grads, :params)}(tuple(getfield(#self#, :grads), getfield(#self#, :params))), Optimisers._update!, tᵢ, xᵢ)
│││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:87 Optimisers.:(var"#_update!#11")(grads, params, _3, ℓ, x)
││││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92  = Optimisers.subtract!(x, x̄′)
│││││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:98 Base.broadcasted(Optimisers.:-, x, x̄)
││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:339 Zygote.map(#91, gs)
│││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:354 Zygote.map!(tuple(f, gsout, gs1), gss...)
││││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 f(%46)
│││││││││││││││ runtime dispatch detected: f::Zygote.var"#91#92"{typeof(-), Vector{Float32}}(%46::Any)::Any
││││││││││││││└────────────────────────────────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92 Optimisers.subtract!(x, %52)
│││││││││││ runtime dispatch detected: Optimisers.subtract!(x::Vector{Float32}, %52::Any)::Any
││││││││││└──────────────────────────────────────────────────────────────────
││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 x′′ = re(Optimisers.map(#9, tree, x′))
│││││┌ @ /Users/me/.julia/packages/Functors/V2McK/src/functor.jl:19 (y[1], y[2], y[3])
││││││┌ @ /Users/me/.julia/dev/Flux/src/layers/basic.jl:165 init(out, in)
│││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:106 Flux.:(var"#glorot_uniform#4")(tuple(pairs(NamedTuple()), #self#), dims...)
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:106 Flux.glorot_uniform(tuple(Flux.default_rng_value()), dims...)
│││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:102 Flux.:(var"#glorot_uniform#3")(tuple(1, #self#, rng), dims...)
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 Flux.sum(%4)
│││││││││││ runtime dispatch detected: Flux.sum(%4::Tuple{Integer, Bool})::Any
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 24.0 Flux.:/ %5
│││││││││││ runtime dispatch detected: (24.0 Flux.:/ %5::Any)::Any
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 Flux.sqrt(%6)
│││││││││││ runtime dispatch detected: Flux.sqrt(%6::Any)::Any
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 %1 Flux.:* %7
│││││││││││ runtime dispatch detected: (%1::Float32 Flux.:* %7::Any)::Any
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Flux.rand([quote], , %9, %10)
│││││││││││ runtime dispatch detected: Flux.rand([quote]::Random.TaskLocalRNG, ::Float32, %9::Bool, %10::Integer)::Matrix{Float32}
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Base.broadcasted(Flux.:*, %15, %8)
│││││││││││ runtime dispatch detected: Base.broadcasted(Flux.:*, %15::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2}, Nothing, typeof(-), Tuple{Matrix{Float32}, Float32}}, %8::Any)::Any
││││││││││└──────────────────────────────────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Base.materialize(%16)
│││││││││││ runtime dispatch detected: Base.materialize(%16::Any)::Any
││││││││││└──────────────────────────────────────────────
││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 #186(%30)
│││││ runtime dispatch detected: #186(%30::NamedTuple{(:weight, :bias, :σ), _A} where _A<:Tuple{Any, Any, Any})::Dense
││││└──────────────────────────────────────────────────────────────────
│┌ @ /Users/me/.julia/dev/Flux/src/train.jl:102 withgradient(#4, model)
││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:125  = Zygote.map(Zygote._project, args, grad)
│││┌ @ tuple.jl:298 f(t[1], s[1])
││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:188 Zygote._project(x, Zygote.reshape(Zygote.collect(dx), Zygote.axes(x)))
│││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:184 Zygote.zygote2differential(dx, x)
││││││┌ @ logging.jl:360 msg = string("primal should be a tuple, not ", primal)
│││││││┌ @ strings/io.jl:185 Base.print_to_string(xs...)
││││││││┌ @ strings/io.jl:144 print(%38, %41)
│││││││││ runtime dispatch detected: print(%38::IOBuffer, %41::Union{String, AbstractArray})::Any
││││││││└─────────────────────

Or without train!:

julia> grads = Flux.gradient(loss3, predict, data[1]...)
((weight = Float32[88.54284;;], bias = Float32[24.615192], σ = nothing), [2.1115810871124268 15.427908897399902 … 55.37689971923828 68.69322204589844], [-0.24469606578350067 -1.7878303527832031 … -6.417233943939209 -7.9603681564331055])

julia> Flux.update!(state, predict, grads[1])
((weight = Leaf(Descent{Float64}(0.1), nothing), bias = Leaf(Descent{Float64}(0.1), nothing), σ = ()), Dense(1 => 1))

julia> @report_opt Flux.update!(state, predict, grads[1])
═════ 39 possible errors found ═════

julia> @report_opt Flux.update!(state, predict, grads[1])
═════ 39 possible errors found ═════
┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:69 Optimisers._grads!(tuple(grads, tree, model, grad), higher...)
│┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:112 Optimisers.foreach(tuple(#15, tree, x′), x̄s′...)
││┌ @ abstractarray.jl:3068 f(z...)
│││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:112 Optimisers._grads!(tuple(getfield(#self#, :dict), tᵢ, xᵢ), x̄sᵢ...)
││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:103 Optimisers.map(Optimisers.:+, x̄s, x̄s₀)
│││││┌ @ abstractarray.jl:3377 collect(Base.Generator(tuple(f), iters...))
││││││┌ @ array.jl:759 et = Base.promote_typejoin_union(T)
│││││││┌ @ promotion.jl:182 Base.promote_typejoin_union(T.a)
││││││││┌ @ promotion.jl:182 Base.promote_typejoin(Base.promote_typejoin_union(T.a), Base.promote_typejoin_union(T.b))
│││││││││┌ @ promotion.jl:171 c = typejoin(Base._promote_typesubtract(a), Base._promote_typesubtract(b))
││││││││││┌ @ promotion.jl:34 typejoin(a.body, b)
│││││││││││┌ @ promotion.jl:36 typejoin(a, b.body)
││││││││││││┌ @ promotion.jl:121 %359.var
│││││││││││││ runtime dispatch detected: (%359::Any).var::Any
││││││││││││└────────────────────
││││││││││││┌ @ promotion.jl:122 %359.body
│││││││││││││ runtime dispatch detected: (%359::Any).body::Any
││││││││││││└────────────────────
││││││││││││┌ @ promotion.jl:126 Base.UnionAll(%426, %428)
│││││││││││││ runtime dispatch detected: Base.UnionAll(%426::Any, %428::Any)::Any
││││││││││││└────────────────────
│││││││││││┌ @ promotion.jl:121 %354.var
││││││││││││ runtime dispatch detected: (%354::Any).var::Any
│││││││││││└────────────────────
│││││││││││┌ @ promotion.jl:122 %354.body
││││││││││││ runtime dispatch detected: (%354::Any).body::Any
│││││││││││└────────────────────
│││││││││││┌ @ promotion.jl:126 Base.UnionAll(%421, %423)
││││││││││││ runtime dispatch detected: Base.UnionAll(%421::Any, %423::Any)::Any
│││││││││││└────────────────────
││││││││││┌ @ promotion.jl:121 %349.var
│││││││││││ runtime dispatch detected: (%349::Any).var::Any
││││││││││└────────────────────
││││││││││┌ @ promotion.jl:122 %349.body
│││││││││││ runtime dispatch detected: (%349::Any).body::Any
││││││││││└────────────────────
││││││││││┌ @ promotion.jl:126 Base.UnionAll(%416, %418)
│││││││││││ runtime dispatch detected: Base.UnionAll(%416::Any, %418::Any)::Any
││││││││││└────────────────────
│││││││┌ @ promotion.jl:184 Base.typejoin_union_tuple(T)
││││││││┌ @ promotion.jl:210 Base.promote_typejoin_union(%45)
│││││││││ runtime dispatch detected: Base.promote_typejoin_union(%45::Any)::Type
││││││││└────────────────────
││││││┌ @ array.jl:779 Base.grow_to!(Core.apply_type(Vector, et)(), itr)
│││││││┌ @ array.jl:859 Base.grow_to!(dest2, itr, y[2])
││││││││┌ @ array.jl:883 new = Base.push_widen(dest, el)
│││││││││┌ @ array.jl:864 Base.promote_typejoin(eltype(dest), Base.typeof(el))
││││││││││┌ @ promotion.jl:171 c = typejoin(Base._promote_typesubtract(a), Base._promote_typesubtract(b))
│││││││││││┌ @ promotion.jl:121 %336.var
││││││││││││ runtime dispatch detected: (%336::Any).var::Any
│││││││││││└────────────────────
│││││││││││┌ @ promotion.jl:122 %336.body
││││││││││││ runtime dispatch detected: (%336::Any).body::Any
│││││││││││└────────────────────
│││││││││││┌ @ promotion.jl:126 Base.UnionAll(%403, %405)
││││││││││││ runtime dispatch detected: Base.UnionAll(%403::Any, %405::Any)::Any
│││││││││││└────────────────────
│││││││││┌ @ array.jl:869 append!(new, dest)
││││││││││┌ @ array.jl:1122 copyto!(a, length(a) - n + 1, items, first(itemindices), n)
│││││││││││┌ @ array.jl:319 Base._copyto_impl!(dest, doffs, src, soffs, n)
││││││││││││┌ @ array.jl:327 unsafe_copyto!(dest, doffs, src, soffs, n)
│││││││││││││┌ @ array.jl:280 destp = pointer(dest, doffs)
││││││││││││││┌ @ abstractarray.jl:1243 Base._memory_offset(x, i)
│││││││││││││││┌ @ abstractarray.jl:1247 Base.elsize(x)
││││││││││││││││┌ @ abstractarray.jl:251 Base.elsize(Base.typeof(A))
│││││││││││││││││┌ @ array.jl:213 Base.aligned_sizeof(T)
││││││││││││││││││┌ @ reflection.jl:392 Base.datatype_alignment(T)
│││││││││││││││││││ runtime dispatch detected: Base.datatype_alignment(T::Type)::Int64
││││││││││││││││││└─────────────────────
││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:103 Optimisers.map(Optimisers.:+, x̄s, %7)
│││││ runtime dispatch detected: Optimisers.map(Optimisers.:+, x̄s::Tuple{Matrix{Float32}}, %7::Any)::Any
││││└───────────────────────────────────────────────────────────────────
││┌ @ abstractarray.jl:3068 f(%8, %9, %10)
│││ runtime dispatch detected: f::Optimisers.var"#15#17"{IdDict{Optimisers.Leaf, Any}}(%8::Union{Tuple{}, Optimisers.Leaf{Optimisers.Descent{Float64}, Nothing}}, %9::Union{typeof(identity), VecOrMat{Float32}}, %10::Union{Nothing, VecOrMat{Float32}})::Any
││└─────────────────────────
┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:72 newmodel = Core.kwcall(NamedTuple{(:grads, :params)}(tuple(grads, Optimisers.IdDict())), Optimisers._update!, tree, model)
│┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:76 Optimisers.:(var"#_update!#8")(grads, params, _3, tree, x)
││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Optimisers.map(#9, tree, x′)
│││┌ @ namedtuple.jl:219 map(tuple(f), map(Base.Tuple, tuple(tuple(nt), nts...))...)
││││┌ @ tuple.jl:302 f(t[1], s[1])
│││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Core.kwcall(NamedTuple{(:grads, :params)}(tuple(getfield(#self#, :grads), getfield(#self#, :params))), Optimisers._update!, tᵢ, xᵢ)
││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:87 Optimisers.:(var"#_update!#11")(grads, params, _3, ℓ, x)
│││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92  = Optimisers.subtract!(x, x̄′)
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:98 Base.broadcasted(Optimisers.:-, x, x̄)
│││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:339 Zygote.map(#91, gs)
││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:354 Zygote.map!(tuple(f, gsout, gs1), gss...)
│││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 Base.Generator(%41, gss)
││││││││││││ runtime dispatch detected: Base.Generator(%41::Zygote.var"#96#98", gss::Tuple{Zygote.Grads})::Base.Generator{Tuple{Zygote.Grads}}
│││││││││││└────────────────────────────────────────────────────────────────────────
│││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 f(%46)
││││││││││││ runtime dispatch detected: f::Zygote.var"#91#92"{typeof(-), Matrix{Float32}}(%46::Any)::Any
│││││││││││└────────────────────────────────────────────────────────────────────────
│││││││││┌ @ broadcast.jl:1317 Base.Broadcast.broadcasted(tuple(Base.Broadcast.combine_styles(tuple(arg1′, arg2′), args′...), f, arg1′, arg2′), args′...)
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/deprecations.jl:22 Base.depwarn("Broadcasting is not safe to use with RNNs, as it does not guarantee an iteration order.\nRe-writing this as a comprehension would be better.", :broadcasted)
│││││││││││┌ @ deprecated.jl:121 Base.:(var"#depwarn#1006")(false, #self#, msg, funcsym)
││││││││││││┌ @ deprecated.jl:132 Base.firstcaller(bt, funcsym)
│││││││││││││┌ @ deprecated.jl:152 Base.firstcaller(bt, tuple(funcsym))
││││││││││││││┌ @ deprecated.jl:157 lkups = Base.StackTraces.lookup(ip)
│││││││││││││││┌ @ stacktraces.jl:114 StackFrame(%63, %66, %69, %70, %73, %76, %4)
││││││││││││││││ runtime dispatch detected: StackFrame(%63::Symbol, %66::Symbol, %69::Int64, %70::Any, %73::Bool, %76::Bool, %4::UInt64)::Base.StackTraces.StackFrame
│││││││││││││││└──────────────────────
│││││││││││││││┌ @ stacktraces.jl:146 Base.IRShow.method_name(lineinfo)
││││││││││││││││┌ @ compiler/ssair/show.jl:178 Base.IRShow.normalize_method_name(%1)
│││││││││││││││││ runtime dispatch detected: Base.IRShow.normalize_method_name(%1::Any)::Any
││││││││││││││││└──────────────────────────────
│││││││││││││││┌ @ stacktraces.jl:142 %56[%59]
││││││││││││││││ runtime dispatch detected: (%56::Any)[%59::Int32]::Any
│││││││││││││││└──────────────────────
│││││││││││││││┌ @ stacktraces.jl:148 %107[%108]
││││││││││││││││ runtime dispatch detected: (%107::Any)[%108::Int32]::Any
│││││││││││││││└──────────────────────
││││││││││││┌ @ logging.jl:376 logger = Base.CoreLogging.current_logger_for_env(std_level, group, _module)
│││││││││││││┌ @ logging.jl:502 Base.CoreLogging.env_override_minlevel(group, _module)
││││││││││││││┌ @ logging.jl:568 Base.moduleroot(_module)
│││││││││││││││┌ @ reflection.jl:45 Base.is_root_module(m)
││││││││││││││││┌ @ lock.jl:265 lock(temp)
│││││││││││││││││┌ @ lock.jl:147 slowlock(rl)
││││││││││││││││││┌ @ lock.jl:156 wait(c)
│││││││││││││││││││┌ @ condition.jl:125 Base.:(var"#wait#624")(false, #self#, c)
││││││││││││││││││││┌ @ condition.jl:132 Base.list_deletefirst!(ct.queue, ct)
│││││││││││││││││││││┌ @ linked_list.jl:145 isequal(h.value, val)
││││││││││││││││││││││┌ @ gcutils.jl:4 isequal(%1, v)
│││││││││││││││││││││││ runtime dispatch detected: isequal(%1::Any, v::Task)::Bool
││││││││││││││││││││││└────────────────
││││││││││││││││││││┌ @ condition.jl:132 Base.list_deletefirst!(%43, %37)
│││││││││││││││││││││ runtime dispatch detected: Base.list_deletefirst!(%43::Any, %37::Task)::Any
││││││││││││││││││││└────────────────────
││││││││││││││││┌ @ lock.jl:269 unlock(temp)
│││││││││││││││││┌ @ lock.jl:177 _unlock(rl)
││││││││││││││││││┌ @ lock.jl:183 notifywaiters(rl)
│││││││││││││││││││┌ @ lock.jl:187  = notify(cond_wait)
││││││││││││││││││││┌ @ condition.jl:148 #self#(c, Base.nothing)
│││││││││││││││││││││┌ @ condition.jl:148 Base.:(var"#notify#625")(true, false, #self#, c, arg)
││││││││││││││││││││││┌ @ condition.jl:148 notify(c, arg, all, error)
│││││││││││││││││││││││┌ @ condition.jl:154 Core.kwcall(NamedTuple{(:error,)}(tuple(error)), schedule, t, arg)
││││││││││││││││││││││││┌ @ task.jl:827 Base.:(var"#schedule#652")(_7, _3, t, arg)
│││││││││││││││││││││││││┌ @ task.jl:831 Base.list_deletefirst!(%10, t)
││││││││││││││││││││││││││ runtime dispatch detected: Base.list_deletefirst!(%10::Any, t::Task)::Any
│││││││││││││││││││││││││└───────────────
││││││││││││┌ @ logging.jl:364 Base.CoreLogging.logging_error(logger, level, _module, group, id, file, line, err, true)
│││││││││││││┌ @ logging.jl:466 Core.kwcall(%54, Base.CoreLogging.handle_message, %37, Base.CoreLogging.Error, %44, %38, :logevent_error, %39, %40, %41)
││││││││││││││ runtime dispatch detected: Core.kwcall(%54::NamedTuple{(:exception,), _A} where _A<:Tuple{Tuple{Any, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}, Base.CoreLogging.handle_message, %37::Base.CoreLogging.AbstractLogger, Base.CoreLogging.Error, %44::Union{LazyString, String}, %38::Module, :logevent_error, %39::Tuple{Union{Ptr{Nothing}, Base.InterpreterIP}, Symbol}, %40::String, %41::Int64)::Any
│││││││││││││└──────────────────
││││││││││┌ @ /Users/me/.julia/dev/Flux/src/deprecations.jl:24 Flux.map(tuple(f), args...)
│││││││││││┌ @ abstractarray.jl:3377 Base.Generator(tuple(f), iters...)
││││││││││││┌ @ generator.jl:36 zip(tuple(I1, I2), Is...)
│││││││││││││┌ @ iterators.jl:370 Base.Iterators.Zip(a)
││││││││││││││ runtime dispatch detected: Base.Iterators.Zip(a::Tuple{typeof(-), Matrix{Float32}, Any})::Base.Iterators.Zip
│││││││││││││└────────────────────
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:98 Base.materialize!(x, Base.broadcasted(Optimisers.:-, x, x̄))
│││││││││┌ @ broadcast.jl:881 Base.Broadcast.materialize!(Base.Broadcast.combine_styles(dest, bc), dest, bc)
││││││││││┌ @ /Users/me/.julia/packages/GPUArrays/fqD8z/src/host/broadcast.jl:41 GPUArrays._copyto!(dest, GPUArrays.instantiate(Core.apply_type(GPUArrays.Broadcasted, Style)(bc.f, bc.args, GPUArrays.axes(dest))))
│││││││││││┌ @ /Users/me/.julia/packages/GPUArrays/fqD8z/src/host/broadcast.jl:49 Base.Broadcast.throwdm(GPUArrays.axes(dest), GPUArrays.axes(bc))
││││││││││││┌ @ broadcast.jl:1050 string("destination axes ", axdest, " are not compatible with source axes ", axsrc)
│││││││││││││┌ @ strings/io.jl:185 Base.print_to_string(xs...)
││││││││││││││┌ @ strings/io.jl:144 print(%74, %77)
│││││││││││││││ runtime dispatch detected: print(%74::IOBuffer, %77::Any)::Any
││││││││││││││└─────────────────────
│││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92 Optimisers.subtract!(x, %52)
││││││││ runtime dispatch detected: Optimisers.subtract!(x::Matrix{Float32}, %52::Any)::Any
│││││││└──────────────────────────────────────────────────────────────────
││││┌ @ tuple.jl:302 map(f, Base.tail(t), Base.tail(s))
│││││┌ @ tuple.jl:299 f(t[1], s[1])
││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 Core.kwcall(NamedTuple{(:grads, :params)}(tuple(getfield(#self#, :grads), getfield(#self#, :params))), Optimisers._update!, tᵢ, xᵢ)
│││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:87 Optimisers.:(var"#_update!#11")(grads, params, _3, ℓ, x)
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92  = Optimisers.subtract!(x, x̄′)
│││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:98 Base.broadcasted(Optimisers.:-, x, x̄)
││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:339 Zygote.map(#91, gs)
│││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:354 Zygote.map!(tuple(f, gsout, gs1), gss...)
││││││││││││┌ @ /Users/me/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:361 f(%46)
│││││││││││││ runtime dispatch detected: f::Zygote.var"#91#92"{typeof(-), Vector{Float32}}(%46::Any)::Any
││││││││││││└────────────────────────────────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:92 Optimisers.subtract!(x, %52)
│││││││││ runtime dispatch detected: Optimisers.subtract!(x::Vector{Float32}, %52::Any)::Any
││││││││└──────────────────────────────────────────────────────────────────
││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 x′′ = re(Optimisers.map(#9, tree, x′))
│││┌ @ /Users/me/.julia/packages/Functors/V2McK/src/functor.jl:19 (y[1], y[2], y[3])
││││┌ @ /Users/me/.julia/dev/Flux/src/layers/basic.jl:165 init(out, in)
│││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:106 Flux.:(var"#glorot_uniform#4")(tuple(pairs(NamedTuple()), #self#), dims...)
││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:106 Flux.glorot_uniform(tuple(Flux.default_rng_value()), dims...)
│││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:102 Flux.:(var"#glorot_uniform#3")(tuple(1, #self#, rng), dims...)
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 Flux.sum(%4)
│││││││││ runtime dispatch detected: Flux.sum(%4::Tuple{Integer, Bool})::Any
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 24.0 Flux.:/ %5
│││││││││ runtime dispatch detected: (24.0 Flux.:/ %5::Any)::Any
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 Flux.sqrt(%6)
│││││││││ runtime dispatch detected: Flux.sqrt(%6::Any)::Any
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:103 %1 Flux.:* %7
│││││││││ runtime dispatch detected: (%1::Float32 Flux.:* %7::Any)::Any
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Flux.rand([quote], , %9, %10)
│││││││││ runtime dispatch detected: Flux.rand([quote]::Random.TaskLocalRNG, ::Float32, %9::Bool, %10::Integer)::Matrix{Float32}
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Base.broadcasted(Flux.:*, %15, %8)
│││││││││ runtime dispatch detected: Base.broadcasted(Flux.:*, %15::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2}, Nothing, typeof(-), Tuple{Matrix{Float32}, Float32}}, %8::Any)::Any
││││││││└──────────────────────────────────────────────
││││││││┌ @ /Users/me/.julia/dev/Flux/src/utils.jl:104 Base.materialize(%16)
│││││││││ runtime dispatch detected: Base.materialize(%16::Any)::Any
││││││││└──────────────────────────────────────────────
││┌ @ /Users/me/.julia/packages/Optimisers/ab3IF/src/interface.jl:80 #186(%30)
│││ runtime dispatch detected: #186(%30::NamedTuple{(:weight, :bias, :σ), _A} where _A<:Tuple{Any, Any, Any})::Dense
││└──────────────────────────────────────────────────────────────────