[Question] Neural CDE for regression

Description

I would like to create a neural CDE for regression. For that, I have taken the example from neural CDE for classification and adapted using the content from neural ODE for regression.

I am encountering some issues which I do not know how to solve. I would appreciate if someone could point me in the right direction. Thank you!

Code

import time

import diffrax
import equinox as eqx  # https://github.com/patrick-kidger/equinox
import jax
import jax.nn as jnn
import jax.numpy as jnp
import jax.random as jr
import jax.scipy as jsp
import matplotlib
import matplotlib.pyplot as plt
import optax

class Func(eqx.Module):
    mlp: eqx.nn.MLP

    def __init__(self, data_size, width_size, depth, *, key, **kwargs):
        super().__init__(**kwargs)
        self.mlp = eqx.nn.MLP(
            in_size=data_size,
            out_size=data_size,
            width_size=width_size,
            depth=depth,
            activation=jnn.tanh,
            final_activation=jnn.tanh,
            key=key,
        )

    def __call__(self, t, y, args):
        return self.mlp(y)

class NeuralCDE(eqx.Module):
    initial: eqx.nn.MLP
    func: Func

    def __init__(self, data_size, width_size, depth, *, key, **kwargs):
        super().__init__(**kwargs)
        ikey, fkey, lkey = jr.split(key, 3)
        self.initial = eqx.nn.MLP(in_size=data_size, out_size=data_size, width_size=width_size, depth=depth, key=ikey)
        self.func = Func(data_size, width_size, depth, key=fkey)
        
    def __call__(self, ts, coeffs, evolving_out=False):
        # Each sample of data consists of some timestamps `ts`, and some `coeffs`
        # parameterising a control path. These are used to produce a continuous-time
        # input path `control`.
        control = diffrax.CubicInterpolation(ts, coeffs)
        term = diffrax.ControlTerm(self.func, control).to_ode()
        solver = diffrax.Tsit5()
        dt0 = ts[1] - ts[0]
        y0 = self.initial(control.evaluate(ts[0]))
        solution = diffrax.diffeqsolve(
            term,
            solver,
            ts[0],
            ts[-1],
            dt0,
            y0,
            stepsize_controller=diffrax.PIDController(rtol=1e-3, atol=1e-6),
            saveat=diffrax.SaveAt(ts=ts),
        )
        return solution.ys

# ============================================================================
def _get_data(ts, *, key):
    y0 = jr.uniform(key, (2,), minval=-0.6, maxval=1)

    def f(t, y, args):
        x = y / (1 + y)
        return jnp.stack([x[1], -x[0]], axis=-1)

    solver = diffrax.Tsit5()
    dt0 = 0.1
    saveat = diffrax.SaveAt(ts=ts)
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(f), solver, ts[0], ts[-1], dt0, y0, saveat=saveat
    )
    ys = sol.ys
    return ys


def get_data(dataset_size, *, key):
    length = 100
    ts = jnp.linspace(0, 10, length)
    key = jr.split(key, dataset_size)
    ys = jax.vmap(lambda key: _get_data(ts, key=key))(key)
    ts_broadcasted = jnp.broadcast_to(ts, (dataset_size, length))
    ys = jnp.concatenate([ts_broadcasted[:, :, None], ys], axis=-1) # time is a channel
    coeffs = jax.vmap(diffrax.backward_hermite_coefficients)(ts_broadcasted, ys)
    return ts_broadcasted, ys, coeffs

# ============================================================================

def dataloader(arrays, batch_size, *, key):
    dataset_size = arrays[0].shape[0]
    assert all(array.shape[0] == dataset_size for array in arrays)
    indices = jnp.arange(dataset_size)
    while True:
        perm = jr.permutation(key, indices)
        (key,) = jr.split(key, 1)
        start = 0
        end = batch_size
        while end < dataset_size:
            batch_perm = perm[start:end]
            yield tuple(array[batch_perm] for array in arrays)
            start = end
            end = start + batch_size

def main(
    dataset_size=256,
    batch_size=32,
    lr_strategy=(3e-3, 3e-3),
    steps_strategy=(500, 500),
    length_strategy=(0.1, 1),
    width_size=64,
    depth=2,
    seed=5678,
    plot=True,
    print_every=100,
):
    key = jr.PRNGKey(seed)
    data_key, model_key, loader_key = jr.split(key, 3)

    ts, ys, coeffs = get_data(dataset_size, key=data_key)
    _, length_size, data_size = ys.shape

    model = NeuralCDE(data_size, width_size, depth, key=model_key)

    # Training loop like normal.
    #
    # Only thing to notice is that up until step 500 we train on only the first 10% of
    # each time series. This is a standard trick to avoid getting caught in a local
    # minimum.

    @eqx.filter_jit # value_and_grad
    def loss(model, ti, yi, coeff_i):
        y_pred = jax.vmap(model, in_axes=(None, 0))(ti[0, :], coeff_i)
        # MSE without time column
        return jnp.mean((yi[:, :, 1:] - y_pred[:, :, 1:]) ** 2)

    grad_loss = eqx.filter_value_and_grad(loss, has_aux=True)

    @eqx.filter_jit
    def make_step(data_i, model, opt_state):
        ti, yi, *coeff_i = data_i
        loss, grads = grad_loss(model, ti, yi, coeff_i)
        updates, opt_state = optim.update(grads, opt_state)
        model = eqx.apply_updates(model, updates)
        return loss, model, opt_state

    for lr, steps, length in zip(lr_strategy, steps_strategy, length_strategy):
        optim = optax.adam(lr)
        opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))
        _ts = ts[:, : int(length_size * length)]
        _ys = ys[:, : int(length_size * length)]
        _coeffs = tuple(arr[:, :int(length_size * length) - 1] for arr in coeffs)
        for step, data_i in zip(
            range(steps), dataloader((_ts, _ys) + _coeffs, batch_size, key=loader_key)
        ):
            start = time.time()
            loss, model, opt_state = make_step(data_i, model, opt_state)
            end = time.time()
            if (step % print_every) == 0 or step == steps - 1:
                print(f"Step: {step}, Loss: {loss}, Computation time: {end - start}")

    if plot:
        plt.plot(ts, ys[0, :, 0], c="dodgerblue", label="Real")
        plt.plot(ts, ys[0, :, 1], c="dodgerblue")
        sample_coeffs = tuple(c[-1] for c in coeffs)
        pred = model(ts, sample_coeffs, evolving_out=True)
        plt.plot(ts, pred[:, 0], c="crimson", label="Model")
        plt.plot(ts, pred[:, 0], c="crimson")
        plt.legend()
        plt.tight_layout()
        plt.savefig("neural_ode.png")
        plt.show()

    return ts, ys, model


ts, ys, model = main()

Error

The error originates at line

return jnp.mean((yi[:, :, 1:] - y_pred[:, :, 1:]) ** 2)

The error message is quite large to write it down here. To replicate the error, please run the code above.
Note: My intention is to compute the MSE between the predicted values and the true values. The variables y and y_pred contain the time series values on the first column. Therefore, for the MSE I only use the last two columns.

Specifications

jax 0.4.35
jaxlib 0.4.35
jaxtyping 0.2.34
diffrax 0.6.0
equinox 0.11.8
numpy 2.1.2
optax 0.2.3

Start by running the supplied main() example with the listed package versions and inspect the shapes at loss(), especially yi and y_pred at the reported MSE line. The issue does not include the full error, so first capture that output and determine whether the prediction and target slices are compatible; done means the example runs through the loss calculation and computes the intended MSE.