How to import jax.numpy at jax.debug.breakpoint()

[HHalva]

I often use jdb.breakpoint() to debug and it's very helpful. However, unlike e.g. pdb, or some other debugger, I'm unable to do something like this:

(jdb) P.sum()
3232.9345616809346
(jdb) P.shape
(6144, 6144)
(jdb) jnp.linalg.norm(P)
*** NameError: name 'jnp' is not defined
(jdb) import jax.numpy as jnp
*** SyntaxError: invalid syntax

Is there some way of doing something like this?

[sharadmv]

Try adding the line import jax.numpy as jnp right before you run the breakpoint. I'd also suggest using numpy not jax.numpy if you're on an accelerator, since the breakpoint uses NumPy arrays, not JAX arrays.

[HHalva]

Awesome -- that worked! Thank you.