{"record":{"id":"ce23f258dea35c5c","repo":"jax-ml/jax","slug":"unknown-cluster-axis-axis-name-available-axes","errorCode":null,"errorMessage":"Unknown cluster axis {axis_name}, available axes: {[*axis_names.cluster]}","messagePattern":"Unknown cluster axis (.+?), available axes: (.+?)","errorType":"exception","errorClass":"LookupError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3388,"sourceCode":"\ndef _block_id(ctx: LoweringRuleContext, dim: gpu_dialect.Dimension) -> ir.Value:\n  result = gpu_dialect.block_id(dim)\n  cluster_size = ctx.launch_ctx.cluster_size\n  if math.prod(cluster_size) == 1 or cluster_size[dim.value] == 1:\n    return result\n  # We scale the grid in the presence of clusters, so we need to scale the\n  # block ID back here.\n  return arith_dialect.divui(result, _as_index(cluster_size[dim.value]))\n\n\ndef _resolve_cluster_axis(axis_names: _AxisNames | None, axis_name: Hashable):\n  if not axis_names:\n    raise LookupError(\n        \"No axis names are available. Make sure you are using `pl.core_map`\"\n        \" with a `plgpu.Mesh`.\"\n    )\n  if not axis_names or axis_name not in axis_names.cluster:\n    raise LookupError(\n        f\"Unknown cluster axis {axis_name}, available axes:\"\n        f\" {[*axis_names.cluster]}\"\n    )\n  return gpu_dialect.Dimension(axis_names.cluster.index(axis_name))\n\n\ndef _is_block_local_scope(collective_axes: CollectiveAxesType,\n                          axis_names: _AxisNames):\n  \"\"\"Returns whether the collective axes represents a block scope.\"\"\"\n  if axis_names.wg is None:\n    return not collective_axes\n  else:\n    return collective_axes == (axis_names.wg,)\n\n\ndef _is_global_scope(collective_axes: CollectiveAxesType,\n                     axis_names: _AxisNames):\n  \"\"\"Returns whether the collective axes represents a GPU global scope.\"\"\"","sourceCodeStart":3370,"sourceCodeEnd":3406,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3370-L3406","documentation":"The axis name passed to a cluster collective does not match any axis declared in the plgpu.Mesh used with pl.core_map. The error lists the available cluster axes so you can correct the name.","triggerScenarios":"Calling a cluster collective with axis_name='x' when the mesh was declared with axis_names=('row','col') — a typo or mismatched naming between kernel body and launch site.","commonSituations":"Renaming mesh axes in one place but not the other; copy-pasting kernels between pipelines with different mesh layouts.","solutions":["Use one of the axes listed in the error message in the collective call","Or rename the mesh axis in plgpu.Mesh to match the kernel's axis_name","Centralize axis names in shared constants to avoid drift"],"exampleFix":"// before\nmesh = plgpu.Mesh(axis_names=('row',))\n# in kernel: mgpu.cluster_id('x')\n// after\nmesh = plgpu.Mesh(axis_names=('row',))\n# in kernel: mgpu.cluster_id('row')","handlingStrategy":"type-guard","validationCode":"assert axis_name in mesh.axis_names, f'{axis_name=} not in {mesh.axis_names}'","typeGuard":"def is_valid_cluster_axis(name: str, mesh) -> bool:\n    return name in mesh.axis_names","tryCatchPattern":null,"preventionTips":["Define axis names as module constants","Validate axis names at kernel-build time"],"tags":["jax","pallas","mosaic-gpu","cluster","axis-name","mesh"],"backgroundTag":"axis-name-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}