jax-ml/jax · error · ValueError

Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma

Error message

Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}

What it means

The TPU Pallas Threefry counter-based PRNG kernel addresses elements with uint32 counters, so threefry_2x32_count rejects shapes whose total element count (np.prod(shape)) exceeds 2**32-1, preventing counter wraparound.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/threefry.py:53

  This function is a fusion of prng.shaped_iota and prng.threefry_2x32 from
  the JAX core library.

  Args:
    key: A threefry key of shape (2,).
    shape: The shape of the output. Must be divisible by `block_size`.
    unpadded_shape: If `shape` is padded, then this is the shape of the
      output tensor if it were not padded. This is important for indexing
      calculations within the kernel. If `shape` is not padded, then this
      should be equal to `shape`.
    block_size: The block size of the kernel.

  Returns:
    A tensor of random bits of shape `shape`.
  """
  shape = tuple(shape)
  if np.prod(shape) > jnp.iinfo(jnp.uint32).max:
    raise ValueError(
        f"Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}")

  if (shape[-2] % block_size[-2] != 0) or (shape[-1] % block_size[-1] != 0):
    raise ValueError(
        f"Shape dimension {shape[-2:]} must be divisible by {block_size}")
  grid_dims = shape[:-2] + (
      shape[-2] // block_size[-2], shape[-1] // block_size[1],)

  def kernel(key_ref, out_ref):
    counts_idx = tuple(pl.program_id(i) for i in range(len(grid_dims)))
    offset = prng_utils.compute_scalar_offset(
        counts_idx, unpadded_shape, block_shape)
    counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)
    counts_lo = counts_lo + offset.astype(jnp.uint32)
    counts_lo = counts_lo.astype(jnp.uint32)
    # TODO(justinfu): Support hi bits on count.
    counts_hi = jnp.zeros_like(counts_lo)
    k1 = jnp.reshape(key_ref[0, 0], (1, 1))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Chunk the request into pieces under 2**32 elements, tracking counter offsets
2. Generate in a loop over the leading axis and concatenate results
3. Use jax.random with sharding to distribute generation across devices

Example fix

// before
bits = plthreefry_random_bits(key, 32, (2**23, 1024))
// after
bits = jnp.concatenate([plthreefry_random_bits(fold_in(key, i), 32, (chunk, 1024)) for i, chunk in enumerate(chunks)])

assert int(np.prod(shape)) <= np.iinfo(np.uint32).max, 'chunk threefry generation below 2**32 elements'

When it happens

Trigger: Calling threefry_2x32_count / plthreefry_random_bits for a buffer with more than ~4.29 billion elements, e.g. (2**23, 1024).

Common situations: One-shot generation of very large random tensors for simulation or initialization on TPU instead of chunked generation.

Related errors

• Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).ma
• Shape dimension {shape[-2:]} must be divisible by {block_siz
• Only 32-bit PRNG supported.
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Bit width must be 32

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