{"record":{"id":"7c6ba2d8d354fad3","repo":"jax-ml/jax","slug":"scan-number-of-arguments-doesn-t-match-the-number","errorCode":null,"errorMessage":"scan number of arguments doesn't match the number of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}","messagePattern":"scan number of arguments doesn't match the number of jaxpr arguments: (.+?) vs (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":735,"sourceCode":"      partitions=(*length_spec, *aval.sharding.spec.partitions)))\n  # TODO(yashkatariya): Replace `lax.empty2` with `lax.empty` once\n  # AllocateBuffer issues are fixed. Also delete `empty2` after this usage is\n  # removed. Basically uncomment the following 2 lines.\n  # lax.empty will also need to take a memory_space argument.\n  # empty = lax.empty((*prefix, *aval.shape), aval.dtype, out_sharding=sharding,\n  #                   memory_space=aval.memory_space)\n  # return core.pvary(empty, tuple(aval.mat.varying))\n  empty = core.pvary(lax.empty2(aval.dtype, memory_space=aval.memory_space),\n                     tuple(aval.mat.varying))\n  with use_abstract_mesh(sharding.mesh):\n    out = lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)\n  return out\n\n\ndef _scan_abstract_eval(*args, reverse, length, ft_in, ft_out, jaxpr,\n                        unroll):\n  if len(args) != len(jaxpr.in_avals):\n    raise ValueError(\"scan number of arguments doesn't match the number \"\n                     \"of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}\")\n  out_carry_avals, y_avals = ft_out.update(jaxpr.out_avals).unpack()\n  _, in_carry_avals, _ = ft_in.update(args).unpack()\n  if ([i.mat for i in in_carry_avals if isinstance(i, core.ShapedArray)] !=\n      [o.mat for o in out_carry_avals if isinstance(o, core.ShapedArray)]):\n    raise ValueError(\n        'Scan carry input and output got mismatched varying manual axes '\n        f'{in_carry_avals} and {out_carry_avals}. Please open an '\n        'issue at https://github.com/jax-ml/jax/issues, and as a '\n        'temporary workaround pass the check_vma=False argument to '\n        '`jax.shard_map`')\n  ys_avals = _map(partial(core.unmapped_leading_aval, length), y_avals)\n  return list(out_carry_avals) + list(ys_avals), core.positional_effects(jaxpr)\n\ndef _scan_jvp(primals, tangents, reverse, length, jaxpr, ft_in, ft_out, unroll):\n  nonzeros = [type(t) is not ad_util.Zero for t in tangents]\n  const_nz, init_nz, xs_nz = ft_in.update(nonzeros).unpack()\n","sourceCodeStart":717,"sourceCodeEnd":753,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L717-L753","documentation":"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.","triggerScenarios":"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.","commonSituations":"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).","solutions":["If using public jax.lax.scan, update/upgrade JAX — this indicates an internal inconsistency or a transformation bug worth reporting","If binding scan_p manually, verify that the flat args list length equals len(jaxpr.in_avals) plus consts handling before bind","Audit custom interpreters/transformations that reconstruct scan equations to forward the original jaxpr and arg count together"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(flat_args) == len(jaxpr.in_avals), (len(flat_args), len(jaxpr.in_avals))  # before scan_p.bind","typeGuard":null,"tryCatchPattern":"try:\n    out = scan_p.bind(*args, **params)\nexcept ValueError as e:\n    if 'number of arguments' in str(e):\n        # re-derive jaxpr from the same args\n        raise","preventionTips":["Never bind scan_p directly; always go through jax.lax.scan","Keep args and jaxpr produced by the same trace when writing custom transforms","Pin a known-good JAX version if writing low-level jaxpr code"],"tags":["jax","internal","scan","jaxpr","argument-count"],"backgroundTag":"jax-internal-invariant-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}