{"record":{"id":"ba759483896c1a6f","repo":"jax-ml/jax","slug":"indexing-barrier-with-transforms-not-supported-i","errorCode":null,"errorMessage":"Indexing barrier with {transforms} not supported in GPU interpret mode","messagePattern":"Indexing barrier with (.+?) not supported in GPU interpret mode","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":130,"sourceCode":"\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:\n    return None\n  if not hasattr(transforms, \"__len__\") or len(transforms) != 1:\n    raise NotImplementedError(\n        f\"Indexing barrier with {transforms} not supported in GPU interpret\"\n        \" mode\"\n    )\n  if not isinstance(transforms[0], indexing.NDIndexer):\n    raise ValueError(f\"Expected an `NDIndexer`, but got {transforms[0]}\")\n  if len(transforms[0].indices) == 1:\n    return transforms[0].indices[0]\n  return tuple(transforms[0].indices)\n\n\ndef _get_barrier_allocation_key_from_inval(\n    inval, transforms_treedef, transforms_leaves\n) -> jax.Array:\n  # `inval` is expected to correspond to a barrier. Since we are interpreting,\n  # `inval` will in fact contain the allocation key (which is a Jax array) for\n  # the barrier.\n  allocation_key_as_array = inval\n","sourceCodeStart":112,"sourceCodeEnd":148,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L112-L148","documentation":"When a barrier is indexed in interpret mode, only a single indexing transform (an NDIndexer) is supported. Multiple transforms or non-index transforms cannot be mapped to a barrier allocation key by the interpreter.","triggerScenarios":"Passing a barrier reference with two or more transforms (e.g. slicing then unswizzling), or a transforms tuple without __len__, into copy_gmem_to_smem, barrier_arrive, barrier_wait, or tcgen05 ops during interpretation.","commonSituations":"Composing multiple reference transformations on a barrier allocated as an array of barriers; chaining slicing with other transforms generated by compiler plumbing.","solutions":["Index the barrier once with a single NDIndexer (slice/index) and pass it directly","Allocate separate scalar barriers instead of an array of barriers when per-element access is needed","Simplify barrier transforms before interpret mode runs"],"exampleFix":"// before\nop(barrier[0:2][1])  # chained transforms\n// after\nop(barrier[1])      # single NDIndexer","handlingStrategy":"type-guard","validationCode":"# ensure barrier refs carry at most one indexing transform\nassert len(ref.transforms or ()) <= 1, 'multiple transforms on barrier unsupported in interpret mode'","typeGuard":"def is_simple_indexed_barrier(ref) -> bool:\n    ts = getattr(ref, 'transforms', ()) or ()\n    return len(ts) <= 1 and all(\n        type(t).__name__ == 'NDIndexer' for t in ts)","tryCatchPattern":null,"preventionTips":["Index barriers exactly once with a plain slice/index","Prefer scalar barriers over barrier arrays in interpret-mode tests"],"tags":["jax","mosaic-gpu","barrier","transforms"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}