{"record":{"id":"b8677f4c523887da","repo":"jax-ml/jax","slug":"scan-got-length-argument-of-which-disagrees-w","errorCode":null,"errorMessage":"scan got `length` argument of {} which disagrees with leading axis sizes {}.","messagePattern":"scan got `length` argument of (.+?) which disagrees with leading axis sizes (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":516,"sourceCode":"\n  xs_shaped_avals = lax_utils.ensure_shaped(*xs_avals)\n  if not all(a.sharding.spec.partitions[0] is None for a in xs_shaped_avals):\n    raise ValueError('0th dimension of all xs should be replicated. Got '\n                     f'{\", \".join(str(a.sharding.spec) for a in xs_shaped_avals)}')\n\n  if length is not None:\n    try:\n      length = int(length)\n    except core.ConcretizationTypeError:\n      msg = ('The `length` argument to `scan` expects a concrete `int` value.'\n             ' For scan-like iteration with a dynamic length, use `while_loop`'\n             ' or `fori_loop`.')\n      raise core.ConcretizationTypeError(length, msg) from None\n    else:\n      if not all(length == l for l in lengths):\n        msg = (\"scan got `length` argument of {} which disagrees with \"\n              \"leading axis sizes {}.\")\n        raise ValueError(msg.format(length, [x.shape[0] for x in xs_flat]))\n      return length\n  else:\n    unique_lengths = set(lengths)\n    if len(unique_lengths) > 1:\n      msg = \"scan got values with different leading axis sizes: {}.\"\n      raise ValueError(msg.format(', '.join(str(x.shape[0]) for x in xs_flat)))\n    elif len(unique_lengths) == 0:\n      msg = \"scan got no values to scan over and `length` not provided.\"\n      raise ValueError(msg)\n    else:\n      return list(unique_lengths)[0]\n\ndef _capitalize(s):\n  # s.capitalize() converts s[1:] to lowercase which we don't want.\n  return s[0].capitalize() + s[1:]\n\ndef _check_carry_type(name, body_fun, in_carry, out_carry):\n  try:","sourceCodeStart":498,"sourceCodeEnd":534,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L498-L534","documentation":"scan cross-checks the explicit length argument against the leading-axis size of every xs element. If they disagree (e.g. length=10 but xs.shape[0]==8), it raises ValueError showing both values.","triggerScenarios":"lax.scan(f, init, xs, length=n) where n != xs.shape[0] for any xs leaf; or xs of differing leading sizes with an explicit length matching only some.","commonSituations":"Off-by-one loop counts; slicing xs after computing length; passing a stale length constant after data preprocessing changed the batch size.","solutions":["Pass length=None and let scan infer from xs","Recompute length from data: length=xs.shape[0]","Fix the off-by-one in how the length is derived (range(len(data)) vs len(data)+1)"],"exampleFix":"// before\nlax.scan(body, init, xs, length=num_steps)  # xs has len num_steps+1\n// after\nlax.scan(body, init, xs[:num_steps], length=num_steps)\n# or simply\nlax.scan(body, init, xs)","handlingStrategy":"validation","validationCode":"flat = jax.tree_util.tree_leaves(xs)\nif length is not None:\n    assert all(x.shape[0] == length for x in flat), \\\n        [(x.shape[0], length) for x in flat]","typeGuard":"def length_matches_xs(xs, length) -> bool:\n    return length is None or all(x.shape[0] == length for x in jax.tree_util.tree_leaves(xs))","tryCatchPattern":"null","preventionTips":["Derive length from the data: length=xs.shape[0]","Omit length and let scan infer","Add a unit assert before scan in data-loading code"],"tags":["jax","scan","length-mismatch","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}