{"record":{"id":"52cba7ff48903449","repo":"jax-ml/jax","slug":"0th-dimension-of-all-xs-should-be-replicated-got","errorCode":null,"errorMessage":"0th dimension of all xs should be replicated. Got {}","messagePattern":"0th dimension of all xs should be replicated\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":501,"sourceCode":"    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):\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)","sourceCodeStart":483,"sourceCodeEnd":519,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L483-L519","documentation":"When xs avals carry sharding annotations, scan requires the leading (scanned) dimension to be replicated across all partitions. A xs value sharded on its 0th axis cannot be iterated in a well-defined way by scan, so ValueError is raised listing the offending sharding specs.","triggerScenarios":"Passing xs with a NamedSharding/GSPMD sharding whose 0th dimension is partitioned, e.g. sharding arrays on the axis you scan over.","commonSituations":"Multi-host or multi-device code where arrays are pre-sharded; mismatch between the sharded data axis and the intended time axis (sharding axis 0 when data is laid out (time, batch)).","solutions":["Shard a different axis: keep the scan axis replicated (e.g. P(None, 'data') for (time, batch))","Use jax.lax.map or pmap-style batching where the leading axis is mapped instead of scanned","Reshard xs to replicated 0th dimension before scan (e.g. via device_get/device_put)"],"exampleFix":"# before: sharding the scan axis\nsh = NamedSharding(mesh, P('devices'))  # shards axis 0\nxs = jax.device_put(xs, sh)\nlax.scan(body, c, xs)\n# after: replicate axis 0, shard axis 1\nsh = NamedSharding(mesh, P(None, 'devices'))\nxs = jax.device_put(xs, sh)\nlax.scan(body, c, xs)","handlingStrategy":"validation","validationCode":"for a in jax.tree_util.tree_leaves(xs):\n    spec = getattr(getattr(a, 'sharding', None), 'spec', None)\n    if spec is not None:\n        assert spec.partitions[0] is None, 'scan axis must be replicated'","typeGuard":"def scan_axis_replicated(xs) -> bool:\n    return all(a.sharding.spec.partitions[0] is None\n               for a in jax.tree_util.tree_leaves(xs)\n               if hasattr(a, 'sharding') and hasattr(a.sharding, 'spec'))","tryCatchPattern":"null","preventionTips":["Shard batch dims, not the time/scan dim (P(None, 'dev'))","Put sharded data on axis 1+ by transposing layouts","Validate shardings of inputs before scan in multi-host code"],"tags":["jax","scan","sharding","multi-device","jaxmesh"],"backgroundTag":"sharded-axis-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}