jax-ml/jax · error · LookupError

No axis names are available. Make sure you are using `pl.cor

Error message

No axis names are available. Make sure you are using `pl.core_map` with a `plgpu.Mesh`.

What it means

Raised when a Pallas kernel calls a collective op with a named cluster axis but the kernel was launched without cluster axis names. Cluster axis names are only populated when the kernel is mapped with pl.core_map using a plgpu.Mesh that declares cluster axes.

Source

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

  else:
    raise NotImplementedError(f"Unsupported dtype {x_aval.dtype}")
  kind = vector_dialect.CombiningKind.MUL
  return _reduce_lowering_rule_wg(ctx, kind, acc, x, axes)


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,)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Launch the kernel with pl.core_map and pass a plgpu.Mesh(axis_names=...) declaring the cluster axes
2. Give the mesh axis the same name you use in the collective call
3. If you don't need cluster collectives, remove the axis-name-based collective from the kernel

Example fix

// before
kernel = p.map(step, in_specs=..., out_specs=...)
// after
mesh = plgpu.Mesh(axis_names=('row','))
kernel = pl.core_map(step, mesh=mesh, in_specs=..., out_specs=...)

assert mesh is not None and mesh.axis_names, 'cluster collectives require plgpu.Mesh with axis_names'

try:
    pl.core_map(step, mesh=mesh, ...)
except LookupError as e:
    raise ValueError(f'Bad launch config: {e}') from e

When it happens

Trigger: Calling e.g. mgpu.cluster_id / collective primitives with an axis_name inside a kernel launched via p.map or without plgpu.Mesh, so ctx.module_ctx.axis_names for the cluster is None/empty.

Common situations: Migrating kernels from single-CTA launches to multi-CTA clusters but forgetting to wrap with pl.core_map and a plgpu.Mesh; using TPU-style mesh APIs on GPU.

Related errors

• Unknown cluster axis {axis_name}, available axes: {[*axis_na
• dims and idxs must have the same length
• cluster= must be at most 3D, got {self}.
• Memory space {self.memory_space} is not supported by mesh {s
• You can't use two different TensorCoreMeshes.

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