jax-ml/jax · error · ValueError

Unmapped values passed to vmap cannot be sharded along the…

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.

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.

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