patrick-kidger / patrick-kidger/diffrax
Vmapped jnp.interp throws error when debugging nan's with disabled jit
Open
Nobody has claimed this yet.
question
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
We've run into a strange error when using jnp.interp and debugging nans with jit disabled. I'm aware that these issues can sometimes arise and enabling only one of these flags is typically recommended. Nevertheless, it might be worth looking into in case it's an easy fix or a sign of something more serious going on.
MWE:
%env JAX_DEBUG_NANS=1
%env JAX_DISABLE_JIT=1
import diffrax
import jax.numpy as jnp
xp = jnp.array([0,1,2])
fp = jnp.array([0.0,1.0,0.0])
def vector_field(t, u, args):
return jnp.interp(u, xp, fp)
result = diffrax.diffeqsolve(diffrax.ODETerm(vector_field),
diffrax.Euler(),
t0=0.0,
t1=1.0,
dt0=0.01,
y0=jnp.array([0.5, 1.5]))
If I disable either or both flags this runs fine. If y0 is scalar or of length 1 it also runs fine. It also runs fine if I use adjoint=diffrax.ForwardMode() or adjoint=diffrax.DirectAdjoint. Furthermore, the following also runs fine with the flags on:
jnp.interp(jnp.array([0.5, 1.5]), xp, fp)
My theory is this has something to do with RecursiveCheckpointAdjoint but it doesn't seem to simply be that jnp.interp is not reverse-mode differentiable as I can calculate jacrev just fine.
Traceback
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation Traceback (most recent call last)
<frozen runpy> in ?()
--> 198 'Could not get source, probably due dynamically evaluated source code.'
<frozen runpy> in ?()
---> 88 'Could not get source, probably due dynamically evaluated source code.'
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel_launcher.py in ?()
---> 18 """Entry point for launching an IPython kernel.
19
~/.virtualenvs/ergodic/lib/python3.13/site-packages/traitlets/config/application.py in ?()
-> 1075 app.start()
1076
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/kernelapp.py in ?()
--> 739 self.io_loop.start()
~/.virtualenvs/ergodic/lib/python3.13/site-packages/tornado/platform/asyncio.py in ?()
--> 205 self.asyncio_loop.run_forever()
/opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/base_events.py in ?()
--> 679 self._run_once()
/opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/base_events.py in ?()
-> 2027 handle._run()
/opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/events.py in ?()
---> 89 self._context.run(self._callback, *self._args)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/kernelbase.py in ?()
--> 545 await self.process_one()
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/kernelbase.py in ?()
--> 534 await dispatch(*args)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/kernelbase.py in ?()
--> 437 await result
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/ipkernel.py in ?()
--> 362 await super().execute_request(stream, ident, parent)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/kernelbase.py in ?()
--> 778 reply_content = await reply_content
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/ipkernel.py in ?()
--> 449 res = shell.run_cell(
~/.virtualenvs/ergodic/lib/python3.13/site-packages/ipykernel/zmqshell.py in ?()
--> 549 return super().run_cell(*args, **kwargs)
/var/folders/7l/46fqh6j56rv1n6xbrntn3bmm0000gn/T/ipykernel_3478/3052998300.py in ?()
---> 13 get_ipython().run_line_magic('env', 'JAX_DEBUG_NANS=1')
14 get_ipython().run_line_magic('env', 'JAX_DISABLE_JIT=1')
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py in ?()
-> 1416 final_state, aux_stats = adjoint.loop(
1417 args=args,
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_adjoint.py in ?()
--> 299 final_state = self._loop(
300 terms=terms,
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py in ?()
--> 638 final_state = outer_while_loop(
639 cond_fun, body_fun, init_state, max_steps=max_steps, buffers=_outer_buffers
/opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/contextlib.py in ?()
---> 85 return func(*args, **kwds)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/loop.py in ?()
--> 107 return checkpointed_while_loop(
~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/checkpointed.py in ?()
--> 247 body_fun_ = filter_closure_convert(body_fun_, init_val_)
248 vjp_arg = (init_val_, body_fun_)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/common.py in ?()
--> 471 buffer_val2 = body_fun(buffer_val)
472 # Needed to work with `disable_jit`, as then we lose the automatic
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py in ?()
--> 635 new_state, _, _ = body_fun_aux(state)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py in ?()
--> 349 (y, y_error, dense_info, solver_state, solver_result) = solver.step(
350 terms,
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_solver/euler.py in ?()
---> 60 y1 = (y0**ω + terms.vf_prod(t0, y0, args, control) ** ω).ω
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_term.py in ?()
--> 756 return self.term.vf_prod(t, y, args, control)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_term.py in ?()
--> 157 return self.prod(self.vf(t, y, args), control)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_term.py in ?()
--> 194 out = self.vector_field(t, y, args)
195 if jtu.tree_structure(out) != jtu.tree_structure(y):
/var/folders/7l/46fqh6j56rv1n6xbrntn3bmm0000gn/T/ipykernel_3478/3052998300.py in ?()
---> 11 return jnp.interp(u, xp, fp)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py in ?()
-> 2752 return jitted_interp(x, xp, fp, left, right, period)
2753
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py in ?()
-> 2668 i = clip(searchsorted(xp_arr, x_arr, side='right'), 1, len(xp_arr) - 1)
2669 df = fp_arr[i] - fp_arr[i - 1]
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py in ?()
-> 9992 return impl(a, v, side, dtype) # type: ignore
9993
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/vectorize.py in ?()
--> 346 result = vectorized_func(*squeezed_args)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/vectorize.py in ?()
--> 144 out = func(*args)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/vectorize.py in ?()
--> 187 return func(*args, **kwargs, **static_kwargs)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py in ?()
-> 9881 carry, _ = lax.scan(body_fun, init, (), length=n_levels,
9882 unroll=n_levels if unrolled else 1)
~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py in ?()
-> 9877 go_left = op(query, sorted_arr[mid])
JaxStackTraceBeforeTransformation: FloatingPointError: invalid value (nan) encountered in broadcast_in_dim
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
FloatingPointError Traceback (most recent call last)
Cell In[1], line 13
10 def vector_field(t, u, args):
11 return jnp.interp(u, xp, fp)
---> 13 result = diffrax.diffeqsolve(diffrax.ODETerm(vector_field),
14 diffrax.Euler(),
15 t0=0.0,
16 t1=1.0,
17 dt0=0.01,
18 y0=jnp.array([0.5, 1.5]))
[... skipping hidden 2 frame]
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/_jit.py:55, in _filter_jit_cache.<locals>.fun_wrapped(dynamic_donate, dynamic_nodonate, static)
53 *args, dummy_arg = (first_arg,) + rest_args
54 assert dummy_arg is None
---> 55 out = fun(*args, **kwargs)
56 dynamic_out, static_out = partition(out, is_array)
57 marker = jnp.array(0)
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py:1416, in diffeqsolve(terms, solver, t0, t1, dt0, y0, args, saveat, stepsize_controller, adjoint, event, max_steps, throw, progress_meter, solver_state, controller_state, made_jump, discrete_terminating_event)
1389 init_state = State(
1390 y=y0,
1391 tprev=tprev,
(...) 1409 event_mask=event_mask,
1410 )
1412 #
1413 # Main loop
1414 #
-> 1416 final_state, aux_stats = adjoint.loop(
1417 args=args,
1418 terms=terms,
1419 solver=solver,
1420 stepsize_controller=stepsize_controller,
1421 event=event,
1422 saveat=saveat,
1423 t0=t0,
1424 t1=t1,
1425 dt0=dt0,
1426 max_steps=max_steps,
1427 init_state=init_state,
1428 throw=throw,
1429 passed_solver_state=passed_solver_state,
1430 passed_controller_state=passed_controller_state,
1431 progress_meter=progress_meter,
1432 )
1434 #
1435 # Finish up
1436 #
1438 progress_meter.close(final_state.progress_meter_state)
[... skipping hidden 1 frame]
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_adjoint.py:299, in RecursiveCheckpointAdjoint.loop(***failed resolving arguments***)
295 outer_while_loop = ft.partial(
296 _outer_loop, kind="checkpointed", checkpoints=self.checkpoints
297 )
298 msg = None
--> 299 final_state = self._loop(
300 terms=terms,
301 saveat=saveat,
302 init_state=init_state,
303 max_steps=max_steps,
304 inner_while_loop=inner_while_loop,
305 outer_while_loop=outer_while_loop,
306 **kwargs,
307 )
308 if msg is not None:
309 final_state = eqxi.nondifferentiable_backward(
310 final_state, msg=msg, symbolic=True
311 )
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/diffrax/_integrate.py:638, in loop(solver, stepsize_controller, event, saveat, t0, t1, dt0, max_steps, terms, args, init_state, inner_while_loop, outer_while_loop, progress_meter)
635 new_state, _, _ = body_fun_aux(state)
636 return new_state
--> 638 final_state = outer_while_loop(
639 cond_fun, body_fun, init_state, max_steps=max_steps, buffers=_outer_buffers
640 )
641 result = final_state.result
643 if event is None or event.root_finder is None:
File /opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/contextlib.py:85, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
82 @wraps(func)
83 def inner(*args, **kwds):
84 with self._recreate_cm():
---> 85 return func(*args, **kwds)
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/loop.py:107, in while_loop(***failed resolving arguments***)
105 elif kind == "checkpointed":
106 del kind, base
--> 107 return checkpointed_while_loop(
108 cond_fun,
109 body_fun,
110 init_val,
111 max_steps=max_steps,
112 buffers=buffers,
113 checkpoints=checkpoints,
114 )
115 elif kind == "bounded":
116 del kind, checkpoints
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/checkpointed.py:249, in checkpointed_while_loop(***failed resolving arguments***)
247 body_fun_ = filter_closure_convert(body_fun_, init_val_)
248 vjp_arg = (init_val_, body_fun_)
--> 249 final_val_ = _checkpointed_while_loop(
250 vjp_arg, cond_fun_, checkpoints, buffers_, max_steps
251 )
252 _, _, _, final_val = _stop_gradient_on_unperturbed(init_val_, final_val_, body_fun_)
253 return final_val
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/_ad.py:1107, in filter_custom_vjp.__call__(self, vjp_arg, *args, **kwargs)
1103 array_args_kwargs, nonarray_args_kwargs = partition((args, kwargs), is_array)
1104 array_args_kwargs = nondifferentiable(
1105 array_args_kwargs, name="`*args` and `**kwargs` to `filter_custom_vjp`"
1106 )
-> 1107 out = self.fn_wrapped(
1108 nonarray_vjp_arg,
1109 nonarray_args_kwargs,
1110 diff_array_vjp_arg,
1111 nondiff_array_vjp_arg,
1112 array_args_kwargs,
1113 )
1114 diff_array_out, nondiff_array_out, nonarray_out = out
1115 return combine(diff_array_out, nondiff_array_out, nonarray_out.value)
[... skipping hidden 11 frame]
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/checkpointed.py:271, in _checkpointed_while_loop(***failed resolving arguments***)
268 while_loop = jax.named_call(lax.while_loop, name="checkpointed-no-vjp")
269 # Hashable wrapper; JAX issue #13554 and
270 # https://github.com/patrick-kidger/equinox/issues/768
--> 271 return while_loop(lambda x: cond_fun(x), lambda x: body_fun(x), init_val)
File /opt/homebrew/Cellar/python@3.13/3.13.0_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/contextlib.py:85, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
82 @wraps(func)
83 def inner(*args, **kwds):
84 with self._recreate_cm():
---> 85 return func(*args, **kwds)
[... skipping hidden 2 frame]
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/equinox/internal/_loop/checkpointed.py:271, in _checkpointed_while_loop.<locals>.<lambda>(x)
268 while_loop = jax.named_call(lax.while_loop, name="checkpointed-no-vjp")
269 # Hashable wrapper; JAX issue #13554 and
270 # https://github.com/patrick-kidger/equinox/issues/768
--> 271 return while_loop(lambda x: cond_fun(x), lambda x: body_fun(x), init_val)
[... skipping hidden 8 frame]
File ~/.virtualenvs/ergodic/lib/python3.13/site-packages/jax/_src/api.py:121, in _nan_check_posthook(fun, args, kwargs, output)
119 f = fun._fun
120 if getattr(f, '_apply_primitive', False):
--> 121 raise FloatingPointError(f"invalid value ({e.ty}) encountered in {f.__qualname__}") from None
122 # compiled_fun can only raise in this case
123 dispatch.maybe_recursive_nan_check(e, f, args, kwargs)
FloatingPointError: invalid value (nan) encountered in broadcast_in_dim
@aidancrilly who first noticed this issue
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the MWE with JAX_DEBUG_NANS and JAX_DISABLE_JIT enabled, then start in diffrax/_adjoint.py at RecursiveCheckpointAdjoint.loop and follow the loop into diffrax/_integrate.py. Compare this path with the working adjoints named in the issue and the jnp.interp traceback. Done means the vector-valued example no longer raises the erroneous FloatingPointError and has regression coverage.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100