jax-ml/jax · error · TypeError

_rbg_random_bits got invalid prng key.

Error message

_rbg_random_bits got invalid prng key.

What it means

The RBG (stateless random-bit generator) implementation stores each key as exactly 4 uint32 words. This TypeError from _rbg_random_bits fires when the key argument does not have shape (4,) and dtype uint32 — typically because a threefry-style (2,)-word key or an ordinary array was passed to an rbg-based operation.

Source

Thrown at jax/_src/random/rbg.py:61

def _rbg_split(key: typing.Array, shape: prng.Shape) -> typing.Array:
  if config.threefry_partitionable.value:
    _threefry_split = threefry2x32._threefry_split_foldlike
  else:
    _threefry_split = threefry2x32._threefry_split_original
  halfkeys = key.reshape(2, 2)
  return api.vmap(
      _threefry_split, (0, None), len(shape))(halfkeys, shape).reshape(
          *shape, 4)

def _rbg_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  assert not data.shape
  return api.vmap(threefry2x32._threefry_fold_in, (0, None), 0)(key.reshape(2, 2), data).reshape(4)

def _rbg_random_bits(key: typing.Array, bit_width: int, shape: Sequence[int]
                     ) -> typing.Array:
  if not key.shape == (4,) and key.dtype == np.dtype('uint32'):
    raise TypeError("_rbg_random_bits got invalid prng key.")
  if bit_width not in (8, 16, 32, 64):
    raise TypeError("requires 8-, 16-, 32- or 64-bit field width.")
  _, bits = lax.rng_bit_generator(key, shape, dtype=prng.UINT_DTYPES[bit_width])
  return bits

rbg_prng_impl = prng.PRNGImpl(
    key_shape=(4,),
    seed=_rbg_seed,
    split=_rbg_split,
    random_bits=_rbg_random_bits,
    fold_in=_rbg_fold_in,
    name='rbg',
    tag='rbg')

prng.register_prng(rbg_prng_impl)


def _unsafe_rbg_split(key: typing.Array, shape: prng.Shape) -> typing.Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the key with matching impl: jax.random.key(seed, impl='rbg')
2. Validate key shape/dtype before use: key.shape == (4,) and key.dtype == jnp.uint32 (via key_data for typed keys)
3. Regenerate any persisted keys after changing implementations

Example fix

// before
key = jax.random.key(0)  # threefry, shape (2,)
bits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32, impl='rbg')

// after
key = jax.random.key(0, impl='rbg')
bits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32, impl='rbg')

import jax, jax.numpy as jnp
data = jax.random.key_data(key)
assert data.shape == (4,) and data.dtype == jnp.uint32, f'bad rbg key: {data.shape} {data.dtype}'

import jax, jax.numpy as jnp
def is_rbg_key(key) -> bool:
    d = jax.random.key_data(key)
    return d.shape == (4,) and d.dtype == jnp.uint32

When it happens

Trigger: Calling jax.random.bits or sampling functions with impl='rbg' on a threefry key; passing lax.rng_bit_generator a key of the wrong family; hand-building rbg key data with 2 words instead of 4.

Common situations: Mixing keys created with jax.random.key(seed) (threefry default) into code configured for rbg; configuration changes of the global PRNG implementation without regenerating stored keys.

Related errors

• JAX encountered invalid PRNG key data: expected key_data.sha
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out

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