{"record":{"id":"d10f7a408ec89d43","repo":"jax-ml/jax","slug":"random-bits-array-of-size-exceeding-2-64","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/philox2x32.py","lineNumber":200,"sourceCode":"\ndef philox2x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Sample uniform random bits using a Philox 2x32 key.\"\"\"\n  if not _is_philox2x32_key(key):\n    raise TypeError(\"philox2x32_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 _philox2x32_random_bits(key, bit_width, shape)\n\n\n@api.jit(static_argnums=(1, 2), inline=True)\ndef _philox2x32_random_bits(\n    key: typing.Array, bit_width: int, shape: tuple[int, ...]\n) -> typing.Array:\n  \"\"\"Internal implementation of philox2x32_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  counts1, counts2 = prng.iota_2x32_shape(shape)\n  out0, out1 = philox2x32_p.bind(key[0], counts1, counts2)\n\n  dtype = prng.UINT_DTYPES[bit_width]\n  if bit_width == 64:\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    return out0 ^ out1\n  else:\n    return lax.convert_element_type(out0 ^ out1, dtype)\n\n\n# -- PRNGImpl registration --\n\nphilox2x32_prng_impl = prng.PRNGImpl(","sourceCodeStart":182,"sourceCodeEnd":218,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/philox2x32.py#L182-L218","documentation":"_philox2x32_random_bits refuses to generate random bits when the requested shape's total element count (math.prod of constant dims) exceeds 2**64, because the internal 64-bit iota-based counter would wrap and produce correlated/duplicated values.","triggerScenarios":"Requesting a shape whose product exceeds 2**64, e.g. (2**33, 2**33), from a single call with static dimensions.","commonSituations":"Accidental huge shapes from mis-multiplied batch*sequence dims or a shape computed from floats (e.g. int(1e19)); symbolically-shaped (non-constant) dims skip the check so tracing code may only fail at runtime with concrete shapes.","solutions":["Split the generation into chunks over multiple keys via jax.random.split and concatenate","Fix the shape computation (verify with math.prod(shape) before calling)","Generate lazily/streamed with scan over batches instead of one giant allocation"],"exampleFix":"# before\nbits = random.philox2x32_random_bits(key, 32, (2**33, 2**33))\n# after\nkeys = jax.random.split(key, 8)\nchunks = [random.philox2x32_random_bits(k, 32, (2**33, 2**30)) for k in keys]\nbits = jnp.concatenate(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":["Chunk large generations across split keys","Sanity-check math.prod(shape) in debug builds"],"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"}