Repository: google/jax

Issue:
[better errors] passing batched key to `jax.random` leads to inscrutible error
```python
import jax

batch_size = 8
sample_shape = (batch_size,)
rng = jax.random.PRNGKey(0)
rng = jax.random.split(rng, batch_size)

jax.random.normal(rng, shape=sample_shape)
```
```pytb
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], line 6
      3 rng = jax.random.PRNGKey(0)
      4 rng = jax.random.split(rng, batch_size)
----> 6 jax.random.normal(rng, shape=sample_shape)

File ~/github/google/jax/jax/_src/random.py:707, in normal(key, shape, dtype)
    705 dtype = dtypes.canonicalize_dtype(dtype)
    706 shape = core.as_named_shape(shape)
--> 707 return _normal(key, shape, dtype)

    [... skipping hidden 12 frame]

File ~/github/google/jax/jax/_src/random.py:720, in _normal(key, shape, dtype)
    718   return (_re + 1j * _im) / sqrt2
    719 else:
--> 720   return _normal_real(key, shape, dtype)

    [... skipping hidden 12 frame]

File ~/github/google/jax/jax/_src/random.py:727, in _normal_real(key, shape, dtype)
    725 lo = np.nextafter(np.array(-1., dtype), np.array(0., dtype), dtype=dtype)
    726 hi = np.array(1., dtype)
--> 727 u = uniform(key, shape, dtype, lo, hi)  # type: ignore[arg-type]
    728 return lax.mul(np.array(np.sqrt(2), dtype), lax.erf_inv(u))

File ~/github/google/jax/jax/_src/random.py:394, in uniform(key, shape, dtype, minval, maxval)
    392 dtype = dtypes.canonicalize_dtype(dtype)
    393 shape = core.as_named_shape(shape)
--> 394 return _uniform(key, shape, dtype, minval, maxval)

    [... skipping hidden 12 frame]

File ~/github/google/jax/jax/_src/random.py:432, in _uniform(key, shape, dtype, minval, maxval)
    425 float_bits = lax.bitwise_or(
    426     lax.shift_right_logical(bits, np.array(rng_bits - nmant, uint_dtype)),
    427     np.array(1.0, dtype).view(uint_dtype),
    428 )
    429 floats = lax.bitcast_convert_type(float_bits, dtype) - np.array(1., dtype)
    430 return lax.max(
    431     minval,
--> 432     lax.reshape(floats * (maxval - minval) + minval, shape.positional))

    [... skipping hidden 8 frame]

File ~/github/google/jax/jax/_src/lax/lax.py:3386, in _reshape_shape_rule(operand, new_sizes, dimensions)
   3383 if (not config.dynamic_shapes.value and
   3384     not math.prod(np.shape(operand)) == math.prod(new_sizes)):
   3385   msg = 'reshape total size must be unchanged, got new_sizes {} for shape {}.'
-> 3386   raise TypeError(msg.format(new_sizes, np.shape(operand)))
   3387 if dimensions is not None:
   3388   if set(dimensions) != set(range(np.ndim(operand))):

TypeError: reshape total size must be unchanged, got new_sizes (8,) for shape (8, 8).
```


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
@@ -69,12 +69,18 @@ def _isnan(x: ArrayLike) -> Array:
   return lax.ne(x, x)
 
 
-def _check_prng_key(key: KeyArrayLike) -> tuple[KeyArray, bool]:
+# TODO(jakevdp) Finalize batched input deprecation by setting error_on_batched=True.
+# FutureWarning Added 2024-01-17
+def _check_prng_key(name: str, key: KeyArrayLike, *,
+                    allow_batched: bool = False,
+                    error_on_batched: bool = False) -> tuple[KeyArray, bool]:
   if isinstance(key, Array) and dtypes.issubdtype(key.dtype, dtypes.prng_key):
-    return key, False
+    wrapped_key = key
+    wrapped = False
   elif _arraylike(key):
     # Call random_wrap here to surface errors for invalid keys.
     wrapped_key = prng.random_wrap(key, impl=default_prng_impl())
+    wrapped = True
     if config.legacy_prng_key.value == 'error':
       raise ValueError(
         'Legacy uint32 key array passed as key to jax.random function. '
@@ -91,10 +97,20 @@ def _check_prng_key(key: KeyArrayLike) -> tuple[KeyArray, bool]:
           'Raw arrays as random keys to jax.random functions are deprecated. '
           'Assuming valid threefry2x32 key for now.',
           FutureWarning)
-    return wrapped_key, True
   else:
     raise TypeError(f'unexpected PRNG key type {type(key)}')
 
+  if (not allow_batched) and wrapped_key.ndim:
+    msg = (f"{name} accepts a single key, but was given a key array of "
+           f"shape {np.shape(key)} != (). Use jax.vmap for batching.")
+    if error_on_batched:
+      raise ValueError(msg)
+    else:
+      warnings.warn(msg + " In a future JAX version, this will be an error.",
+                    FutureWarning, stacklevel=3)
+
+  return wrapped_key, wrapped
+
 
 def _return_prng_keys(was_wrapped, key):
   # TODO(frostig): remove once we always enable_custom_prng
@@ -245,10 +261,7 @@ def fold_in(key: KeyArrayLike, data: IntegerArray) -> KeyArray:
     A new PRNG key that is a deterministic function of the inputs and is
     statistically safe for producing a stream of new pseudo-random values.
   """
-  key, wrapped = _check_prng_key(key)
-  if np.ndim(key):
-    raise TypeError("fold_in accepts a single key, but was given a key array of"
-                    f"shape {np.shape(key)} != (). Use jax.vmap for batching.")
+  key, wrapped = _check_prng_key("fold_in", key, error_on_batched=True)
   if np.ndim(data):
     raise TypeError("fold_in accepts a scalar, but was given an array of"
                     f"shape {np.shape(data)} != (). Use jax.vmap for batching.")
@@ -262,7 +275,7 @@ def _split(key: KeyArray, num: int | tuple[int, ...] = 2) -> KeyArray:
   # to always enable_custom_prng
   assert jnp.issubdtype(key.dtype, dtypes.prng_key)
   if key.ndim:
-    raise TypeError("split accepts a single key, but was given a key array of"
+    raise TypeError("split accepts a single key, but was given a key array of "
                     f"shape {key.shape} != (). Use jax.vmap for batching.")
   shape = tuple(num) if isinstance(num, Sequence) else (num,)
   return prng.random_split(key, shape=shape)
@@ -278,7 +291,7 @@ def split(key: KeyArrayLike, num: int | tuple[int, ...] = 2) -> KeyArray:
   Returns:
     An array-like object of `num` new PRNG keys.
   """
-  typed_key, wrapped = _check_prng_key(key)
+  typed_key, wrapped = _check_prng_key("split", key, error_on_batched=True)
   return _return_prng_keys(wrapped, _split(typed_key, num))
 
 
@@ -288,7 +301,7 @@ def _key_impl(keys: KeyArray) -> PRNGImpl:
   return keys_dtype._impl
 
 def key_impl(keys: KeyArrayLike) -> Hashable:
-  typed_keys, _ = _check_prng_key(keys)
+  typed_keys, _ = _check_prng_key("key_impl", keys, allow_batched=True)
   return PRNGSpec(_key_impl(typed_keys))
 
 
@@ -298,7 +311,7 @@ def _key_data(keys: KeyArray) -> Array:
 
 def key_data(keys: KeyArrayLike) -> Array:
   """Recover the bits of key data underlying a PRNG key array."""
-  keys, _ = _check_prng_key(keys)
+  keys, _ = _check_prng_key("key_data", keys, allow_batched=True)
   return _key_data(keys)
 
 
@@ -350,7 +363,7 @@ def bits(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("bits", key)
   if dtype is None:
     dtype = dtypes.canonicalize_dtype(jnp.uint)
   else:
@@ -383,7 +396,7 @@ def uniform(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("uniform", key)
   dtypes.check_user_dtype_supported(dtype)
 
   if not dtypes.issubdtype(dtype, np.floating):
@@ -452,7 +465,7 @@ def randint(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("randint", key)
   dtypes.check_user_dtype_supported(dtype)
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = core.canonicalize_shape(shape)
@@ -535,7 +548,7 @@ def shuffle(key: KeyArrayLike, x: ArrayLike, axis: int = 0) -> Array:
   msg = ("jax.random.shuffle is deprecated and will be removed in a future release. "
          "Use jax.random.permutation with independent=True.")
   warnings.warn(msg, FutureWarning)
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("shuffle", key)
   return _shuffle(key, x, axis)  # type: ignore
 
 
@@ -556,7 +569,7 @@ def permutation(key: KeyArrayLike,
   Returns:
     A shuffled version of x or array range
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("permutation", key)
   check_arraylike("permutation", x)
   axis = canonicalize_axis(axis, np.ndim(x) or 1)
   if not np.ndim(x):
@@ -630,7 +643,7 @@ def choice(key: KeyArrayLike,
   Returns:
     An array of shape `shape` containing samples from `a`.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("choice", key)
   if not isinstance(shape, Sequence):
     raise TypeError("shape argument of jax.random.choice must be a sequence, "
                     f"got {shape}")
@@ -697,7 +710,7 @@ def normal(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("normal", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.inexact):
     raise ValueError(f"dtype argument to `normal` must be a float or complex dtype, "
@@ -764,7 +777,7 @@ def multivariate_normal(key: KeyArrayLike,
     ``shape + mean.shape[-1:]`` if ``shape`` is not None, or else
     ``broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + mean.shape[-1:]``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("multivariate_normal", key)
   dtypes.check_user_dtype_supported(dtype)
   mean, cov = promote_dtypes_inexact(mean, cov)
   if method not in {'svd', 'eigh', 'cholesky'}:
@@ -843,7 +856,7 @@ def truncated_normal(key: KeyArrayLike,
     ``shape`` is not None, or else by broadcasting ``lower`` and ``upper``.
     Returns values in the open interval ``(lower, upper)``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("truncated_normal", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `truncated_normal` must be a float "
@@ -901,7 +914,7 @@ def bernoulli(key: KeyArrayLike,
     A random array with boolean dtype and shape given by ``shape`` if ``shape``
     is not None, or else ``p.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("bernoulli", key)
   dtype = dtypes.canonicalize_dtype(lax.dtype(p))
   if shape is not None:
     shape = core.as_named_shape(shape)
@@ -952,7 +965,7 @@ def beta(key: KeyArrayLike,
     A random array with the specified dtype and shape given by ``shape`` if
     ``shape`` is not None, or else by broadcasting ``a`` and ``b``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("beta", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `beta` must be a float "
@@ -1005,7 +1018,7 @@ def cauchy(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("cauchy", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `cauchy` must be a float "
@@ -1057,7 +1070,7 @@ def dirichlet(key: KeyArrayLike,
     ``shape + (alpha.shape[-1],)`` if ``shape`` is not None, or else
     ``alpha.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("dirichlet", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `dirichlet` must be a float "
@@ -1116,7 +1129,7 @@ def exponential(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("exponential", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `exponential` must be a float "
@@ -1297,7 +1310,7 @@ def gamma(key: KeyArrayLike,
     loggamma : sample gamma values in log-space, which can provide improved
       accuracy for small values of ``a``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("gamma", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `gamma` must be a float "
@@ -1339,7 +1352,7 @@ def loggamma(key: KeyArrayLike,
   See Also:
     gamma : standard gamma sampler.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("loggamma", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `gamma` must be a float "
@@ -1475,7 +1488,7 @@ def poisson(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape is not None, or else by ``lam.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("poisson", key)
   dtypes.check_user_dtype_supported(dtype)
   # TODO(frostig): generalize underlying poisson implementation and
   # remove this check
@@ -1515,7 +1528,7 @@ def gumbel(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("gumbel", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `gumbel` must be a float "
@@ -1550,7 +1563,7 @@ def categorical(key: KeyArrayLike,
     A random array with int dtype and shape given by ``shape`` if ``shape``
     is not None, or else ``np.delete(logits.shape, axis)``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("categorical", key)
   check_arraylike("categorical", logits)
   logits_arr = jnp.asarray(logits)
 
@@ -1593,7 +1606,7 @@ def laplace(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("laplace", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `laplace` must be a float "
@@ -1630,7 +1643,7 @@ def logistic(key: KeyArrayLike,
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("logistic", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `logistic` must be a float "
@@ -1673,7 +1686,7 @@ def pareto(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``b.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("pareto", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `pareto` must be a float "
@@ -1722,7 +1735,7 @@ def t(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``df.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("t", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `t` must be a float "
@@ -1775,7 +1788,7 @@ def chisquare(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``df.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("chisquare", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError("dtype argument to `chisquare` must be a float "
@@ -1833,7 +1846,7 @@ def f(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``df.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("f", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError("dtype argument to `f` must be a float "
@@ -1885,7 +1898,7 @@ def rademacher(key: KeyArrayLike,
     a 50% change of being 1 or -1.
 
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("rademacher", key)
   dtypes.check_user_dtype_supported(dtype)
   dtype = dtypes.canonicalize_dtype(dtype)
   shape = core.canonicalize_shape(shape)
@@ -1921,7 +1934,7 @@ def maxwell(key: KeyArrayLike,
   """
   # Generate samples using:
   # sqrt(X^2 + Y^2 + Z^2), X,Y,Z ~N(0,1)
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("maxwell", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `maxwell` must be a float "
@@ -1964,7 +1977,7 @@ def double_sided_maxwell(key: KeyArrayLike,
     A jnp.array of samples.
 
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("double_sided_maxwell", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `double_sided_maxwell` must be a float"
@@ -2016,7 +2029,7 @@ def weibull_min(key: KeyArrayLike,
     A jnp.array of samples.
 
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("weibull_min", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError(f"dtype argument to `weibull_min` must be a float "
@@ -2055,7 +2068,7 @@ def orthogonal(
   Returns:
     A random array of shape `(*shape, n, n)` and specified dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("orthogonal", key)
   dtypes.check_user_dtype_supported(dtype)
   _check_shape("orthogonal", shape)
   n = core.concrete_or_error(index, n, "The error occurred in jax.random.orthogonal()")
@@ -2090,7 +2103,7 @@ def generalized_normal(
   Returns:
     A random array with the specified shape and dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("generalized_normal", key)
   dtypes.check_user_dtype_supported(dtype)
   _check_shape("generalized_normal", shape)
   keys = split(key)
@@ -2120,7 +2133,7 @@ def ball(
   Returns:
     A random array of shape `(*shape, d)` and specified dtype.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("ball", key)
   dtypes.check_user_dtype_supported(dtype)
   _check_shape("ball", shape)
   d = core.concrete_or_error(index, d, "The error occurred in jax.random.ball()")
@@ -2158,7 +2171,7 @@ def rayleigh(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``scale.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("rayleigh", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError("dtype argument to `rayleigh` must be a float "
@@ -2212,7 +2225,7 @@ def wald(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``mean.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("wald", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError("dtype argument to `wald` must be a float "
@@ -2268,7 +2281,7 @@ def geometric(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``p.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("geometric", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.integer):
     raise ValueError("dtype argument to `geometric` must be an int "
@@ -2330,7 +2343,7 @@ def triangular(key: KeyArrayLike,
     A random array with the specified dtype and with shape given by ``shape`` if
     ``shape`` is not None, or else by ``left.shape``, ``mode.shape`` and ``right.shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("triangular", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):
     raise ValueError("dtype argument to `triangular` must be a float "
@@ -2384,7 +2397,7 @@ def lognormal(key: KeyArrayLike,
   Returns:
     A random array with the specified dtype and with shape given by ``shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("lognormal", key)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.inexact):
     raise ValueError(f"dtype argument to `lognormal` must be a float or complex dtype, "
@@ -2597,7 +2610,7 @@ def binomial(
     A random array with the specified dtype and with shape given by
     ``np.broadcast(n, p).shape``.
   """
-  key, _ = _check_prng_key(key)
+  key, _ = _check_prng_key("binomial", key)
   check_arraylike("binomial", n, p)
   dtypes.check_user_dtype_supported(dtype)
   if not dtypes.issubdtype(dtype, np.floating):


Test patch:
diff --git a/jax/experimental/jax2tf/tests/shape_poly_test.py b/jax/experimental/jax2tf/tests/shape_poly_test.py
--- a/jax/experimental/jax2tf/tests/shape_poly_test.py
+++ b/jax/experimental/jax2tf/tests/shape_poly_test.py
@@ -2108,7 +2108,7 @@ def f2(z, w):  # z: f32[a, 5] w: f32[a + b, 5] -> f32[2*a + b, 10]
       # non-partitionable), and unsafe_rbg.
       [
         PolyHarness("random_gamma", f"{flags_name}",
-                    lambda key, a: jax.random.gamma(key, a),
+                    lambda key, a: jax.vmap(jax.random.gamma)(key, a),
                     arg_descriptors=[RandArg((3, key_size), np.uint32), RandArg((3, 4, 5), _f32)],
                     polymorphic_shapes=["b, ...", "b, w, ..."], tol=1E-5,
                     override_jax_config_flags=override_jax_config_flags),  # type: ignore
diff --git a/tests/key_reuse_test.py b/tests/key_reuse_test.py
--- a/tests/key_reuse_test.py
+++ b/tests/key_reuse_test.py
@@ -616,7 +616,7 @@ def test_reuse_after_broadcast(self):
     def f():
       key = jax.random.key(0)
       key2 = key[None]
-      return jax.random.bits(key) + jax.random.bits(key2)
+      return jax.random.bits(key) + jax.vmap(jax.random.bits)(key2)
 
     with self.assertRaisesRegex(KeyReuseError, self.random_bits_error):
       self.check_key_reuse(f)
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -1247,6 +1247,34 @@ def testBinomialCornerCases(self):
     self.assertArraysAllClose(samples2, jnp.array([jnp.nan, 0., jnp.nan, jnp.nan]), check_dtypes=False)
     self.assertArraysAllClose(samples3, jnp.array([jnp.nan, jnp.nan, jnp.nan]), check_dtypes=False)
 
+  def test_batched_key_warnings(self):
+    keys = jax.random.split(self.make_key(0))
+    msg = "{} accepts a single key, but was given a key array of shape.*"
+
+    # Check a handful of functions that are expected to warn.
+    with self.assertWarnsRegex(FutureWarning, msg.format('bits')):
+      jax.random.bits(keys, shape=(2,))
+    with self.assertWarnsRegex(FutureWarning, msg.format('chisquare')):
+      jax.random.chisquare(keys, 1.0, shape=(2,))
+    with self.assertWarnsRegex(FutureWarning, msg.format('dirichlet')):
+      jax.random.dirichlet(keys, jnp.arange(2.0), shape=(2,))
+    with self.assertWarnsRegex(FutureWarning, msg.format('gamma')):
+      jax.random.gamma(keys, 1.0, shape=(2,))
+    with self.assertWarnsRegex(FutureWarning, msg.format('loggamma')):
+      jax.random.loggamma(keys, 1.0, shape=(2,))
+
+    # Other functions should error; test a few cases.
+    with self.assertRaisesRegex(ValueError, msg.format('fold_in')):
+      jax.random.fold_in(keys, 0)
+    with self.assertRaisesRegex(ValueError, msg.format('split')):
+      jax.random.split(keys)
+
+    # Some shouldn't error or warn
+    with self.assertNoWarnings():
+      jax.random.key_data(keys)
+      jax.random.key_impl(keys)
+
+
 threefry_seed = prng_internal.threefry_seed
 threefry_split = prng_internal.threefry_split
 threefry_random_bits = prng_internal.threefry_random_bits
diff --git a/tests/shape_poly_test.py b/tests/shape_poly_test.py
--- a/tests/shape_poly_test.py
+++ b/tests/shape_poly_test.py
@@ -1917,7 +1917,7 @@ def test_vmap_error(self):
       # non-partitionable), and unsafe_rbg.
       [
         PolyHarness("random_gamma", f"{flags_name}",
-                    lambda key, a: jax.random.gamma(
+                    lambda key, a: jax.vmap(jax.random.gamma)(
                       jax.random.wrap_key_data(key), a),
                     arg_descriptors=[RandArg((3, key_size), np.uint32),
                                      RandArg((3, 4, 5), _f32)],