internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/jax

Issue:
exceptions that reach the user should have public type (e.g. ConcretizationTypeError)
Consider `except jax._src.errors.ConcretizationTypeError: # pylint: disable=protected-access`. Lame!


Discussion:
Actually, these already have a public name, and are [documented](https://jax.readthedocs.io/en/latest/errors.html), e.g. `jax.errors.ConcretizationTypeError` after we `import jax`.
We can improve this situation by making the repr of the errors print their public path; #15698 should take care of that

Resolving patch:
diff --git a/jax/_src/errors.py b/jax/_src/errors.py
--- a/jax/_src/errors.py
+++ b/jax/_src/errors.py
@@ -14,6 +14,8 @@
 from __future__ import annotations
 
 from jax._src import core
+from jax._src.util import set_module
+
 
 class _JAXErrorMixin:
   """Mixin for JAX-specific errors"""
@@ -29,14 +31,17 @@ def __init__(self, message: str):
     super().__init__(error_msg)  # type: ignore
 
 
+@set_module('jax.errors')
 class JAXTypeError(_JAXErrorMixin, TypeError):
   pass
 
 
+@set_module('jax.errors')
 class JAXIndexError(_JAXErrorMixin, IndexError):
   pass
 
 
+@set_module('jax.errors')
 class ConcretizationTypeError(JAXTypeError):
   """
   This error occurs when a JAX Tracer object is used in a context where a
@@ -180,6 +185,7 @@ def __init__(self, tracer: core.Tracer, context: str = ""):
         f"{tracer}\n{context}{tracer._origin_msg()}\n")
 
 
+@set_module('jax.errors')
 class NonConcreteBooleanIndexError(JAXIndexError):
   """
   This error occurs when a program attempts to use non-concrete boolean indices
@@ -271,6 +277,7 @@ def __init__(self, tracer: core.Tracer):
         f"Array boolean indices must be concrete; got {tracer}\n")
 
 
+@set_module('jax.errors')
 class TracerArrayConversionError(JAXTypeError):
   """
   This error occurs when a program attempts to convert a JAX Tracer object into
@@ -351,6 +358,7 @@ def __init__(self, tracer: core.Tracer):
         f"the JAX Tracer object {tracer}{tracer._origin_msg()}")
 
 
+@set_module('jax.errors')
 class TracerIntegerConversionError(JAXTypeError):
   """
   This error can occur when a JAX Tracer object is used in a context where a
@@ -443,6 +451,7 @@ def __init__(self, tracer: core.Tracer):
         f"The __index__() method was called on the JAX Tracer object {tracer}")
 
 
+@set_module('jax.errors')
 class UnexpectedTracerError(JAXTypeError):
   """
   This error occurs when you use a JAX value that has leaked out of a function.
diff --git a/jax/_src/util.py b/jax/_src/util.py
--- a/jax/_src/util.py
+++ b/jax/_src/util.py
@@ -547,6 +547,14 @@ def _original_func(f):
   return f
 
 
+def set_module(module):
+  def wrapper(func):
+    if module is not None:
+      func.__module__ = module
+    return func
+  return wrapper
+
+
 if TYPE_CHECKING:
   def use_cpp_class(cpp_cls: Any) -> Callable[[T], T]:
     def wrapper(cls: T) -> T:


Test patch:
Repository: google/jax

Issue:
Missing in reference docs: `jax.lax.with_sharding_constraint`
This page should have everything in `jax.lax`: https://jax.readthedocs.io/en/latest/jax.lax.html#

but it's missing `jax.lax.with_sharding_constraint`, which is part of the public API, because it's mentioned [here](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code)


Discussion:

Resolving patch:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1757,6 +1757,21 @@ def _resolve_wsc_args(axis_resources, shardings):
 # period is finished. The deprecation period expires 3 months from Feb 13, 2023.
 def with_sharding_constraint(x, shardings=UNSPECIFIED,
                              axis_resources=UNSPECIFIED):
+  """Mechanism to constrain the sharding of an Array inside a jitted computation
+
+  This is a strict constraint for the GSPMD partitioner and not a hint. For examples
+  of how to use this function, see `Distributed arrays and automatic parallelization`_.
+
+  Args:
+    x: PyTree of jax.Arrays which will have their shardings constrainted
+    shardings: PyTree of sharding specifications. Valid values are the same as for
+      the ``in_shardings`` argument of :func:`jax.experimental.pjit`.
+    axis_resources: (deprecated) use shardings instead.
+  Returns:
+    x_with_shardings: PyTree of jax.Arrays with specified sharding constraints.
+
+  .. _Distributed arrays and automatic parallelization: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
+  """
   final_shardings = _resolve_wsc_args(axis_resources, shardings)
   x_flat, tree = tree_flatten(x)
   user_shardings, _, _ = prepare_axis_resources(
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -368,7 +368,7 @@
 from jax._src.ad_util import stop_gradient_p as stop_gradient_p
 from jax.lax import linalg as linalg
 
-from jax._src.pjit import with_sharding_constraint
-from jax._src.dispatch import device_put_p
+from jax._src.pjit import with_sharding_constraint as with_sharding_constraint
+from jax._src.dispatch import device_put_p as device_put_p
 
 from math import prod  # TODO(phawkins): remove this accidental export


Test patch:
Repository: google/jax

Issue:
`jit` + `linear_transpose` raises error due to missing DCE  call
### Description

```python
import jax

def f1(x):
    return 2 * x, x > 0

def g1(x):
    return f1(x)[0]

jax.linear_transpose(g1, 1.)(1.)  # works

@jax.jit
def f2(x):
    return 2 * x, x > 0

def g2(x):
    return f2(x)[0]

jax.linear_transpose(g2, 1.)(1.)  # error!
# jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'gt' not implemented
```

It looks like `jax.linear_transpose` is missing a call to `pe.dce_jaxpr` prior to attempting the transpose.

I checked that `jax.grad` doesn't have the same issue. I assume it has a DCE call in there already somewhere.

### What jax/jaxlib version are you using?

JAX 0.4.8


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2250,6 +2250,7 @@ def linear_transpose(fun: Callable, *primals, reduce_axes=()) -> Callable:
   in_pvals = map(pe.PartialVal.unknown, in_avals)
   jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,
                                                       instantiate=True)
+  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)
   out_avals, _ = unzip2(out_pvals)
   out_dtypes = map(dtypes.dtype, out_avals)
   if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2076,6 +2076,12 @@ def test_linear_transpose_integer(self):
     expected = 6
     self.assertEqual(actual, expected)
 
+  def test_linear_transpose_dce(self):
+    # https://github.com/google/jax/issues/15660
+    f = jit(lambda x: (2 * x, x > 0))
+    g = lambda x: f(x)[0]
+    api.linear_transpose(g, 1.)(1.)
+
   def test_linear_transpose_error(self):
     with self.assertRaisesRegex(
         TypeError, "linear_transpose only supports"):
Repository: google/jax

Issue:
Inconsistent dtype output in jax.random.choice for double precision enabled
### Description

When using double precision, `jax.random.choice` returns inconsistent dtypes. When using a weight vector, the output dtype is set using `searchsorted` rules, i.e., based on the size of the array of weights at hand, while when none is used (uniform weights), the dtype is set using `jax.random.randint` rules (defaulting to `jnp.int_`). I believe the behaviour should be homogenous, no matter what the choice ends up being.

```python
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)
key = jax.random.PRNGKey(42)
w = jnp.array([1/3, 1/3, 1/3])

print(jax.random.choice(key, 3).dtype)
print(jax.random.choice(key, 3, p=w).dtype)
```

In practice, I believe that the implementation choice made in https://github.com/google/jax/blob/c22302ae2e0913efd5be23e136f2ff14b520d9cf/jax/_src/numpy/lax_numpy.py#L4824 is perhaps a bit redundant. Indeed, if the config has not enables double precision, `int64` will default to `int32` anyway, while, if it has, you were happy to pay the price in the first place.

WDYT?



### What jax/jaxlib version are you using?

jax v0.4.10 

### Which accelerator(s) are you using?

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -567,7 +567,7 @@ def choice(key: KeyArray,
     if replace:
       p_cuml = jnp.cumsum(p_arr)
       r = p_cuml[-1] * (1 - uniform(key, shape, dtype=p_cuml.dtype))
-      ind = jnp.searchsorted(p_cuml, r)
+      ind = jnp.searchsorted(p_cuml, r).astype(int)
     else:
       # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
       g = -gumbel(key, (n_inputs,), dtype=p_arr.dtype) - jnp.log(p_arr)


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -754,8 +754,10 @@ def testChoice(self, dtype, input_range_or_shape, shape, replace, weighted, axis
     sample = rand(key, x)
     if not is_range:
       self.assertEqual(dtype, sample.dtype)
-    np_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
-    self.assertEqual(np_shape, sample.shape)
+    expected_shape = np.shape(np_choice(x, shape or None, replace, p, axis))
+    self.assertEqual(expected_shape, sample.shape)
+    expected_dtype = dtypes.result_type(int if is_range else x)
+    self.assertEqual(expected_dtype, sample.dtype)
     if not replace and shape:
       def lsort(x):
         if not math.prod(x.shape): return x
Repository: google/jax

Issue:
Potential typo for `jax.log_compiles`
### Description

`jax.log_compiles()` currently sets the flag for `jax_log_checkpoint_residuals` instead of `jax_log_compiles`, which looks like a typo.

This looks like a copy-paste error introduced by 6b4262d9f6ee36aea16ba7cfac639904b499661c.

https://github.com/google/jax/blob/9da52e890537c507b4d559ee6e3d8e8358e77f3f/jax/_src/config.py#L737-L750

### What jax/jaxlib version are you using?

jax v0.4.10, jaxlib v0.4.10

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/config.py b/jax/_src/config.py
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -742,7 +742,7 @@ def update_thread_local_jit_state(**kw):
           'option is set, the log level is WARNING; otherwise the level is '
           'DEBUG.'))
 
-log_compiles = config.define_bool_state(
+log_checkpoint_residuals = config.define_bool_state(
     name='jax_log_checkpoint_residuals',
     default=False,
     help=('Log a message every time jax.checkpoint (aka jax.remat) is '


Test patch:
Repository: google/jax

Issue:
[shard-map] Check failed using unsafe_rbg_key
Based on code from @jheek :

```python
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'

from functools import partial

import numpy as np

import jax
import jax.numpy as jnp
from jax import random
from jax.experimental.shard_map import shard_map, PartitionSpec as P, Mesh

from flax import linen

jax.config.update('jax_enable_custom_prng', True)  # doesn't crash if we exclude this line!

dense = linen.Dense(128)

mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('model', 'data'))

@jax.jit
def f(rng):
  @partial(shard_map, mesh=mesh, in_specs=P("model"), out_specs=P("model"), check_rep=False)
  def g(rng):
    rng = rng[0]
    return dense.init(rng, jnp.zeros((8, 128)))

  return g(random.split(rng, 4))

f(random.PRNGKey(0)) # ==> crashes VM on ASSERT(IsManual(sharding))
```


Discussion:

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
@@ -1500,7 +1500,7 @@ def _wrap_with_spmd_op(name: str,
   else:
     backend_config = ""
   result_type = aval_to_ir_type(aval_out)
-  out_shape = aval_out.shape  # type: ignore
+  out_shape = core.physical_aval(aval_out).shape  # type: ignore
   if core.is_constant_shape(out_shape):
     result_shapes = None
   else:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -470,7 +470,8 @@ def physical_hlo_sharding(aval, op_sharding_proto):
       op_sharding_proto = op_sharding_proto.to_proto()
     new_op_sharding = op_sharding_proto.clone()
     tad = list(new_op_sharding.tile_assignment_dimensions)
-    tad.extend([1] * len(key_shape))
+    suffix = [tad.pop()] if op_sharding_proto.replicate_on_last_tile_dim else []
+    tad.extend([1] * len(key_shape) + suffix)
     new_op_sharding.tile_assignment_dimensions = tad
     return xc.HloSharding.from_proto(new_op_sharding)
 


Test patch:
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -740,6 +740,20 @@ def g(x):
     # error!
     jax.jit(g)(x)  # doesn't crash
 
+  def test_key_array_with_replicated_last_tile_dim(self):
+    # See https://github.com/google/jax/issues/16137
+
+    mesh = jtu.create_global_mesh((2, 4), ('i', 'j'))
+
+    def f(rng):
+      @partial(shard_map, mesh=mesh, in_specs=P('i'), out_specs=P('i'),
+               check_rep=False)
+      def g(rng):
+        return jnp.array([jax.random.normal(rng[0])])
+      return g(jax.random.split(rng, 4))
+
+    jax.jit(f)(jax.random.key(0))  # doesn't crash
+
   # same method appears in api_test.py:DCETest
   # TODO(mattjj): consider moving this method to be a helper in jtu
   def assert_dce_result(self, jaxpr: core.Jaxpr, used_outputs: List[bool],
Repository: google/jax

Issue:
Interaction between `vmap` and new PRNG
### Description

Hi, I'm having trouble seeing how the following MWE using [Diffrax](https://github.com/patrick-kidger/diffrax) errors:
```py
import jax
from jax import numpy as jnp
from diffrax import diffeqsolve, ControlTerm, Euler, MultiTerm, ODETerm, SaveAt, VirtualBrownianTree

B = 20
D = 5

tol = 1e-3
dt0 = 0.05
T0 = 0.0
Tf = 1.0

def f(x0, key):
    drift = lambda t, x, args: jnp.zeros((D, ))
    diffusion = lambda t, x, args: jnp.zeros((D, D))
    brownian_motion = VirtualBrownianTree(T0, Tf, tol=tol, shape=x0.shape, key=key)
    terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))
    solver = Euler()
    saveat = SaveAt(dense=True)
    sol = diffeqsolve(terms, solver, T0, Tf, dt0=dt0, y0=x0, saveat=saveat)
    return sol.evaluate(Tf)

x0 = jnp.zeros((D, ))
key = jax.random.key(1)
print(f(x0, key))

x0 = jnp.zeros((B, D))
key = jax.random.split(key, B)
print(jax.vmap(f)(x0, key))
```
In particular, the last three lines (with `vmap`) cause an error, while the previous three lines seem to work fine.
I raised [this issue in Diffrax](https://github.com/patrick-kidger/diffrax/issues/267) and was told that it was an issue only with the new PRNG, so the bug is in Jax rather than Diffrax. Hence I'm raising it here.

Here are links for:
- [The stack trace.](https://pastebin.com/46hbkCKb)
- [The output of my `conda list`.](https://pastebin.com/3xjrfYYz)

### What jax/jaxlib version are you using?

0.4.12/0.4.12

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10.11, WSL 2 -> Ubuntu 22.04.2 LTS

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3488,7 +3488,6 @@ def _select_hlo_lowering_opaque(ctx, which, *cases):
   avals_in = ctx.avals_in
   aval_out, = ctx.avals_out
   assert all(aval_case == aval_out for aval_case in avals_in[1:])
-  assert avals_in[0].ndim == aval_out.ndim
   select_lower = _select_hlo_lowering
 
   physical_aval_out = core.physical_aval(aval_out)


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1860,6 +1860,15 @@ def test_select(self):
     self.assertIsInstance(ys, random.KeyArray)
     self.assertEqual(ys.shape, (3, 2))
 
+  @skipIf(not config.jax_enable_custom_prng,
+          'requires config.jax_enable_custom_prng')
+  def test_select_scalar_cond(self):
+    # regression test for https://github.com/google/jax/issues/16422
+    ks = self.make_keys(3)
+    ys = lax.select(True, ks, ks)
+    self.assertIsInstance(ys, random.KeyArray)
+    self.assertEqual(ys.shape, (3,))
+
   @skipIf(not config.jax_enable_custom_prng,
           'requires config.jax_enable_custom_prng')
   def test_select2(self):
Repository: google/jax

Issue:
ModuleNotFoundError: No module named 'requests' when executing `jax.distributed.initialize()` on TPU Pods
### Description

```python
import jax
jax.distributed.initialize()
```

Error:

```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/distributed.py", line 168, in initialize
    global_state.initialize(coordinator_address, num_processes, process_id, local_device_ids)
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/distributed.py", line 47, in initialize
    clusters.ClusterEnv.auto_detect_unset_distributed_params(
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cluster.py", line 55, in auto_detect_unset_distributed_params
    coordinator_address = env.get_coordinator_address()
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 54, in get_coordinator_address
    return cls._get_worker_endpoints()[0].split(':')[2] + ':8476'
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 74, in _get_worker_endpoints
    return get_metadata('worker-network-endpoints').split(',')
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 23, in get_metadata
    import requests  # pytype: disable=import-error
    ^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'requests'
```

I think `requests` should be automatically installed when installing JAX.

### What jax/jaxlib version are you using?

jax 0.4.14

### Which accelerator(s) are you using?

TPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks - I think that makes sense. It looks like requests is only imported in `jax/_src/clusters/cloud_tpu_cluster.py`, so it would probably be better as a TPU-only requirement. Perhaps we could include it in in the optional requirements for `jax[tpu]`?

Tagging @skye who knows a lot about cloud TPU runtimes – what do you think?
I agree it should be required for `jax[tpu]`. I think it used to be at some point actually, but was removed because we didn't need it anymore (but now we do!).

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -85,8 +85,11 @@ def generate_proto(source):
 
         # Cloud TPU VM jaxlib can be installed via:
         # $ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-        'tpu': [f'jaxlib=={_current_jaxlib_version}',
-                f'libtpu-nightly=={_libtpu_version}'],
+        'tpu': [
+          f'jaxlib=={_current_jaxlib_version}',
+          f'libtpu-nightly=={_libtpu_version}',
+          'requests',  # necessary for jax.distributed.initialize
+        ],
 
         # $ pip install jax[australis]
         'australis': ['protobuf>=3.13,<4'],


Test patch:
Repository: google/jax

Issue:
jvp of jnp.power(arr, scalar) fails
Repro:
```python
>>> import jax
>>> jax.jacrev(jax.jit(jax.lax.pow))(jax.numpy.arange(4.0), 2)
...
TypeError: select cases must have the same shapes, got [(4,), ()].
```
It looks like this was introduced by https://github.com/google/jax/pull/16419.


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1989,6 +1989,10 @@ def _pow_jvp_lhs(g, ans, x, y):
   y_dtype = dtypes.dtype(y)
   x, y = jax._src.numpy.util.promote_dtypes_numeric(x, y)  # TODO replace this
   if dtypes.issubdtype(y_dtype, np.integer):
+    if x.shape != y.shape:
+      shape = broadcast_shapes(x.shape, y.shape)
+      x = _maybe_broadcast(shape, x)
+      y = _maybe_broadcast(shape, y)
     jac = select(eq(y, _const(y, 0)), _ones(y),
                  mul(_replace_zero(y), pow(x, sub(y, _ones(y)))))
   else:


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -1140,6 +1140,14 @@ def f(x):
     with self.assertRaises(NotImplementedError):
       jax.jacrev(f)(x)
 
+  def testPowShapeMismatch(self):
+    # Regression test for https://github.com/google/jax/issues/17294
+    x = lax.iota('float32', 4)
+    y = 2
+    actual = jax.jacrev(jax.jit(jax.lax.pow))(x, y)  # no error
+    expected = jax.numpy.diag(y * x ** (y - 1))
+    self.assertArraysEqual(actual, expected)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax

Issue:
jax.random.key(0).itemsize crashes
### Description

I would expect the following to work, as it works for numpy arrays.

```
In [13]: jax.random.key(3).itemsize
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[13], line 1
----> 1 jax.random.key(3).itemsize

File ~/Documents/pythonenvs/mpi4jax/python-3.11.1/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:779, in _make_abstract_method.<locals>.method(*args, **kwargs)
    776 @abc.abstractmethod
    777 @wraps(func)
    778 def method(*args, **kwargs):
--> 779   raise NotImplementedError(f"Cannot call abstract method {name}")

NotImplementedError: Cannot call abstract method itemsize
```

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -167,6 +167,10 @@ def size(self) -> int: ...
   @abc.abstractmethod
   def dtype(self): ...
 
+  @property
+  @abc.abstractmethod
+  def itemsize(self): ...
+
   @property
   @abc.abstractmethod
   def sharding(self): ...
@@ -280,6 +284,10 @@ def ndim(self):
   def dtype(self):
     return KeyTy(self.impl)
 
+  @property
+  def itemsize(self):
+    return self.dtype.itemsize
+
   _device = property(op.attrgetter('_base_array._device'))
   _committed = property(op.attrgetter('_base_array._committed'))
   device = property(op.attrgetter('_base_array.device'))  # type: ignore[assignment]


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1805,6 +1805,12 @@ def test_key_dtype_attributes(self):
     self.assertEqual(key.size * key.dtype.itemsize,
                      key_raw.size * key_raw.dtype.itemsize)
 
+  def test_key_attributes(self):
+    key = self.make_keys()
+    self.assertEqual(key.itemsize, key.dtype.itemsize)
+    self.assertEqual(key.size, math.prod(key.shape))
+    self.assertEqual(key.ndim, len(key.shape))
+
   def test_isinstance(self):
     @jax.jit
     def f(k):