jax-ml/jax · error · ValueError

Bit width must be 32

Error message

Bit width must be 32

What it means

The Pallas TPU random-bit generator is hardwired to 32-bit output, so _random_bits raises ValueError when bit_width != 32. This is called through the PRNGImpl when random_bits/normal/etc. request a different width inside a Pallas kernel.

Source

Thrown at jax/_src/pallas/mosaic/random.py:71

  if vmapped_key:
    pallas_key_data = jax.vmap(generate_key)(key)
  else:
    pallas_key_data = generate_key(key)
  return jax_api_random.wrap_key_data(pallas_key_data, impl="pallas_tpu")

def is_pallas_impl(impl: jax_prng.PRNGImpl) -> bool:
  """Returns True if the PRNGImpl is a Pallas-specific implementation."""
  return impl == tpu_key_impl or impl == tpu_internal_stateful_impl


def _seed_func(seed: jnp.int32):
  seed_data = jnp.zeros(tpu_key_impl.key_shape, dtype=jnp.int32)
  return (seed_data + seed).astype(jnp.uint32)  # Broadcast the seed.

def _random_bits(key: typing.Array, bit_width: int, shape: Shape):
  if bit_width != 32:
    raise ValueError("Bit width must be 32")
  prng_seed(key)
  return prng_random_bits(shape)

def _fold_in(key: jax_prng.PRNGKeyArray, data: typing.Array):
  key0, key1 = unwrap_pallas_seed(key)
  # Perform a cheap mixing of data into the key.
  key1 = key1 + data
  [key0, key1] = threefry2x32.apply_round([key0, key1], 13)
  return wrap_pallas_seed(key0, key1, impl="pallas_tpu")

def _split(key: typing.Array, shape: Shape):
  del key, shape
  raise NotImplementedError(
      "Cannot split a Pallas key. Use fold_in instead to generate new keys."
  )

tpu_key_impl = jax_prng.PRNGImpl(
    key_shape=(1, 2),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Request 32-bit random values inside Pallas kernels (uint32/float32 samples)
2. Generate non-32-bit values on the host or derive them from 32-bit samples

Example fix

# before
bits = random_bits(key, bit_width=8, shape=shape)
# after
bits = (random_bits(key, bit_width=32, shape=shape) & 0xFF).astype(jnp.uint8)

assert bit_width == 32, 'Pallas TPU PRNG only supports 32-bit'

When it happens

Trigger: Using jax.random bits APIs inside a Pallas TPU kernel with the pallas_tpu PRNG impl where the requested bit width is not 32 (e.g. random bits of width 8 or 64).

Common situations: Calling random-level APIs that default to the dtype width (uint16/uint64 keys) inside a kernel; porting host-side jax.random code into Pallas.

Related errors

• PRNG keys must be loaded from SMEM. Did you set the memory s
• Cannot split a Pallas key. Use fold_in instead to generate n
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.
• Can only load a single key per load.

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