{"record":{"id":"2f2737cf1862db00","repo":"jax-ml/jax","slug":"scoped-allocation-must-have-the-thread-axis-in-its","errorCode":null,"errorMessage":"Scoped allocation must have the thread axis in its collective axes.","messagePattern":"Scoped allocation must have the thread axis in its collective axes\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":109,"sourceCode":"  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:\n    return None\n  transforms = jax.tree.unflatten(transforms_treedef, transforms_leaves)\n\n  if not transforms:","sourceCodeStart":91,"sourceCodeEnd":127,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L91-L127","documentation":"Mosaic GPU interpret mode requires scoped allocations to be collective along the thread axis. Since the interpreter executes one thread's view of the block, a scoped allocation that is not collective on the thread axis would have undefined semantics, so it is rejected.","triggerScenarios":"Calling run_scoped with collective_axes that omit the thread axis name of the current WarpMesh (for non-REGS memory spaces, which route through _raise_if_unsupported_collective_axes).","commonSituations":"Passing an empty collective_axes tuple or only cluster axis names; refactoring kernels and accidentally dropping the thread axis from collective_axes.","solutions":["Add the thread axis name (from mesh.axis_names[-1]) to collective_axes","Pass collective_axes=None/empty only if the API permits non-collective scope, otherwise include thread axis","Check WarpMesh construction to confirm the intended thread axis name"],"exampleFix":"// before\nrun_scoped(buf, collective_axes=())\n// after\nrun_scoped(buf, collective_axes=('thread',))","handlingStrategy":"validation","validationCode":"thread_axis = mesh.axis_names[-1] if mesh is not None else None\nif thread_axis is not None:\n    assert thread_axis in (collective_axes or ()), 'thread axis must be collective'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always derive collective_axes from the mesh rather than hardcoding","Include the thread axis name whenever scoping SMEM/TMEM buffers"],"tags":["jax","mosaic-gpu","run-scoped","collective-axes"],"backgroundTag":"invalid-parameter-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}