JIT

[pnkraemer]

Maybe we can write some of the code in a way that it is out-of-the-box jittable.

https://jax.readthedocs.io/en/latest/jax.html#jax.jit

[pnkraemer]

Inputs and outputs of jittable functions are jnp.arrays, scalars, or tuples/lists thereof.

To JIT the solver, there are a few speedbumps:

• the assert ... statements in EK1
• The if step_change < min_change stuff in adaptive steps.

The following works after commenting out the EK1 asserts:

@jax.jit
def solve(y0):
    ivp = tornado.ivp.vanderpol(y0=y0)

    solution, _ = tornado.ivpsolve.solve(
        ivp,
        method="ek1_diag",
        solver_order=5,
        adaptive=False,
        dt=0.1,
        benchmark_mode=True,
    )
    return solution.y.mean

y0 = jnp.array([2.0, 0.0])
solve(y0)
solve(y0)

Note: it is pointless to test this on such high level, because compilation takes ages.