Nested AD #Zygote #ForwardDiff

ad.jl

using Zygote
using Flux

x0 = rand(2)
dz0 = rand(2)  # test data

model = Chain(Dense(2, 5, relu), Dense(5, 1)) |> f64
f(x) = model(x)

df(x) = gradient(z->f(z)[1], x)[1]
df(x0)  # works :)
## or use forward_jacobian
# df(x) = Zygote.forward_jacobian(f, x)[2]
# df(x0) # works :)

p = params(model)
loss(x,y) = sum(df(x) .- y)  # some loss function 
loss(x0, dz0)

gradient(()->loss(x0, dz0), p)  # ERROR: Mutating arrays is not supported :( I don't know the reason here
# gradient(()->loss(x0, dz0), p)  # ERROR: Mutating arrays is not supported :(
# I think this is because forward_jacobian will need to change dtype of x to Dual !