Improve memory performance of ODE solvers via diffrax.diffrax.SubSaveAt

[arik-shurygin]

diffrax.SubSaveAt allows you to define a function to apply to each timestep saved, this way we do not need to load our large compartments into memory and can instead pre-process them as they are solved by the diffrax.diffeqsolve()

See example described here: https://docs.kidger.site/diffrax/api/saveat/

I remember hearing about some out of memory errors from @kokbent, this may help if that issue ever comes back.