# Mutation error in Zygote hessian

Hi there everyone ,

I am new to the Julia ML ecosystem and would like some help on a specific problem. I am operating with Julia v1.8, and Flux v0.13.16 .

So I am minimizing a function , in fact this function is the ground state energy.

The method I am developing requires the evaluation of a custom gradient, it looks like (the following code is an example and is pretty sloppy) :

``````
N=5
π«=rand(Float32,N)

struct NeuralAnsatz
chain::Chain
end

function (m::NeuralAnsatz)(x)
return exp.(m.chain(x)).^2
end

Flux.@functor NeuralAnsatz

βΒ²(g::NeuralAnsatz,π±::Vector)=sum(Diagonal(hessian(π±->sum(g(π±)),π±)))

## example
chain = Chain(Dense(1, 1,relu))
Ξ¨ = NeuralAnsatz(chain)

Ο(π±::Vector)=1/2*π±[1]^2
HΜ(π±::Vector, Ο::NeuralAnsatz)=-βΒ²(Ο,π±)/2 .+Ο(π±)*Ο(π±)
Ξ΅β(π±::Vector,Ο::NeuralAnsatz)=Ο(π±).^-1 .*HΜ(π±,Ο) ## loss function

##

for i in 2:N:
end

g=g./N
# ##
``````

The problem is that the code throws up a typical error about mutating arrays, I know that the mutation comes from :

`HΜ(π±::Vector, Ο::NeuralAnsatz)=-βΒ²(Ο,π±)/2 .+Ο(π±)*Ο(π±)`

I am guessing that even though there is no explicit mutation, that somewhere under the hood a mutation occurs to vector π± through hessian function. How can I deal with this? Will I have to do some crazy stuff like define my own chain rule?

There is another way of evaluating the gradient of the loss and maybe if there is no good way here I can make another post for that.

Here is the stack trace for :

``````gradient(()->sum(βΒ²(Ο,[1.0])/2),Flux.params(Ξ¨))
``````
``````ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Float64}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)

Possible fixes:
- avoid mutating operations (preferred)
- or read the documentation and solutions for this error
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
[1] error(s::String)
@ Base ./error.jl:35
[2] _throw_mutation_error(f::Function, args::Matrix{Float64})
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/lib/array.jl:88
[3] (::Zygote.var"#550#551"{Matrix{Float64}})(#unused#::Nothing)
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/lib/array.jl:100
[4] (::Zygote.var"#2620#back#552"{Zygote.var"#550#551"{Matrix{Float64}}})(Ξ::Nothing)
[5] Pullback
@ ~/.julia/packages/Zygote/SuKWp/src/lib/forward.jl:31 [inlined]
[6] (::Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}, Val{1}}, Any})(Ξ::Tuple{Nothing, Diagonal{Float64, FillArrays.Fill{Float64, 1, Tuple{Base.OneTo{Int64}}}}})
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface2.jl:0
[7] Pullback
@ ~/.julia/packages/Zygote/SuKWp/src/lib/forward.jl:44 [inlined]
[8] Pullback
@ ~/.julia/packages/Zygote/SuKWp/src/lib/forward.jl:42 [inlined]
[9] Pullback
[10] (::Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1473", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}, Val{1}}, Any}, Zygote.var"#1891#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Nothing, false}}, Zygote.var"#1998#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float64}}}}})(Ξ::Diagonal{Float64, FillArrays.Fill{Float64, 1, Tuple{Base.OneTo{Int64}}}})
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface2.jl:0
[11] Pullback
[12] Pullback
@ ./REPL[8]:1 [inlined]
[13] (::Zygote.Pullback{Tuple{typeof(βΒ²), NeuralAnsatz, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(hessian), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1473", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}, Val{1}}, Any}, Zygote.var"#1891#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Nothing, false}}, Zygote.var"#1998#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float64}}}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#17#18"{NeuralAnsatz}, Nothing, false}}, Zygote.var"#2987#back#777"{Zygote.var"#771#775"{Diagonal{Float64, Vector{Float64}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2050"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}}})(Ξ::Float64)
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface2.jl:0
[14] Pullback
@ ./REPL[863]:1 [inlined]
[15] (::Zygote.Pullback{Tuple{var"#389#390"}, Tuple{Zygote.Pullback{Tuple{typeof(sum), Float64}, Tuple{Zygote.Pullback{Tuple{Base.var"##sum#267", Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, typeof(sum), Float64}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.var"#2331#back#414"{Zygote.var"#pairs_namedtuple_pullback#413"{(), NamedTuple{(), Tuple{}}}}}}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float64, Float64, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, NeuralAnsatz}}, Zygote.ZBack{ChainRules.var"#vect_pullback#1369"{1, Tuple{ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, Zygote.Pullback{Tuple{typeof(βΒ²), NeuralAnsatz, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(hessian), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1473", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}, Val{1}}, Any}, Zygote.var"#1891#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Nothing, false}}, Zygote.var"#1998#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float64}}}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#17#18"{NeuralAnsatz}, Nothing, false}}, Zygote.var"#2987#back#777"{Zygote.var"#771#775"{Diagonal{Float64, Vector{Float64}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2050"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}}}}})(Ξ::Float64)
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface2.jl:0
[16] (::Zygote.var"#118#119"{Params{Zygote.Buffer{Any, Vector{Any}}}, Zygote.Pullback{Tuple{var"#389#390"}, Tuple{Zygote.Pullback{Tuple{typeof(sum), Float64}, Tuple{Zygote.Pullback{Tuple{Base.var"##sum#267", Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, typeof(sum), Float64}, Any}, Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.var"#2331#back#414"{Zygote.var"#pairs_namedtuple_pullback#413"{(), NamedTuple{(), Tuple{}}}}}}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float64, Float64, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, NeuralAnsatz}}, Zygote.ZBack{ChainRules.var"#vect_pullback#1369"{1, Tuple{ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, Zygote.Pullback{Tuple{typeof(βΒ²), NeuralAnsatz, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(hessian), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#17#18"{NeuralAnsatz}, Vector{Float64}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.var"#1955#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1473", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Vector{Float64}, Val{1}}, Any}, Zygote.var"#1891#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#17#18"{NeuralAnsatz}}, Nothing, false}}, Zygote.var"#1998#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float64}}}}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#17#18"{NeuralAnsatz}, Nothing, false}}, Zygote.var"#2987#back#777"{Zygote.var"#771#775"{Diagonal{Float64, Vector{Float64}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float64}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2050"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}}}}}, Zygote.Context{true}})(Ξ::Float64)
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface.jl:389
@ Zygote ~/.julia/packages/Zygote/SuKWp/src/compiler/interface.jl:97
[18] top-level scope
@ REPL[863]:1
``````

I edited my comment per the response in order to provide all necessary information.

I think the problem here is actually nesting Zygote over itself, which should work but sometimes doesnβt. Here is a simplified version that shows a similar error:

``````using Flux
using Zygote
using LinearAlgebra

# nest struct definition in a submodule so we can edit it during development by
# changing the submodule name.
module sub1
using ..Flux: Chain
struct NeuralAnsatz
chain::Chain
end
end

# Not good for performance since this is an untyped global but fine for now
NeuralAnsatz = sub1.NeuralAnsatz

function (m::NeuralAnsatz)(x)
return exp.(m.chain(x)).^2
end

Flux.@functor NeuralAnsatz

# All the other methods were not called in this example (and I suspect defining
# β(::Function, ...) is a mistake).
βΒ²(g::NeuralAnsatz,π±::Vector)=sum(Diagonal(hessian(π±->sum(g(π±)),π±)))

chain = Chain(Dense(1, 1,relu))
Ξ¨ = NeuralAnsatz(chain)
Ξ = Flux.params(Ξ¨)

Ο(π±::Vector)=1/2*π±[1]^2
HΜ(π±::Vector, Ο::NeuralAnsatz)=-βΒ²(Ο,π±)/2 .+Ο(π±)*Ο(π±)
Ξ΅β(π±::Vector,Ο::NeuralAnsatz)=Ο(π±).^-1 .*HΜ(π±,Ο) ## loss function

# You didn't define either N or metropolis_hastings, but this is enough to get
# going
N = 10
π« = randn(Float32, N);

``````
Stacktrace

ERROR: Mutating arrays is not supported β called setindex!(Matrix{Float32}, β¦)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= β¦)

Possible fixes:

• avoid mutating operations (preferred)
• or read the documentation and solutions for this error
Limitations Β· Zygote

Stacktrace:
[1] error(s::String)
@ Base ./error.jl:35
[2] _throw_mutation_error(f::Function, args::Matrix{Float32})
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/lib/array.jl:88
[3] (::Zygote.var"#551#552"{Matrix{Float32}})(#unused#::Nothing)
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/lib/array.jl:100
[4] (::Zygote.var"#2659#back#553"{Zygote.var"#551#552"{Matrix{Float32}}})(Ξ::Nothing)
[5] Pullback
@ ~/.julia/packages/Zygote/JeHtr/src/lib/forward.jl:31 [inlined]
[6] (::Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}, Val{1}}, Any})(Ξ::Tuple{Nothing, Diagonal{Float32, FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}}}})
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/compiler/interface2.jl:0
[7] Pullback
@ ~/.julia/packages/Zygote/JeHtr/src/lib/forward.jl:44 [inlined]
[8] Pullback
@ ~/.julia/packages/Zygote/JeHtr/src/lib/forward.jl:42 [inlined]
[9] Pullback
[10] (::Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1990#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1464", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}, Val{1}}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2033#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float32}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Nothing, false}}}})(Ξ::Diagonal{Float32, FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}}})
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/compiler/interface2.jl:0
[11] Pullback
[12] Pullback
@ ./REPL[16]:1 [inlined]
[13] (::Zygote.Pullback{Tuple{typeof(βΒ²), Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.var"#3027#back#778"{Zygote.var"#772#776"{Diagonal{Float32, Vector{Float32}}}}, Zygote.Pullback{Tuple{typeof(hessian), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1990#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1464", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}, Val{1}}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2033#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float32}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Nothing, false}}}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float32}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2049"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{var"#11#12"{Main.sub1.NeuralAnsatz}, Nothing, false}}}})(Ξ::Float32)
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/compiler/interface2.jl:0
[14] Pullback
@ ./REPL[23]:1 [inlined]
[15] (::Zygote.Pullback{Tuple{typeof(Δ€), Vector{Float32}, Main.sub1.NeuralAnsatz}, Tuple{Zygote.Pullback{Tuple{typeof(Ο), Vector{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#literal_pow_pullback#327"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback2#1346"{Float64, Float32}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float64, Float64, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.literal_getindex), Vector{Float32}, Val{1}}, Tuple{Zygote.var"#2610#back#529"{Zygote.var"#539#541"{1, Float32, Vector{Float32}, Tuple{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float64}}, Tuple{}}, Zygote.Pullback{Tuple{Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Vector{Float32}}, Tuple{Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#1162"{Tuple{Nothing, Nothing}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Vector{Float32}}, Tuple{}}, Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#4123#back#1356"{Zygote.var"#bc_fwd_back#1394"{Vector{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Vector{Float32}}, Val{1}}}}}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2881#back#684"{Zygote.var"#map_back#678"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:chain, Zygote.Context{true}, Main.sub1.NeuralAnsatz, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:layers, Zygote.Context{true}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#2033#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#61"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#size_pullback#919"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:Ο, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.var"#3754#back#1177"{Zygote.var"#1171#1175"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Tuple{Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Pair{Int64, Int64}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.ZBack{Flux.var"#177#178"}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:bias, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.ZBack{Flux.var"#_size_check_pullback#204"{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}, Pair{Int64, Int64}}}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#1481"{Matrix{Float32}, Vector{Float32}}}}}}}}}, Zygote.var"#3882#back#1235"{Zygote.var"#1231#1234"{2, Vector{Float32}}}}}, Zygote.ZBack{ChainRules.var"#times_pullback#1498"{Float64, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float32, Float32, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(βΒ²), Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.var"#3027#back#778"{Zygote.var"#772#776"{Diagonal{Float32, Vector{Float32}}}}, Zygote.Pullback{Tuple{typeof(hessian), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1990#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1464", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}, Val{1}}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2033#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float32}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Nothing, false}}}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float32}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2049"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{var"#11#12"{Main.sub1.NeuralAnsatz}, Nothing, false}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#1344"{Int64, ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}}}, Zygote.var"#3754#back#1177"{Zygote.var"#1171#1175"{Tuple{Float32, Vector{Float64}}}}}})(Ξ::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/compiler/interface2.jl:0
[16] Pullback
@ ./REPL[24]:1 [inlined]
[17] (::Zygote.Pullback{Tuple{typeof(Ξ΅β), Vector{Float32}, Main.sub1.NeuralAnsatz}, Tuple{Zygote.Pullback{Tuple{typeof(Δ€), Vector{Float32}, Main.sub1.NeuralAnsatz}, Tuple{Zygote.Pullback{Tuple{typeof(Ο), Vector{Float32}}, Tuple{Zygote.ZBack{Zygote.var"#literal_pow_pullback#327"{2, Float32}}, Zygote.ZBack{ChainRules.var"#times_pullback2#1346"{Float64, Float32}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float64, Float64, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.literal_getindex), Vector{Float32}, Val{1}}, Tuple{Zygote.var"#2610#back#529"{Zygote.var"#539#541"{1, Float32, Vector{Float32}, Tuple{Int64}}}}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float64}}, Tuple{}}, Zygote.Pullback{Tuple{Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Vector{Float32}}, Tuple{Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#1162"{Tuple{Nothing, Nothing}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Vector{Float32}}, Tuple{}}, Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#4123#back#1356"{Zygote.var"#bc_fwd_back#1394"{Vector{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Vector{Float32}}, Val{1}}}}}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2881#back#684"{Zygote.var"#map_back#678"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:chain, Zygote.Context{true}, Main.sub1.NeuralAnsatz, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:layers, Zygote.Context{true}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#2033#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#61"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#size_pullback#919"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:Ο, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.var"#3754#back#1177"{Zygote.var"#1171#1175"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Tuple{Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Pair{Int64, Int64}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.ZBack{Flux.var"#177#178"}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:bias, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.ZBack{Flux.var"#_size_check_pullback#204"{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}, Pair{Int64, Int64}}}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#1481"{Matrix{Float32}, Vector{Float32}}}}}}}}}, Zygote.var"#3882#back#1235"{Zygote.var"#1231#1234"{2, Vector{Float32}}}}}, Zygote.ZBack{ChainRules.var"#times_pullback#1498"{Float64, Vector{Float32}, ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.ZBack{ChainRules.var"#/_pullback#1336"{Float32, Float32, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(βΒ²), Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.var"#3027#back#778"{Zygote.var"#772#776"{Diagonal{Float32, Vector{Float32}}}}, Zygote.Pullback{Tuple{typeof(hessian), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.hessian_dual), var"#11#12"{Main.sub1.NeuralAnsatz}, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1990#back#190"{Zygote.var"#186#189"{Zygote.Context{true}, GlobalRef, Int64}}, Zygote.Pullback{Tuple{Zygote.var"##forward_jacobian#1464", Int64, typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}}, Tuple{Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.ZBack{ChainRules.var"#min_pullback#1342"{Bool, Bool, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}, Zygote.Pullback{Tuple{typeof(Zygote.forward_jacobian), Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Vector{Float32}, Val{1}}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#747"}}}}}, Zygote.var"#2033#back#209"{Zygote.var"#back#207"{2, 2, Zygote.Context{true}, Matrix{Float32}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Zygote.var"#121#122"{var"#11#12"{Main.sub1.NeuralAnsatz}}, Nothing, false}}}}}}, Zygote.Pullback{Tuple{Type{Diagonal}, Matrix{Float32}}, Tuple{Zygote.ZBack{ChainRules.var"#diag_pullback#2049"}, Zygote.ZBack{typeof(ChainRules._Diagonal_pullback)}}}, Zygote.var"#2214#back#309"{Zygote.Jnew{var"#11#12"{Main.sub1.NeuralAnsatz}, Nothing, false}}}}, Zygote.ZBack{ChainRules.var"#-_pullback#1344"{Int64, ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}}}, Zygote.var"#3754#back#1177"{Zygote.var"#1171#1175"{Tuple{Float32, Vector{Float64}}}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float64}}, Tuple{}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.var"#3882#back#1235"{Zygote.var"#1231#1234"{-1, Vector{Float32}}}, Zygote.var"#3802#back#1201"{Zygote.var"#1197#1200"{Vector{Float32}, Vector{Float64}}}, Zygote.Pullback{Tuple{Main.sub1.NeuralAnsatz, Vector{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcasted), typeof(exp), Vector{Float32}}, Tuple{Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.var"#combine_styles_pullback#1162"{Tuple{Nothing, Nothing}}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.broadcastable), Vector{Float32}}, Tuple{}}, Zygote.var"#2173#back#289"{Zygote.var"#287#288"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{}}, Zygote.var"#4123#back#1356"{Zygote.var"#bc_fwd_back#1394"{Vector{ForwardDiff.Dual{Nothing, Float32, 1}}, Tuple{Vector{Float32}}, Val{1}}}}}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2017#back#200"{typeof(identity)}, Zygote.var"#2881#back#684"{Zygote.var"#map_back#678"{typeof(Base.Broadcast.broadcastable), 1, Tuple{Tuple{}}, Tuple{Val{0}}, Tuple{}}}}}, Zygote.var"#1926#back#157"{Zygote.var"#153#156"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:chain, Zygote.Context{true}, Main.sub1.NeuralAnsatz, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}}, Zygote.Pullback{Tuple{Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:layers, Zygote.Context{true}, Chain{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}}, Zygote.Pullback{Tuple{typeof(Flux._applychain), Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}, Vector{Float32}}, Tuple{Zygote.var"#2033#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}, Tuple{Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{NNlib.var"#broadcasted_relu_pullback#61"{Vector{Float32}}}, Zygote.ZBack{ChainRules.var"#size_pullback#919"}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:Ο, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, typeof(relu)}}, Zygote.var"#3754#back#1177"{Zygote.var"#1171#1175"{Tuple{Vector{Float32}, Vector{Float32}}}}, Zygote.Pullback{Tuple{Type{Pair}, Int64, Int64}, Tuple{Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.Pullback{Tuple{typeof(Core.convert), Type{Int64}, Int64}, Tuple{}}, Zygote.var"#2214#back#309"{Zygote.Jnew{Pair{Int64, Int64}, Nothing, false}}, Zygote.ZBack{ChainRules.var"#fieldtype_pullback#421"}}}, Zygote.Pullback{Tuple{typeof(Base.Broadcast.materialize), Vector{Float32}}, Tuple{}}, Zygote.ZBack{Flux.var"#177#178"}, Zygote.Pullback{Tuple{typeof(NNlib.fast_act), typeof(relu), Vector{Float32}}, Tuple{}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:bias, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}}}, Zygote.ZBack{Flux.var"#_size_check_pullback#204"{Tuple{Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Vector{Float32}, Pair{Int64, Int64}}}}, Zygote.var"#2184#back#299"{Zygote.var"#back#298"{:weight, Zygote.Context{true}, Dense{typeof(relu), Matrix{Float32}, Vector{Float32}}, Matrix{Float32}}}, Zygote.ZBack{ChainRules.var"#times_pullback#1481"{Matrix{Float32}, Vector{Float32}}}}}}}}}, Zygote.var"#3882#back#1235"{Zygote.var"#1231#1234"{2, Vector{Float32}}}}}}})(Ξ::FillArrays.Fill{Float64, 1, Tuple{Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/compiler/interface2.jl:0
[18] Pullback
@ ./REPL[27]:1 [inlined]

It looks like the error occurs while taking the pullback of `Zygote.hessian_dual`, hopefully someone who has a better understanding of Zygote can make a suggestion about how to fix it. Have you tried any other autodiff tools? For third derivatives like this it might be worth looking at `TaylorDiff.jl`

Hmm I have not, but I would like to stick to zygote if possible. Also the loss function is a hessian over the position variables, but the gradient is over the implicit parameters of the neural network. So this is not just a third derivative tout court. It is I admit a weird loss function, but it is theoretically well motivated.