Accelerate ODE solver [What did I miss?]

Hi,

I am playing around with diffrax's ODE solving functionality. In a nutshell, I define a simple feedforward MLP with random initialization and benchmark the runtime of using it as the temporal derivatives of an ODE. I wrote the following code to record the run-time of ODE solving and got run-time around 3.7 sec, which seems much slower compared to other ODE solver frameworks.

I am new to jax and diffrax. What did I miss in my code implemenation?

import equinox as eqx
import jax
import diffrax
import jax.numpy as jnp
import time


class MLPeqx(eqx.Module):
    layers: list
    activation: callable = eqx.static_field()

    def __init__(self, hidden_dims):
        super().__init__()
        tmp_key = jax.random.split(jax.random.PRNGKey(0), len(hidden_dims) - 1)
        self.layers = [eqx.nn.Linear(hidden_dims[i], hidden_dims[i + 1], key=tmp_key[i]) for i in
                       range(len(hidden_dims) - 1)]
        self.activation = jax.nn.relu

    def __call__(self, x):
        for i in range(len(self.layers) - 1):
            x = self.activation(self.layers[i](x))
        x = self.layers[-1](x)
        return x


class ODEjax(eqx.Module):
    func: MLPeqx

    def __init__(self, hidden_dims):
        super().__init__()
        self.func = MLPeqx(hidden_dims)

    def __call__(self, t, y, args=None):
        return self.func(y)


def solve_ode(input_x, t, func, cfg):
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(func),
        cfg['method'],
        t0=t[0],
        t1=t[-1],
        y0=input_x,
        dt0=None,
        saveat=diffrax.SaveAt(ts=t),
        stepsize_controller=diffrax.PIDController(atol=cfg['atol'], rtol=cfg['rtol']),
    )
    return sol.ys


def run_diffrax(hidden_dims, input_x, t, num_t, cfg):
    t = jnp.linspace(t[0], t[1], num_t)
    func = ODEjax(hidden_dims)
    y = jax.vmap(solve_ode, in_axes=(0, None, None, None))(input_x, t, func, cfg)
    return y


if __name__ == '__main__':
    batch_size = 128
    hidden_dims = [100, 100, 100]
    input_x = jax.random.normal(jax.random.PRNGKey(0), (128, 100))

    start_time = time.time()
    run_diffrax(hidden_dims, input_x, [0.0, 1.0], 100, {
        'method': diffrax.Dopri5(),
        'atol': 1e-5,
        'rtol': 1e-5})
    end_time = time.time()

    print(f"run time = {end_time - start_time:.3f} (sec)")

Start with the provided run_diffrax and solve_ode benchmark, checking how the vmap, compilation, timing, solver configuration, and MLP setup affect the reported runtime. Compare a correctly measured run against the supplied 3.7-second result; done means identifying the missing factor or reproducing a confirmed performance issue with an actionable explanation.