{"record":{"id":"5bcea337855c6708","repo":"jax-ml/jax","slug":"no-axis-names-are-available-make-sure-you-are-usi-5bcea3","errorCode":null,"errorMessage":"No axis names are available. Make sure you are using `pl.core_map` with a `plgpu.Mesh` or an appropriate JAX device mesh.","messagePattern":"No axis names are available\\. Make sure you are using `pl\\.core_map` with a `plgpu\\.Mesh` or an appropriate JAX device mesh\\.","errorType":"exception","errorClass":"LookupError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3455,"sourceCode":"    assert axis_name in axis_names.grid\n    idx = axis_names.grid.index(axis_name)\n    return block_ids[gpu_dialect.Dimension(idx)]\n\n\n@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(","sourceCodeStart":3437,"sourceCodeEnd":3473,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3437-L3473","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nout = kernel(x)  # kernel uses p.sum(x, 'i')\n// after\nmesh = plgpu.Mesh(axis_names=('i',))\nout = pl.core_map(kernel, mesh=mesh, ...)(x)","handlingStrategy":"validation","validationCode":"if not (gpu_mesh_axes or jax_mesh_axes):\n    raise ValueError('Named collectives need a plgpu.Mesh or jax.sharding.Mesh')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Set up meshes before launching kernels with collectives","Name axes consistently across codebase"],"tags":["jax","pallas","mesh","collective","axis-name"],"backgroundTag":"missing-mesh-axis-names","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}