jax-ml/jax · error · ValueError

Only 32-bit PRNG supported.

Error message

Only 32-bit PRNG supported.

What it means

The TPU Pallas Philox implementation only generates 32-bit random values: philox_random_bits raises ValueError for any bit_width other than 32 because the kernel produces uint32 words only.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/philox.py:193

        key, padded_shape, shape,
        block_size=BLOCK_SIZE, offset=offset,
        fuse_output=fuse_output)
    return padded_result[..., :shape[-2], :shape[-1]]
  else:
    return philox_4x32_kernel(key, shape, shape,
                              block_size=BLOCK_SIZE, offset=offset,
                              fuse_output=fuse_output)


def philox_split(key, shape: Shape):
  """Splits the key into two keys of the same shape."""
  bits1, bits2 = philox_4x32_count(key, shape, fuse_output=False)
  return jnp.stack([bits1, bits2], axis=bits1.ndim)


def philox_random_bits(key, bit_width: int, shape: Shape):
  if bit_width != 32:
    raise ValueError("Only 32-bit PRNG supported.")
  return philox_4x32_count(key, shape, fuse_output=True)


def philox_fold_in(key, data):
  assert data.ndim == 0
  return philox_4x32_count(key, (), offset=data, fuse_output=False)


plphilox_prng_impl = prng.PRNGImpl(
    key_shape=(2,),
    seed=threefry2x32.threefry_seed,
    split=philox_split,
    random_bits=philox_random_bits,
    fold_in=philox_fold_in,
    name="pallas_philox4x32",
    tag="pllox")

prng.register_prng(plphilox_prng_impl)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Request 32 bits and truncate/shift yourself for smaller widths: (bits >> 24).astype(jnp.uint8)
2. For 64-bit, generate two 32-bit halves and combine: (hi.astype(uint64) << 32) | lo
3. Route non-32-bit requests to the standard jax.random implementation instead

Example fix

// before
u8 = philox_random_bits(key, 8, shape)
// after
u8 = (philox_random_bits(key, 32, shape) >> 24).astype(jnp.uint8)

assert bit_width == 32, 'TPU pallas philox only supports 32-bit; derive smaller widths by shifting'

When it happens

Trigger: Calling philox_random_bits(key, 8, shape) or (key, 16, shape) or (key, 64, shape).

Common situations: Plugging this backend into a generic PRNG interface (e.g. jax.random.bits with dtype uint8/uint16/uint64) where bit width varies by dtype.

Related errors

• Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
• Shape dimension {shape[-2:]} must be divisible by {block_siz
• Offset must be scalar, got {offset.shape}
• Only 32-bit PRNG supported.
• 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/60d35695ae8b927d. Report an issue: GitHub.