Mutation error in Zygote hessian

Specific Domains Machine Learning
,
1

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 

## 
g=gradient(()->sum(Ξ΅β‚€([[𝐫[1]],Ξ¨)),Flux.params(Ξ¨)) 

for i in 2:N: 
     g.+=gradient(()->sum(Ξ΅β‚€([[𝐫[i]],Ξ¨)),Flux.params(Ξ¨)) ## somehow need to accumulate this gradient
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)
    @ Zygote ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:71
  [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
    @ ~/.julia/packages/Zygote/SuKWp/src/lib/grad.jl:64 [inlined]
 [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
    @ ~/.julia/packages/Zygote/SuKWp/src/lib/grad.jl:62 [inlined]
 [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
 [17] gradient(f::Function, args::Params{Zygote.Buffer{Any, Vector{Any}}})
    @ 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.

2

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);

gradient(()->sum(Ξ΅β‚€([𝐫[1]],Ξ¨)),Θ)
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)
@ Zygote ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:71
[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
@ ~/.julia/packages/Zygote/JeHtr/src/lib/grad.jl:64 [inlined]
[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
@ ~/.julia/packages/Zygote/JeHtr/src/lib/grad.jl:62 [inlined]
[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

3

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.