jax-ml/jax · error · ValueError

Unmapped values passed to vmap cannot be sharded along the m

Error message

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}

What it means

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.

Source

Thrown at jax/_src/api.py:1291

      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"
            f" axis you are vmapping over. Got type: {aval.str_short(True)},"
            f" in_axes: {i} and vmapped mesh axis: {ema}")

def _mapped_axis_size(fn, tree, vals, dims, name, axis_size=None):
  if not vals:
    if axis_size is not None:
      return axis_size
    args, kwargs = tree_unflatten(tree, vals)
    raise ValueError(
        f"{name} wrapped function must be passed at least one argument "
        "containing an array or axis_size must be specified, got empty "
        f"*args={args} and **kwargs={kwargs}"
    )

  def _get_axis_size(name: str, x, axis: int) -> core.AxisSize | None:
    shape: tuple[core.AxisSize, ...] = ()
    try:

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
bias = jax.device_put(bias, NamedSharding(mesh, P('data')))
jax.vmap(f, in_axes=(0, None))(x, bias)
// after
bias = jax.device_put(bias, NamedSharding(mesh, P()))
jax.vmap(f, in_axes=(0, None))(x, bias)

Defensive patterns

Strategy: validation

Validate before calling

mesh_axes = set(flatten_spec(jax.typeof(y).sharding.spec)) if in_axis is None else set()
assert not (mesh_axes & set(ema)), 'unmapped arg sharded on vmapped mesh axis'

Type guard

def unmapped_is_replicated(x, ema):
    sh = getattr(jax.typeof(x), 'sharding', None)
    spec = set(flatten_spec(sh.spec)) if sh else set()
    return not (spec & set(ema))

Try / catch

try:
    jax.vmap(f, in_axes=in_axes)(*args)
except ValueError as e:
    if 'cannot be sharded along the mesh axis' in str(e):
        args = [jax.device_put(a, NamedSharding(mesh, P())) if unmapped else a for ...]
        jax.vmap(f, in_axes=in_axes)(*args)
    else: raise

Prevention

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

When it happens

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

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

Related errors

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