We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/_src/callback.py`
Content:
```
1 # Copyright 2022 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Module for JAX callbacks."""
15 from __future__ import annotations
16 
17 from typing import Any, Callable, Sequence
18 
19 from jax import core
20 from jax import tree_util
21 from jax._src import dtypes
22 from jax._src import lib as jaxlib
23 from jax._src import util
24 from jax.interpreters import ad
25 from jax.interpreters import batching
26 from jax.interpreters import mlir
27 import numpy as np
28 
29 # `pure_callback_p` is the main primitive for staging out Python pure callbacks.
30 pure_callback_p = core.Primitive("pure_callback")
31 pure_callback_p.multiple_results = True
32 
33 map, unsafe_map = util.safe_map, map
34 
35 
36 @pure_callback_p.def_impl
37 def pure_callback_impl(*args, result_avals, callback: Callable[..., Any],
38                        vectorized: bool):
39   del vectorized, result_avals
40   return callback(*args)
41 
42 
43 @pure_callback_p.def_abstract_eval
44 def pure_callback_abstract_eval(*avals, callback: Callable[..., Any],
45                                 result_avals, vectorized: bool):
46   del avals, callback, vectorized
47   return result_avals
48 
49 
50 def pure_callback_jvp_rule(*args, **kwargs):
51   del args, kwargs
52   raise ValueError(
53       "Pure callbacks do not support JVP. "
54       "Please use `jax.custom_jvp` to use callbacks while taking gradients.")
55 
56 
57 ad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule
58 
59 
60 def pure_callback_transpose_rule(*args, **kwargs):
61   del args, kwargs
62   raise ValueError(
63       "Pure callbacks do not support transpose. "
64       "Please use `jax.custom_vjp` to use callbacks while taking gradients.")
65 
66 ad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule
67 
68 
69 def pure_callback_batching_rule(args, dims, *, callback, vectorized: bool,
70                                 result_avals: Sequence[core.ShapedArray]):
71   axis_size = next(a.shape[0] for a, d in zip(args, dims)
72                    if d is not batching.not_mapped)
73   new_args = [arg if dim is batching.not_mapped else
74               batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)]
75   if vectorized:
76     result_avals = tuple(
77         core.unmapped_aval(axis_size, core.no_axis_name, 0, aval)  # type: ignore
78         for aval in result_avals)
79     outvals = pure_callback_p.bind(
80         *new_args, callback=callback, vectorized=vectorized,
81         result_avals=result_avals)
82   else:
83     is_batched = [d is not batching.not_mapped for d in dims]
84     unbatched_args, batched_args = util.partition_list(is_batched, new_args)
85     def _batch_fun(*batched_args):
86       merged_args = util.merge_lists(is_batched, unbatched_args, batched_args)
87       return pure_callback_p.bind(
88           *merged_args, callback=callback, result_avals=result_avals,
89           vectorized=vectorized)
90     from jax._src.lax.control_flow import map as lax_map
91     outvals = lax_map(_batch_fun, *batched_args)
92   return tuple(outvals), (0,) * len(outvals)
93 
94 
95 batching.primitive_batchers[pure_callback_p] = pure_callback_batching_rule
96 
97 
98 def pure_callback_lowering(ctx, *args, callback, **params):
99 
100   if ctx.module_context.platform == "TPU" and jaxlib.version < (0, 3, 15):
101     raise NotImplementedError("Pure callbacks on TPU not supported. "
102                               "Please upgrade to a jaxlib >= 0.3.15.")
103 
104   def _callback(*flat_args):
105     return tuple(pure_callback_p.impl(*flat_args, callback=callback, **params))
106 
107   result, _, keepalive = mlir.emit_python_callback(
108       ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,
109       sharding=None)
110   ctx.module_context.add_keepalive(keepalive)
111   return result
112 
113 mlir.register_lowering(pure_callback_p, pure_callback_lowering)
114 
115 def _check_shape_dtype(shape_dtype):
116   dt = np.dtype(shape_dtype.dtype)
117   if dtypes.canonicalize_dtype(dt) != dt:
118     raise ValueError(
119         "Cannot return 64-bit values when `jax_enable_x64` is disabled")
120 
121 def pure_callback(callback: Callable[..., Any], result_shape_dtypes: Any,
122                   *args: Any, vectorized: bool = False, **kwargs: Any):
123   def _flat_callback(*flat_args):
124     args, kwargs = tree_util.tree_unflatten(in_tree, flat_args)
125     return tree_util.tree_leaves(callback(*args, **kwargs))
126 
127   flat_args, in_tree = tree_util.tree_flatten((args, kwargs))
128   tree_util.tree_map(_check_shape_dtype, result_shape_dtypes)
129   result_avals = tree_util.tree_map(
130       lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)
131   flat_result_avals, out_tree = tree_util.tree_flatten(result_avals)
132   out_flat = pure_callback_p.bind(
133       *flat_args, callback=_flat_callback,
134       result_avals=tuple(flat_result_avals), vectorized=vectorized)
135   return tree_util.tree_unflatten(out_tree, out_flat)
136
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/jax/_src/callback.py b/jax/_src/callback.py
--- 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._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 @@
 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 @@
                               "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,

{"golden_diff": "diff --git a/jax/_src/callback.py b/jax/_src/callback.py\n--- a/jax/_src/callback.py\n+++ b/jax/_src/callback.py\n@@ -14,6 +14,8 @@\n \"\"\"Module for JAX callbacks.\"\"\"\n from __future__ import annotations\n \n+import functools\n+\n from typing import Any, Callable, Sequence\n \n from jax import core\n@@ -21,6 +23,7 @@\n from jax._src import dtypes\n from jax._src import lib as jaxlib\n from jax._src import util\n+from jax._src import dispatch\n from jax.interpreters import ad\n from jax.interpreters import batching\n from jax.interpreters import mlir\n@@ -33,11 +36,12 @@\n map, unsafe_map = util.safe_map, map\n \n \n-@pure_callback_p.def_impl\n def pure_callback_impl(*args, result_avals, callback: Callable[..., Any],\n                        vectorized: bool):\n   del vectorized, result_avals\n   return callback(*args)\n+pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,\n+                                           pure_callback_p))\n \n \n @pure_callback_p.def_abstract_eval\n@@ -102,7 +106,7 @@\n                               \"Please upgrade to a jaxlib >= 0.3.15.\")\n \n   def _callback(*flat_args):\n-    return tuple(pure_callback_p.impl(*flat_args, callback=callback, **params))\n+    return tuple(pure_callback_impl(*flat_args, callback=callback, **params))\n \n   result, _, keepalive = mlir.emit_python_callback(\n       ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,\n", "issue": "pure_callback passes jax.DeviceArray to the callback on CPU when not jitted\n### Description\r\n\r\nAs title. \r\nThis only happens when the callback is executed outside of jit boundaries.\r\n\r\nThis is different from the documented behaviour.\r\n\r\n```python\r\n\u279c python\r\nPython 3.10.6 (main, Aug 23 2022, 11:35:18) [Clang 13.1.6 (clang-1316.0.21.2.5)] on darwin\r\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\r\n>>> import jax\r\n>>> def test(x):\r\n...   print(type(x))\r\n...   return x\r\n...\r\n>>> def f(x):\r\n...   return jax.pure_callback(test, x, x)\r\n...\r\n>>> x= jax.numpy.ones(3)\r\n>>> f(x)\r\n<class 'jaxlib.xla_extension.DeviceArray'>\r\nDeviceArray([1., 1., 1.], dtype=float32)\r\n>>> jax.jit(f)(x)\r\n<class 'numpy.ndarray'>\r\nDeviceArray([1., 1., 1.], dtype=float32)\r\n```\r\n\r\n### What jax/jaxlib version are you using?\r\n\r\njax 0.3.17\r\n\r\n### Which accelerator(s) are you using?\r\n\r\nCPU\r\n\r\n### Additional System Info\r\n\r\nMacOs\n", "before_files": [{"content": "# Copyright 2022 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Module for JAX callbacks.\"\"\"\nfrom __future__ import annotations\n\nfrom typing import Any, Callable, Sequence\n\nfrom jax import core\nfrom jax import tree_util\nfrom jax._src import dtypes\nfrom jax._src import lib as jaxlib\nfrom jax._src import util\nfrom jax.interpreters import ad\nfrom jax.interpreters import batching\nfrom jax.interpreters import mlir\nimport numpy as np\n\n# `pure_callback_p` is the main primitive for staging out Python pure callbacks.\npure_callback_p = core.Primitive(\"pure_callback\")\npure_callback_p.multiple_results = True\n\nmap, unsafe_map = util.safe_map, map\n\n\n@pure_callback_p.def_impl\ndef pure_callback_impl(*args, result_avals, callback: Callable[..., Any],\n                       vectorized: bool):\n  del vectorized, result_avals\n  return callback(*args)\n\n\n@pure_callback_p.def_abstract_eval\ndef pure_callback_abstract_eval(*avals, callback: Callable[..., Any],\n                                result_avals, vectorized: bool):\n  del avals, callback, vectorized\n  return result_avals\n\n\ndef pure_callback_jvp_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support JVP. \"\n      \"Please use `jax.custom_jvp` to use callbacks while taking gradients.\")\n\n\nad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule\n\n\ndef pure_callback_transpose_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support transpose. \"\n      \"Please use `jax.custom_vjp` to use callbacks while taking gradients.\")\n\nad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule\n\n\ndef pure_callback_batching_rule(args, dims, *, callback, vectorized: bool,\n                                result_avals: Sequence[core.ShapedArray]):\n  axis_size = next(a.shape[0] for a, d in zip(args, dims)\n                   if d is not batching.not_mapped)\n  new_args = [arg if dim is batching.not_mapped else\n              batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)]\n  if vectorized:\n    result_avals = tuple(\n        core.unmapped_aval(axis_size, core.no_axis_name, 0, aval)  # type: ignore\n        for aval in result_avals)\n    outvals = pure_callback_p.bind(\n        *new_args, callback=callback, vectorized=vectorized,\n        result_avals=result_avals)\n  else:\n    is_batched = [d is not batching.not_mapped for d in dims]\n    unbatched_args, batched_args = util.partition_list(is_batched, new_args)\n    def _batch_fun(*batched_args):\n      merged_args = util.merge_lists(is_batched, unbatched_args, batched_args)\n      return pure_callback_p.bind(\n          *merged_args, callback=callback, result_avals=result_avals,\n          vectorized=vectorized)\n    from jax._src.lax.control_flow import map as lax_map\n    outvals = lax_map(_batch_fun, *batched_args)\n  return tuple(outvals), (0,) * len(outvals)\n\n\nbatching.primitive_batchers[pure_callback_p] = pure_callback_batching_rule\n\n\ndef pure_callback_lowering(ctx, *args, callback, **params):\n\n  if ctx.module_context.platform == \"TPU\" and jaxlib.version < (0, 3, 15):\n    raise NotImplementedError(\"Pure callbacks on TPU not supported. \"\n                              \"Please upgrade to a jaxlib >= 0.3.15.\")\n\n  def _callback(*flat_args):\n    return tuple(pure_callback_p.impl(*flat_args, callback=callback, **params))\n\n  result, _, keepalive = mlir.emit_python_callback(\n      ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,\n      sharding=None)\n  ctx.module_context.add_keepalive(keepalive)\n  return result\n\nmlir.register_lowering(pure_callback_p, pure_callback_lowering)\n\ndef _check_shape_dtype(shape_dtype):\n  dt = np.dtype(shape_dtype.dtype)\n  if dtypes.canonicalize_dtype(dt) != dt:\n    raise ValueError(\n        \"Cannot return 64-bit values when `jax_enable_x64` is disabled\")\n\ndef pure_callback(callback: Callable[..., Any], result_shape_dtypes: Any,\n                  *args: Any, vectorized: bool = False, **kwargs: Any):\n  def _flat_callback(*flat_args):\n    args, kwargs = tree_util.tree_unflatten(in_tree, flat_args)\n    return tree_util.tree_leaves(callback(*args, **kwargs))\n\n  flat_args, in_tree = tree_util.tree_flatten((args, kwargs))\n  tree_util.tree_map(_check_shape_dtype, result_shape_dtypes)\n  result_avals = tree_util.tree_map(\n      lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)\n  flat_result_avals, out_tree = tree_util.tree_flatten(result_avals)\n  out_flat = pure_callback_p.bind(\n      *flat_args, callback=_flat_callback,\n      result_avals=tuple(flat_result_avals), vectorized=vectorized)\n  return tree_util.tree_unflatten(out_tree, out_flat)\n", "path": "jax/_src/callback.py"}], "after_files": [{"content": "# Copyright 2022 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Module for JAX callbacks.\"\"\"\nfrom __future__ import annotations\n\nimport functools\n\nfrom typing import Any, Callable, Sequence\n\nfrom jax import core\nfrom jax import tree_util\nfrom jax._src import dtypes\nfrom jax._src import lib as jaxlib\nfrom jax._src import util\nfrom jax._src import dispatch\nfrom jax.interpreters import ad\nfrom jax.interpreters import batching\nfrom jax.interpreters import mlir\nimport numpy as np\n\n# `pure_callback_p` is the main primitive for staging out Python pure callbacks.\npure_callback_p = core.Primitive(\"pure_callback\")\npure_callback_p.multiple_results = True\n\nmap, unsafe_map = util.safe_map, map\n\n\ndef pure_callback_impl(*args, result_avals, callback: Callable[..., Any],\n                       vectorized: bool):\n  del vectorized, result_avals\n  return callback(*args)\npure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,\n                                           pure_callback_p))\n\n\n@pure_callback_p.def_abstract_eval\ndef pure_callback_abstract_eval(*avals, callback: Callable[..., Any],\n                                result_avals, vectorized: bool):\n  del avals, callback, vectorized\n  return result_avals\n\n\ndef pure_callback_jvp_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support JVP. \"\n      \"Please use `jax.custom_jvp` to use callbacks while taking gradients.\")\n\n\nad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule\n\n\ndef pure_callback_transpose_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support transpose. \"\n      \"Please use `jax.custom_vjp` to use callbacks while taking gradients.\")\n\nad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule\n\n\ndef pure_callback_batching_rule(args, dims, *, callback, vectorized: bool,\n                                result_avals: Sequence[core.ShapedArray]):\n  axis_size = next(a.shape[0] for a, d in zip(args, dims)\n                   if d is not batching.not_mapped)\n  new_args = [arg if dim is batching.not_mapped else\n              batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)]\n  if vectorized:\n    result_avals = tuple(\n        core.unmapped_aval(axis_size, core.no_axis_name, 0, aval)  # type: ignore\n        for aval in result_avals)\n    outvals = pure_callback_p.bind(\n        *new_args, callback=callback, vectorized=vectorized,\n        result_avals=result_avals)\n  else:\n    is_batched = [d is not batching.not_mapped for d in dims]\n    unbatched_args, batched_args = util.partition_list(is_batched, new_args)\n    def _batch_fun(*batched_args):\n      merged_args = util.merge_lists(is_batched, unbatched_args, batched_args)\n      return pure_callback_p.bind(\n          *merged_args, callback=callback, result_avals=result_avals,\n          vectorized=vectorized)\n    from jax._src.lax.control_flow import map as lax_map\n    outvals = lax_map(_batch_fun, *batched_args)\n  return tuple(outvals), (0,) * len(outvals)\n\n\nbatching.primitive_batchers[pure_callback_p] = pure_callback_batching_rule\n\n\ndef pure_callback_lowering(ctx, *args, callback, **params):\n\n  if ctx.module_context.platform == \"TPU\" and jaxlib.version < (0, 3, 15):\n    raise NotImplementedError(\"Pure callbacks on TPU not supported. \"\n                              \"Please upgrade to a jaxlib >= 0.3.15.\")\n\n  def _callback(*flat_args):\n    return tuple(pure_callback_impl(*flat_args, callback=callback, **params))\n\n  result, _, keepalive = mlir.emit_python_callback(\n      ctx, _callback, None, list(args), ctx.avals_in, ctx.avals_out, False,\n      sharding=None)\n  ctx.module_context.add_keepalive(keepalive)\n  return result\n\nmlir.register_lowering(pure_callback_p, pure_callback_lowering)\n\ndef _check_shape_dtype(shape_dtype):\n  dt = np.dtype(shape_dtype.dtype)\n  if dtypes.canonicalize_dtype(dt) != dt:\n    raise ValueError(\n        \"Cannot return 64-bit values when `jax_enable_x64` is disabled\")\n\ndef pure_callback(callback: Callable[..., Any], result_shape_dtypes: Any,\n                  *args: Any, vectorized: bool = False, **kwargs: Any):\n  def _flat_callback(*flat_args):\n    args, kwargs = tree_util.tree_unflatten(in_tree, flat_args)\n    return tree_util.tree_leaves(callback(*args, **kwargs))\n\n  flat_args, in_tree = tree_util.tree_flatten((args, kwargs))\n  tree_util.tree_map(_check_shape_dtype, result_shape_dtypes)\n  result_avals = tree_util.tree_map(\n      lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)\n  flat_result_avals, out_tree = tree_util.tree_flatten(result_avals)\n  out_flat = pure_callback_p.bind(\n      *flat_args, callback=_flat_callback,\n      result_avals=tuple(flat_result_avals), vectorized=vectorized)\n  return tree_util.tree_unflatten(out_tree, out_flat)\n", "path": "jax/_src/callback.py"}]}