{"record":{"id":"749cd77485835fed","repo":"jax-ml/jax","slug":"barrier-does-not-support-slice-with-stride-1","errorCode":null,"errorMessage":"Barrier does not support slice with `stride != 1`","messagePattern":"Barrier does not support slice with `stride != 1`","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1489,"sourceCode":"      collective_axes=collective_axes,\n      leader_tracked=leader_tracked,\n  )\n  return None\n\n\ndef _get_barrier_base_index(aval, transforms) -> ir.Value | None:\n  if not transforms:\n    return None\n  strides = list(pallas_utils.strides_from_shape(aval.shape))\n  base_index: ir.Value | None = None\n  while transforms:\n    match transforms:\n      case [indexing.NDIndexer() as indexer, *transforms]:\n        num_int_idxs = 0\n        for i, (idx, stride) in enumerate(zip(indexer.indices, strides[:])):\n          if isinstance(idx, indexing.Slice):\n            if idx.stride != 1:\n              raise NotImplementedError(\n                  \"Barrier does not support slice with `stride != 1`\"\n              )\n            idx = idx.start\n          else:\n            # This dimension is absent for any corresponding `NDIndexer`s, so\n            # we remove the corresponding stride.\n            strides.pop(i - num_int_idxs)\n            num_int_idxs += 1\n\n          if isinstance(\n              idx, (int, ir.Value, mgpu.FragmentedArray, literals.TypedNdArray)\n          ):\n            idx = lowering._as_index(idx)  # pylint: disable=protected-access\n          else:\n            raise ValueError(\n                \"Barrier can only be indexed with integers or slices, got\"\n                f\" {idx}\"\n            )","sourceCodeStart":1471,"sourceCodeEnd":1507,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1471-L1507","documentation":"When computing a barrier's base address, the indexing slice must have stride 1 because the lowering computes a single linear base offset. Slices with a step (e.g. barrier_ref[0::2]) cannot be represented.","triggerScenarios":"Indexing a barrier/mma accumulator ref with a strided slice, e.g. barrier[0:4:2], in any op that lowers through _get_barrier_base_index (barrier.arrive/wait/test, async_store_smem, tcgen05_mma, copies).","commonSituations":"Using Python slice syntax with a step on barrier refs while trying to select every other barrier in an array of barriers.","solutions":["Index barriers individually or with contiguous slices (stride 1), e.g. barrier[i] or barrier[0:4].","Restructure so you don't need strided barrier selection (allocate separate barriers per subgroup)."],"exampleFix":"# before\nbarriers[0:8:2].arrive()\n# after\nfor i in range(0, 8, 2):\n    barriers[i].arrive()","handlingStrategy":"validation","validationCode":"def check_slice(idx):\n    if isinstance(idx, slice):\n        assert idx.step in (None, 1), f'stride {idx.step} unsupported on barriers'\nfor i in indices: check_slice(i)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never use stepped slices on barrier/accumulator refs.","Loop over individual barriers instead of strided selection."],"tags":["mosaic-gpu","pallas","barrier","indexing","slice"],"backgroundTag":"unsupported-slice-stride","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}