jax-ml/jax · error · NameError

Found an unbound axis name: {name}. To fix this, please call

Error message

Found an unbound axis name: {name}. To fix this, please call {api_name} under `jax.shard_map`.

What it means

Collectives like psum, all_gather, ppermute, all_to_all require a bound mesh axis name. _check_axis_names looks the name up in the current axis environment and raises NameError if no vmap/shard_map frame introduced it.

Source

Thrown at jax/_src/lax/parallel.py:1033

  check_unreduced_args([aval], axes, 'psum')
  out_aval = ShapedArray(
      lax._reduce_op_shape_rule(aval, axes=pos_axes), aval.dtype,
      sharding=lax._reduce_op_sharding_rule(aval, axes=pos_axes))
  return out_aval, {core.NamedAxisEffect(axis) for axis in named_axes}

# TODO(yashkatariya): Replace this with _psum_invariant_abstract_eval
def _pmin_pmax_abstract_eval(name, aval, *, axes, axis_index_groups):
  if not config._check_vma.value:
    return _allreduce_effectful_abstract_eval(
        aval, axes=axes, axis_index_groups=axis_index_groups)
  return _psum_invariant_abstract_eval(name, aval, axes=axes)

def _check_axis_names(axes, api_name):
  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))
  axis_env = core.get_axis_env()
  for name in named_axes:
    if not axis_env.axis_exists(name):
      raise NameError(
          f"Found an unbound axis name: {name}. To fix this, please call"
          f" {api_name} under `jax.shard_map`.")

def _lower_reducer_into_block(ctx, prim, scalar_aval, block):
  with ir.InsertionPoint(block):
    lower_reducer = mlir.lower_fun(prim.bind, multiple_results=False)
    reducer_ctx = ctx.replace(
        primitive=None, avals_in=[scalar_aval] * 2, avals_out=[scalar_aval]
    )
    out_nodes = lower_reducer(reducer_ctx, *block.arguments)
    flat_out_nodes, _ = mlir.ir_tree_registry.flatten(out_nodes)
    if isinstance(block.owner, func_dialect.FuncOp):
      func_dialect.return_(flat_out_nodes)
    else:
      hlo.return_(flat_out_nodes)


def _build_reducer_func_op(ctx, prim, aval_in):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the call in jax.shard_map with a mesh that declares the axis name
2. Check for a typo in the axis name against Mesh(axis_names=...)
3. If using vmap-only, wrap with jax.vmap(..., axis_name='i')

Example fix

// before
y = jax.lax.psum(x, 'i')
// after
y = jax.shard_map(lambda b: jax.lax.psum(b, 'i'), mesh)(x)

def assert_axis_bound(mesh, name):
    assert name in mesh.axis_names, f'{name} not in mesh axes {mesh.axis_names}'

# and always call collectives inside shard_map

try:
    y = jax.shard_map(f, mesh)(x)
except NameError as e:
    if 'unbound axis name' in str(e): raise RuntimeError(f'mesh misconfigured: {e}')
    raise

When it happens

Trigger: Calling jax.lax.psum(x, 'i') (or all_to_all/all_gather/psend) outside jax.shard_map or without the axis being declared on the mesh/shard_map in_axes.

Common situations: Running distributed code eagerly (no shard_map) during debugging; axis name typo ('dp' vs 'data'); mesh declared but function not wrapped in shard_map.

Related errors

• Mesh must be provided for shard_map with checkify.
• When `check_vma=True` on `jax.shard_map`, `manual_axis_type`
• The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
• The context mesh {ctx_mesh} should match the mesh passed to
• shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.

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