jax-ml/jax · error · TypeError

JAX encountered invalid PRNG key data: expected key_data.dty

Error message

JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; got dtype={key_data.dtype}

What it means

JAX PRNG key data must be stored as uint32 (or the internal float0 sentinel). This TypeError fires when key_data has any other dtype (int64, float32, etc.), because the RNG bit-manipulation kernels only operate on 32-bit unsigned words. The check lives in _check_prng_key_data and runs from PRNGKeyArray.__init__ and jax.random.random_wrap.

Source

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

    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``).
  """
  # TODO(jakevdp): potentially add tolist(), tobytes(),
  #    device_buffer, device_buffers, __cuda_interface__()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast explicitly: jax.random.random_wrap(arr.astype(jnp.uint32), impl=...)
2. Prefer creating keys from seeds via jax.random.key(seed)
3. When x64 mode is on, be aware asarray produces int64 and always cast to uint32 before wrapping

Example fix

// before
key = jax.random.random_wrap(jnp.array([1, 2]), impl='threefry2x32')

// after
key = jax.random.random_wrap(jnp.array([1, 2], dtype=jnp.uint32), impl='threefry2x32')

import jax.numpy as jnp
data = jnp.asarray(data)
if data.dtype != jnp.uint32:
    data = data.astype(jnp.uint32)

import jax.numpy as jnp
def is_uint32_key_data(x) -> bool:
    return x.dtype == jnp.uint32

When it happens

Trigger: Calling jax.random.random_wrap on an int64 array (common when x64 mode is enabled), a float array, or a numpy array of dtype int32; converting seeds with jnp.array(seed) (default int32/int64) and wrapping directly.

Common situations: JAX_ENABLE_X64=1 environments where asarray defaults to int64; loading key data from files where it was stored as int64 or float; numpy interop where np.array([1,2]) yields int64 on Linux.

Related errors

• threefry_2x32 requires uint32 arguments, got {}
• random_unwrap takes key array operand, got {keys.dtype=}
• requires 8-, 16-, 32- or 64-bit field width.
• PRNG key seed must be an integer; got {seed!r}
• requires 8-, 16-, 32- or 64-bit field width.

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