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 Philox counter-based PRNG Pallas kernel on TPU addresses output elements with uint32 offsets, so the total number of elements (np.prod(shape)) must not exceed 2**32-1. Larger requests are rejected up front to avoid silent wraparound/correlated random numbers.

Source

Thrown at jax/experimental/pallas/ops/tpu/random/philox.py:103

  Args:
    key: A Philox 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.
    offset: An optional offset to the counts.
    fuse_output: Whether to fuse the output bits into a single value.

  Returns:
    A tensor of random bits of shape `shape` if fuse_output=True. Otherwise,
    this will return a tensor of shape (2, *shape) with the first channel being
    the high bits and the second channel being the low bits.
  """
  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],)
  offset = jnp.array(offset, dtype=jnp.uint32)
  if offset.ndim != 0:
    raise ValueError(f"Offset must be scalar, got {offset.shape}")
  offset = jnp.reshape(offset, (1,))

  def kernel(offset_ref, 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) + offset_ref[0]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Split the generation into chunks each under 2**32 elements and use the offset parameter to keep streams distinct
Generate along a leading axis in a loop and concatenate
Reconsider whether the full buffer must be materialized at once (use lazy/streamed generation)

Example fix

// before
bits = philox_random_bits(key, 32, (2**24, 512))  # > uint32 max
// after
outs = [philox_4x32_count(key, (chunk, 512), offset=i*chunk*512) for i, chunk in enumerate(chunks)]
bits = jnp.concatenate(outs)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np, jax.numpy as jnp
prod = int(np.prod(shape))
assert prod <= np.iinfo(np.uint32).max, f'{prod} elements too large; chunk the request'

Prevention

Chunk large random generation into pieces < 2**32 elements
Use offsets to keep chunked streams non-overlapping

When it happens

Trigger: Calling philox_4x32_count / philox_random_bits with a shape whose product exceeds 4294967295, e.g. (2**24, 512) on a very large buffer.

Common situations: Generating huge random bit buffers in one call for simulations or weight initialization; forgetting to chunk large random generation requests.

Related errors

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