{"record":{"id":"86364d918b2e0f75","repo":"jax-ml/jax","slug":"axis-axis-name-does-not-refer-to-a-gpu-mesh-axis","errorCode":null,"errorMessage":"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]})","messagePattern":"Axis (.+?) does not refer to a GPU mesh axis \\(available axes: (.+?)\\) or a JAX mesh axis \\(available axes: (.+?)\\)","errorType":"exception","errorClass":"LookupError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3460,"sourceCode":"@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(lax.axis_index_p, *gpu_core.LANExWARP_SEMANTICS)\n@register_lowering_rule(lax.axis_index_p, mgpu.LoweringSemantics.Warpgroup)\n@register_lowering_rule(lax.axis_index_p, *gpu_core.WGxWARP_SEMANTICS)\ndef _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    if axis_name == ctx.module_ctx.warp_axis_name:\n      w_idx = mgpu.warp_idx(sync=True)\n      i32 = ir.IntegerType.get_signless(32)\n      return arith_dialect.remui(w_idx, _ir_constant(4, i32))\n  gpu_axis_names = ctx.module_ctx.axis_names\n  jax_axis_names = getattr(ctx.module_ctx.mesh_info, \"axis_names\", ())\n  if gpu_axis_names is None and not jax_axis_names:\n    raise LookupError(\n        \"No axis names are available. Make sure you are using `pl.core_map`\"\n        \" with a `plgpu.Mesh` or an appropriate JAX device mesh.\"\n    )\n  if axis_name not in itertools.chain(gpu_axis_names or (), jax_axis_names):\n    raise LookupError(\n        f\"Axis {axis_name} does not refer to a GPU mesh axis (available axes:\"\n        f\" {[*gpu_axis_names]}) or a JAX mesh axis (available axes:\"\n        f\" {[*jax_axis_names]})\"\n    )\n  if axis_name in jax_axis_names:\n    jax_mesh = ctx.module_ctx.mesh_info\n    assert jax_mesh is not None\n    device_id = ctx.launch_ctx.device_id()\n    jax_mesh_shape = jax_mesh.mesh_shape\n    axis_index = jax_axis_names.index(axis_name)\n    i32 = ir.IntegerType.get_signless(32)\n    axis_size = _ir_constant(jax_mesh_shape[axis_index], i32)\n    minor_divisor = _ir_constant(\n        np.prod(jax_mesh_shape[axis_index + 1 :], dtype=np.int32), i32\n    )\n    return arith_dialect.remsi(arith_dialect.divsi(device_id, minor_divisor), axis_size)\n\n  # We already checked that the axis is in scope and it wasn't a JAX mesh axis.","sourceCodeStart":3442,"sourceCodeEnd":3478,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3442-L3478","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use a name from either list in the error message","Add the axis to the plgpu.Mesh or the jax.sharding.Mesh as appropriate","Keep axis names in shared enums/constants used by both launch and kernel code"],"exampleFix":"// before\nwith jax.sharding.Mesh(devs, ('row',)):\n  kernel uses p.sum(x, 'col')\n// after\nwith jax.sharding.Mesh(devs, ('row',)):\n  kernel uses p.sum(x, 'row')","handlingStrategy":"type-guard","validationCode":"assert axis_name in set(gpu_axis_names or ()) | set(jax_axis_names or ())","typeGuard":"def axis_available(name, gpu_axes, jax_axes) -> bool:\n    return name in set(gpu_axes or ()) | set(jax_axes or ())","tryCatchPattern":null,"preventionTips":["Use enums for axis names","Fail fast in kernel constructors when axis unknown"],"tags":["jax","pallas","mesh","collective","axis-name"],"backgroundTag":"axis-name-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}