Repository: google/jax
Issue: pure_callback passes jax.DeviceArray to the callback on CPU when not jitted

### Description

As title. 
This only happens when the callback is executed outside of jit boundaries.

This is different from the documented behaviour.

```python
➜ python
Python 3.10.6 (main, Aug 23 2022, 11:35:18) [Clang 13.1.6 (clang-1316.0.21.2.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> def test(x):
...   print(type(x))
...   return x
...
>>> def f(x):
...   return jax.pure_callback(test, x, x)
...
>>> x= jax.numpy.ones(3)
>>> f(x)
<class 'jaxlib.xla_extension.DeviceArray'>
DeviceArray([1., 1., 1.], dtype=float32)
>>> jax.jit(f)(x)
<class 'numpy.ndarray'>
DeviceArray([1., 1., 1.], dtype=float32)
```

### What jax/jaxlib version are you using?

jax 0.3.17

### Which accelerator(s) are you using?

CPU

### Additional System Info

MacOs

Changed files:
tests/python_callback_test.py, jax/_src/callback.py

diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index 955ad776b..ced20d119 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -14,6 +14,8 @@
 """Module for JAX callbacks."""
 from __future__ import annotations
 
+import functools
+
 from typing import Any, Callable, Sequence
 
 from jax import core
@@ -21,6 +23,7 @@ from jax import tree_util
 from jax._src import dtypes
 from jax._src import lib as jaxlib
 from jax._src import util
+from jax._src import dispatch
 from jax.interpreters import ad
 from jax.interpreters import batching
 from jax.interpreters import mlir
@@ -33,11 +36,12 @@ pure_callback_p.multiple_results = True
 map, unsafe_map = util.safe_map, map
 
 
-@pure_callback_p.def_impl
 def pure_callback_impl(*args, result_avals, callback: Callable[..., Any],
                        vectorized: bool):
   del vectorized, result_avals
   return callback(*args)
+pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,
+                                           pure_callback_p))
 
 
 @pure_callback_p.def_abstract_eval
@@ -102,7 +106,7 @@ def pure_callback_lowering(ctx, *args, callback, **params):
                               "Please upgrade to a jaxlib >= 0.3.15.")
 
   def _callback(*flat_args):
-    return tuple(pure_callback_p.impl(*flat_args, callback=callback, **params))
+    return tuple(pure_callback_impl(*flat_args, callback=callback, **params))
 
   result, _, keepalive = mlir.emit_python_callback(
       ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index 2afad9058..7f88d8912 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -469,6 +469,17 @@ class PurePythonCallbackTest(jtu.JaxTestCase):
     super().tearDown()
     dispatch.runtime_tokens.clear()
 
+  @jtu.skip_on_devices(*disabled_backends)
+  def test_pure_callback_passes_ndarrays_without_jit(self):
+
+    def cb(x):
+      self.assertIs(type(x), np.ndarray)
+      return x
+
+    def f(x):
+      return jax.pure_callback(cb, x, x)
+    f(jnp.array(2.))
+
   @jtu.skip_on_devices(*disabled_backends)
   def test_simple_pure_callback(self):