{"record":{"id":"10fd12f0f7d7026b","repo":"jax-ml/jax","slug":"shape-too-large-np-prod-shape-np-iinfo-jnp-10fd12","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/threefry.py","lineNumber":53,"sourceCode":"\n  This function is a fusion of prng.shaped_iota and prng.threefry_2x32 from\n  the JAX core library.\n\n  Args:\n    key: A threefry 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\n  Returns:\n    A tensor of random bits of shape `shape`.\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\n  def kernel(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)\n    counts_lo = counts_lo.astype(jnp.uint32)\n    # TODO(justinfu): Support hi bits on count.\n    counts_hi = jnp.zeros_like(counts_lo)\n    k1 = jnp.reshape(key_ref[0, 0], (1, 1))","sourceCodeStart":35,"sourceCodeEnd":71,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/random/threefry.py#L35-L71","documentation":"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.","triggerScenarios":"Calling threefry_2x32_count / plthreefry_random_bits for a buffer with more than ~4.29 billion elements, e.g. (2**23, 1024).","commonSituations":"One-shot generation of very large random tensors for simulation or initialization on TPU instead of chunked generation.","solutions":["Chunk the request into pieces under 2**32 elements, tracking counter offsets","Generate in a loop over the leading axis and concatenate results","Use jax.random with sharding to distribute generation across devices"],"exampleFix":"// before\nbits = plthreefry_random_bits(key, 32, (2**23, 1024))\n// after\nbits = jnp.concatenate([plthreefry_random_bits(fold_in(key, i), 32, (chunk, 1024)) for i, chunk in enumerate(chunks)])","handlingStrategy":"validation","validationCode":"assert int(np.prod(shape)) <= np.iinfo(np.uint32).max, 'chunk threefry generation below 2**32 elements'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Chunk big random tensors and use fold_in/key-splitting per chunk","Consider sharded generation across TPU devices"],"tags":["jax","pallas","tpu","random","threefry","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"}