jax-ml/jax · error · TypeError

requires 8-, 16-, 32- or 64-bit field width.

Error message

requires 8-, 16-, 32- or 64-bit field width.

What it means

philox4x32_random_bits only supports bit widths 8, 16, 32, and 64; any other bit_width raises this TypeError since output packing is defined only for those widths.

Source

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

@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(
      k0, k1, counts1, counts2, zeros, zeros
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use one of 8, 16, 32, 64
2. Mask a wider draw for odd widths: bits & ((1 << w) - 1)
3. Use jax.random.bits(key, shape, dtype=jnp.uintN) for the same effect

Example fix

# before
b = random.philox4x32_random_bits(key, 24, (100,))
# after
b = random.philox4x32_random_bits(key, 32, (100,)) & 0xFFFFFF

assert bit_width in (8, 16, 32, 64), f'invalid bit_width {bit_width}'

def is_valid_bit_width(w) -> bool:
    return w in (8, 16, 32, 64)

When it happens

Trigger: Calling random.philox4x32_random_bits(key, 24, shape) or any width outside {8,16,32,64}.

Common situations: Generating custom-width random integers; typos or off-by-one widths in config-driven pipelines.

Related errors

• requires 8-, 16-, 32- or 64-bit field width.
• PRNG key seed must be a scalar; got {seed!r}.
• philox2x32_random_bits got invalid prng key.
• 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/cea9dcc52d059f5a. Report an issue: GitHub.