{"record":{"id":"36c7b889744372db","repo":"jax-ml/jax","slug":"attempting-to-operate-on-barrier-without-indexing","errorCode":null,"errorMessage":"Attempting to operate on barrier without indexing, but `num_barriers = {num_barriers}`","messagePattern":"Attempting to operate on barrier without indexing, but `num_barriers = (.+?)`","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py","lineNumber":165,"sourceCode":"  allocation_key_as_array = inval\n\n  # Assert to check internal consistency: `allocation_key_as_array` should be\n  # at least a 2D array, and the size of the last dimension is 5 (which matches the\n  # fields count of HostAllocationKey).\n  assert len(allocation_key_as_array.shape) >= 2\n  assert (\n      allocation_key_as_array.shape[-1:]\n      == gpu_callbacks.HostAllocationKey.shape_and_dtype().shape\n  )\n  num_barriers = math.prod(allocation_key_as_array.shape[:-1])\n\n  index = _get_index_for_barrier_allocation_key(\n      transforms_treedef, transforms_leaves\n  )\n\n  if index is None:\n    if num_barriers != 1:\n      raise ValueError(\n          \"Attempting to operate on barrier without indexing, but\"\n          f\" `num_barriers = {num_barriers}`\"\n      )\n    idx = (0,) * (len(allocation_key_as_array.shape) - 1)\n    return allocation_key_as_array[idx]\n  else:\n    return allocation_key_as_array[index]\n\n\ndef _get_num_threads_sharing_collective_allocation(\n    axes_dims: tuple[int, ...],\n    is_last_thread_axis_collective: bool,\n) -> int:\n  \"\"\"Returns the number of threads that share a collective allocation.\"\"\"\n  if is_last_thread_axis_collective:\n    return axes_dims[-1]\n  else:\n    return 1","sourceCodeStart":147,"sourceCodeEnd":183,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py#L147-L183","documentation":"In interpret mode, barriers allocated as arrays (num_barriers > 1) must be indexed when used. Operating on a multi-barrier allocation without an index is ambiguous and therefore rejected.","triggerScenarios":"Calling barrier_arrive, barrier_wait, copy_gmem_to_smem, tcgen05_mma, or tcgen05_commit_arrive with an un-indexed barrier ref where the barrier was allocated with num_barriers > 1.","commonSituations":"Allocating a barrier array for pipelined copies (multiple buffers) but forgetting to index it per pipeline stage; refactoring from a single barrier to multiple barriers.","solutions":["Index the barrier: barrier_arrive(barrier[i], ...) with the stage index","Allocate with num_barriers=1 if only one barrier is needed","Check loop bodies that use barrier arrays to ensure the loop induction variable indexes the barrier"],"exampleFix":"// before\nbarrier_arrive(barrier, sem)\n// after\nbarrier_arrive(barrier[i], sem)","handlingStrategy":"validation","validationCode":"if num_barriers > 1:\n    assert barrier_transforms_present, f'index barrier: num_barriers={num_barriers}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always index barrier arrays per pipeline stage","Use num_barriers=1 unless pipelining truly requires multiple"],"tags":["jax","mosaic-gpu","barrier","indexing"],"backgroundTag":"missing-required-index","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}