Leave some elements of `jax` array unchanged during `diffeqsolve` without computation

[BenjaminDAnjou]

I have a differential system where the solution ys is a simple jax.Array. I need to repeatedly evolve this array using a sequence of diffeqsolve, each time evolving a different subarray based on a condition. The simplest way is to set the vector field to zero when I don't want an element to change. Here is an example:

ys = jnp.ones((10,))
thresh = jnp.linspace(0, 1, 10)

def vf(t, ys, args):
    thresh = args
    return jnp.where(ys > thresh, -ys, jnp.zeros_like(ys))

term = dx.ODETerm(vf)
solver = dx.Tsit5()
stepsize_controller = dx.PIDController(rtol=1e-8, atol=1e-8)

sol = dx.diffeqsolve(term, solver, t0=0., t1=1., dt0=0.01, y0=ys, stepsize_controller=stepsize_controller, args=thresh)

sol = dx.diffeqsolve(term, solver, t0=1., t1=2., dt0=0.01, y0=sol.ys, stepsize_controller=stepsize_controller, args=jnp.flip(thresh))

This gives the desired result. However, it seems wasteful to evolve the elements for which the vector field is zero.

I could try to split the relevant subarray and only evolve that part, but the size of that subarray could vary between iterations, so I'd rather operate on the whole array.

Is there an intelligent way to operate diffeqolve on a subarray without actually propagating the rest, or at least propagating the rest at marginal cost? Or is the best way to set the vector field to zero?

Another option is to convert the array into a list and use equinox.partition. But I would like to avoid changing the structure if I can. Basically, my question is: is it possible to do something like equinox.partition at the jax.Array level?

[patrick-kidger]

If the threshold is known at compile time, and you're willing to compile each diffeqsolve separately (which will increase compilation time proportionally), then just writing out each diffeqsolve with a different vector field is probably the optimal approach.

If the theshold is not known at compile time and/or you're not willing to pay the compilation overhead of each separate diffeqsolve, then probably a lax.scan over your diffeqsolves would be optimal, with your vector field looking like the one you have here.

(Either way I wouldn't suggest using a single vector field but then writing out multiple diffeqsolves, as you have it here, as then you're compiling each one separately without the benefit of sizing each vector field appropriately.)