{"record":{"id":"79e0d391bd1e95c0","repo":"jax-ml/jax","slug":"expected-a-single-barrier-got-a-barrier-reference","errorCode":null,"errorMessage":"Expected a single barrier, got a barrier reference with shape {transformed_barrier.shape}","messagePattern":"Expected a single barrier, got a barrier reference with shape (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":586,"sourceCode":"      flat_ref_transforms_avals\n  )\n  barrier_transform_avals = barrier_transforms_treedef.unflatten(\n      flat_barrier_transforms_avals\n  )\n  transformed_ref = pallas_core.TransformedRef(ref, ref_transform_avals)\n  if src.shape != transformed_ref.shape:\n    raise TypeError(\n        f\"The stored value has shape {src.shape}, but the target reference has\"\n        f\" shape {transformed_ref.shape}\"\n    )\n  if src.dtype != transformed_ref.dtype:\n    raise TypeError(\n        f\"The stored value has dtype {src.dtype}, but the target reference has\"\n        f\" dtype {transformed_ref.dtype}\"\n    )\n  transformed_barrier = pallas_core.TransformedRef(barrier, barrier_transform_avals)\n  if transformed_barrier.size != 1:\n    raise TypeError(\n        \"Expected a single barrier, got a barrier reference with shape\"\n        f\" {transformed_barrier.shape}\"\n    )\n\n  effs = {gpu_core._memory_effect, state.WriteEffect(1)}\n  return (), effs\n\n\n@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Lane)\n@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Warpgroup)\ndef _async_store_smem_lowering(\n    ctx: lowering.LoweringRuleContext,\n    src,\n    ref,\n    barrier,\n    cluster_idx,\n    *flat_transforms,\n    ref_transforms_treedef,","sourceCodeStart":568,"sourceCodeEnd":604,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L568-L604","documentation":"async_store_smem takes a barrier reference used to signal transaction completion; the abstract eval requires the transformed barrier to hold exactly one element (transformed_barrier.size == 1). A barrier with a non-scalar shape cannot be passed to the arrive_expect_tx operation.","triggerScenarios":"Passing an array-shaped barrier (e.g. a (4,) barrier vector) or slicing the barrier so more than one element remains, to async_store_smem.","commonSituations":"Allocating one barrier per warp/iteration as a vector and passing the whole array instead of a single element; mis-indexing barrier buffers with block indices meant for the value.","solutions":["Index the barrier down to a single element, e.g. barrier[0] or barrier[i], before passing it","Allocate barriers as scalar-shaped SMEM (pl.SMEM((), plint.barrier_dtype))","Use the dedicated async_barrier/barrier APIs if per-warp barriers are needed"],"exampleFix":"# before\nasync_store_smem(smem, x, barriers)\n# after\nasync_store_smem(smem, x, barriers[0])","handlingStrategy":"validation","validationCode":"assert barrier_ref.size == 1, 'async_store_smem needs a scalar barrier'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate scalar barriers; index barrier arrays to a single element before passing"],"tags":["jax","pallas","mosaic-gpu","barrier","async-store"],"backgroundTag":"invalid-argument-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}