{"record":{"id":"ec2ffadcff686c5a","repo":"jax-ml/jax","slug":"scan-with-num-extensive-inputs-extensive-argumen","errorCode":null,"errorMessage":"Scan with {num_extensive_inputs} extensive argument(s) is not supported. Found {num_consts} consts and {num_carry} carry arguments.","messagePattern":"Scan with (.+?) extensive argument\\(s\\) is not supported\\. Found (.+?) consts and (.+?) carry arguments\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/utils.py","lineNumber":122,"sourceCode":"def next_power_of_2(x: int) -> int:\n  \"\"\"Returns the next power of two greater than or equal to `x`.\"\"\"\n  if x < 0:\n    raise ValueError(\"`next_power_of_2` requires a non-negative integer.\")\n  return 1 if x == 0 else 2 ** (x - 1).bit_length()\n\n\ndef pattern_match_scan_to_fori_loop(\n    jaxpr: jax_core.Jaxpr, num_consts: int, num_carry: int\n) -> tuple[jax_core.Jaxpr, bool]:\n  num_extensive_inputs = len(jaxpr.invars) - num_consts - num_carry\n  num_extensive_outputs = len(jaxpr.outvars) - num_carry\n  if num_extensive_outputs:\n    raise ValueError(\n        f\"Scan with {num_extensive_outputs} extensive output(s) is not\"\n        \" supported.\"\n    )\n  if num_extensive_inputs:\n    raise ValueError(\n        f\"Scan with {num_extensive_inputs} extensive argument(s) is not\"\n        f\" supported. Found {num_consts} consts and {num_carry} carry\"\n        \" arguments.\"\n    )\n  if num_carry > 0:\n    # Pattern match onto fori_loop:\n    # We expect the first carry argument to the jaxpr to be the loop index and\n    # for the loop index + 1 to be returned as the first value out of the loop.\n    in_index_var = jaxpr.invars[num_consts]\n    out_index_var = jaxpr.outvars[0]\n    assert isinstance(in_index_var.aval, jax_core.ShapedArray)\n    # Check that the loop index argument is an int32 scalar\n    if (in_index_var.aval.shape or\n        in_index_var.aval.dtype not in (jnp.int32, jnp.int64)):\n      # The loop index is not an int32 scalar so we assume that the loop index\n      # has been DCEd and the body does *not* expect a loop index as an\n      # argument.\n      return jaxpr, False","sourceCodeStart":104,"sourceCodeEnd":140,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/utils.py#L104-L140","documentation":"When pattern-matching a lax.scan to a fori_loop inside Pallas, the scanned jaxpr may only take consts and carry arguments. Extra (extensive) inputs — len(invars) - num_consts - num_carry > 0 — cannot be lowered to the Triton loop.","triggerScenarios":"A pallas kernel body calls lax.scan on a closure that captures or takes additional per-iteration array inputs beyond consts and carry, e.g. scan(f)(x_stacked) where the stacked input is sliced per iteration.","commonSituations":"Using scan over a leading axis of an input array instead of folding that data into the carry or consts; closures capturing block arrays inside scan.","solutions":["Move per-iteration data into the carry tuple and update it each iteration","Hoist the data out as a const (compute outside the kernel or before scan)","Convert the scan to an explicit Python/fori loop indexing into a Ref"],"exampleFix":"# before\n_, out = lax.scan(lambda c, xs: f(c, xs), c0, stacked)  # xs is extensive\n# after\nc, _ = lax.scan(lambda c, _: f(c, stacked_const), c0, None, length=n)","handlingStrategy":"validation","validationCode":"# all scan inputs must be consts or carry; fold per-iter data into carry\nx, _ = lax.scan(lambda c, _: step(c, const_data), c0, None, length=n)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't scan over stacked inputs in pallas kernels","Put static data in consts, dynamic state in carry"],"tags":["jax","pallas","triton","scan","control-flow"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}