Speeding Up Evaluation of Padded Time Series Using Events?

I am currently training an online NCDE on a set of time series that are irregularly sampled and have varying lengths. To batch these time series, I pad them with the last (t_{N_i}, x_{N_i}) pair such that each time series has the same number of observations N_max, where N_i and N_max are the number of observations in the i'th and most observed time series, respectively.

I need to evaluate the NCDE at each observation in a time series due to the online nature of the model. Since in most cases N_i << N_max, this results in a large amount of wasted calculation. My thinking was to use an Event to stop the integration early for each time series. I've implemented a MWE below, where I've replaced the NCDE with an ODE.

import jax
jax.config.update("jax_enable_x64", True)

import diffrax
import numpy as np
from time import time

def f(t, y, args):
    return -y

N = 50
step = 0.2
cols = np.arange(N+1) * step
rows = np.arange(N+1)[:, None]
ts = np.clip(cols, 0, rows)[1:]

@jax.jit
def solve_ode_with_event(ts, idx_max):

    def cond_fn(t, y, args, **kwargs):

        return t >= jax.lax.dynamic_index_in_dim(ts, idx_max, keepdims=False)

    term = diffrax.ODETerm(f)
    solver = diffrax.Tsit5()
    saveat = diffrax.SaveAt(ts=ts)
    event = diffrax.Event(cond_fn)

    y0 = 1
    solution = diffrax.diffeqsolve(
        term,
        solver,
        t0=ts[0],
        t1=ts[-1],
        dt0=0.1,
        y0=y0,
        saveat=saveat,
        event=event
    )

    return solution

@jax.jit
def solve_ode(ts):

    term = diffrax.ODETerm(f)
    solver = diffrax.Tsit5()
    saveat = diffrax.SaveAt(ts=ts)
    stepsize_controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)

    y0 = 1
    solution = diffrax.diffeqsolve(
        term,
        solver,
        t0=ts[0],
        t1=ts[-1],
        dt0=0.1,
        y0=y0,
        saveat=saveat,
        stepsize_controller=stepsize_controller
    )

    return solution

_ = solve_ode(ts[0])
_ = solve_ode_with_event(ts[0], 5)

n_runs = int(1e4)

t_odes = np.zeros(N)
for i in range(N):
    print(i)
    t0 = time()
    for j in range(n_runs):
        _ = solve_ode(ts[i])
    t_ode = (time() - t0) / n_runs
    t_odes[i] = t_ode

t_events = np.zeros(N)
for i in range(N):
    t0 = time()
    for j in range(n_runs):
        _ = solve_ode_with_event(ts[i], i+1)
    t_event = (time() - t0) / n_runs
    t_events[i-1] = t_event

speed_ups = t_odes / t_events

Running this example, I find that the event-based approach is slower across the board compared to simply just running the diffeqsolve over the padded time series and does not seem to vary significantly with the choice of time to stop integration at.

As far as I understand, in the case where an event is detected then the remaining steps of the solver are calculated using the value of the ode at t_event to ensure static array shapes, after which these steps are masked out. I'd think the observed slow down when using Events is then simply due to the additional overhead of evaluating events at each step.

Is this understanding correct? Or would we be able to get speed ups using this approach for more costly ODE evaluations? This relates somewhat to #601 .

Start with the supplied MWE and compare the compiled behavior of diffrax.diffeqsolve with and without diffrax.Event across different stopping indices. Trace the event and padded-save paths, then establish whether event overhead or masked solver work dominates; done means a supported conclusion about the performance tradeoff, ideally with a reproducible benchmark.