Description

jax.debug.breakpoint() give strange results.
Some of the inputs are tensors of 1s with shape of 2048 (bfloat16). But if we do x.sum() is return 256.
But if I do that in the python script, I get the good results (2048):

# python test.py
2048 <class 'jax.numpy.bfloat16'> <class 'jaxlib.xla_extension.ArrayImpl'>
Entering jdb:
(jdb) x.sum()
256
(jdb) c
2048.0 <class 'jax.numpy.float16'> <class 'jaxlib.xla_extension.ArrayImpl'>
Entering jdb:
(jdb) x.sum()
2048.0
(jdb) c
2048.0 <class 'jax.numpy.float32'> <class 'jaxlib.xla_extension.ArrayImpl'>
Entering jdb:
(jdb) x.sum()
2048.0

To prevent all that category of issues, breakpoint should give ArrayImpl object and not numpy object.
@sharadmv

What jax/jaxlib version are you using?

upstream

Which accelerator(s) are you using?

Independent of backend.

Additional system info

No response

NVIDIA GPU info

No response