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` or an appropriate JAX device mesh.

What it means

Raised when a kernel uses a named axis (e.g. in a collective like p.sum over an axis_name) but no axis names exist: neither a plgpu.Mesh from pl.core_map nor a JAX device mesh is attached to the module context.

Source

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

    assert axis_name in axis_names.grid
    idx = axis_names.grid.index(axis_name)
    return block_ids[gpu_dialect.Dimension(idx)]


@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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Wrap the kernel launch in pl.core_map with a plgpu.Mesh declaring the axis
Alternatively run under a jax.sharding.Mesh with the matching axis name
Double-check the axis name string matches the mesh declaration exactly

Example fix

// before
out = kernel(x)  # kernel uses p.sum(x, 'i')
// after
mesh = plgpu.Mesh(axis_names=('i',))
out = pl.core_map(kernel, mesh=mesh, ...)(x)

Defensive patterns

Strategy: validation

Validate before calling

if not (gpu_mesh_axes or jax_mesh_axes):
    raise ValueError('Named collectives need a plgpu.Mesh or jax.sharding.Mesh')

Prevention

Set up meshes before launching kernels with collectives
Name axes consistently across codebase

When it happens

Trigger: Calling axis-named collectives (p.sum axis_name=..., ppermute, etc.) in a kernel launched without pl.core_map+plgpu.Mesh and outside any jax.sharding.Mesh context.

Common situations: Converting single-device kernels to sharded/multi-device ones without setting up a Mesh; forgetting that GPU mesh axes require plgpu.Mesh, not just a JAX Mesh in some paths.

Related errors

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