jax-ml/jax · error · ValueError

0th dimension of leaf passed to `jax.lax.map` should be repl

Error message

0th dimension of leaf passed to `jax.lax.map` should be replicated. Got {}

What it means

jax.lax.map splits an input leaf into (num_batches, batch_size, ...) via _scan_leaf; this requires the leading axis to be replicated across devices (its sharding spec entry must be None). If the leaf's 0th dimension is sharded, reshaping across the mapped axis is not valid under the array's sharding, so JAX raises a ValueError naming the offending aval.

Source

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

  if lower_dtype != dtype:
    lower = lax.convert_element_type(lower, dtype)
  if upper_dtype != dtype:
    upper = lax.convert_element_type(upper, dtype)
  while_body_fun = _fori_body_fun(body_fun, body_fun_dbg)
  _, _, result = while_loop(_fori_cond_fun, while_body_fun,
                            (lower, upper, init_val))
  return result

### map and miscellaneous rules

def _scan_leaf(leaf, batch_elems, num_batches, batch_size):
  def f(l):
    return l[:batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:])

  aval = core.typeof(leaf)
  if aval.sharding.spec[0] is not None:
    raise ValueError(
        '0th dimension of leaf passed to `jax.lax.map` should be replicated.'
        f' Got {aval.str_short(True, True)}')

  out_s = aval.sharding.update(spec=P(None, None, *aval.sharding.spec[1:]))
  out_s = canonicalize_sharding(out_s, 'lax.map')
  if out_s is not None and out_s.mesh._any_axis_explicit:
    return auto_axes(f, out_sharding=out_s, axes=out_s.mesh.explicit_axes)(leaf)
  return f(leaf)

def _remainder_leaf(leaf, batch_elems):
  def f(l):
    return l[batch_elems:]
  sharding = canonicalize_sharding(core.typeof(leaf).sharding, 'lax.map')
  if sharding is not None and sharding.mesh._any_axis_explicit:
    return auto_axes(
        f, out_sharding=sharding, axes=sharding.mesh.explicit_axes
    )(leaf)
  return f(leaf)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replicate the input's leading axis before lax.map: e.g. re-shard with jax.lax.with_sharding_constraint(x, jax.sharding.NamedSharding(mesh, P(None, ...))) so axis 0 is replicated
2. Map over a different, replicated axis by transposing so the sharded axis is not the mapped one
3. Replace lax.map with vmap or scan over a replicated copy (x = jax.device_get then re-shard, or use jax.make_array_from_process_to_device_index with replicated dim 0)
4. Compute in batches whose size divides the axis so the _batch_and_remainder path (which triggers _scan_leaf) is avoided

Example fix

// before
out = jax.lax.map(f, x)  # x has dim0 sharded across mesh
// after
from jax.sharding import NamedSharding, P
x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(None, 'data')))
out = jax.lax.map(f, x)  # mapped axis replicated

import jax
def leading_axis_replicated(x):
    sh = getattr(x, 'sharding', None)
    spec = getattr(sh, 'spec', None)
    return spec is None or spec[0] is None

if not leading_axis_replicated(x):
    raise ValueError('replicate axis 0 before lax.map')

def map_safe_leaf(leaf) -> bool:
    s = getattr(leaf, 'sharding', None)
    spec = getattr(s, 'spec', None)
    return spec is None or spec[0] is None

When it happens

Trigger: Calling jax.lax.map on arrays whose first axis is sharded (e.g. NamedSharding with the 0th axis mapped to a mesh axis, or arrays produced from sharded pjit/jit computations), when lax.map internally batches for scan — i.e. when the mapped axis length times batch does not divide evenly / the batching path with _batch_and_remainder triggers.

Common situations: Using jax.lax.map on sharded inputs under multi-device jax.jit with sharding constraints (GSPMD); large datasets laid out with the batch dimension sharded and then mapped over; combining jax.experimental.mesh_utils sharding with lax.map.

Related errors

• Query, key and value should have same sharding.
• to_dlpack can only pack a dlpack tensor from an array on a s
• from_dlpack can only unpack a dlpack tensor onto a singular
• 0th dimension of all xs should be replicated. Got {}
• unreduced rule for {name} is not implemented. Please file an

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