Saving metrics during solve

Hi, thanks again for the library.

we have a numerical integration we perform using lax.scan. We’d like to port it over to diffrax. I wanted to get your thoughts on the shortest path to this implementation.

From https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html

def scan(f, init, xs, length=None):
  if xs is None:
    xs = [None] * length
  carry = init
  ys = []
  for x in xs:
    carry, y = f(carry, x)
    ys.append(y)
  return carry, np.stack(ys)

we rely quite a bit on the implicit append and stack calls that happen inside lax.scan. I think the core of my question is, is there a way to avoid the pre-allocation as you suggest in https://github.com/patrick-kidger/diffrax/issues/60#issuecomment-1034768386 . My guess is the answer is no but figured I’d see if you had any ideas.

TY in advance!

[patrick-kidger]

Sounds good! I’ve adjusted the title of this issue to make this a feature request. (Which if you ever feel like digging into the guts of Diffrax I’d be happy to shepard a PR on.)

[patrick-kidger]

Yep! It’s this big block of code. The idea will be to change from state.ys – which is where we save our output – to state.out.

Now the bad news is that this part of the code is pretty complicated to read, as it needs to carefully make sure it gets good performance. Fortunately the good news is that you shouldn’t need to change any of that – just carefully figure out what’s what, and replace saving ys with saving save_fn(ys).

I’d suggest that the saving-function should be specified as part of the diffeqsolve(..., saveat=...) argument, and that when called it should take state as an argument.

As a nice bonus, I’ve realised that this makes the solver_state, controller_state, made_jump arguments to SaveAt superfluous: these can be obtained simply by using the appropriate metric-saving function.