jax-ml/jax · error · ValueError

0th dimension of all xs should be replicated. Got {}

Error message

0th dimension of all xs should be replicated. Got {}

What it means

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.

Source

Thrown at jax/_src/lax/control_flow/loops.py:501

    if length is None:
      raise ValueError(
          "must provide `length` to `scan`, since the leading-axis size of "
          "non-array (hijax) types cannot be inferred")
    return length
  xs_flat = [x for x, h in zip(xs_flat, is_hi) if not h]
  xs_avals = [a for a, h in zip(xs_avals, is_hi) if not h]

  try:
    lengths: list[int] = [x.shape[0] for x in xs_flat]
  except AttributeError as err:
    msg = "scan got value with no leading axis to scan over: {}."
    raise ValueError(
      msg.format(', '.join(str(x) for x in xs_flat
                           if not hasattr(x, 'shape')))) from err

  xs_shaped_avals = lax_utils.ensure_shaped(*xs_avals)
  if not all(a.sharding.spec.partitions[0] is None for a in xs_shaped_avals):
    raise ValueError('0th dimension of all xs should be replicated. Got '
                     f'{", ".join(str(a.sharding.spec) for a in xs_shaped_avals)}')

  if length is not None:
    try:
      length = int(length)
    except core.ConcretizationTypeError:
      msg = ('The `length` argument to `scan` expects a concrete `int` value.'
             ' For scan-like iteration with a dynamic length, use `while_loop`'
             ' or `fori_loop`.')
      raise core.ConcretizationTypeError(length, msg) from None
    else:
      if not all(length == l for l in lengths):
        msg = ("scan got `length` argument of {} which disagrees with "
              "leading axis sizes {}.")
        raise ValueError(msg.format(length, [x.shape[0] for x in xs_flat]))
      return length
  else:
    unique_lengths = set(lengths)

View on GitHub (pinned to 1e1c6a8fc0)

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)

Example fix

# before: sharding the scan axis
sh = NamedSharding(mesh, P('devices'))  # shards axis 0
xs = jax.device_put(xs, sh)
lax.scan(body, c, xs)
# after: replicate axis 0, shard axis 1
sh = NamedSharding(mesh, P(None, 'devices'))
xs = jax.device_put(xs, sh)
lax.scan(body, c, xs)

Defensive patterns

Strategy: validation

Validate before calling

for a in jax.tree_util.tree_leaves(xs):
    spec = getattr(getattr(a, 'sharding', None), 'spec', None)
    if spec is not None:
        assert spec.partitions[0] is None, 'scan axis must be replicated'

Type guard

def scan_axis_replicated(xs) -> bool:
    return all(a.sharding.spec.partitions[0] is None
               for a in jax.tree_util.tree_leaves(xs)
               if hasattr(a, 'sharding') and hasattr(a.sharding, 'spec'))

Try / catch

null

Prevention

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

When it happens

Trigger: Passing xs with a NamedSharding/GSPMD sharding whose 0th dimension is partitioned, e.g. sharding arrays on the axis you scan over.

Common situations: 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)).

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/52cba7ff48903449. Report an issue: GitHub.