jax-ml/jax · error · ValueError

scan number of arguments doesn't match the number of jaxpr a

Error message

scan number of arguments doesn't match the number of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}

What it means

This internal abstract-evaluation check fires when the number of flat arguments presented to the scan primitive does not equal the number of input avars in its jaxpr. It should never be reachable from normal user code — scan's Python wrapper flattens arguments itself — and it also contains a formatting bug: the message lacks an f-prefix, so it literally prints '{len(args)} vs {len(jaxpr.in_avals)}' without interpolating values.

Source

Thrown at jax/_src/lax/control_flow/loops.py:735

      partitions=(*length_spec, *aval.sharding.spec.partitions)))
  # TODO(yashkatariya): Replace `lax.empty2` with `lax.empty` once
  # AllocateBuffer issues are fixed. Also delete `empty2` after this usage is
  # removed. Basically uncomment the following 2 lines.
  # lax.empty will also need to take a memory_space argument.
  # empty = lax.empty((*prefix, *aval.shape), aval.dtype, out_sharding=sharding,
  #                   memory_space=aval.memory_space)
  # return core.pvary(empty, tuple(aval.mat.varying))
  empty = core.pvary(lax.empty2(aval.dtype, memory_space=aval.memory_space),
                     tuple(aval.mat.varying))
  with use_abstract_mesh(sharding.mesh):
    out = lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)
  return out


def _scan_abstract_eval(*args, reverse, length, ft_in, ft_out, jaxpr,
                        unroll):
  if len(args) != len(jaxpr.in_avals):
    raise ValueError("scan number of arguments doesn't match the number "
                     "of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}")
  out_carry_avals, y_avals = ft_out.update(jaxpr.out_avals).unpack()
  _, in_carry_avals, _ = ft_in.update(args).unpack()
  if ([i.mat for i in in_carry_avals if isinstance(i, core.ShapedArray)] !=
      [o.mat for o in out_carry_avals if isinstance(o, core.ShapedArray)]):
    raise ValueError(
        'Scan carry input and output got mismatched varying manual axes '
        f'{in_carry_avals} and {out_carry_avals}. Please open an '
        'issue at https://github.com/jax-ml/jax/issues, and as a '
        'temporary workaround pass the check_vma=False argument to '
        '`jax.shard_map`')
  ys_avals = _map(partial(core.unmapped_leading_aval, length), y_avals)
  return list(out_carry_avals) + list(ys_avals), core.positional_effects(jaxpr)

def _scan_jvp(primals, tangents, reverse, length, jaxpr, ft_in, ft_out, unroll):
  nonzeros = [type(t) is not ad_util.Zero for t in tangents]
  const_nz, init_nz, xs_nz = ft_in.update(nonzeros).unpack()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. If using public jax.lax.scan, update/upgrade JAX — this indicates an internal inconsistency or a transformation bug worth reporting
2. If binding scan_p manually, verify that the flat args list length equals len(jaxpr.in_avals) plus consts handling before bind
3. Audit custom interpreters/transformations that reconstruct scan equations to forward the original jaxpr and arg count together

assert len(flat_args) == len(jaxpr.in_avals), (len(flat_args), len(jaxpr.in_avals))  # before scan_p.bind

try:
    out = scan_p.bind(*args, **params)
except ValueError as e:
    if 'number of arguments' in str(e):
        # re-derive jaxpr from the same args
        raise

When it happens

Trigger: Calling the low-level scan primitive directly (jax.lax.scan_p.bind or constructing a JaxprEqn by hand) with mismatched arg counts; transformations (custom vmap/pmap/interpreters) that re-bind scan_p with a stale jaxpr; essentially never triggered via public jax.lax.scan.

Common situations: Writing a custom JAX primitive or interpreter that rebinds scan_p; programmatic jaxpr construction (jax.core) with a hand-built scan equation; version upgrades that change scan_p parameters (e.g. ft_in/ft_out addition).

Related errors

• {_vjp_too_many_args(name, len(extra_args) + 1)}
• `JaxprInputEffect` {eff} does not have a corresponding jaxpr
• Effects not supported in `scan`: {disallowed_effects}
• `unroll` must be a `bool` or a non-negative `int`.
• lax.scan: f argument should be a callable.

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