jax-ml/jax · error · ValueError

Only 32-bit PRNG supported.

Error message

Only 32-bit PRNG supported.

What it means

The Pallas TPU Threefry backend only supports 32-bit random output; plthreefry_random_bits raises ValueError for bit_width != 32 because the 2x32 kernel emits uint32 words only.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/threefry.py:92

        k1, k2, counts_hi, counts_lo)
    out_bits = o1 ^ o2
    out_ref[...] = out_bits.reshape(out_ref.shape)

  key = key.reshape((1, 2))
  out = jax.ShapeDtypeStruct(shape, dtype=jnp.uint32)
  block_shape = (1,) * (len(shape)-2) + block_size
  result = pl.pallas_call(
      kernel,
      in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)],
      out_specs=pl.BlockSpec(block_shape, lambda *idxs: idxs),
      grid=grid_dims,
      out_shape=out,
  )(key)
  return result

def plthreefry_random_bits(key, bit_width: int, shape: Shape):
  if bit_width != 32:
    raise ValueError("Only 32-bit PRNG supported.")
  if len(shape) == 0:
    return plthreefry_random_bits(key, bit_width, (1, 1))[0, 0]
  elif len(shape) == 1:
    return plthreefry_random_bits(key, bit_width, (1, *shape))[0]

  requires_pad = (
      shape[-2] % BLOCK_SIZE[-2] != 0) or (shape[-1] % BLOCK_SIZE[-1] != 0)
  if requires_pad:
    padded_shape = tuple(shape[:-2]) + (
        prng_utils.round_up(shape[-2], BLOCK_SIZE[-2]),
        prng_utils.round_up(shape[-1], BLOCK_SIZE[-1]),
    )
    padded_result = threefry_2x32_count(
        key, padded_shape, shape, block_size=BLOCK_SIZE)
    return padded_result[..., :shape[-2], :shape[-1]]
  else:
    return threefry_2x32_count(key, shape, shape, block_size=BLOCK_SIZE)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Generate 32 bits and downcast: (bits >> 24).astype(jnp.uint8)
2. For 64-bit, combine two 32-bit draws
3. Dispatch non-32-bit widths to the default threefry PRNG implementation

Example fix

// before
u16 = plthreefry_random_bits(key, 16, shape)
// after
u16 = (plthreefry_random_bits(key, 32, shape) >> 16).astype(jnp.uint16)

assert bit_width == 32, 'pallas threefry only emits uint32; downcast manually for smaller widths'

When it happens

Trigger: Calling plthreefry_random_bits(key, 16, shape) or any width other than 32.

Common situations: A generic PRNG-dispatch layer (like jax.random.bits' internal backend selection) forwarding dtypes such as uint8/uint16/uint64 to this TPU backend.

Related errors

• Only 32-bit PRNG supported.
• Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
• Shape dimension {shape[-2:]} must be divisible by {block_siz
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Bit width must be 32

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