jax-ml/jax · error · NotImplementedError

Non-decrementing wait is not supported.

Error message

Non-decrementing wait is not supported.

What it means

The Pallas TPU interpreter only supports semaphore waits that decrement the semaphore's value; semaphore_wait with decrement=False (a non-destructive read of the semaphore) is not implemented and raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1736

            ctx.axis_indices)
        token = callback.io_callback(
            functools.partial(semaphore_signal, source_info=eqn.source_info),
            TOKEN_SHAPE_DTYPE,
            token,
            ctx.device_id,
            ctx.local_core_id,
            state_discharge.transform_array(sem, sem_transforms),
            inc,
            target_device_id,
            core_index,
        )
        out = []

      elif prim is primitives.semaphore_wait_p:
        sem, sem_transforms, value, decrement = (
            jax.tree.unflatten(eqn.params['args_tree'], deferred_invals()))
        if not decrement:
          raise NotImplementedError('Non-decrementing wait is not supported.')
        token = callback.io_callback(
            semaphore_wait,
            TOKEN_SHAPE_DTYPE,
            token,
            ctx.device_id,
            ctx.local_core_id,
            state_discharge.transform_array(sem, sem_transforms),
            value,
        )
        out = []

      else:
        if ctx.interpret_params.skip_floating_point_ops and all(
            interpret_utils.is_float(ovar.aval.dtype) for ovar in eqn.outvars
        ):
          # Skip `prim.bind` since `prim` only produces floating-point values.
          # It is safe to populate `out` with avals since mapping `write` over
          #  `out` below only relies on the shape and dtype (for writing

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use decrementing waits (decrement=True), ensuring each wait pairs with a signal/lock post
2. Restructure spin-wait patterns to signal/wait discipline instead of non-destructive polling
3. Fall back to real TPU execution if non-decrementing waits are required

Example fix

# before
pltpu.semaphore_wait(sem, value=0, decrement=False)
# after
pltpu.semaphore_wait(sem, value=0)  # decrement=True default

assert decrement is not False, 'non-decrementing semaphore wait unsupported in TPU interpret mode'

null

try:
    interpret_run(kernel)
except NotImplementedError as e:
    if 'Non-decrementing wait' in str(e):
        run_on_tpu_hardware(kernel)

When it happens

Trigger: Calling primitives.semaphore_wait / pltpu.semaphore_wait with decrement=False inside a kernel run in TPU interpret mode.

Common situations: Polling a semaphore without consuming it (common in spin-wait patterns); GPU-style atomic semantics ported to TPU; probing semaphore state for debugging.

Related errors

• masked load_p
• run_scoped_p with collective axes is not supported
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• Aliasing of scalar prefetch arguments is not currently suppo
• Kernel input {j} in HBM but does not have trivial BlockSpec.

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