{"record":{"id":"75322030d465b4bf","repo":"jax-ml/jax","slug":"no-cluster-axes-found","errorCode":null,"errorMessage":"No cluster axes found.","messagePattern":"No cluster axes found\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3193,"sourceCode":"    collective = False\n\n  with predicate_ctx:\n    mgpu.dialect.tcgen05_commit_arrive(\n        barrier_ref.as_barrier_memref(), collective=collective\n    )\n  return []\n\n\ndef _collective_mma_predicate(ctx: lowering.LoweringRuleContext,\n                              collective_axis: str) -> ir.Value:\n  \"\"\"Computes a predicate to run only on the leader block.\"\"\"\n  cluster_axis = lowering._resolve_cluster_axis(\n      ctx.module_ctx.axis_names, collective_axis)\n  if cluster_axis != gpu_dialect.Dimension(0):\n    # Note: resolve_cluster_axis checks if axis_names exists.\n    assert ctx.module_ctx.axis_names is not None\n    if len(ctx.module_ctx.axis_names.cluster) <= 1:\n      raise ValueError(\"No cluster axes found.\")\n    minormost_cluster_axis = ctx.module_ctx.axis_names.cluster[0]\n    raise ValueError(\n        \"Can only perform collective MMA along minormost cluster axis. \"\n        f\"Got {collective_axis}, expected {minormost_cluster_axis}.\")\n  index = ir.IndexType.get()\n  is_leader_block = arith_dialect.cmpi(\n      arith_dialect.CmpIPredicate.eq,\n      mgpu.utils.cluster_idx(cluster_axis), mgpu.c(0, index))\n  return is_leader_block\n\n\ncommit_tmem_p = jax_core.Primitive(\"commit_tmem\")\ncommit_tmem_p.multiple_results = True\n\n\n@commit_tmem_p.def_effectful_abstract_eval\ndef _commit_tmem_abstract_eval():\n  return (), {gpu_core._memory_effect}","sourceCodeStart":3175,"sourceCodeEnd":3211,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3175-L3211","documentation":"For collective MMA (multicast across a cluster), the kernel must be launched with more than one cluster axis; the lowering needs a minormost cluster axis to elect a leader block. If ctx.module_ctx.axis_names.cluster has length <= 1 there is nothing to multicast over, so it errors.","triggerScenarios":"Calling collective MMA (or async copy to TMEM) with a collective_axis while the kernel's cluster (num_clusters/cluster axes) has 1 or fewer axes, e.g. running without ClusterSpec/collective axes.","commonSituations":"Developing a single-block kernel then enabling collective_axis; forgetting to declare the cluster dimension in the Pallas grid interpretation.","solutions":["Declare at least 2 cluster axes in the kernel launch (e.g. set the cluster/collective dimension so axis_names.cluster has length > 1)","Only pass collective_axis when the kernel is actually launched as a multi-block cluster"],"exampleFix":"// before\nkernel = pl.make_interp... # no cluster\ntcgen05_mma(..., collective_axis='io')\n// after\n# launch with cluster axes, e.g. Interpretation with num clusters > 1 along an axis\nkernel = pl.pallas_call(..., in_specs=[BlockSpec(..., collective_axis='io')])","handlingStrategy":"validation","validationCode":"assert collective_axis is None or len(kernel_cluster_axes) > 1","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Declare multi-axis clusters before enabling collective MMA","Gate collective_axis on launch configuration"],"tags":["jax","pallas","tcgen05","cluster","collective-mma"],"backgroundTag":"missing-required-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}