jax-ml/jax · error · ValueError

Mapped away dimension of inputs passed to vmap should be sha

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.

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)

Solutions

1. Make all mapped arguments share the same sharding spec (device_put them with the same NamedSharding before vmap)
2. Replicate (don't shard) inputs along the mapped axis, or move the sharding outside the vmap
3. 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)

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}'

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

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

• Unmapped values passed to vmap cannot be sharded along the m
• The error code state and the predicate must be on the same m
• Reduced axes can only refer to mesh axes that is of type `Ex
• pallas_call with a mesh does not support batching
• callbacks are only supported in spmd computations when all m

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