Return value of the vector field at t0 without recomputation

[BenjaminDAnjou]

So I know that it's not possible to return the values of the vector field at the save times without recomputation, because in general diffrax interpolates to obtain the solution.

But shouldn't it at least possible to save the vector field at t0 without recomputation, either at the level of diffeqsolve() or at the level of step()? I am not aware of a differential equation solver that does not compute the vector field at t0.

For now, I can get around this for FSAL solvers, which contain the vector field in the solver_state. But that requires implementing a bit of annoying logic, and doesn't work for non-FSAL solvers.

So I was wondering if it was conceivable to return some output vector_field_at_t0 in the solution, or to output the value of the vector field at t0 at each step()? Or, more generally, to store the output of AbstractTerm methods at t0?

I suspect that this would be annoying to implement. But the vector field is an important quantity that I'm sure many people would like to access. So if it's available at the t0 of every step, and there is a clean way to output it, it might be worth it, and so I just wanted to put the idea out there.

[patrick-kidger]

So if I had to guess then this should probably work as you intend simply by calling the vector field yourself, and letting the XLA compiler deduplicate the call with the one inside diffeqsolve.

diffrax/diffrax/_integrate.py

Line 1268 in 5ea5fcb

solver_state = solver.init(terms, t0, tnext, y0, args)

If I had to guess then the detail that needs making sure of here is that it really is the same t0 etc., for example that it's the same dtype:

diffrax/diffrax/_integrate.py

Line 1076 in 5ea5fcb

t0 = jnp.asarray(t0, dtype=time_dtype)

Alternatively I think you could add a SubSaveAt(t0=..., fn=...); I've not checked but again that may be deduplicated. (And if not it ideally would be, so potentially there could be a PR modifying that.)

I think adding a specialized API for this is probably unlikely I'm afraid, as I suspect this is a fairly niche use-case.

[BenjaminDAnjou]

Thanks I'll look into it. If deduplication can be guaranteed at compilation, then that would be enough for me. The internals of jax compilation are still a bit opaque for me; I will read up on it.