Repository: google/jax

Issue:
id_tap crashes with new custom PRNG key
### Description

```python
from jax import enable_custom_prng, jit
from jax.experimental.host_callback import id_tap
from jax.random import KeyArray, PRNGKey

with enable_custom_prng():
    key = PRNGKey(123)
    @jit
    def f(x: KeyArray) -> KeyArray:
        def tap(tap_x, transforms):
            print(tap_x)
        return id_tap(tap, x, result=x)

    f(key)
```
produces
```shell
Traceback (most recent call last):
  File "/home/neil/src/tjax/a.py", line -1, in <module>
  File "/home/neil/src/tjax/a.py", line -1, in f
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/experimental/host_callback.py", line -1, in id_tap
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/experimental/host_callback.py", line -1, in _call
jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: array_shape(): incompatible function arguments. The following argument types are supported:
    1. (type: jaxlib.xla_extension.PrimitiveType, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]
    2. (type: dtype, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]

Invoked with: key<fry>, ()

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/neil/src/tjax/a.py", line 13, in <module>
    f(key)
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/api.py", line 622, in cache_miss
    execute = dispatch._xla_call_impl_lazy(fun_, *tracers, **params)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/dispatch.py", line 241, in _xla_call_impl_lazy
    return xla_callable(fun, device, backend, name, donated_invars, keep_unused,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/linear_util.py", line 303, in memoized_fun
    ans = call(fun, *args)
          ^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/dispatch.py", line 357, in _xla_callable_uncached
    computation = sharded_lowering(fun, device, backend, name, donated_invars,
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/dispatch.py", line 348, in sharded_lowering
    return pxla.lower_sharding_computation(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/pxla.py", line 2936, in lower_sharding_computation
    lowering_result = mlir.lower_jaxpr_to_module(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/mlir.py", line 707, in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/mlir.py", line 988, in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/mlir.py", line 1122, in jaxpr_subcomp
    ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/experimental/host_callback.py", line 1186, in _outside_call_lowering
    results, next_token, keep_alive = mlir.emit_python_callback(ctx,
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/mlir.py", line 1747, in emit_python_callback
    [xla.aval_to_xla_shapes(op_aval) for op_aval in operand_avals])
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/mlir.py", line 1747, in <listcomp>
    [xla.aval_to_xla_shapes(op_aval) for op_aval in operand_avals])
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/xla.py", line 222, in aval_to_xla_shapes
    return xla_shape_handlers[type(aval)](aval)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/interpreters/xla.py", line 91, in _make_array_shape
    return (xc.Shape.array_shape(a.dtype, a.shape),)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jax._src.traceback_util.UnfilteredStackTrace: TypeError: array_shape(): incompatible function arguments. The following argument types are supported:
    1. (type: jaxlib.xla_extension.PrimitiveType, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]
    2. (type: dtype, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]

Invoked with: key<fry>, ()

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/neil/src/tjax/a.py", line 13, in <module>
    f(key)
  File "/home/neil/.cache/pypoetry/virtualenvs/tjax-DRMj3ydM-py3.11/lib/python3.11/site-packages/jax/experimental/host_callback.py", line 1186, in _outside_call_lowering
    results, next_token, keep_alive = mlir.emit_python_callback(ctx,
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: array_shape(): incompatible function arguments. The following argument types are supported:
    1. (type: jaxlib.xla_extension.PrimitiveType, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]
    2. (type: dtype, dims: object, layout: Optional[object] = None, dynamic_dimensions: Optional[List[bool]] = None) -> StatusOr[jaxlib.xla_extension.Shape]

Invoked with: key<fry>, ()
```

### What jax/jaxlib version are you using?

0.4.1


Discussion:
Incidentally, there are a couple stack frames that should probably be marked as belonging Jax.
Assinging @sharadmv as the callback expert, and @froystig as the custom PRNG expert

Resolving patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1671,9 +1671,13 @@ def _layout_to_mlir_layout(minor_to_major: Optional[Sequence[int]]):
     layout = np.array(minor_to_major, dtype="int64")
   return ir.DenseIntElementsAttr.get(layout, type=ir.IndexType.get())
 
-def _aval_to_default_layout(aval):
+def _aval_to_default_layouts(aval):
+  if core.is_opaque_dtype(aval.dtype):
+    avals = aval.dtype._rules.physical_avals(aval)
+  else:
+    avals = [aval]
   # Row major order is default for `NumPy`.
-  return list(range(aval.ndim - 1, -1, -1))
+  return [list(range(aval.ndim - 1, -1, -1)) for aval in avals]
 
 def emit_python_callback(
     ctx: LoweringRuleContext, callback, token: Optional[Any],
@@ -1695,17 +1699,12 @@ def emit_python_callback(
       [xla.aval_to_xla_shapes(op_aval) for op_aval in operand_avals])
   # Handling layouts
   if operand_layouts is None:
-    operand_layouts = map(_aval_to_default_layout, operand_avals)
-  operand_mlir_layouts = [
-      _layout_to_mlir_layout(_aval_to_default_layout(layout)) if layout is None
-      else _layout_to_mlir_layout(layout) for layout, aval
-      in zip(operand_layouts, operand_avals)]
+    operand_layouts = util.concatenate(
+        map(_aval_to_default_layouts, operand_avals))
+  operand_mlir_layouts = map(_layout_to_mlir_layout, operand_layouts)
   if result_layouts is None:
-    result_layouts = map(_aval_to_default_layout, result_avals)
-  result_mlir_layouts = [
-      _layout_to_mlir_layout(_aval_to_default_layout(aval)) if layout is None
-      else _layout_to_mlir_layout(layout) for layout, aval
-      in zip(result_layouts, result_avals)]
+    result_layouts = util.concatenate(map(_aval_to_default_layouts, result_avals))
+  result_mlir_layouts = map(_layout_to_mlir_layout, result_layouts)
 
   # First we apply checks to ensure output shapes and dtypes match the expected
   # ones.
diff --git a/jax/_src/interpreters/xla.py b/jax/_src/interpreters/xla.py
--- a/jax/_src/interpreters/xla.py
+++ b/jax/_src/interpreters/xla.py
@@ -54,11 +54,15 @@ def identity(x): return x
 
 _scalar_types = dtypes.python_scalar_dtypes.keys()
 
-def _make_array_shape(a: ShapedArray) -> Sequence[xc.Shape]:
-  if a.dtype == dtypes.float0:
-    return (xc.Shape.array_shape(np.dtype('bool'), a.shape),)
+def _make_array_shape(aval: ShapedArray) -> Sequence[xc.Shape]:
+  def dt(aval):
+    return np.dtype('bool') if aval.dtype == dtypes.float0 else aval.dtype
+
+  if core.is_opaque_dtype(aval.dtype):
+    avals = aval.dtype._rules.physical_avals(aval)
   else:
-    return (xc.Shape.array_shape(a.dtype, a.shape),)
+    avals = [aval]
+  return tuple(xc.Shape.array_shape(dt(a), a.shape) for a in avals)
 
 def get_canonical_source_file(frame: source_info_util.Frame):
   source_file = frame.file_name


Test patch:
diff --git a/tests/host_callback_test.py b/tests/host_callback_test.py
--- a/tests/host_callback_test.py
+++ b/tests/host_callback_test.py
@@ -1908,6 +1908,15 @@ def func(x, transforms, y):
     with self.assertRaisesRegex(TypeError, r"Support for \*\*kwargs in ``id_tap``"):
       hcb.id_tap(func, 1, y=2)
 
+  def test_tap_id_tap_random_key(self):
+    # See https://github.com/google/jax/issues/13949
+    with jax.enable_custom_prng():
+      @jax.jit
+      def f(x):
+        def tap(tap_x, _): pass
+        return hcb.id_tap(tap, x, result=x)
+      f(jax.random.PRNGKey(123))
+
   def test_tap_odeint(self):
     # TODO: find a smaller repro for bug #4015
     # Seems to be xla_call(scan(xla_call)), all under grad.