jax-ml/jax · error · ValueError

Scan with {num_extensive_outputs} extensive output(s) is not

Error message

Scan with {num_extensive_outputs} extensive output(s) is not supported.

What it means

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.

Source

Thrown at jax/_src/pallas/utils.py:117

    size = size // s
    strides.append(int(size))
  return tuple(strides)


def next_power_of_2(x: int) -> int:
  """Returns the next power of two greater than or equal to `x`."""
  if x < 0:
    raise ValueError("`next_power_of_2` requires a non-negative integer.")
  return 1 if x == 0 else 2 ** (x - 1).bit_length()


def pattern_match_scan_to_fori_loop(
    jaxpr: jax_core.Jaxpr, num_consts: int, num_carry: int
) -> tuple[jax_core.Jaxpr, bool]:
  num_extensive_inputs = len(jaxpr.invars) - num_consts - num_carry
  num_extensive_outputs = len(jaxpr.outvars) - num_carry
  if num_extensive_outputs:
    raise ValueError(
        f"Scan with {num_extensive_outputs} extensive output(s) is not"
        " supported."
    )
  if num_extensive_inputs:
    raise ValueError(
        f"Scan with {num_extensive_inputs} extensive argument(s) is not"
        f" supported. Found {num_consts} consts and {num_carry} carry"
        " arguments."
    )
  if num_carry > 0:
    # Pattern match onto fori_loop:
    # We expect the first carry argument to the jaxpr to be the loop index and
    # for the loop index + 1 to be returned as the first value out of the loop.
    in_index_var = jaxpr.invars[num_consts]
    out_index_var = jaxpr.outvars[0]
    assert isinstance(in_index_var.aval, jax_core.ShapedArray)
    # Check that the loop index argument is an int32 scalar
    if (in_index_var.aval.shape or

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure the scan so every output is a carry (fold results into the carry tuple)
2. Write per-iteration results into an output Ref inside the scan body instead of returning them
3. Replace scan with an explicit fori_loop/while_loop over Refs

Example fix

# before
def body(c, _):
  return c, c * 2  # second output is extensive
# after
def body(c, _):
  return (c * 2,), None  # everything is carry

# ensure scan returns only carry: len(outvars) == num_carry
def body(carry, _):
    ...
    return (new_carry,), None  # all outputs are carry

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Scan with {num_extensive_inputs} extensive argument(s) is no
• Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• end must be greater than start, but got: {end} <= {start}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/220f3b2615e767ec. Report an issue: GitHub.