jax-ml/jax · error · LookupError

Unknown cluster axis {axis_name}, available axes: {[*axis_na

Error message

Unknown cluster axis {axis_name}, available axes: {[*axis_names.cluster]}

What it means

The axis name passed to a cluster collective does not match any axis declared in the plgpu.Mesh used with pl.core_map. The error lists the available cluster axes so you can correct the name.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3388

def _block_id(ctx: LoweringRuleContext, dim: gpu_dialect.Dimension) -> ir.Value:
  result = gpu_dialect.block_id(dim)
  cluster_size = ctx.launch_ctx.cluster_size
  if math.prod(cluster_size) == 1 or cluster_size[dim.value] == 1:
    return result
  # We scale the grid in the presence of clusters, so we need to scale the
  # block ID back here.
  return arith_dialect.divui(result, _as_index(cluster_size[dim.value]))


def _resolve_cluster_axis(axis_names: _AxisNames | None, axis_name: Hashable):
  if not axis_names:
    raise LookupError(
        "No axis names are available. Make sure you are using `pl.core_map`"
        " with a `plgpu.Mesh`."
    )
  if not axis_names or axis_name not in axis_names.cluster:
    raise LookupError(
        f"Unknown cluster axis {axis_name}, available axes:"
        f" {[*axis_names.cluster]}"
    )
  return gpu_dialect.Dimension(axis_names.cluster.index(axis_name))


def _is_block_local_scope(collective_axes: CollectiveAxesType,
                          axis_names: _AxisNames):
  """Returns whether the collective axes represents a block scope."""
  if axis_names.wg is None:
    return not collective_axes
  else:
    return collective_axes == (axis_names.wg,)


def _is_global_scope(collective_axes: CollectiveAxesType,
                     axis_names: _AxisNames):
  """Returns whether the collective axes represents a GPU global scope."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use one of the axes listed in the error message in the collective call
2. Or rename the mesh axis in plgpu.Mesh to match the kernel's axis_name
3. Centralize axis names in shared constants to avoid drift

Example fix

// before
mesh = plgpu.Mesh(axis_names=('row',))
# in kernel: mgpu.cluster_id('x')
// after
mesh = plgpu.Mesh(axis_names=('row',))
# in kernel: mgpu.cluster_id('row')

assert axis_name in mesh.axis_names, f'{axis_name=} not in {mesh.axis_names}'

def is_valid_cluster_axis(name: str, mesh) -> bool:
    return name in mesh.axis_names

When it happens

Trigger: Calling a cluster collective with axis_name='x' when the mesh was declared with axis_names=('row','col') — a typo or mismatched naming between kernel body and launch site.

Common situations: Renaming mesh axes in one place but not the other; copy-pasting kernels between pipelines with different mesh layouts.

Related errors

• No axis names are available. Make sure you are using `pl.cor
• {non_mesh_axes}
• dims and idxs must have the same length
• cluster= must be at most 3D, got {self}.
• No axis names are available. Make sure you are using `pl.cor

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