jax-ml/jax · error · NameError

unbound axis name: {axis_name}

Error message

unbound axis name: {axis_name}

What it means

Inside sharded computations with a Mesh, collective operations reference axes by name (e.g. axis_index('data'), psum(..., axis_name='data')). _axis_read looks up the name among the context mesh's axis names and raises NameError if absent. The `max(...)` makes repeated axis names use the last (innermost) occurrence.

Source

Thrown at jax/_src/interpreters/pxla.py:450

    if self.mut is None:
      return out
    else:
      out_ = []
      for i, o in zip(self.mut.out_mut, out):
        if i is not None:
          try: args[i]._refs._buf._replace_with(o)
          except AttributeError: pass  # TODO(mattjj): remove float0
        else:
          out_.append(o)
      return out_


def _axis_read(axis_names, axis_name):
  try:
    return max(i for i, name in enumerate(axis_names) if name == axis_name)
  except ValueError:
    raise NameError(f"unbound axis name: {axis_name}") from None

def axis_groups(axis_ctx, name) -> tuple[tuple[int, ...]]:
  assert not isinstance(axis_ctx, sharding_impls.ShardingContext)
  size = axis_ctx.mesh.size
  axis_names = axis_ctx.mesh.axis_names
  axis_sizes = axis_ctx.mesh.axis_sizes
  if not isinstance(name, (list, tuple)):
    name = (name,)
  mesh_axes = tuple(unsafe_map(partial(_axis_read, axis_names), name))
  trailing_size, ragged = divmod(size, math.prod(axis_sizes))
  assert not ragged
  mesh_spec = axis_sizes + (trailing_size,)
  return _axis_groups(mesh_spec, mesh_axes)

def _axis_groups(mesh_spec, mesh_axes):
  """Computes replica group ids for a collective performed over a subset of the mesh.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use an axis name that exists in the Mesh: check mesh.axis_names and pass one of those to the collective
2. Fix typos/differences between the declared Mesh axis names and the axis_name argument
3. If the axis genuinely should exist, verify the collective is evaluated under the right mesh context (e.g. within the sharded computation, not outside)

Example fix

# before
mesh = Mesh(devices, axis_names=('data',))
... psum(x, axis_name='model')

# after
mesh = Mesh(devices, axis_names=('data',))
... psum(x, axis_name='data')

def check_axis(mesh, axis_name):
    assert axis_name in mesh.axis_names, \
        f'{axis_name!r} not in mesh axes {mesh.axis_names}'

def axis_is_bound(mesh, axis_name: str) -> bool:
    return axis_name in mesh.axis_names

When it happens

Trigger: Using a collective with an axis_name not present in the active mesh: psum(x, 'model') under a Mesh with only a 'data' axis; using xmap/pjit with axis names mismatched between the annotation and the mesh; typos or renaming mesh axes without updating collectives.

Common situations: Renaming mesh axes; copy-pasting collective code between models with different mesh layouts; nested/sharded_jit contexts where the axis context isn't the expected mesh.

Related errors

• {name} is a Unreduced -> Invariant collective. This means th
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {name} only accepts inputs that are unreduced. Got {aval.str
• {name}'s input cannot be varying across the axis_name provi

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