jax-ml/jax · error · NotImplementedError

Unhandled transforms for semaphore_read: {transforms}

Error message

Unhandled transforms for semaphore_read: {transforms}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4555

      else:
        current_indices[i] = current_start_index + _ensure_idx_fa(dim_indexer)
        removed_dimensions.add(i)
  return indexing.NDIndexer(
      indices=tuple(current_indices),
      shape=root_shape,
      int_indexer_shape=(),
  )


@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Warpgroup)
def _semaphore_read_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):
  sem, transforms = tree_util.tree_unflatten(args_tree, args)
  sem_aval, transform_avals = tree_util.tree_unflatten(args_tree, ctx.avals_in)
  assert isinstance(sem_aval, state_types.AbstractRef)
  sem, _, transforms = _handle_transforms(ctx, sem_aval, sem, transform_avals, transforms)
  if transforms:
    raise NotImplementedError(f"Unhandled transforms for semaphore_read: {transforms}")
  sem_ptr = mgpu.utils.memref_ptr(sem)
  i32_ty = ir.IntegerType.get_signless(32)
  result = llvm_dialect.inline_asm(
      i32_ty,
      [sem_ptr],
      "ld.acquire.sys.u32 $0,[$1];",
      "=r,l",
      has_side_effects=True,
  )
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    return _ensure_fa(result, jnp.int32)
  return result


@contextlib.contextmanager
def _wrap_in_custom_primitive_if_wg(
    ctx: LoweringRuleContext, operands: Sequence[ir.Value]
):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder operations so semaphore_read operates on the untransformed semaphore ref
2. Perform transforms on data buffers, not the semaphore reference
3. Check JAX version — transform handling for semaphores may be added later

assert not transforms, f"semaphore_read with transforms: {transforms}"

When it happens

Trigger: Calling semaphore_read on a ref that was discharges/scattered or otherwise wrapped in transforms that were not fully consumed before the read.

Common situations: Using semaphores together with swizzled or untiled buffers inside a plgpu kernel.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Non-decrementing wait is not supported.
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4e3460092af6ae90. Report an issue: GitHub.