jax-ml/jax · error · TypeError

JAX encountered invalid PRNG key data: expected key_data.sha

Error message

JAX encountered invalid PRNG key data: expected key_data.shape to end with {impl.key_shape}; got shape={key_data.shape} for {impl=}

What it means

Every JAX RNG implementation defines a fixed key_shape (e.g. (2,) for threefry2x32, (4,) for rbg). This TypeError means the key_data's trailing dimensions do not match that shape, so the data cannot be interpreted as keys of the requested implementation. It is raised by _check_prng_key_data when key_data.shape[-len(impl.key_shape):] != impl.key_shape.

Source

Thrown at jax/_src/random/prng.py:135

def register_prng(impl: PRNGImpl):
  if impl.name in prngs:
    raise ValueError(f'PRNG with name {impl.name} already registered: {impl}')
  prngs[impl.name] = impl


# -- PRNG key arrays

def _check_prng_key_data(impl, key_data: typing.Array):
  ndim = len(impl.key_shape)
  if not all(hasattr(key_data, attr) for attr in ['ndim', 'shape', 'dtype']):
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data "
                    f"to have ndim, shape, and dtype attributes. Got {key_data}")
  if key_data.ndim < 1:
    raise TypeError("JAX encountered invalid PRNG key data: expected "
                    f"key_data.ndim >= 1; got ndim={key_data.ndim}")
  if key_data.shape[-ndim:] != impl.key_shape:
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data.shape to "
                    f"end with {impl.key_shape}; got shape={key_data.shape} for {impl=}")
  if key_data.dtype not in [np.uint32, float0]:
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; "
                    f"got dtype={key_data.dtype}")


class PRNGKeyArray(Array):
  """An array of PRNG keys backed by an RNG implementation.

  This class lifts the definition of a PRNG, provided in the form of a
  ``PRNGImpl``, into an array-like pytree class. Instances of this
  class behave like an array whose base elements are keys, hiding the
  fact that keys are typically arrays (of ``uint32`` dtype) themselves.

  PRNGKeyArrays are also restricted relative to JAX arrays in that
  they do not expose arithmetic operations. They instead expose
  wrapper methods around the PRNG implementation functions (``split``,
  ``random_bits``, ``fold_in``).

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match the data to the impl: use 2 uint32 words for threefry2x32, 4 for rbg
2. Re-derive keys with jax.random.key(seed, impl=impl) instead of transplanting raw data between implementations
3. Print impl.key_shape for your implementation and reshape/extend the data accordingly

Example fix

// before
impl = jax.random.rbg_prng_impl
key = jax.random.random_wrap(jnp.uint32([1, 2]), impl=impl)  # rbg needs 4 words

// after
key = jax.random.key(42, impl='rbg')

import jax.numpy as jnp
KEY_SHAPES = {'threefry2x32': (2,), 'rbg': (4,), 'unsafe_rbg': (4,)}
assert data.shape[-len(KEY_SHAPES[impl_name]):] == KEY_SHAPES[impl_name]

import jax.numpy as jnp
def matches_impl_shape(data, impl) -> bool:
    ks = impl.key_shape
    return data.shape[-len(ks):] == ks

When it happens

Trigger: Wrapping a length-4 uint32 array as a 'threefry2x32' key (expects trailing shape (2,)); wrapping a (2,)-shaped array as an 'rbg' key (expects (4,)); wrapping data with an extra partial dimension such as shape (3,) or (5,).

Common situations: Switching the impl argument when migrating from legacy PRNGKey arrays to new-style keys or to rbg; mixing key data from different generator families (threefry vs rbg vs unsafe_rbg) in the same codebase.

Related errors

• JAX encountered invalid PRNG key data: expected key_data to
• JAX encountered invalid PRNG key data: expected key_data.ndi
• _rbg_random_bits got invalid prng key.
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0f9ceea7104f8cdc. Report an issue: GitHub.