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

philox2x32_random_bits validates bit_width and only accepts 8, 16, 32, or 64 — the widths for which it packs uint outputs. Any other integer (or non-integer) width raises this TypeError.

Source

Thrown at jax/_src/random/philox2x32.py:190

  assert not data.shape
  return _philox2x32_fold_in(key, jnp.asarray(data, dtype="uint32"))


@api.jit
def _philox2x32_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  """Internal implementation of philox2x32_fold_in."""
  out0, _ = philox2x32_p.bind(key[0], np.uint32(0), data)
  return jnp.array([out0], dtype=np.uint32)


def philox2x32_random_bits(
    key: typing.Array, bit_width: int, shape: tuple[int, ...]
) -> typing.Array:
  """Sample uniform random bits using a Philox 2x32 key."""
  if not _is_philox2x32_key(key):
    raise TypeError("philox2x32_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 _philox2x32_random_bits(key, bit_width, shape)


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

  counts1, counts2 = prng.iota_2x32_shape(shape)
  out0, out1 = philox2x32_p.bind(key[0], counts1, counts2)

  dtype = prng.UINT_DTYPES[bit_width]
  if bit_width == 64:
    bits_hi = lax.convert_element_type(out0, dtype)
    bits_lo = lax.convert_element_type(out1, dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use one of 8, 16, 32, or 64
  2. For narrower widths, generate 8/16/32 bits and mask: bits & ((1 << w) - 1)
  3. Prefer jax.random.bits(key, shape, dtype) with dtype of uint8/16/32/64

Example fix

# before
b = random.philox2x32_random_bits(key, 12, (1000,))
# after
b = random.philox2x32_random_bits(key, 16, (1000,)) & 0xFFF
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

When it happens

Trigger: Calling random.philox2x32_random_bits(key, 12, shape) or bit_width=128, or passing a width stored in a config as an arbitrary int.

Common situations: Trying to generate arbitrary-width random integers; porting code that assumed any bit count works; typos like 33 for 32.

Related errors


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