{"record":{"id":"0b7bf738c9573a5b","repo":"jax-ml/jax","slug":"collective-allocations-along-cluster-axes-are-not","errorCode":null,"errorMessage":"Collective allocations along cluster axes are not supported.","messagePattern":"Collective allocations along cluster axes are not supported\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":105,"sourceCode":"def _raise_if_unsupported_collective_axes(\n    mesh: mosaic_gpu_core.Mesh | None,\n    is_collective_by_thread_cluster_axis: tuple[bool, ...],\n):\n  if not mesh or not mesh.thread_name:\n    if any(is_collective_by_thread_cluster_axis):\n      raise ValueError(\n          \"Requesting collective allocations, but no explicit thread axis\"\n          \" specified.\"\n      )\n  else:\n    # Note that the leading entries in `is_collective_by_thread__cluster_axis`\n    # correspond to the cluster axes, while the last entry corresponds to the\n    # thread axis within a block.\n    *is_collective_by_cluster_axis, is_thread_axis_collective = (\n        is_collective_by_thread_cluster_axis\n    )\n    if any(is_collective_by_cluster_axis):\n      raise ValueError(\n          \"Collective allocations along cluster axes are not supported.\"\n      )\n    if not is_thread_axis_collective:\n      raise ValueError(\n          \"Scoped allocation must have the thread axis in its collective axes.\"\n      )\n\n\n# TODO(nrink): Try unifying this function with `_extract_barrier_slice_base`\n# from `jax._src.pallas.mosaic_gpu.primitives`.\ndef _get_index_for_barrier_allocation_key(\n    transforms_treedef, transforms_leaves,\n) -> indexing.DimIndexer | None:\n  # TODO(nrink): The working out of `transforms` and the returned index below\n  # may need tidying up. Specifically, GPU interpret mode should correctly\n  # support legal ways to index into barriers. (Here, 'legal' is to be read as\n  # 'allowed by the Pallas GPU semantics'.)\n  if transforms_treedef is None:","sourceCodeStart":87,"sourceCodeEnd":123,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L87-L123","documentation":"In Mosaic GPU interpret mode, scoped memory allocations (run_scoped) may only be collective along the thread axis, not along cluster axes of the warp mesh. The interpreter simulates each thread sequentially and cannot model allocations shared across cluster dimensions, so it raises this ValueError when any cluster axis is marked collective.","triggerScenarios":"Calling run_scoped with a WarpMesh whose collective_axes include one or more cluster axes (i.e. anything other than the single thread axis), while interpreting an MGPU kernel.","commonSituations":"Porting TPU-style collective SMEM allocations to Mosaic GPU kernels; using collective_axes computed from a multi-dimension mesh without filtering to the thread axis.","solutions":["Restrict collective_axes in run_scoped to only the thread axis name of your WarpMesh","Redesign the kernel to allocate per-warp and synchronize with barriers instead of cluster-wide scoped allocations","Run the kernel on real hardware compilation instead of interpret mode if cluster-collective allocation is required"],"exampleFix":"// before\nrun_scoped(alloc, collective_axes=('row', 'col'))  # mesh cluster axes\n// after\nrun_scoped(alloc, collective_axes=('thread',))    # only thread axis","handlingStrategy":"validation","validationCode":"mesh = kernel_launch_mesh\ncollective = set(collective_axes or ())\ncluster_axes = set(mesh.axis_names[:-1]) if mesh else set()\nassert not (collective & cluster_axes), 'cluster-axis collective allocations unsupported'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate collective_axes against mesh.axis_names before run_scoped","Keep a lint/test that runs kernels in interpret mode in CI to catch unsupported patterns early"],"tags":["jax","mosaic-gpu","interpret-mode","collective-allocations"],"backgroundTag":"unsupported-feature-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}