{"record":{"id":"7464a90b2a444f6e","repo":"jax-ml/jax","slug":"scan-got-values-with-different-leading-axis-sizes","errorCode":null,"errorMessage":"scan got values with different leading axis sizes: {}.","messagePattern":"scan got values with different leading axis sizes: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":522,"sourceCode":"  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:\n    sig = inspect.signature(body_fun)\n  except (ValueError, TypeError):\n    sig = None\n  carry_name = sig and list(sig.parameters)[0]\n  if carry_name:\n    component = lambda p: (f'the input carry component {carry_name}{keystr(p)}'","sourceCodeStart":504,"sourceCodeEnd":540,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L504-L540","documentation":"When length is not given, scan infers it from the leading axes of xs and requires all leaves to agree. Mixed leading axis sizes (e.g. xs=(a of shape (10,...), b of shape (12,...))) raise this ValueError.","triggerScenarios":"Passing a pytree of xs whose arrays have different first-dimension sizes; accidentally including a broadcastable array of different length.","commonSituations":"Tuple of sequences misaligned in preprocessing; concatenating features along the wrong axis so lengths drift; including a static array in xs that has another size.","solutions":["Align all xs leading axes before scan (truncate/pad to common length)","Move non-conforming leaves out of xs into constants closed over by f","Pass explicit length and slice every leaf to it"],"exampleFix":"# before\nxs = (a, b)  # a.shape==(10,..), b.shape==(12,..)\nlax.scan(body, init, xs)\n# after\nxs = (a[:10], b[:10])\nlax.scan(body, init, xs)","handlingStrategy":"validation","validationCode":"sizes = {x.shape[0] for x in jax.tree_util.tree_leaves(xs)}\nassert len(sizes) <= 1, f'inconsistent leading axes: {sizes}'","typeGuard":"def consistent_leading_axes(xs) -> bool:\n    return len({x.shape[0] for x in jax.tree_util.tree_leaves(xs)}) <= 1","tryCatchPattern":"null","preventionTips":["Slice/pad all xs leaves to a common length in preprocessing","Keep per-step data in one aligned struct","Assert alignment after any concatenation or resampling"],"tags":["jax","scan","shape-mismatch","length-inference"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}