jax-ml/jax · error · TypeError

threefry_random_bits got invalid prng key.

Error message

threefry_random_bits got invalid prng key.

What it means

threefry_random_bits requires its key argument to be a valid threefry PRNG key: raw uint32 data of shape (2,) (or batched equivalents). This TypeError fires when the key is from another implementation (e.g. rbg with 4 words), already-consumed raw data of the wrong shape, or a non-key array.

Source

Thrown at jax/_src/random/threefry2x32.py:324

  k1, k2 = key
  counts1, counts2 = prng.iota_2x32_shape(shape)
  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)
  return jnp.stack([bits1, bits2], axis=bits1.ndim)


def threefry_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  assert not data.shape
  return _threefry_fold_in(key, jnp.asarray(data, dtype='uint32'))

@api.jit
def _threefry_fold_in(key, data):
  return threefry_2x32(key, threefry_seed(data))


def threefry_random_bits(key: typing.Array, bit_width, shape):
  """Sample uniform random bits of given width and shape using PRNG key."""
  if not _is_threefry_prng_key(key):
    raise TypeError("threefry_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.")

  if config.threefry_partitionable.value:
    return _threefry_random_bits_partitionable(key, bit_width, shape)
  else:
    return _threefry_random_bits_original(key, bit_width, shape)

def _threefry_random_bits_partitionable(key: typing.Array, bit_width, shape):
  if all(core.is_constant_dim(d) for d in shape) and math.prod(shape) > 2 ** 64:
    raise NotImplementedError('random bits array of size exceeding 2 ** 64')

  k1, k2 = key
  counts1, counts2 = prng.iota_2x32_shape(shape)
  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)

  dtype = prng.UINT_DTYPES[bit_width]
  if bit_width == 64:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Generate keys with the matching implementation: jax.random.key(seed, impl='threefry2x32')
2. Validate raw key data shape (2,) and uint32 dtype before calling
3. Use public jax.random.bits/rand APIs, which dispatch on the key's impl

Example fix

// before
key = jax.random.key(0, impl='rbg')  # 4-word key
bits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32)  # routed to threefry path

// after
key = jax.random.key(0)  # threefry2x32
bits = jax.random.bits(key, shape=(8,), dtype=jnp.uint32)

import jax, jax.numpy as jnp
data = jax.random.key_data(key)
assert data.shape[-2:] == (2,) and data.dtype == jnp.uint32, 'not a threefry key'

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

When it happens

Trigger: Passing an rbg key to threefry-based threefry_random_bits (reachable via jax.random.bits with mismatched impls); calling the internal API with arbitrary uint32 arrays; mixing keys across generator families after config changes.

Common situations: Global PRNG implementation switched to/from rbg without regenerating keys; code paths that conditionally use _double_threefry_random_bits with keys of the wrong family; custom samplers calling internal bit-generators directly.

Related errors

• philox2x32_random_bits got invalid prng key.
• philox4x32_random_bits got invalid prng key.
• threefry_2x32 requires uint32 arguments, got {}
• requires 8-, 16-, 32- or 64-bit field width.
• PRNG keys must be loaded from SMEM. Did you set the memory s

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