jax-ml/jax · error · NameError

unbound axis name: {axis_name}

Error message

unbound axis name: {axis_name}

What it means

JAX NameError thrown by AxisEnv.axis_size when a named axis (from pmap/xmap/sharding annotations) is looked up but is not bound in the current axis environment. JAX tracks named logical axes explicitly, and referencing one that was never introduced into the trace's axis environment is an error.

Source

Thrown at jax/_src/core.py:1376

  spmd_axis_names : frozenset[AxisName]
  explicit_mesh_axis_names: frozenset[AxisName]

  @staticmethod
  @weak_value_interner
  def _create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names):
    obj = object.__new__(AxisEnv)
    object.__setattr__(obj, 'axis_sizes', axis_sizes)
    object.__setattr__(obj, 'spmd_axis_names', spmd_axis_names)
    object.__setattr__(obj, 'explicit_mesh_axis_names', explicit_mesh_axis_names)
    return obj

  def __new__(cls, axis_sizes=FrozenDict({}), spmd_axis_names=frozenset(),
              explicit_mesh_axis_names=frozenset()):
    return cls._create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names)

  def axis_size(self, axis_name):
    if axis_name not in self.axis_sizes:
      raise NameError(f"unbound axis name: {axis_name}")
    else:
      return self.axis_sizes[axis_name]

  def axis_exists(self, axis_name):
    return axis_name in self.axis_sizes

  def axis_names(self):
    return tuple(k for k in self.axis_sizes)

  def pop_pure(self, axis_name):
    new_sizes = dict(self.axis_sizes)
    new_sizes.pop(axis_name)
    return AxisEnv(FrozenDict(new_sizes), self.spmd_axis_names,
                   self.explicit_mesh_axis_names)

  def extend_pure(self, name_size_pairs):
    new_sizes = dict(self.axis_sizes)
    new_sizes.update((name, size) for name, size in name_size_pairs

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the spelling/case of axis_name against the name passed to the enclosing vmap/pmap/xmap
2. Ensure the collective (psum/pmax/all_reduce...) is called inside the transformation that declared that axis name
3. Nest transformations so that every axis_name referenced is declared by an enclosing vmap(name=...) or pmap(axis_name=...)
4. Print/inspect the AxisEnv (axis_sizes keys) in a minimal repro to see which names are bound

Example fix

// before
out = jax.vmap(fn)(x)  # fn uses lax.psum(x, 'batch')

// after
out = jax.vmap(fn, axis_name='batch')(x)

def check_axis_bound(axis_name, declared_names):
    assert axis_name in declared_names, f"axis {axis_name!r} not in {declared_names}"

When it happens

Trigger: Calling axis_size(axis_name) on an AxisEnv lacking that name; using vmap/xmap/pmap with an out_axes/in_axes axis_name not declared; using an axis name in collectives (e.g., psum, all_gather) that is not among the mapped axis names of the enclosing transformation.

Common situations: Typos in axis_name strings; using psum(..., axis_name='i') inside vmap without axis_name='i'; refactoring renamed axes in one place but not the collective call; mixing pmap axis_name with vmap-only contexts.

Related errors

• unbound axis name: {axis_name}
• axis_index_groups not supported in vmap collectives. Please
• all_gather_reduced is a Varying -> Reduced collective. This
• Please file an issue at https://github.com/jax-ml/jax/issues
• unreduced_psum_scatter is a Unreduced -> Varying collective.

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