Hello all. Yes, this is another question about avoiding mutation when using Zygote for AD. I am working on a problem where I will eventually need to differentiate through a time dependent simulation. As a toy problem consider
### Simulate forward Euler from x0 and record the trace.
x = zeros(T,length(tplot)) # Create simulation vector of same size as Tvals
x[1] = x0
for i = 2:length(tplot)
dt = tplot[i] - tplot[i-1]
x[i] = x[i-1] + dt * (-param[1]*x[i-1])
end
cost = sum(abs2,x)This obviously can’t be differentiated with zygote since the array that stores the simulation is mutated. Of course, an easy fix in this or similar cases would be to constantly create new vectors
### Simulate forward Euler from x0 and record the trace.
x = [x0]
for i = 2:length(tplot)
dt = tplot[i] - tplot[i-1]
xn = x[end] + dt * (-param[1]*x[end])
x = vcat(x,[xn])
endThis however generates lots of allocations.
I suspect everything related to this topic may be very problem specific. But I wanted to see if there is any generalizable advice here? Namely, when you need to 1) iteratively simulate with a for loop or similar and 2) store the results of each loop iteration for later use, is there a way to do this without mutating and without excess allocations?
On a related note, since GPU kernels cannot return values (and thus have to mutate something), how does zygote interact with GPUs? This is not a direct concern for me at the moment, just a curiosity in the same vein.
I’ve looked through a number of threads on this, but they seem to be tackling more complex mutation than I have here, and usually have complicated solutions often involving lots of allocations. Any advice you have is greatly appreciated.
Also, I’m relatively new to Julia, so apologies if this is a settled issue and I just didn’t find the right info.