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

_philox2x32_random_bits refuses to generate random bits when the requested shape's total element count (math.prod of constant dims) exceeds 2**64, because the internal 64-bit iota-based counter would wrap and produce correlated/duplicated values.

Source

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

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)
    return lax.shift_left(bits_hi, jnp.asarray(32, dtype=dtype)) | bits_lo
  elif bit_width == 32:
    return out0 ^ out1
  else:
    return lax.convert_element_type(out0 ^ out1, dtype)


# -- PRNGImpl registration --

philox2x32_prng_impl = prng.PRNGImpl(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the generation into chunks over multiple keys via jax.random.split and concatenate
2. Fix the shape computation (verify with math.prod(shape) before calling)
3. Generate lazily/streamed with scan over batches instead of one giant allocation

Example fix

# before
bits = random.philox2x32_random_bits(key, 32, (2**33, 2**33))
# after
keys = jax.random.split(key, 8)
chunks = [random.philox2x32_random_bits(k, 32, (2**33, 2**30)) for k in keys]
bits = jnp.concatenate(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 shape whose product exceeds 2**64, e.g. (2**33, 2**33), from a single call with static dimensions.

Common situations: Accidental huge shapes from mis-multiplied batch*sequence dims or a shape computed from floats (e.g. int(1e19)); symbolically-shaped (non-constant) dims skip the check so tracing code may only fail at runtime 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/d10f7a408ec89d43. Report an issue: GitHub.