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

threefry_random_bits only produces unsigned integer output of width 8, 16, 32, or 64 bits, matching the UINT_DTYPES the kernels can emit. This TypeError fires when bit_width is outside that set — e.g. requesting float widths or other integer sizes through the internal API (the public equivalent is jax.random.bits with a bad dtype).

Source

Thrown at jax/_src/random/threefry2x32.py:326

  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)
  return jnp.stack([bits1, bits2], axis=bits1.ndim)


def threefry_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  assert not data.shape
  return _threefry_fold_in(key, jnp.asarray(data, dtype='uint32'))

@api.jit
def _threefry_fold_in(key, data):
  return threefry_2x32(key, threefry_seed(data))


def threefry_random_bits(key: typing.Array, bit_width, shape):
  """Sample uniform random bits of given width and shape using PRNG key."""
  if not _is_threefry_prng_key(key):
    raise TypeError("threefry_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.")

  if config.threefry_partitionable.value:
    return _threefry_random_bits_partitionable(key, bit_width, shape)
  else:
    return _threefry_random_bits_original(key, bit_width, shape)

def _threefry_random_bits_partitionable(key: typing.Array, bit_width, shape):
  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')

  k1, k2 = key
  counts1, counts2 = prng.iota_2x32_shape(shape)
  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use uint8/uint16/uint32/uint64 only, then convert: bits.astype(jnp.float32)/2**32 for floats
2. Use jax.random.uniform/normal for real-valued randomness
3. Validate bit_width against (8,16,32,64) in dynamic pipelines

Example fix

// before
u = jax.random.bits(key, shape=(4,), dtype=jnp.float32)

// after
u = jax.random.bits(key, shape=(4,), dtype=jnp.uint32).astype(jnp.float32) / 2**32

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

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

When it happens

Trigger: Calling threefry_random_bits with bit_width in {4, 128} or derived from a float dtype; jax.random.bits(..., dtype=jnp.float32); dynamic width computation landing on an unsupported value.

Common situations: Users expecting float uniform values from a bits API (should use jax.random.uniform); dtype-driven generic code passing arbitrary dtypes; porting NumPy randint-style code with unusual widths.

Related errors

• requires 8-, 16-, 32- or 64-bit field width.
• threefry_2x32 requires uint32 arguments, got {}
• 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}

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