jax.debug.breakpoint() for jax.experimental.ode.odeint?

[virajpandya]

I am trying to debug an ODE integration code that I ported from pure python (using scipy.integrate.solve_ivp) to JAX. Whereas scipy.integrate.solve_ivp works very well for different ODE solver choices, my JAX port based on jax.experimental.ode.odeint is failing and I'm trying to figure out why. Currently I have put a bunch of jax.debug.print() statements into my integrator function (the one that is fed into the solve_ivp / odeint), but this is tedious and I have not yet found my problem. It's possible I have a bug in one of the other (jitted) functions I've defined that gets called by my main integrator function since the "ys" array returned by odeint has a bunch of NaN's.

Is there a way to use jax.debug.breakpoint() to interactively inspect the values of intermediate variables computed within the integrator function fed to odeint? I tried putting jax.debug.breakpoint() in my integrator function but I'm getting an XLA RuntimeError (see below). My guess is that breakpoint() doesn't make sense for odeint but wanted to double check.

ys = odeint(integrator,initial_conditions,t_eval,parameters,rtol=rtol,atol=atol)


File /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/experimental/ode.py:179, in odeint(func, y0, t, rtol, atol, mxstep, hmax, *args)
    176   raise TypeError(f"t must be an array of floats, but got {t}.")
    178 converted, consts = custom_derivatives.closure_convert(func, y0, t[0], *args)
--> 179 return _odeint_wrapper(converted, rtol, atol, mxstep, hmax, y0, t, *args, *consts)

    [... skipping hidden 4 frame]

File /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py:2095, in ExecuteReplicated._call_with_tokens(self, input_bufs)
   2092   input_bufs = [*tokens, *input_bufs]
   2093 num_output_tokens = len(self.ordered_effects)
   2094 out_bufs, sharded_token = (
-> 2095     self.xla_executable.execute_sharded_on_local_devices_with_tokens(
   2096       input_bufs))
   2097 token_bufs, out_bufs = util.split_list(out_bufs, [num_output_tokens])
   2098 for i, device in enumerate(self._local_devices):

XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: SystemExit: 0

At:
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugger/cli_debugger.py(147): do_quit
  /opt/anaconda3/envs/jax/lib/python3.10/cmd.py(217): onecmd
  /opt/anaconda3/envs/jax/lib/python3.10/cmd.py(142): cmdloop
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugger/cli_debugger.py(163): run
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugger/cli_debugger.py(168): run_debugger
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugger/core.py(208): _breakpoint_callback
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugging.py(213): _flat_callback
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugging.py(81): debug_callback_impl
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/debugging.py(132): _callback
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/mlir.py(1798): _wrapped_callback
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py(2095): _call_with_tokens
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/interpreters/pxla.py(2111): __call__
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/profiler.py(314): wrapper
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/api.py(565): cache_miss
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/traceback_util.py(162): reraise_with_filtered_traceback
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/jax/experimental/ode.py(179): odeint
  /var/folders/8c/m_h30b8j6r3cvxgz8fj0zff00000gn/T/ipykernel_45663/2158904990.py(42): <module>
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3442): run_code
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3382): run_ast_nodes
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3203): run_cell_async
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/async_helpers.py(129): _pseudo_sync_runner
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3000): _run_cell
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/IPython/core/interactiveshell.py(2945): run_cell
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/zmqshell.py(540): run_cell
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/ipkernel.py(423): do_execute
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/kernelbase.py(729): execute_request
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/kernelbase.py(409): dispatch_shell
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/kernelbase.py(502): process_one
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/kernelbase.py(513): dispatch_queue
  /opt/anaconda3/envs/jax/lib/python3.10/asyncio/events.py(80): _run
  /opt/anaconda3/envs/jax/lib/python3.10/asyncio/base_events.py(1906): _run_once
  /opt/anaconda3/envs/jax/lib/python3.10/asyncio/base_events.py(603): run_forever
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/tornado/platform/asyncio.py(215): start
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel/kernelapp.py(728): start
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/traitlets/config/application.py(1043): launch_instance
  /opt/anaconda3/envs/jax/lib/python3.10/site-packages/ipykernel_launcher.py(17): <module>
  /opt/anaconda3/envs/jax/lib/python3.10/runpy.py(86): _run_code
  /opt/anaconda3/envs/jax/lib/python3.10/runpy.py(196): _run_module_as_main