{"record":{"id":"faa00e96f70716ae","repo":"jax-ml/jax","slug":"scan-got-value-with-no-leading-axis-to-scan-over","errorCode":null,"errorMessage":"scan got value with no leading axis to scan over: {}.","messagePattern":"scan got value with no leading axis to scan over: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":495,"sourceCode":"    length: Any | None) -> int:\n\n  # TODO(dougalm): put this in some sort of `scannable` typeclass\n  from jax._src.hijax import HiType\n  is_hi = [isinstance(a, HiType) for a in xs_avals]\n  if xs_flat and all(is_hi):\n    if length is None:\n      raise ValueError(\n          \"must provide `length` to `scan`, since the leading-axis size of \"\n          \"non-array (hijax) types cannot be inferred\")\n    return length\n  xs_flat = [x for x, h in zip(xs_flat, is_hi) if not h]\n  xs_avals = [a for a, h in zip(xs_avals, is_hi) if not h]\n\n  try:\n    lengths: list[int] = [x.shape[0] for x in xs_flat]\n  except AttributeError as err:\n    msg = \"scan got value with no leading axis to scan over: {}.\"\n    raise ValueError(\n      msg.format(', '.join(str(x) for x in xs_flat\n                           if not hasattr(x, 'shape')))) from err\n\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):","sourceCodeStart":477,"sourceCodeEnd":513,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L477-L513","documentation":"scan infers the loop length from x.shape[0] of each xs element. If an element has no .shape (a scalar, a Python number, or a non-array object mixed into the xs pytree), inference fails and raises ValueError naming the offending values.","triggerScenarios":"lax.scan(f, init, xs=(1, array)) — a Python scalar leaf inside xs; or xs being a plain scalar; or an object without .shape in the pytree.","commonSituations":"Passing hyperparameters or scalars inside xs by accident; mixing a per-step constant into the scanned sequence instead of closing over it.","solutions":["Close over constants in f instead of putting them in xs","Convert scalars to at least 1-D arrays with shape (length,) via jnp.broadcast_to or [c]*n wrapped in jnp.asarray","Ensure every leaf of xs is an ndarray with a leading axis"],"exampleFix":"// before\nconst = 2.0\nlax.scan(lambda c, t: (c, t[0] + const), 0.0, (xs, const))\n# const has no shape -> error\n// after\ndef body(c, x):\n  return c, x + 2.0\ncarry, ys = lax.scan(body, 0.0, xs)","handlingStrategy":"validation","validationCode":"flat, _ = jax.tree_util.tree_flatten(xs)\nassert all(hasattr(x, 'shape') and x.ndim >= 1 for x in flat), \\\n    f'leaves without leading axis: {[x for x in flat if not hasattr(x, \"shape\")]}'","typeGuard":"def all_xs_have_leading_axis(xs) -> bool:\n    return all(hasattr(x, 'shape') and getattr(x, 'ndim', 0) >= 1\n               for x in jax.tree_util.tree_leaves(xs))","tryCatchPattern":"null","preventionTips":["Close over scalars in the body instead of putting them in xs","Convert constants with jnp.asarray([c]*n) when a per-step value is needed","Lint scan calls where xs leaves include Python numbers"],"tags":["jax","scan","shape-inference","scalar-operand","pytree"],"backgroundTag":"missing-leading-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}