Import issue

[anvvalade]

In ode_util.py:

from jax._src import linear_util as lu

should be replaced by

from jax.extend import linear_util as lu

to avoid issues with debug_info in wrap_init:

https://github.com/jax-ml/jax/blob/main/jax/extend/linear_util.py