{"record":{"id":"275695d024aebfa2","repo":"jax-ml/jax","slug":"requesting-collective-allocations-but-no-explicit","errorCode":null,"errorMessage":"Requesting collective allocations, but no explicit thread axis specified.","messagePattern":"Requesting collective allocations, but no explicit thread axis specified\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":93,"sourceCode":"    space: mosaic_gpu_core.MemorySpace | None,\n):\n  # TODO(nrink): Support more memory spaces.\n  if space is not None and space not in [\n      mosaic_gpu_core.MemorySpace.GMEM,\n      mosaic_gpu_core.MemorySpace.SMEM,\n      mosaic_gpu_core.MemorySpace.TMEM,\n      mosaic_gpu_core.MemorySpace.REGS,\n  ]:\n    raise NotImplementedError(f\"Unsupported memory space: {space}\")\n\n\ndef _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      )","sourceCodeStart":75,"sourceCodeEnd":111,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L75-L111","documentation":"The interpreter validates collective (cross-thread) allocations: if any axis is marked collective but the mesh has no named thread axis (mesh.thread_name is None/empty), there is no defined group of threads to make the allocation collective over, so it raises ValueError.","triggerScenarios":"Creating a Block/scratch with collective axes (is_collective_by_thread_cluster_axis contains True) while the kernel's Mesh lacks thread_name, in interpret mode.","commonSituations":"Sharing scratch across warps without declaring the thread axis in the Mesh; meshes built only with grid/cluster dims; refactor removing thread_name while keeping collective buffers.","solutions":["Add thread_name (and num_threads) to the Mesh: Mesh(grid=..., thread_name='tid', num_threads=N)","If the allocation shouldn't be shared, drop the collective axis from the allocation","Align collective axis flags with the actual mesh axis layout (cluster axes first, thread axis last)"],"exampleFix":"# before\nmesh = Mesh(grid=(1,), num_threads=128)  # no thread_name\n# after\nmesh = Mesh(grid=(1,), thread_name='tid', num_threads=128)","handlingStrategy":"validation","validationCode":"if any(is_collective_by_thread_cluster_axis):\n    assert mesh is not None and mesh.thread_name, 'collective alloc needs a named thread axis'","typeGuard":"def mesh_has_thread_axis(mesh) -> bool:\n    return mesh is not None and bool(mesh.thread_name)","tryCatchPattern":null,"preventionTips":["Always set thread_name/num_threads on Mesh when sharing scratch across threads","Audit collective flags when refactoring meshes"],"tags":["pallas","mosaic-gpu","collective","mesh","thread-axis","config"],"backgroundTag":"grid-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}