Batch training of NODE with varying external input (forcing) per batch element

Hi,

sry for my slightly uninformed question, but I am new to the Jax ecosystem.

I have different data sets of measurements with different excitations u(t) for one dynamic system, which dynamics I want to learn. So, excitation changes, but the system (ODE->NODE) is the same.

I want to use equinox+diffrax to train a neural ODE via batching, which has an external input u, meaning the ODE is described by xdot = f(x,u(t)). The dependency u(t) from time is not known explicitly (interpolation from data has to be used) and varies per batch element.

Looking in the docu I found the forcing term and the batch training of NODEs.
My problem is how to combine both. My first hack was to map each u(t) of every batch element to non-overlapping time periods to get a unique mapping from time to the correct input time series. Then I am able to use vmap directly via

@eqx.filter_value_and_grad
def grad_loss(model, ti, yi):
    y_pred = jax.vmap(model, in_axes=(0, 0, None))(ti, yi[:, 0], input_concatenated)
    return jnp.mean((yi - y_pred) ** 2)

Are there any better options to handle this? Note, that the gradient should not be calculated wrt parameters of the interpolation object representing u(t).

Thanks. If there are any questions, let me know.

Start by comparing the forcing-term and batch-training examples in the Diffrax documentation, then inspect the grad_loss and jax.vmap usage shown in the issue. Determine whether a supported approach can batch distinct interpolated u(t) inputs while excluding interpolation parameters from differentiation; completion would require a documented, concrete approach or an explicitly scoped implementation request.