jax-ml/jax · error · ValueError

Mapped away dimension of inputs passed to vmap should be…

Error message

Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: {out_spec} vs {spec}

What it means

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.

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

Example fix

// before
x = jax.device_put(x, NamedSharding(mesh, P('data')))
y = jax.device_put(y, NamedSharding(mesh, P('model')))
jax.vmap(f)(x, y)
// after
sh = NamedSharding(mesh, P('data'))
x = jax.device_put(x, sh); y = jax.device_put(y, sh)
jax.vmap(f)(x, y)

Defensive patterns

Strategy: validation

Validate before calling

specs = {jax.typeof(a).sharding.spec for a in tree_leaves(args) if hasattr(jax.typeof(a), 'sharding')}
assert len(specs) <= 1, f'mapped args sharded differently: {specs}'

Type guard

def same_sharding(args):
    specs = [getattr(getattr(jax.typeof(a), 'sharding', None), 'spec', None) for a in tree_leaves(args)]
    specs = [s for s in specs if s is not None]
    return len(set(map(str, specs))) <= 1

Try / catch

try:
    jax.vmap(f)(*args)
except ValueError as e:
    if 'inconsistent axis specs' in str(e):
        args = jax.device_put(args, replicate_sharding); jax.vmap(f)(*args)
    else: raise

Prevention

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

When it happens

Trigger: 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.

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

Related errors

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

Appendix: source

Thrown at jax/_src/api.py:1274

      return tree_unflatten(out_tree(), out_flat)

  return cast(F, vmap_f)

def _mapped_axis_spec(args_flat, in_axes_flat):
  def _get_spec(arg, i):
    try:
      # Duck type arrays like BCOO arrays can be passed to vmap.
      return shaped_abstractify(arg).sharding.spec[i]
    except (IndexError, TypeError, AttributeError):
      return None

  out_spec = None
  non_none_count = 0
  for arg, i in zip(args_flat, in_axes_flat):
    if i is not None:
      spec = _get_spec(arg, i)
      if non_none_count != 0 and out_spec != spec:
        raise ValueError(
            "Mapped away dimension of inputs passed to vmap should be sharded"
            f" the same. Got inconsistent axis specs: {out_spec} vs {spec}")
      out_spec = spec
      non_none_count += 1
  if out_spec is not None and not isinstance(out_spec, tuple):
    out_spec = (out_spec,)
  return out_spec

def _check_ema_unmapped_args(ema, args_flat, in_axes_flat):
  if ema is None:
    return
  for a, i in zip(args_flat, in_axes_flat):
    if i is None:
      aval = core.typeof(a)
      spec = set(sharding_impls.flatten_spec(aval.sharding.spec))
      if any(e in spec for e in ema):
        raise ValueError(
            "Unmapped values passed to vmap cannot be sharded along the mesh"

View on GitHub (pinned to 1e1c6a8fc0)