jax-ml/jax · error · ValueError

Must {name} semaphores of the following types: {allowed_sema

Error message

Must {name} semaphores of the following types: {allowed_semaphore_types}. Got {sem_dtype}.

What it means

Semaphore dtype validation: the semaphore Ref must have a dtype compatible with the allowed semaphore types (pallas semaphore, barrier_semaphore, or the interpret dtype).

Source

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

    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."""
  if ref_value.shape == ref_aval.shape:
    return state_discharge.transform_array(ref_value, transforms)
  elif len(ref_value.shape) == 0:
    return ref_value
  else:
    raise ValueError(
        f"Semaphore value shape {ref_value.shape} does not match aval shape"
        f" {ref_aval.shape}"
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Allocate the semaphore with the proper Pallas semaphore dtype (e.g., pallas_core.semaphore / SEMAPHORE_INTERPRET_DTYPE in interpret mode)
Use the library's semaphore allocation helpers rather than raw alloc with a custom dtype

Example fix

// before
sem = alloc((), jnp.float32)
// after
sem = alloc((), semaphore_dtype)  # pallas semaphore dtype

Defensive patterns

Strategy: validation

Validate before calling

import jnp
from jax.experimental.pallas import pallas_core
assert any(jnp.issubdtype(sem.dtype, t) for t in (pallas_core.semaphore, pallas_core.barrier_semaphore))

Prevention

Use library semaphore allocation helpers
Never allocate semaphores with arbitrary dtypes

When it happens

Trigger: Allocating a Ref with an arbitrary dtype (e.g., float32) and using it as a semaphore in signal/wait/read.

Common situations: Creating semaphores with jnp.zeros(..., dtype=jnp.float32) instead of the proper semaphore dtype constant; mixing up buffer and semaphore allocation.

Related errors

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