jax-ml/jax · error · ValueError

shard_map out_specs vma error (msg from _inout_vma_error)

Error message

shard_map out_specs vma error (msg from _inout_vma_error)

What it means

When shard_map's output reassembly encounters a value that is replicated across a mesh axis named in out_specs (i.e. the same data on every shard, so concatenation would duplicate it), a _RepError is raised internally and this ValueError with a message from _inout_vma_error is surfaced. It typically indicates a Virtual Mesh Axis (vmap-introduced axis) conflict or a replicated output being mapped over a named axis.

Source

Thrown at jax/_src/shard_map.py:353

      return ans_ft.with_aux(out_specs_flat)

    try:
      newly_manual_axes = axis_names - set(mesh.manual_axes)
      out_ft = shard_map_p.bind(
          *dyn_args, subfuns=(f_wrapped,), mesh=mesh, in_specs=in_specs_flat,
          check_vma=check_vma, newly_manual_axes=newly_manual_axes, debug_info=dbg)
    except _SpecError as e:
      fails, out_tree = e.args
      msg = _spec_rank_error(SpecErrorType.out, f, out_tree, out_specs, fails)
      if any(fail is not no_fail and not fail.shape for fail in fails):
        msg += (" In particular, for rank 0 outputs which are not constant "
                "over the mesh, add at least one (singleton) axis to them so "
                "that they can be concatenated using out_specs.")
      raise ValueError(msg) from None
    except _RepError as e:
      fails, out_tree, = e.args
      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)
      raise ValueError(msg) from None
    return out_ft.unflatten()
  return cast(F, wrapped)


def _axes_to_pspec(axis_name, axis):
  if axis is None:
    return P()
  return P(*[None] * axis + [axis_name])


def _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):
  if mesh is None:
    mesh = get_abstract_mesh()
    if mesh.empty:
      raise ValueError(
          "The context mesh cannot be empty. Use"
          " `jax.set_mesh(mesh)` to enter into a mesh context")
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Restructure the function so the output genuinely varies along every mesh axis named in its out_specs entry, or drop that axis from the spec
Use jax.lax.axis_index / explicit per-shard values instead of broadcasting a constant
If the value is intentionally replicated, use out_specs that omit the replicated axis (or P() for full replication)

Example fix

// before
def f(x): return jnp.zeros_like(x)[0]  # same value on every shard
jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

// after
def f(x): return x.sum(0)  # varies per shard is fine; or:
jax.shard_map(lambda x: jnp.zeros(1), mesh=mesh, in_specs=P('i'), out_specs=P())(x)

Defensive patterns

Strategy: validation

Validate before calling

# ensure outputs vary along every mesh axis named in out_specs
def f(x):
    per_shard = x.sum()          # constant across shards -> bad for P('i')
    return per_shard + jnp.zeros(x.shape[0])  # shape tied to shard -> varies

Try / catch

try:
    out = shmapped(x)
except ValueError:
    out = jax.shard_map(lambda x: jnp.zeros(1), mesh=mesh,
                        in_specs=P('i'), out_specs=P())(x)  # replicated fallback

Prevention

Never broadcast constants across a named out_specs axis
Compute outputs from shard-local data or lax.axis_index
Prefer P() for intentionally replicated outputs

When it happens

Trigger: A mapped function returns a value that does not vary along a mesh axis that out_specs names for it — e.g. broadcasting a constant across shards while out_specs=P('i'), often when using smap/vmap composition with shard_map.

Common situations: Mixing vmap/smap with shard_map where automatic spmd batching introduces implicit axes; returning mesh-broadcast constants; partially-updated JAX versions refining VMA error reporting.

Related errors

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