jax-ml/jax · error · TypeError

in_axes was not specified when axis_name={axis_names} was of

Error message

in_axes was not specified when axis_name={axis_names} was of type {axis_types} / shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but it was `None` when {axis_names=} are of type {axis_types}

What it means

shard_map requires either explicit PartitionSpec in_specs/in_axes or an explicit-mesh context when axis_names include non-Explicit axis types (e.g. autodiff/vmap-introduced axes). The message has two variants: for smap it says in_axes was unspecified; for shard_map it says in_specs was None while axis_names were of a non-Explicit type. In short, you can't infer specs when the named axes aren't concrete mesh axes.

Source

Thrown at jax/_src/shard_map.py:412

    axis_names = frozenset(axis_names)
  if not axis_names:
    axis_names = mesh_axis_names_wo_vmap
  if not axis_names.issubset(mesh_axis_names_wo_vmap):
    raise ValueError(
        f"jax.shard_map requires axis_names={axis_names} to be a subset of "
        f"mesh.axis_names={mesh_axis_names_wo_vmap}")

  if (in_specs is Infer and
      not all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):
    axis_types = ', '.join(str(mesh._name_to_type[a]) for a in axis_names)
    if _smap:
      msg = (f"in_axes was not specified when axis_name={axis_names} was of"
             f" type {axis_types}")
    else:
      msg = ("shard_map in_specs argument must be a pytree of"
             " `jax.sharding.PartitionSpec` instances, but it was `None` when"
             f" {axis_names=} are of type {axis_types}")
    raise TypeError(msg)

  if in_specs is not Infer and in_specs is not None:
    _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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Specify in_axes/in_specs explicitly so JAX needn't infer specs from non-Explicit axes
2. Restrict axis_names to Explicit mesh axes, or mark the mesh axes as explicit when constructing the abstract mesh

Example fix

// before
jax.experimental.smap(f, mesh=mesh, out_axes=0, axis_name='i')  # 'i' not Explicit, no in_axes

// after
jax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes=0, axis_name='i')

# always specify in_axes/in_specs when non-Explicit axes may be present
if any(mesh._name_to_type[a] != AxisType.Explicit for a in axis_names):
    assert in_axes is not None and in_axes is not Infer, 'specify in_axes explicitly'

When it happens

Trigger: Calling smap without in_axes while axis_name refers to a non-Explicit axis type, or calling shard_map with in_specs=None when any named axis is not an Explicit mesh axis (e.g. inside vmap/grad where implicit axes exist).

Common situations: Composing vmap or grad with shard_map/smap; using abstract/multi-stage meshes where axes have types like Autodiff or Replicated; migrating code that relied on inference that now fails under typed mesh axes.

Related errors

• smap axis_name should be a `str` or a `Hashable`, but got {a
• smap in_axes must be an int, None, jax.sharding.Infer, or a
• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• Primitive {prim_name} requires varying manual axes to match,

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