{"record":{"id":"4e3460092af6ae90","repo":"jax-ml/jax","slug":"unhandled-transforms-for-semaphore-read-transfor","errorCode":null,"errorMessage":"Unhandled transforms for semaphore_read: {transforms}","messagePattern":"Unhandled transforms for semaphore_read: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4555,"sourceCode":"      else:\n        current_indices[i] = current_start_index + _ensure_idx_fa(dim_indexer)\n        removed_dimensions.add(i)\n  return indexing.NDIndexer(\n      indices=tuple(current_indices),\n      shape=root_shape,\n      int_indexer_shape=(),\n  )\n\n\n@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Warpgroup)\ndef _semaphore_read_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):\n  sem, transforms = tree_util.tree_unflatten(args_tree, args)\n  sem_aval, transform_avals = tree_util.tree_unflatten(args_tree, ctx.avals_in)\n  assert isinstance(sem_aval, state_types.AbstractRef)\n  sem, _, transforms = _handle_transforms(ctx, sem_aval, sem, transform_avals, transforms)\n  if transforms:\n    raise NotImplementedError(f\"Unhandled transforms for semaphore_read: {transforms}\")\n  sem_ptr = mgpu.utils.memref_ptr(sem)\n  i32_ty = ir.IntegerType.get_signless(32)\n  result = llvm_dialect.inline_asm(\n      i32_ty,\n      [sem_ptr],\n      \"ld.acquire.sys.u32 $0,[$1];\",\n      \"=r,l\",\n      has_side_effects=True,\n  )\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    return _ensure_fa(result, jnp.int32)\n  return result\n\n\n@contextlib.contextmanager\ndef _wrap_in_custom_primitive_if_wg(\n    ctx: LoweringRuleContext, operands: Sequence[ir.Value]\n):","sourceCodeStart":4537,"sourceCodeEnd":4573,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4537-L4573","documentation":"Raised when lowering a semaphore_read operation that still carries leftover transforms (e.g. UntilingTransform, UnswizzleRef) after transform handling. The Mosaic GPU backend cannot read semaphores through transformed/swizzled references.","triggerScenarios":"Calling semaphore_read on a ref that was discharges/scattered or otherwise wrapped in transforms that were not fully consumed before the read.","commonSituations":"Using semaphores together with swizzled or untiled buffers inside a plgpu kernel.","solutions":["Reorder operations so semaphore_read operates on the untransformed semaphore ref","Perform transforms on data buffers, not the semaphore reference","Check JAX version — transform handling for semaphores may be added later"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert not transforms, f\"semaphore_read with transforms: {transforms}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Read semaphores only from plain, untransformed refs"],"tags":["pallas","mosaic-gpu","semaphore","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}