While updating my python (3.11) environment from diffrax v0.5.1 and jax v0.4.29 to the latest versions I've ran into some problems regarding times for jit compilation of the diffeqsolve function and the solve itself. By backtracking the individual releases, I have noticed two major increases in compile times, one when updating from diffrax v0.6.0 to v0.6.1 and one when updating jax from v0.4.38 to v0.5.3 (I did not check all intermediate releases yet). The duration for the ODE solve itself increases significantly when updating from v0.6.1 to v0.6.2.
Below are the jit and solve times for integration of a neural ode model with time-dependent inputs. The solver used is Tsit5 with a PI step size controller. All runs report the exact same amount of solver steps taken. I'm using double precision due to issues with very small integration step sizes.
********************************
[Compiling module jit_diffeqsolve] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
Scanning through the changelog I could not find any changes that may cause this behaviour. Do you have any idea what might cause these problems?
The code used for testing is below. I can also provide the data and test model if needed. If you are wondering why I'm vmapping the ODEs rhs instead of the solve, I have found that this is significantly faster for my use case (atleast it used to be, maybe it's time to re-evaluate...). However, as only one (very long) input trajectory is used in the test case below, this shouldn't be relevant to this issue.
import json
from time import time
import diffrax as dfx
import equinox as eqx
import jax
import jax.nn as jnn
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
class ODEFunc(eqx.Module):
mlp: eqx.nn.MLP
def __init__(self, state_size, input_size, width_size, depth, *, key, **kwargs):
super().__init__(**kwargs)
mlp = eqx.nn.MLP(
in_size=state_size + input_size,
out_size=state_size,
width_size=width_size,
depth=depth,
activation=jnn.elu,
key=key,
)
self.mlp = mlp
def _forward(self, y, u):
yu = jnp.concatenate((y, u), axis=0)
return self.mlp(yu)
def __call__(self, t, y, u_t):
return jax.vmap(self._forward)(y, u_t(t))
class NeuralODE(eqx.Module):
func: ODEFunc
def __init__(self, state_size, input_size, width_size, depth, *, key, **kwargs):
super().__init__(**kwargs)
self.func = ODEFunc(state_size, input_size, width_size, depth, key=key)
def __call__(self, ts, y0, u_t):
solution = dfx.diffeqsolve(
dfx.ODETerm(self.func),
dfx.Tsit5(),
t0=ts[0],
t1=ts[-1],
dt0=None,
y0=y0,
args=u_t,
stepsize_controller=dfx.PIDController(rtol=1e-4, atol=1e-6, pcoeff=0.4, icoeff=0.3),
saveat=dfx.SaveAt(ts=ts),
max_steps=None
)
return solution.ys, solution.stats
def save(filename, hyperparams, model):
with open(filename, "wb") as f:
hyperparam_str = json.dumps(hyperparams)
f.write((hyperparam_str + "\n").encode())
eqx.tree_serialise_leaves(f, model)
def load(filename, constructor):
with open(filename, "rb") as f:
hyperparams = json.loads(f.readline().decode())
model = constructor(**hyperparams, key=jax.random.PRNGKey(0))
return eqx.tree_deserialise_leaves(f, model)
data = jnp.load("work/inputs.npz")
model = load("work/testmodel.eqx", NeuralODE)
ts, u, y = data["t"] - data["t"][0], data["u"], data["y"]
spline_coeffs = dfx.backward_hermite_coefficients(
ts, u[:, jnp.newaxis]
)
f_ui = dfx.CubicInterpolation(ts, spline_coeffs)
t_jit_start = time()
y_pred, stats = model(ts, y[:1], f_ui.evaluate)
t_jit_stop = time()
t_start = time()
y_pred, stats = model(ts, y[:1], f_ui.evaluate)
t_stop = time()
print(
f"Jax: {jax.__version__}, Equinox: {eqx.__version__}, Diffrax: {dfx.__version__}\n"
f"Jit + Integration time: {t_jit_stop - t_jit_start:.4f}\nIntegration time: {t_stop - t_start:.4f}\nNo. of steps: {stats['num_steps']}"
)
While updating my python (3.11) environment from diffrax v0.5.1 and jax v0.4.29 to the latest versions I've ran into some problems regarding times for jit compilation of the diffeqsolve function and the solve itself. By backtracking the individual releases, I have noticed two major increases in compile times, one when updating from diffrax v0.6.0 to v0.6.1 and one when updating jax from v0.4.38 to v0.5.3 (I did not check all intermediate releases yet). The duration for the ODE solve itself increases significantly when updating from v0.6.1 to v0.6.2.
Below are the jit and solve times for integration of a neural ode model with time-dependent inputs. The solver used is Tsit5 with a PI step size controller. All runs report the exact same amount of solver steps taken. I'm using double precision due to issues with very small integration step sizes.
For jax v0.5.3 following message is printed after some time:
Scanning through the changelog I could not find any changes that may cause this behaviour. Do you have any idea what might cause these problems?
The code used for testing is below. I can also provide the data and test model if needed. If you are wondering why I'm vmapping the ODEs rhs instead of the solve, I have found that this is significantly faster for my use case (atleast it used to be, maybe it's time to re-evaluate...). However, as only one (very long) input trajectory is used in the test case below, this shouldn't be relevant to this issue.