Hi, I’m trying to translate some of my code in Jax over to Julia and running into some errors when trying to use Zygote with Flux neural networks.

I have two problems:

- I get an error when trying to compute the Hessian of a neural network function, and this just may be my misunderstanding of the syntax, although I’ve tried it multiple ways.
- I get an error when trying to compute the jacobian of another function that involves the jacobian of yet another function.

Here is my code:

using Flux, DiffEqFlux, Zygote

using OrdinaryDiffEq, Random

using LinearAlgebra

import Trackerstruct CustomModel

chain::Chain

endfunction(m::CustomModel)(x,dx)

input = vcat(x,dx)

return m.chain(input)

endf = Chain(

Dense(4, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 1)

)g = Chain(

Dense(4, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 1)

)h = Chain(

Dense(4, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 8, swish),

Dense(8, 1)

)Flux.@functor CustomModel

T = CustomModel(f)

U = CustomModel(g)

R = CustomModel(h)L(x,dx) = T(x,dx) .- U(x,dx)

pₖ(x,dx) = jacobian(L,x,dx)[2]function H(x, dx)

hamilt = dot(pₖ(x,dx), dx) .- L(x,dx)

endM(x,dx) = Zygote.hessian(T(x,dx), [x, dx])

function dp(x,dx)

dp = jacobian(H,x,dx)[1] .- jacobian(R,x,dx)[2]

endprintln(M(rand(2),rand(2))) # error 1

println(dp(rand(2),rand(2))) # error 2

And the errors that I get:

- ArgumentError: Cannot create a dual over scalar type Any. If the type behaves as a scalar, define ForwardDiff.can_dual(::Type{Any}) = true.
- Mutating arrays is not supported – called copyto!(::SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, _…)

I’d appreciate any advice. I’m brand new to Julia so I apologize if this is a silly question.

For reference, in jax I am using the jax.jacfwd and jax.hessian functions to accomplish the same thing.

Thank you!