jax-ml/jax · error · LookupError

Axis {axis_name} does not refer to a GPU mesh axis (availabl

Error message

Axis {axis_name} does not refer to a GPU mesh axis (available axes: {[*gpu_axis_names]}) or a JAX mesh axis (available axes: {[*jax_axis_names]})

What it means

The axis name used in a collective is not found among the GPU mesh axes (from plgpu.Mesh) nor the JAX mesh axes (from the enclosing jax.sharding.Mesh). The message lists both sets of valid names.

Source

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

@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(lax.axis_index_p, *gpu_core.LANExWARP_SEMANTICS)
@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(lax.axis_index_p, *gpu_core.WGxWARP_SEMANTICS)
def _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if axis_name == ctx.module_ctx.warp_axis_name:
      w_idx = mgpu.warp_idx(sync=True)
      i32 = ir.IntegerType.get_signless(32)
      return arith_dialect.remui(w_idx, _ir_constant(4, i32))
  gpu_axis_names = ctx.module_ctx.axis_names
  jax_axis_names = getattr(ctx.module_ctx.mesh_info, "axis_names", ())
  if gpu_axis_names is None and not jax_axis_names:
    raise LookupError(
        "No axis names are available. Make sure you are using `pl.core_map`"
        " with a `plgpu.Mesh` or an appropriate JAX device mesh."
    )
  if axis_name not in itertools.chain(gpu_axis_names or (), jax_axis_names):
    raise LookupError(
        f"Axis {axis_name} does not refer to a GPU mesh axis (available axes:"
        f" {[*gpu_axis_names]}) or a JAX mesh axis (available axes:"
        f" {[*jax_axis_names]})"
    )
  if axis_name in jax_axis_names:
    jax_mesh = ctx.module_ctx.mesh_info
    assert jax_mesh is not None
    device_id = ctx.launch_ctx.device_id()
    jax_mesh_shape = jax_mesh.mesh_shape
    axis_index = jax_axis_names.index(axis_name)
    i32 = ir.IntegerType.get_signless(32)
    axis_size = _ir_constant(jax_mesh_shape[axis_index], i32)
    minor_divisor = _ir_constant(
        np.prod(jax_mesh_shape[axis_index + 1 :], dtype=np.int32), i32
    )
    return arith_dialect.remsi(arith_dialect.divsi(device_id, minor_divisor), axis_size)

  # We already checked that the axis is in scope and it wasn't a JAX mesh axis.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a name from either list in the error message
2. Add the axis to the plgpu.Mesh or the jax.sharding.Mesh as appropriate
3. Keep axis names in shared enums/constants used by both launch and kernel code

Example fix

// before
with jax.sharding.Mesh(devs, ('row',)):
  kernel uses p.sum(x, 'col')
// after
with jax.sharding.Mesh(devs, ('row',)):
  kernel uses p.sum(x, 'row')

assert axis_name in set(gpu_axis_names or ()) | set(jax_axis_names or ())

def axis_available(name, gpu_axes, jax_axes) -> bool:
    return name in set(gpu_axes or ()) | set(jax_axes or ())

When it happens

Trigger: p.sum(x, axis_name='data') inside a kernel where the plgpu.Mesh declares ('row',) and no JAX mesh with 'data' is active — any name not in either list triggers this.

Common situations: Typos in axis names; assuming a JAX Mesh axis is automatically visible to GPU collectives when the kernel was mapped with a differently-named plgpu.Mesh.

Related errors

• No axis names are available. Make sure you are using `pl.cor
• {non_mesh_axes}
• Unknown cluster axis {axis_name}, available axes: {[*axis_na
• {axis_name} mixes JAX mesh and Pallas mesh grid axes
• unbound axis name: {axis_name}

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