Error in diffeqsolve when right-hand side is not defined at time zero

[FloList]

Hi Patrick,

I noticed that there are situations in which diffeqsolve raises an error message when the right-hand side of the ODE is undefined at time 0, even though time 0 lies outside the time interval of interest.

Here is a minimal working example for an ODE whose right-hand side is well-defined for t0 <= t <= t1 with t0 = 1 and t1 = 3, but diverges at t = 0:

import jax
from jax import numpy as jnp
from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController

# @jax.jit
def model(t, y, args):
    arr = jnp.ones((1,))
    arr = arr.at[0].set(-1.0 / t)
    return -y * arr[0]

term = ODETerm(model)
solver = Dopri5()
saveat = SaveAt(ts=[1., 2., 3.])
stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)

sol = diffeqsolve(term, solver, t0=1.0, t1=3.0, dt0=0.1, y0=1, saveat=saveat, stepsize_controller=stepsize_controller)

This raises the error message "ValueError: Terms are not compatible with solver!", which can be tracked down to come from a ZeroDivisionError: float division by zero because _check tries to evaluate the right-hand side at t = 0.0, see

diffrax/diffrax/_integrate.py

Line 169 in 4a308b8

vf_type = eqx.filter_eval_shape(term.vf, 0.0, yi, args)

When jitting the model function such that t becomes a tracer, the division by zero goes unnoticed and the code runs without problems. What should probably happen is that _check evaluates the term at time t0 instead of time 0.

Thanks!

[patrick-kidger]

Agreed, and thanks for noticing!

Would you be interested in contributing a PR?

[patrick-kidger]

Fixed in #566.