High order derivative of Flux.Chain network

using Zygote

model = Chain(
Dense(2, 8, σ),
Dense(8, 1)
)

xs = [1,1]

g(x) = model(x)[1]

g_p(x,m) = Zygote.gradient(m → sum(m(x)), m)[1]

g_x(x,m) = Zygote.gradient(x → sum(m(x)), x)[1]

g_xx(x,m) = Zygote.diaghessian(x → sum(m(x)), x)[1]

g_xp(x,m) = Zygote.gradient(m → sum(g_x(x,m)), m)

g_xxp(x,m) = Zygote.gradient(m → sum(g_xx(x,m)), m)

Where the last derivative g_xxp() fails.

What is the right way to do this higher order derivative in Zygote?

ERROR: Mutating arrays is not supported – called setindex!(Vector{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::Vector{Float32})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\array.jl:88
[3] (::Zygote.var"#550#551"{Vector{Float32}})(#unused#::Nothing)
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\array.jl:100
[4] (::Zygote.var"#2620#back#552"{Zygote.var"#550#551"{Vector{Float32}}})(Δ::Nothing)
@ Zygote C:\Users\ASUS.julia\packages\ZygoteRules\OgCVT\src\adjoint.jl:71
[5] Pullback
@ C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\forward.jl:56 [inlined]
[6] (::Zygote.Pullback{Tuple{typeof(Zygote.extract_diag!), Vector{Float32}, Int64, Vector{ForwardDiff.Dual{Nothing, Float32, 2}}}, Any})(Δ::Nothing)
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[7] Pullback
@ C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\forward.jl:64 [inlined]
[8] (::Zygote.Pullback{Tuple{typeof(Zygote.forward_diag), Zygote.var"#134#137"{Int64, Val{1}, var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Vector{Int64}, Val{2}}, Any})(Δ::Tuple{Nothing, FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[9] Pullback
@ C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\forward.jl:75 [inlined]
[10] (::Zygote.Pullback{Tuple{typeof(Zygote.forward_diag), Zygote.var"#134#137"{Int64, Val{1}, var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Vector{Int64}}, Any})(Δ::Tuple{Nothing, FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[11] Pullback
@ C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\grad.jl:260 [inlined]
[12] (::Zygote.Pullback{Tuple{Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Int64}, Any})(Δ::FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[13] Pullback
@ .\ntuple.jl:19 [inlined]
[14] Pullback
@ C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\lib\grad.jl:257 [inlined]
[15] (::Zygote.Pullback{Tuple{typeof(diaghessian), var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Vector{Int64}}, Tuple{Zygote.Pullback{Tuple{typeof(ntuple), Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Int64}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#746"}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Nothing, false}}}})(Δ::Tuple{FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[16] Pullback
@ f:\Julia_codes\Project_EEGsolving\main02.jl:112 [inlined]
[17] (::Zygote.Pullback{Tuple{typeof(g_xx), Vector{Int64}, Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Zygote.var"#1998#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{false}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(diaghessian), var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Vector{Int64}}, Tuple{Zygote.Pullback{Tuple{typeof(ntuple), Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Int64}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#746"}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Nothing, false}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Nothing, false}}}})(Δ::FillArrays.Fill{Float32, 1, Tuple{Base.OneTo{Int64}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[18] Pullback
@ f:\Julia_codes\Project_EEGsolving\main02.jl:118 [inlined]
[19] (::Zygote.Pullback{Tuple{var"#484#485"{Vector{Int64}}, Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Zygote.Pullback{Tuple{typeof(g_xx), Vector{Int64}, Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Zygote.var"#1998#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{false}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(diaghessian), var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Vector{Int64}}, Tuple{Zygote.Pullback{Tuple{typeof(ntuple), Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Int64}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#746"}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Nothing, false}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Nothing, false}}}}, Zygote.var"#2149#back#299"{Zygote.var"#back#298"{:x, Zygote.Context{false}, var"#484#485"{Vector{Int64}}, Vector{Int64}}}, Zygote.var"#2987#back#777"{Zygote.var"#771#775"{Vector{Float32}}}}})(Δ::Float32)
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface2.jl:0
[20] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{var"#484#485"{Vector{Int64}}, Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Zygote.Pullback{Tuple{typeof(g_xx), Vector{Int64}, Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Zygote.var"#1998#back#209"{Zygote.var"#back#207"{1, 1, Zygote.Context{false}, Vector{Float32}}}, Zygote.Pullback{Tuple{typeof(diaghessian), var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Vector{Int64}}, Tuple{Zygote.Pullback{Tuple{typeof(ntuple), Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Int64}, Any}, Zygote.ZBack{ChainRules.var"#length_pullback#746"}, Zygote.var"#2176#back#309"{Zygote.Jnew{Zygote.var"#133#136"{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Tuple{Vector{Int64}}}, Nothing, false}}}}, Zygote.var"#2176#back#309"{Zygote.Jnew{var"#456#457"{Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, Nothing, false}}}}, Zygote.var"#2149#back#299"{Zygote.var"#back#298"{:x, Zygote.Context{false}, var"#484#485"{Vector{Int64}}, Vector{Int64}}}, Zygote.var"#2987#back#777"{Zygote.var"#771#775"{Vector{Float32}}}}}})(Δ::Float32)
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface.jl:45
[21] gradient(f::Function, args::Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
@ Zygote C:\Users\ASUS.julia\packages\Zygote\TSj5C\src\compiler\interface.jl:97
[22] g_xxp(x::Vector{Int64}, m::Chain{Tuple{Dense{typeof(σ), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
@ Main f:\Julia_codes\Project_EEGsolving\main02.jl:118
[23] top-level scope
@ REPL[503]:1