Mutating arrays is not supported #814
Description
It should work, did I miss something?
code:
using DiffEqFlux, Lux, ComponentArrays
using DifferentialEquations, SciMLSensitivity
using Zygote
using Random
nn = Lux.Chain(Lux.Dense(1 => 4, tanh), Lux.Dense(4 => 1, tanh))
rng = Random.default_rng()
ps, st = Lux.setup(Random.default_rng(), nn)
ffjord = FFJORD(nn, (0.0f0, 1.0f0))
r = rand(Float32, 1, 128)
Zygote.jacobian(x -> sum(ffjord(r, x)[1]), ComponentArray(ps))