Repository: google/jax

Issue:
jax.debug.callback changes array type
### Description

```python
import jax.numpy as jnp
from jax import Array, debug, jit

def callback(x) -> None:
    assert isinstance(x, Array)  # Fail!

@jit
def f(x):
    assert isinstance(x, Array)  # Good.
    debug.callback(callback, x)

f(jnp.ones(1))
```

For some reason, `debug.callback` rematerializes the array as a numpy array instead of `jaxlib.xla_extension.ArrayImpl` (the thing created by `jnp.asarray`).  This can cause problems if the array in the callback is mixed with jax arrays when passed to libraries that implement the Array API.

### System info (python version, jaxlib version, accelerator, etc.)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
jax: 0.4.26
jaxlib: 0.4.26
numpy: 1.26.4
python: 3.11.8 (main, Feb 22 2024, 17:25:49) [GCC 11.4.0]


Discussion:
Thanks for the report! This is working as intended: callback functions will convert JAX arrays to NumPy arrays, however this will change soon (see #20325)
@jakevdp Great to hear!   So, in an upcoming release then?   Feel free to close this if so :smile: 

Resolving patch:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -18,6 +18,7 @@
 import importlib.util
 from collections.abc import Sequence
 import functools
+import logging
 import string
 import sys
 from typing import Any, Callable, Union
@@ -25,6 +26,7 @@
 
 import numpy as np
 
+import jax
 import jax.numpy as jnp
 from jax import lax
 
@@ -45,6 +47,8 @@
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
 
+logger = logging.getLogger(__name__)
+
 class DebugEffect(effects.Effect):
   __str__ = lambda self: "Debug"
 debug_effect = DebugEffect()
@@ -73,7 +77,14 @@ class OrderedDebugEffect(effects.Effect):
 def debug_callback_impl(*args, callback: Callable[..., Any],
                         effect: DebugEffect):
   del effect
-  callback(*args)
+  cpu_device, *_ = jax.local_devices(backend="cpu")
+  args = jax.device_put(args, cpu_device)
+  with jax.default_device(cpu_device):
+    try:
+      callback(*args)
+    except BaseException:
+      logger.exception("jax.debug_callback failed")
+      raise
   return ()
 
 @debug_callback_p.def_effectful_abstract_eval


Test patch:
diff --git a/tests/debugger_test.py b/tests/debugger_test.py
--- a/tests/debugger_test.py
+++ b/tests/debugger_test.py
@@ -110,7 +110,7 @@ def f(x):
       return y
     expected = _format_multiline(r"""
     Entering jdb:
-    (jdb) array(2., dtype=float32)
+    (jdb) Array(2., dtype=float32)
     (jdb) """)
     f(jnp.array(2., jnp.float32))
     jax.effects_barrier()
@@ -126,7 +126,7 @@ def f(x):
       return y
     expected = _format_multiline(r"""
     Entering jdb:
-    (jdb) (array(2., dtype=float32), array(3., dtype=float32))
+    (jdb) (Array(2., dtype=float32), Array(3., dtype=float32))
     (jdb) """)
     f(jnp.array(2., jnp.float32))
     jax.effects_barrier()
@@ -196,7 +196,7 @@ def g\(x\):
     ->        y = f\(x\)
               return jnp\.exp\(y\)
     .*
-    \(jdb\) array\(2\., dtype=float32\)
+    \(jdb\) Array\(2\., dtype=float32\)
     \(jdb\) > .*debugger_test\.py\([0-9]+\)
             def f\(x\):
               y = jnp\.sin\(x\)
@@ -225,9 +225,9 @@ def g(x):
       return jnp.exp(y)
     expected = _format_multiline(r"""
     Entering jdb:
-    (jdb) array(3., dtype=float32)
+    (jdb) Array(3., dtype=float32)
     (jdb) Entering jdb:
-    (jdb) array(6., dtype=float32)
+    (jdb) Array(6., dtype=float32)
     (jdb) """)
     g(jnp.array(2., jnp.float32))
     jax.effects_barrier()
@@ -249,9 +249,9 @@ def g(x):
       return jnp.exp(y)
     expected = _format_multiline(r"""
     Entering jdb:
-    (jdb) array(1., dtype=float32)
+    (jdb) Array(1., dtype=float32)
     (jdb) Entering jdb:
-    (jdb) array(2., dtype=float32)
+    (jdb) Array(2., dtype=float32)
     (jdb) """)
     g(jnp.arange(2., dtype=jnp.float32))
     jax.effects_barrier()
@@ -274,9 +274,9 @@ def g(x):
       return jnp.exp(y)
     expected = _format_multiline(r"""
     Entering jdb:
-    \(jdb\) array\(.*, dtype=float32\)
+    \(jdb\) Array\(.*, dtype=float32\)
     \(jdb\) Entering jdb:
-    \(jdb\) array\(.*, dtype=float32\)
+    \(jdb\) Array\(.*, dtype=float32\)
     \(jdb\) """)
     g(jnp.arange(2., dtype=jnp.float32))
     jax.effects_barrier()
@@ -302,7 +302,7 @@ def g(x):
         out_shardings=jax.sharding.PartitionSpec("dev"),
     )
     with jax.sharding.Mesh(np.array(jax.devices()), ["dev"]):
-      arr = (1 + np.arange(8)).astype(np.int32)
+      arr = (1 + jnp.arange(8)).astype(np.int32)
       expected = _format_multiline(r"""
       Entering jdb:
       \(jdb\) {}
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -170,7 +170,7 @@ def f(x):
     with jtu.capture_stdout() as output:
       f(np.array(2, np.int32))
       jax.effects_barrier()
-    self.assertEqual(output(), f"x: {str(dict(foo=np.array(2, np.int32)))}\n")
+    self.assertEqual(output(), f"x: {str(dict(foo=jnp.array(2, np.int32)))}\n")
 
   def test_debug_print_should_use_default_layout(self):
     data = np.array(