jax-ml/jax · error · NotImplementedError

random bits array of size exceeding 2 ** 64

Error message

random bits array of size exceeding 2 ** 64

What it means

_philox4x32_random_bits rejects shapes whose total element count exceeds 2**64 (when all dims are static), since the counter iotas would overflow 64 bits and outputs would repeat/correlate.

Source

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

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
  )

  dtype = prng.UINT_DTYPES[bit_width]
  if bit_width == 64:
    # Combine two 32-bit outputs into one 64-bit value.
    bits_hi = lax.convert_element_type(out0, dtype)
    bits_lo = lax.convert_element_type(out1, dtype)
    return lax.shift_left(bits_hi, jnp.asarray(32, dtype=dtype)) | bits_lo
  elif bit_width == 32:
    # XOR all four outputs for maximum mixing.
    return out0 ^ out1 ^ out2 ^ out3

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify math.prod(shape) <= 2**64 in a precondition
2. Chunk generation across split keys and concatenate
3. Fix the upstream shape computation bug

Example fix

# before
bits = random.philox4x32_random_bits(key, 64, huge_shape)
# after
keys = jax.random.split(key, num_chunks)
bits = jnp.concatenate([random.philox4x32_random_bits(k, 64, chunk) for k, chunk in zip(keys, chunks)])

import math
assert math.prod(shape) <= 2**64, 'requested random bits exceed 2**64 elements'

def size_within_limit(shape) -> bool:
    import math
    return math.prod(shape) <= 2**64

When it happens

Trigger: Requesting a static shape with math.prod(shape) > 2**64 in one call, e.g. from miscomputed batch*dim products.

Common situations: Shape arithmetic errors (int overflow of intended dims), accidental shapes from broadcasting bugs; dynamic dims bypass the check so failures may appear only with concrete shapes.

Related errors

• random bits array of size exceeding 2 ** 64
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.
• PRNG key seed must be a scalar; got {seed!r}.
• philox2x32_random_bits got invalid prng key.

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