{"record":{"id":"414861352f3308fc","repo":"jax-ml/jax","slug":"mapped-away-dimension-of-inputs-passed-to-vmap-sho","errorCode":null,"errorMessage":"Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: {out_spec} vs {spec}","messagePattern":"Mapped away dimension of inputs passed to vmap should be sharded the same\\. Got inconsistent axis specs: (.+?) vs (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1274,"sourceCode":"      return tree_unflatten(out_tree(), out_flat)\n\n  return cast(F, vmap_f)\n\ndef _mapped_axis_spec(args_flat, in_axes_flat):\n  def _get_spec(arg, i):\n    try:\n      # Duck type arrays like BCOO arrays can be passed to vmap.\n      return shaped_abstractify(arg).sharding.spec[i]\n    except (IndexError, TypeError, AttributeError):\n      return None\n\n  out_spec = None\n  non_none_count = 0\n  for arg, i in zip(args_flat, in_axes_flat):\n    if i is not None:\n      spec = _get_spec(arg, i)\n      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\"","sourceCodeStart":1256,"sourceCodeEnd":1292,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1256-L1292","documentation":"Raised by jax.vmap when multiple arguments that are mapped (in_axes not None) are sharded across mesh axes differently, i.e. their mapped-away dimensions have inconsistent named-sharding specs. JAX requires all mapped inputs to be sharded identically along the axis being vmapped so the computation can be coherently distributed.","triggerScenarios":"Calling jax.vmap(f)(x, y) where x and y are arrays sharded over a NamedSharding/mesh with different specs (e.g. x sharded on 'data' and y replicated or sharded on 'model') while both are mapped over in_axes.","commonSituations":"Mixing SPMD-style mesh shardings with vmap; creating inputs from different device_put calls with different shardings; upgrading code that used pmap to vmap over pre-sharded arrays.","solutions":["Make all mapped arguments share the same sharding spec (device_put them with the same NamedSharding before vmap)","Replicate (don't shard) inputs along the mapped axis, or move the sharding outside the vmap","Use pmap or jax.lax.map / SPMD instead of vmap over already-sharded arrays"],"exampleFix":"// before\nx = jax.device_put(x, NamedSharding(mesh, P('data')))\ny = jax.device_put(y, NamedSharding(mesh, P('model')))\njax.vmap(f)(x, y)\n// after\nsh = NamedSharding(mesh, P('data'))\nx = jax.device_put(x, sh); y = jax.device_put(y, sh)\njax.vmap(f)(x, y)","handlingStrategy":"validation","validationCode":"specs = {jax.typeof(a).sharding.spec for a in tree_leaves(args) if hasattr(jax.typeof(a), 'sharding')}\nassert len(specs) <= 1, f'mapped args sharded differently: {specs}'","typeGuard":"def same_sharding(args):\n    specs = [getattr(getattr(jax.typeof(a), 'sharding', None), 'spec', None) for a in tree_leaves(args)]\n    specs = [s for s in specs if s is not None]\n    return len(set(map(str, specs))) <= 1","tryCatchPattern":"try:\n    jax.vmap(f)(*args)\nexcept ValueError as e:\n    if 'inconsistent axis specs' in str(e):\n        args = jax.device_put(args, replicate_sharding); jax.vmap(f)(*args)\n    else: raise","preventionTips":["device_put all mapped inputs with one shared NamedSharding before vmap","Keep SPMD sharding logic outside vmap boundaries","Add a pre-call assert that mapped inputs share one sharding spec"],"tags":["jax","vmap","sharding","mesh","distributed"],"backgroundTag":"sharding-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}