You can definitely do this with Enzyme, but I suggest using DIfferentiationInterface.jl with the Enzyme backend. I find that APi to be more intuitive and you also get access to use other backends for free.
1 Like