EnzymeAD / Reactant.jl

MIT License
66 stars 6 forks source link

`setindex!` doesn't update arguments inplace (optimisers.jl) #212

Open avik-pal opened 2 weeks ago

avik-pal commented 2 weeks ago
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
wsmoses commented 2 weeks ago

@avik-pal can you dive into this one?