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

Random-bit generation in JAX is only defined for unsigned integer widths 8, 16, 32, and 64, because the underlying kernels emit chunks of 32 or 64 random bits and split them evenly. _rbg_random_bits raises this TypeError when bit_width is anything else (e.g. bool, float widths, or arbitrary integers).

Source

Thrown at jax/_src/random/rbg.py:63

  if config.threefry_partitionable.value:
    _threefry_split = threefry2x32._threefry_split_foldlike
  else:
    _threefry_split = threefry2x32._threefry_split_original
  halfkeys = key.reshape(2, 2)
  return api.vmap(
      _threefry_split, (0, None), len(shape))(halfkeys, shape).reshape(
          *shape, 4)

def _rbg_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  assert not data.shape
  return api.vmap(threefry2x32._threefry_fold_in, (0, None), 0)(key.reshape(2, 2), data).reshape(4)

def _rbg_random_bits(key: typing.Array, bit_width: int, shape: Sequence[int]
                     ) -> typing.Array:
  if not key.shape == (4,) and key.dtype == np.dtype('uint32'):
    raise TypeError("_rbg_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.")
  _, bits = lax.rng_bit_generator(key, shape, dtype=prng.UINT_DTYPES[bit_width])
  return bits

rbg_prng_impl = prng.PRNGImpl(
    key_shape=(4,),
    seed=_rbg_seed,
    split=_rbg_split,
    random_bits=_rbg_random_bits,
    fold_in=_rbg_fold_in,
    name='rbg',
    tag='rbg')

prng.register_prng(rbg_prng_impl)


def _unsafe_rbg_split(key: typing.Array, shape: prng.Shape) -> typing.Array:
  # treat 10 iterations of random bits as a 'hash function'
  num = math.prod(shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use one of uint8/uint16/uint32/uint64 with jax.random.bits
2. For booleans, generate uint bits and compare: jax.random.bits(key, shape, dtype=jnp.uint8, impl=...) % 2 == 0 or use jax.random.bernoulli
3. Validate dtype before calling in dynamic pipelines

Example fix

// before
mask = jax.random.bits(key, shape=(4,), dtype=jnp.bool_, impl='rbg')

// after
mask = jax.random.bits(key, shape=(4,), dtype=jnp.uint8, impl='rbg') < 128

VALID = {8, 16, 32, 64}
assert bit_width in VALID, f'bit_width must be one of {VALID}'

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

When it happens

Trigger: Calling jax.random.bits with dtype=bool (bit_width interpreted oddly) or with a float dtype; calling lax.rng_bit_generator via rbg with an unsupported dtype; passing a computed width like 4 or 128.

Common situations: Trying to generate random booleans with jax.random.bits instead of comparisons (e.g. jax.random.uniform(...) < 0.5); dynamic dtype selection that yields non-uint kinds; assuming arbitrary bit widths are supported.

Related errors

• requires 8-, 16-, 32- or 64-bit field width.
• JAX encountered invalid PRNG key data: expected key_data.dty
• random_unwrap takes key array operand, got {keys.dtype=}
• PRNG key seed must be an integer; got {seed!r}
• threefry_2x32 requires uint32 arguments, got {}

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