jax-ml/jax · error · TypeError

JAX encountered invalid PRNG key data: expected key_data to

Error message

JAX encountered invalid PRNG key data: expected key_data to have ndim, shape, and dtype attributes. Got {key_data}

What it means

JAX's typed PRNG key arrays (PRNGKeyArray) wrap an underlying 'key_data' array whose implementation is validated at construction. This TypeError means the object passed as key_data is not array-like from JAX's perspective: it lacks basic ndim/shape/dtype attributes (e.g. a plain Python list, an int, or None was passed where a key data array was expected). JAX throws it as a fail-fast check inside _check_prng_key_data, which runs from PRNGKeyArray.__init__ and jax.random.random_wrap.

Source

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

            pp.nest(2, pp.group(pp.brk() + pp.join(pp.brk(), [
              pp.text(f"{k} = {v}") for k, v in self._asdict().items()
            ]))))


prngs: dict[str, PRNGImpl] = {}

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the data to a JAX array first: jax.random.random_wrap(jnp.asarray(data, dtype=jnp.uint32), impl=...)
2. If starting from a seed, create a key with jax.random.key(seed) or jax.random.PRNGKey(seed) instead of hand-building key data
3. Ensure the wrapped data is uint32 with trailing shape matching the impl's key_shape (e.g. (2,) for threefry2x32)

Example fix

// before
raw = [123, 456]
key = jax.random.random_wrap(raw, impl='threefry2x32')

// after
import jax.numpy as jnp
key = jax.random.random_wrap(jnp.asarray(raw, dtype=jnp.uint32), impl='threefry2x32')

import jax.numpy as jnp
def is_valid_key_data(x):
    return (all(hasattr(x, a) for a in ('ndim', 'shape', 'dtype'))
            and x.ndim >= 1
            and x.shape[-2:] == (2,)
            and x.dtype == jnp.uint32)

import jax.numpy as jnp
from typing import Any
def is_wrappable_key_data(x: Any) -> bool:
    return (all(hasattr(x, a) for a in ('ndim', 'shape', 'dtype'))
            and x.ndim >= 1 and x.shape[-2:] == (2,)
            and x.dtype == jnp.uint32)

When it happens

Trigger: Constructing a PRNGKeyArray/PRNGImpl with a non-array key_data (e.g. PRNGKeyArray(impl, [1,2,3])), or calling jax.random.random_wrap with a Python list/tuple/scalar instead of a uint32 array; also passing objects of array-like libraries that don't expose the trio of attributes.

Common situations: Migrating legacy code that stored raw key data as Python lists; deserializing keys from JSON/config and passing the decoded list straight back; wrapping third-party array objects that lack JAX's array API surface.

Related errors

• JAX encountered invalid PRNG key data: expected key_data.ndi
• JAX encountered invalid PRNG key data: expected key_data.sha
• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error

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