{"record":{"id":"0269733db88031ab","repo":"jax-ml/jax","slug":"shape-too-large-np-prod-shape-np-iinfo-jnp","errorCode":null,"errorMessage":"Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}","messagePattern":"Shape too large: (.+?) > (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/random/philox.py","lineNumber":103,"sourceCode":"  Args:\n    key: A Philox key of shape (2,).\n    shape: The shape of the output. Must be divisible by `block_size`.\n    unpadded_shape: If `shape` is padded, then this is the shape of the\n      output tensor if it were not padded. This is important for indexing\n      calculations within the kernel. If `shape` is not padded, then this\n      should be equal to `shape`.\n    block_size: The block size of the kernel.\n    offset: An optional offset to the counts.\n    fuse_output: Whether to fuse the output bits into a single value.\n\n  Returns:\n    A tensor of random bits of shape `shape` if fuse_output=True. Otherwise,\n    this will return a tensor of shape (2, *shape) with the first channel being\n    the high bits and the second channel being the low bits.\n  \"\"\"\n  shape = tuple(shape)\n  if np.prod(shape) > jnp.iinfo(jnp.uint32).max:\n    raise ValueError(\n        f\"Shape too large: {np.prod(shape)} > {np.iinfo(jnp.uint32).max}\")\n\n  if (shape[-2] % block_size[-2] != 0) or (shape[-1] % block_size[-1] != 0):\n    raise ValueError(\n        f\"Shape dimension {shape[-2:]} must be divisible by {block_size}\")\n  grid_dims = shape[:-2] + (\n      shape[-2] // block_size[-2], shape[-1] // block_size[1],)\n  offset = jnp.array(offset, dtype=jnp.uint32)\n  if offset.ndim != 0:\n    raise ValueError(f\"Offset must be scalar, got {offset.shape}\")\n  offset = jnp.reshape(offset, (1,))\n\n  def kernel(offset_ref, key_ref, out_ref):\n    counts_idx = tuple(pl.program_id(i) for i in range(len(grid_dims)))\n    offset = prng_utils.compute_scalar_offset(\n        counts_idx, unpadded_shape, block_shape)\n    counts_lo = prng_utils.blocked_iota(block_size, unpadded_shape)\n    counts_lo = counts_lo + offset.astype(jnp.uint32) + offset_ref[0]","sourceCodeStart":85,"sourceCodeEnd":121,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/philox.py#L85-L121","documentation":"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.","triggerScenarios":"Calling philox_4x32_count / philox_random_bits with a shape whose product exceeds 4294967295, e.g. (2**24, 512) on a very large buffer.","commonSituations":"Generating huge random bit buffers in one call for simulations or weight initialization; forgetting to chunk large random generation requests.","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)"],"exampleFix":"// before\nbits = philox_random_bits(key, 32, (2**24, 512))  # > uint32 max\n// after\nouts = [philox_4x32_count(key, (chunk, 512), offset=i*chunk*512) for i, chunk in enumerate(chunks)]\nbits = jnp.concatenate(outs)","handlingStrategy":"validation","validationCode":"import numpy as np, jax.numpy as jnp\nprod = int(np.prod(shape))\nassert prod <= np.iinfo(np.uint32).max, f'{prod} elements too large; chunk the request'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Chunk large random generation into pieces < 2**32 elements","Use offsets to keep chunked streams non-overlapping"],"tags":["jax","pallas","tpu","random","philox","shape-limit"],"backgroundTag":"tensor-shape-too-large","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}