{"record":{"id":"97ff444193465a29","repo":"jax-ml/jax","slug":"0th-dimension-of-leaf-passed-to-jax-lax-map-shou","errorCode":null,"errorMessage":"0th dimension of leaf passed to `jax.lax.map` should be replicated. Got {}","messagePattern":"0th dimension of leaf passed to `jax\\.lax\\.map` should be replicated\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2680,"sourceCode":"\n  if lower_dtype != dtype:\n    lower = lax.convert_element_type(lower, dtype)\n  if upper_dtype != dtype:\n    upper = lax.convert_element_type(upper, dtype)\n  while_body_fun = _fori_body_fun(body_fun, body_fun_dbg)\n  _, _, result = while_loop(_fori_cond_fun, while_body_fun,\n                            (lower, upper, init_val))\n  return result\n\n### map and miscellaneous rules\n\ndef _scan_leaf(leaf, batch_elems, num_batches, batch_size):\n  def f(l):\n    return l[:batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:])\n\n  aval = core.typeof(leaf)\n  if aval.sharding.spec[0] is not None:\n    raise ValueError(\n        '0th dimension of leaf passed to `jax.lax.map` should be replicated.'\n        f' Got {aval.str_short(True, True)}')\n\n  out_s = aval.sharding.update(spec=P(None, None, *aval.sharding.spec[1:]))\n  out_s = canonicalize_sharding(out_s, 'lax.map')\n  if out_s is not None and out_s.mesh._any_axis_explicit:\n    return auto_axes(f, out_sharding=out_s, axes=out_s.mesh.explicit_axes)(leaf)\n  return f(leaf)\n\ndef _remainder_leaf(leaf, batch_elems):\n  def f(l):\n    return l[batch_elems:]\n  sharding = canonicalize_sharding(core.typeof(leaf).sharding, 'lax.map')\n  if sharding is not None and sharding.mesh._any_axis_explicit:\n    return auto_axes(\n        f, out_sharding=sharding, axes=sharding.mesh.explicit_axes\n    )(leaf)\n  return f(leaf)","sourceCodeStart":2662,"sourceCodeEnd":2698,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2662-L2698","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Map over a different, replicated axis by transposing so the sharded axis is not the mapped one","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)","Compute in batches whose size divides the axis so the _batch_and_remainder path (which triggers _scan_leaf) is avoided"],"exampleFix":"// before\nout = jax.lax.map(f, x)  # x has dim0 sharded across mesh\n// after\nfrom jax.sharding import NamedSharding, P\nx = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(None, 'data')))\nout = jax.lax.map(f, x)  # mapped axis replicated","handlingStrategy":"validation","validationCode":"import jax\ndef leading_axis_replicated(x):\n    sh = getattr(x, 'sharding', None)\n    spec = getattr(sh, 'spec', None)\n    return spec is None or spec[0] is None\n\nif not leading_axis_replicated(x):\n    raise ValueError('replicate axis 0 before lax.map')","typeGuard":"def map_safe_leaf(leaf) -> bool:\n    s = getattr(leaf, 'sharding', None)\n    spec = getattr(s, 'spec', None)\n    return spec is None or spec[0] is None","tryCatchPattern":null,"preventionTips":["Replicate the mapped axis before jax.lax.map in sharded pipelines","Prefer vmap/scan for sharded inputs","Add sharding-constraint assertions in tests for multi-device code"],"tags":["jax","lax-map","sharding","multi-device","gspmd"],"backgroundTag":"jax-sharded-axis-not-replicated","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}