{"record":{"id":"11ff86ce1aeba24f","repo":"jax-ml/jax","slug":"barrier-does-not-support-arbitrary-transforms","errorCode":null,"errorMessage":"Barrier does not support arbitrary transforms","messagePattern":"Barrier does not support arbitrary transforms","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1515,"sourceCode":"            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            )\n\n          idx = arith_dialect.muli(idx, lowering._as_index(stride))  # pylint: disable=protected-access\n          if base_index is None:\n            base_index = idx\n          else:\n            base_index = arith_dialect.addi(base_index, idx)\n      case _:\n        raise ValueError(\"Barrier does not support arbitrary transforms\")\n  return base_index\n\n\nbarrier_arrive_p = jax_core.Primitive(\"barrier_arrive\")\nbarrier_arrive_p.multiple_results = True\n\n\n@barrier_arrive_p.def_effectful_abstract_eval\ndef _barrier_arrive_abstract_eval(barrier, *args, **params):\n  del args, params  # Unused.\n  _check_ref(barrier, \"barrier\", gpu_core.SMEM)\n  return (), {gpu_core._memory_effect}\n\n\ndef _barrier_arrive_pp_eqn(\n    eqn: jax_core.JaxprEqn,\n    context: jax_core.JaxprPpContext,\n    settings: jax_core.JaxprPpSettings,","sourceCodeStart":1497,"sourceCodeEnd":1533,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1497-L1533","documentation":"_get_barrier_base_index only handles transform lists that begin with an NDIndexer (plain indexing). Any other transform (swizzles, tilings, transposes, compositions) applied to a barrier ref cannot be lowered to an address, so it is rejected.","triggerScenarios":"Applying non-indexing transforms (e.g. via TransformedRef with swizzle/transpose) to a barrier ref and then arriving/waiting on it.","commonSituations":"Reusing generic ref-transformation helpers that add layout transforms to all refs in a kernel, accidentally including barriers/accumulators.","solutions":["Remove transforms from barrier refs; keep them plain SMEM-allocated buffers.","Apply transforms only to data refs, using separate code paths for barriers."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu import indexing\nassert all(isinstance(t, indexing.NDIndexer) for t in transforms), 'barriers take only indexing transforms'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep barrier refs free of layout transforms; transform only data refs."],"tags":["mosaic-gpu","pallas","barrier","transforms"],"backgroundTag":"unsupported-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}