{"record":{"id":"220f3b2615e767ec","repo":"jax-ml/jax","slug":"scan-with-num-extensive-outputs-extensive-output","errorCode":null,"errorMessage":"Scan with {num_extensive_outputs} extensive output(s) is not supported.","messagePattern":"Scan with (.+?) extensive output\\(s\\) is not supported\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/utils.py","lineNumber":117,"sourceCode":"    size = size // s\n    strides.append(int(size))\n  return tuple(strides)\n\n\ndef 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","sourceCodeStart":99,"sourceCodeEnd":135,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/utils.py#L99-L135","documentation":"When Pallas lowers a lax.scan to a Triton loop, the scanned jaxpr's outputs must all be carry values; outputs that are newly-created (extensive) per iteration are not supported because the loop has no way to accumulate them. num_extensive_outputs = len(outvars) - num_carry must be zero.","triggerScenarios":"Writing a pallas kernel whose body uses lax.scan where the scan returns values not derived from the carry, e.g. return carry, i*2 as second output; or yielding new arrays from the scanned function.","commonSituations":"Refactoring loop code into scan with extra outputs; using scan to emit per-iteration results instead of accumulating into carry or writing to a Ref.","solutions":["Restructure the scan so every output is a carry (fold results into the carry tuple)","Write per-iteration results into an output Ref inside the scan body instead of returning them","Replace scan with an explicit fori_loop/while_loop over Refs"],"exampleFix":"# before\ndef body(c, _):\n  return c, c * 2  # second output is extensive\n# after\ndef body(c, _):\n  return (c * 2,), None  # everything is carry","handlingStrategy":"validation","validationCode":"# ensure scan returns only carry: len(outvars) == num_carry\ndef body(carry, _):\n    ...\n    return (new_carry,), None  # all outputs are carry","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Write per-iteration outputs into Refs, not scan outputs","Every output of the scanned function must be part of the 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"}