{"record":{"id":"dadbbd4253c71950","repo":"jax-ml/jax","slug":"unmapped-values-passed-to-vmap-cannot-be-sharded-a","errorCode":null,"errorMessage":"Unmapped values passed to vmap cannot be sharded along the mesh axis you are vmapping over. Got type: {aval.str_short(True)}, in_axes: {i} and vmapped mesh axis: {ema}","messagePattern":"Unmapped values passed to vmap cannot be sharded along the mesh axis you are vmapping over\\. Got type: (.+?), in_axes: (.+?) and vmapped mesh axis: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1291,"sourceCode":"      if non_none_count != 0 and out_spec != spec:\n        raise ValueError(\n            \"Mapped away dimension of inputs passed to vmap should be sharded\"\n            f\" the same. Got inconsistent axis specs: {out_spec} vs {spec}\")\n      out_spec = spec\n      non_none_count += 1\n  if out_spec is not None and not isinstance(out_spec, tuple):\n    out_spec = (out_spec,)\n  return out_spec\n\ndef _check_ema_unmapped_args(ema, args_flat, in_axes_flat):\n  if ema is None:\n    return\n  for a, i in zip(args_flat, in_axes_flat):\n    if i is None:\n      aval = core.typeof(a)\n      spec = set(sharding_impls.flatten_spec(aval.sharding.spec))\n      if any(e in spec for e in ema):\n        raise ValueError(\n            \"Unmapped values passed to vmap cannot be sharded along the mesh\"\n            f\" axis you are vmapping over. Got type: {aval.str_short(True)},\"\n            f\" in_axes: {i} and vmapped mesh axis: {ema}\")\n\ndef _mapped_axis_size(fn, tree, vals, dims, name, axis_size=None):\n  if not vals:\n    if axis_size is not None:\n      return axis_size\n    args, kwargs = tree_unflatten(tree, vals)\n    raise ValueError(\n        f\"{name} wrapped function must be passed at least one argument \"\n        \"containing an array or axis_size must be specified, got empty \"\n        f\"*args={args} and **kwargs={kwargs}\"\n    )\n\n  def _get_axis_size(name: str, x, axis: int) -> core.AxisSize | None:\n    shape: tuple[core.AxisSize, ...] = ()\n    try:","sourceCodeStart":1273,"sourceCodeEnd":1309,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1273-L1309","documentation":"Raised by jax.vmap when an argument with in_axes=None (unmapped) is nevertheless sharded along a mesh axis that vmap is mapping over. Since vmap would have to replicate/handle that axis, JAX forbids unmapped inputs sharded on the vmapped mesh axis.","triggerScenarios":"jax.vmap(f, in_axes=(0, None))(x, y) where y is a NamedSharding array whose spec includes the mesh axis being mapped (the axis x is sharded/mapped along).","commonSituations":"Passing global constants (bias vectors, weights) that were device_put with a sharding covering the whole mesh into a vmap over that mesh; converting pmap code where broadcast was implicit.","solutions":["Replicate the unmapped argument (shard it on no axes: P() or P(None))","Change its in_axes so it is mapped along that axis too","Move the sharding of that constant outside/after the vmap call"],"exampleFix":"// before\nbias = jax.device_put(bias, NamedSharding(mesh, P('data')))\njax.vmap(f, in_axes=(0, None))(x, bias)\n// after\nbias = jax.device_put(bias, NamedSharding(mesh, P()))\njax.vmap(f, in_axes=(0, None))(x, bias)","handlingStrategy":"validation","validationCode":"mesh_axes = set(flatten_spec(jax.typeof(y).sharding.spec)) if in_axis is None else set()\nassert not (mesh_axes & set(ema)), 'unmapped arg sharded on vmapped mesh axis'","typeGuard":"def unmapped_is_replicated(x, ema):\n    sh = getattr(jax.typeof(x), 'sharding', None)\n    spec = set(flatten_spec(sh.spec)) if sh else set()\n    return not (spec & set(ema))","tryCatchPattern":"try:\n    jax.vmap(f, in_axes=in_axes)(*args)\nexcept ValueError as e:\n    if 'cannot be sharded along the mesh axis' in str(e):\n        args = [jax.device_put(a, NamedSharding(mesh, P())) if unmapped else a for ...]\n        jax.vmap(f, in_axes=in_axes)(*args)\n    else: raise","preventionTips":["Keep constants replicated (P()) when used as unmapped vmap inputs","Document which mesh axis a vmap maps over next to the sharding setup","Test vmap+mesh combos in unit tests"],"tags":["jax","vmap","sharding","mesh"],"backgroundTag":"sharding-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}