using Optimisers, Reactant
ps = (; a = randn(10),)
ps_ra = Reactant.to_rarray(ps)
st_opt = Optimisers.setup(Optimisers.Adam(0.1), ps_ra)
gs_ra = Reactant.to_rarray((; a = randn(10),))
function run_opt!(ps, gs, st)
return last(Optimisers.update!(st, ps, gs))
end
ps_ra_copy = deepcopy(ps_ra)
run_opt!_compiled = @compile run_opt!(ps_ra, gs_ra, st_opt)
ps2 = run_opt!_compiled(ps_ra, gs_ra, st_opt)
ps_ra_copy.a .- ps2.a # Should be different
ps_ra.a .- ps2.a # Should be the same but isn't