What causes "Cannot {name} on a non-semaphore Ref: {sem_aval}" in jax-ml/jax?

Allocate the semaphore with the Pallas allocation API and pass the Ref itself, not a read value

jax-ml/jax · error · ValueError

Cannot {name} on a non-semaphore Ref: {sem_aval}

Error message

Cannot {name} on a non-semaphore Ref: {sem_aval}

What it means

Semaphore operations (signal, wait, read) in Pallas require the target to be a state.AbstractRef holding a semaphore. This error fires when the value passed is not a Ref at all (e.g., a plain array or a different reference type).

Source

Thrown at jax/_src/pallas/primitives.py:890


class DeviceIdType(enum.Enum):
  MESH = "mesh"
  LOGICAL = "logical"


def check_sem_avals(
    sem_aval, sem_transforms_avals, name, allowed_semaphore_types=None
):
  if allowed_semaphore_types is None:
    allowed_semaphore_types = {
        pallas_core.semaphore,
        pallas_core.barrier_semaphore,
        # For interpret mode.
        pallas_core.SEMAPHORE_INTERPRET_DTYPE,
    }
  if not isinstance(sem_aval, state.AbstractRef):
    raise ValueError(f"Cannot {name} on a non-semaphore Ref: {sem_aval}")
  sem_shape = sem_aval.shape
  if sem_transforms_avals:
    sem_shape = sem_transforms_avals[-1].get_indexer_shape()
  if sem_shape:
    raise ValueError(f"Cannot {name} on a non-()-shaped semaphore: {sem_shape}")
  sem_dtype = sem_aval.dtype
  if not any(
      jnp.issubdtype(sem_dtype, sem_type)
      for sem_type in allowed_semaphore_types
  ):
    raise ValueError(
        f"Must {name} semaphores of the following types:"
        f" {allowed_semaphore_types}. Got {sem_dtype}."
    )


def _transform_semaphore(ref_value, transforms, ref_aval):
  """Helper function for indexing into a semaphore during state_discharge."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Allocate the semaphore with the Pallas allocation API and pass the Ref itself, not a read value
Check that the argument is a state.Ref (isinstance(ref, state.Ref)) before calling semaphore ops

Example fix

// before
semaphore_wait(sem[...], 1)  # passed value not Ref
// after
semaphore_wait(sem, 1)  # pass the Ref

Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src import state
assert isinstance(sem, state.Ref), f"expected a Ref semaphore, got {type(sem)}"

Type guard

def is_semaphore_ref(sem) -> bool:
    from jax._src import state
    return isinstance(sem, state.Ref)

Try / catch

try:
    semaphore_wait(sem, 1)
except ValueError as e:
    if "non-semaphore Ref" in str(e):
        raise TypeError("pass the semaphore Ref, not its value") from e
    raise

Prevention

Always pass the Ref itself to semaphore ops
Wrap semaphore helpers with Ref assertions in tests

When it happens

Trigger: Passing a non-Ref value (jnp array, TracedArray, or a Ref of non-semaphore abstraction) to semaphore_signal / semaphore_wait / semaphore_read.

Common situations: Forgetting to allocate the semaphore via pallas state and passing its current value instead of the Ref; unwrapping refs too early in kernel code.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

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