jax-ml/jax · error · LookupError

Axis does not refer to a GPU mesh axis (available axes: )…

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

Raised in the lax.axis_index lowering for Mosaic GPU: the requested axis_name is neither one of the GPU mesh axes nor one of the enclosing JAX mesh axes visible to the kernel, so no index can be produced and a ValueError is raised. The input at fault is the axis_name argument of jnp/lax axis_index.

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…
• mixes JAX mesh and Pallas mesh grid axes
• {non_mesh_axes}
• Unknown cluster axis
• all_gather_reduced is a Varying -> Reduced collective. This…

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

Appendix: 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)