jax-ml/jax · error · TypeError

philox2x32_random_bits got invalid prng key.

Error message

philox2x32_random_bits got invalid prng key.

What it means

philox2x32_random_bits checks its key with _is_philox2x32_key (dtype uint32 and trailing shape (2,)); any other key (threefry, a plain array, or an upgraded legacy key) raises this TypeError. The function is the low-level bit-generator behind the philox2x32 implementation.

Source

Thrown at jax/_src/random/philox2x32.py:188

def philox2x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  """Fold-in an integer value to create a new Philox2x32 key."""
  assert not data.shape
  return _philox2x32_fold_in(key, jnp.asarray(data, dtype="uint32"))


@api.jit
def _philox2x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  """Internal implementation of philox2x32_fold_in."""
  out0, _ = philox2x32_p.bind(key[0], np.uint32(0), data)
  return jnp.array([out0], dtype=np.uint32)


def philox2x32_random_bits(
    key: typing.Array, bit_width: int, shape: tuple[int, ...]
) -> typing.Array:
  """Sample uniform random bits using a Philox 2x32 key."""
  if not _is_philox2x32_key(key):
    raise TypeError("philox2x32_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.")
  return _philox2x32_random_bits(key, bit_width, shape)


@api.jit(static_argnums=(1, 2), inline=True)
def _philox2x32_random_bits(
    key: typing.Array, bit_width: int, shape: tuple[int, ...]
) -> typing.Array:
  """Internal implementation of philox2x32_random_bits."""
  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")

  counts1, counts2 = prng.iota_2x32_shape(shape)
  out0, out1 = philox2x32_p.bind(key[0], counts1, counts2)

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the key with the matching impl: jax.random.PRNGKey(seed, impl='philox2x32')
2. Convert typed keys with jax.random.key_data / raw keys via jax.random.wrap_key_data if interoperating with legacy uint32 arrays
3. Prefer high-level jax.random.bits API which dispatches on the key's impl

Example fix

# before
key = jax.random.PRNGKey(0, impl='threefry2x32')
bits = random.philox2x32_random_bits(key, 32, (4,))
# after
key = jax.random.PRNGKey(0, impl='philox2x32')
bits = random.philox2x32_random_bits(key, 32, (4,))

key = jax.random.PRNGKey(0, impl='philox2x32')
assert _is_philox2x32_key(key)  # or key.impl == 'philox2x32' for typed keys

def is_philox2x32(k) -> bool:
    import jax.numpy as jnp
    d = jnp.asarray(k)
    return d.dtype == jnp.uint32 and d.shape[-1:] == (2,)

When it happens

Trigger: Calling random.philox2x32_random_bits(threefry_key, 32, shape), passing a jax.random.KeyArray of a different impl, or passing a raw uint32 array whose last axis is not size 2.

Common situations: Mixing PRNG implementations after jax.random.use_prng_impl or upgrading old code that manipulated raw key arrays instead of typed keys.

Related errors

• philox4x32_random_bits got invalid prng key.
• PRNG key seed must be a scalar; got {seed!r}.
• requires 8-, 16-, 32- or 64-bit field width.
• random bits array of size exceeding 2 ** 64
• PRNG key seed must be a scalar; got {seed!r}.

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