{"record":{"id":"c5362d0b6b4cd34e","repo":"jax-ml/jax","slug":"random-bits-array-of-size-exceeding-2-64-c5362d","errorCode":null,"errorMessage":"random bits array of size exceeding 2 ** 64","messagePattern":"random bits array of size exceeding 2 \\*\\* 64","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/philox4x32.py","lineNumber":228,"sourceCode":"\ndef philox4x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Sample uniform random bits using a Philox 4x32 key.\"\"\"\n  if not _is_philox4x32_key(key):\n    raise TypeError(\"philox4x32_random_bits got invalid prng key.\")\n  if bit_width not in (8, 16, 32, 64):\n    raise TypeError(\"requires 8-, 16-, 32- or 64-bit field width.\")\n  return _philox4x32_random_bits(key, bit_width, shape)\n\n\n@api.jit(static_argnums=(1, 2), inline=True)\ndef _philox4x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Internal implementation of philox4x32_random_bits.\"\"\"\n  if all(core.is_constant_dim(d) for d in shape) and math.prod(shape) > 2**64:\n    raise NotImplementedError(\"random bits array of size exceeding 2 ** 64\")\n\n  k0, k1 = key[0], key[1]\n  counts1, counts2 = prng.iota_2x32_shape(shape)\n  zeros = jnp.zeros(shape, dtype=np.uint32)\n\n  out0, out1, out2, out3 = philox4x32_p.bind(\n      k0, k1, counts1, counts2, zeros, zeros\n  )\n\n  dtype = prng.UINT_DTYPES[bit_width]\n  if bit_width == 64:\n    # Combine two 32-bit outputs into one 64-bit value.\n    bits_hi = lax.convert_element_type(out0, dtype)\n    bits_lo = lax.convert_element_type(out1, dtype)\n    return lax.shift_left(bits_hi, jnp.asarray(32, dtype=dtype)) | bits_lo\n  elif bit_width == 32:\n    # XOR all four outputs for maximum mixing.\n    return out0 ^ out1 ^ out2 ^ out3","sourceCodeStart":210,"sourceCodeEnd":246,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/philox4x32.py#L210-L246","documentation":"_philox4x32_random_bits rejects shapes whose total element count exceeds 2**64 (when all dims are static), since the counter iotas would overflow 64 bits and outputs would repeat/correlate.","triggerScenarios":"Requesting a static shape with math.prod(shape) > 2**64 in one call, e.g. from miscomputed batch*dim products.","commonSituations":"Shape arithmetic errors (int overflow of intended dims), accidental shapes from broadcasting bugs; dynamic dims bypass the check so failures may appear only with concrete shapes.","solutions":["Verify math.prod(shape) <= 2**64 in a precondition","Chunk generation across split keys and concatenate","Fix the upstream shape computation bug"],"exampleFix":"# before\nbits = random.philox4x32_random_bits(key, 64, huge_shape)\n# after\nkeys = jax.random.split(key, num_chunks)\nbits = jnp.concatenate([random.philox4x32_random_bits(k, 64, chunk) for k, chunk in zip(keys, chunks)])","handlingStrategy":"validation","validationCode":"import math\nassert math.prod(shape) <= 2**64, 'requested random bits exceed 2**64 elements'","typeGuard":"def size_within_limit(shape) -> bool:\n    import math\n    return math.prod(shape) <= 2**64","tryCatchPattern":null,"preventionTips":["Unit-test shape arithmetic to catch broadcast bugs","Stream large draws with scan instead of one call"],"tags":["jax","prng","philox","shape-validation","resource-limits"],"backgroundTag":"output-size-limit-exceeded","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}