jax-ml/jax · error · ValueError

Invalid spec: {spec}

Error message

Invalid spec: {spec}

What it means

In full-manual shard_map mode, _manual_spec converts a PartitionSpec into a manual-mode spec by keeping only entries in manual_axes. If a tuple spec contains None in a position that isn't trailing (after stripping trailing Nones), the spec is malformed for manual conversion and this ValueError 'Invalid spec: {spec}' is raised.

Source

Thrown at jax/_src/shard_map.py:433

    _check_specs(SpecErrorType.input, in_specs, axis_names)
    _check_unreduced(SpecErrorType.input, mesh, axis_names, in_specs)
  _check_specs(SpecErrorType.out, out_specs, axis_names)
  _check_unreduced(SpecErrorType.out, mesh, axis_names, out_specs)
  return mesh, axis_names


def _manual_spec(manual_axes, spec: P, mesh) -> P:
  out: list[str | tuple[str | None, ...] | None] = []
  s: str | None | tuple[str, ...]
  for s in spec.partitions:
    if s is None:
      out.append(s)
    elif isinstance(s, tuple):
      temp = [p if p in manual_axes else None for p in s]
      while temp and temp[-1] is None:
        temp.pop()
      if None in temp:
        raise ValueError(f"Invalid spec: {spec}")
      out.append(None if len(temp) == 0 else tuple(temp))
    else:
      out.append(s if s in manual_axes else None)
  _check_unreduced(SpecErrorType.input, mesh, manual_axes, spec)
  return spec.update(partitions=tuple(out))


# Error checking and messages

SpecErrorType = enum.Enum('SpecErrorType', ['input', 'out'])

def _check_unreduced(error_type, mesh, manual_axes, specs):
  from jax._src.hijax import HiPspec
  prefix = 'in' if error_type == SpecErrorType.input else 'out'
  full_manual = frozenset(mesh.axis_names) == manual_axes
  specs_flat, _ = tree_flatten(specs)
  for s in specs_flat:
    if isinstance(s, HiPspec):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rewrite the spec so non-manual entries are only trailing: P('i','j') or P('i','j',None)
2. Use unreduced/reduced spec syntax or explicit full-manual specs appropriate to manual_axes mode instead of mixing inner Nones

Example fix

// before
jax.shard_map(f, mesh=mesh, in_specs=P('i', None, 'j'), manual_axes=('i','j'), ...)

// after
jax.shard_map(f, mesh=mesh, in_specs=P('i', 'j'), manual_axes=('i','j'), ...)

def valid_manual_spec(spec, manual_axes):
    if isinstance(spec, tuple):
        core = [p for p in spec if p is not None]
        return all(p in manual_axes for p in core)
    return True
assert valid_manual_spec(P('i', None, 'j'), manual_axes), 'inner None not allowed'

When it happens

Trigger: Using shard_map in manual_axes mode with a PartitionSpec like P('i', None, 'j') where an inner None remains after trailing-None removal while the surrounding entries are manual axes — the conversion can't map the hole to a single manual axis.

Common situations: Writing manual-mode sharding code (FSDP/pipeline-style) where specs mix mesh axes and unconstrained dims; porting spec patterns from automatic mode into full manual mode.

Related errors

• in_specs passed to shard_map: {s} does not match the specs o
• unreduced/reduced can only be passed to {prefix}_specs when
• unreduced in {prefix}_specs {s} can only be used when the me
• reduced in {prefix}_specs {s} can only be used when the mesh
• shard_map in_specs argument must be a pytree of `jax.shardin

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