jax-ml/jax · error · ValueError

shard_map {prefix}_specs argument must refer to an axis mark

Error message

shard_map {prefix}_specs argument must refer to an axis marked as manual ({manual_axes}), but:

{msgs}

Check the {prefix}_specs values passed to shard_map.

What it means

In manual/partial-manual shard_map mode, specs may only reference mesh axes that are marked manual (via manual_axes). This ValueError lists each spec path that references a non-manual axis name.

Source

Thrown at jax/_src/shard_map.py:506

      names = (names,) if not isinstance(names, tuple) else names
      for name in names:
        if name is not None and name not in manual_axes:
          return False
    return True

  if all(check_spec(p) for p in tree_leaves(specs)):
    return
  prefix = 'in' if error_type == SpecErrorType.input else 'out'
  msgs = [f"  {prefix}_specs{keystr(key)} is {x} of type {type(x).__name__}, "
          for key, x in generate_key_paths(specs) if not isinstance(x, P)]
  if not msgs:
    for key, p in generate_key_paths(specs):
      for names in p:
        names = (names,) if not isinstance(names, tuple) else names
        for name in names:
          if name is not None and name not in manual_axes:
            msgs.append(f"  {prefix}_specs{keystr(key)} refers to {repr(name)}")
    raise ValueError(
        f"shard_map {prefix}_specs argument must refer to an axis "
        f"marked as manual ({manual_axes}), but:\n\n"
        + '\n\n'.join(msgs) + '\n\n'
        f"Check the {prefix}_specs values passed to shard_map.")
  raise TypeError(
      f"shard_map {prefix}_specs argument must be a pytree of "
      f"`jax.sharding.PartitionSpec` instances, but:\n\n"
      + '\n\n'.join(msgs) + '\n\n'
      f"Check the {prefix}_specs values passed to shard_map.")

class NoFail:
  def __repr__(self):
    return "NoFail()"

no_fail = NoFail()

def _check_specs_vs_args(
    f: Callable, mesh: Mesh | AbstractMesh, in_tree: PyTreeDef, in_specs: Specs,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Add the referenced axis name to manual_axes in the shard_map call
Change the spec to only use axis names already in manual_axes (or use None to replicate)
Re-check which axes the surrounding computation actually treats as manual

Example fix

// before
shard_map(f, mesh, manual_axes=('m',), in_specs=P('m','x'))
// after
shard_map(f, mesh, manual_axes=('m','x'), in_specs=P('m','x'))

Defensive patterns

Strategy: validation

Validate before calling

def specs_within_manual(specs, manual_axes):
    bad = [n for p in jax.tree.leaves(specs) if isinstance(p, P)
           for names in p for n in ((names,) if not isinstance(names, tuple) else names)
           if n is not None and n not in manual_axes]
    return not bad

Type guard

def manual_safe(specs, manual_axes) -> bool:
    ok = lambda p: all(n is None or n in manual_axes for n in (x for part in p for x in (part if isinstance(part, tuple) else (part,))))
    return all(ok(p) for p in jax.tree.leaves(specs) if isinstance(p, P))

Try / catch

try: shard_map(...) except ValueError as e: if 'marked as manual' in str(e): parse listed names and add to manual_axes; else: raise

Prevention

Derive manual_axes from the axes actually named in specs
Keep spec axis-name sets in one constants module

When it happens

Trigger: Calling shard_map(..., manual_axes=('m',)) (or checking a spec in a partially-manual context) with in_specs/out_specs naming an axis not included in manual_axes, e.g. out_specs=P('data') when 'data' is not manual.

Common situations: Migrating code to manual sharding mode and forgetting to update specs; using axis names from the global mesh that were not declared as manual axes.

Related errors

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