Hi all, I’m quite new (on and off) to Julia and I’m trying to play more consistently with it.

I have an issue with AD : I can compute the gradient with respect to the weights of a very simple neural network with Zygote but it fails with Enzyme. I do not have the same problem when I compute the gradient with respect to the input of the neural network. I guess this is related to ComponentArrays.

Would you have any clue about that ? Thanks.

Running the code below, I get the following error :

MethodError: no method matching EnzymeCore.Duplicated(::Vector{Float32}, ::Vector{Float64})

Closest candidates are:

EnzymeCore.Duplicated(::T, !Matched::T) where T at ~/data/.julia/packages/EnzymeCore/cwKkU/src/EnzymeCore.jl:64

Here is the code :

using Lux, ComponentArrays, Random, Enzyme

import Zygote

rng = Random.MersenneTwister(1234)

# Define a basic neural network structure

NN = Lux.Chain( Lux.Dense(5 => 5, tanh),

Lux.Dense(5 => 1) )

# Setup the network

ps, st = Lux.setup(rng, NN)

# Test the intialized network with some input values

xtest = [0.1, 0.2, 0.3, 0.4, 0.5]

NN(xtest, ps, st)[1][1]

# No problem to compute the gradient with respect to the network input with Enzyme…

dx = zeros(size(xtest)[1])

autodiff(Reverse, xt → NN(xt, ps, st)[1][1], Active, Duplicated(xtest, dx))

dx

#…or Zygote

Zygote.gradient(xt → NN(xt, ps, st)[1][1], xtest)

############################################################

# Now let’s try to differentiate wrt the weights of the neural network

ax_test = getaxes( ComponentArray(ps) )

θ_test = getdata( ComponentArray(ps) )

# That works fine with Zygote…

Zygote.gradient( θ → NN(xtest, ComponentArray(θ, ax_test), st)[1][1], θ_test )

# …but not with Enzyme

dθ = zeros(size(θ_test))

autodiff(Reverse, θ → NN(xtest, ComponentArray(θ, ax_init), st)[1][1], Active, Duplicated(θ_test, dθ))