jax-ml/jax · error · TypeError

philox4x32_random_bits got invalid prng key.

Error message

philox4x32_random_bits got invalid prng key.

What it means

philox4x32_random_bits validates that its key is a Philox 4x32 key (uint32 dtype with trailing shape (4,)); keys of any other PRNG implementation or raw arrays raise this TypeError.

Source

Thrown at jax/_src/random/philox4x32.py:216

@api.jit
def _philox4x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  """Internal implementation of philox4x32_fold_in."""
  # Hash the key with the data used as part of the counter.
  k0, k1 = key[0], key[1]
  out0, out1, _, _ = philox4x32_p.bind(
      k0, k1, np.uint32(0), np.uint32(0), np.uint32(0), data
  )
  return jnp.array([out0, out1], dtype=np.uint32)


def philox4x32_random_bits(
    key: typing.Array, bit_width: int, shape: tuple[int, ...]
) -> typing.Array:
  """Sample uniform random bits using a Philox 4x32 key."""
  if not _is_philox4x32_key(key):
    raise TypeError("philox4x32_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 _philox4x32_random_bits(key, bit_width, shape)


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

  k0, k1 = key[0], key[1]
  counts1, counts2 = prng.iota_2x32_shape(shape)
  zeros = jnp.zeros(shape, dtype=np.uint32)

  out0, out1, out2, out3 = philox4x32_p.bind(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create the key with impl='philox4x32' or use jax.random.wrap_key_data on correctly-shaped uint32 data
2. Use impl-generic jax.random.bits instead
3. Check the key's impl with key.impl before calling the low-level API

Example fix

# before
bits = random.philox4x32_random_bits(threefry_key, 32, (8,))
# after
key = jax.random.PRNGKey(0, impl='philox4x32')
bits = random.philox4x32_random_bits(key, 32, (8,))

key = jax.random.PRNGKey(0, impl='philox4x32')
assert jnp.asarray(key).shape[-1:] == (4,) and jnp.asarray(key).dtype == jnp.uint32

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

When it happens

Trigger: Passing a threefry or typed jax.random.KeyArray of a different impl, or a raw uint32 array with last dim != 4.

Common situations: Interoperability code that passes whichever key it has; upgrading legacy raw-key manipulation to typed keys.

Related errors

• philox2x32_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/e6f19a62d50654f2. Report an issue: GitHub.