jax-ml/jax · error · ValueError

Cannot on a non-()-shaped semaphore

Error message

Cannot {name} on a non-()-shaped semaphore: {sem_shape}

What it means

Guard in check_sem_avals, shared by semaphore signal and wait: semaphores are scalar synchronization primitives, so after applying any transforms the effective semaphore shape must be (). A semaphore ref (or view) with a non-empty shape fails, and the offending shape is reported.

Solutions

1. Allocate semaphores as scalar Refs: one Ref per semaphore, shape ()
2. If multiple semaphores are needed, allocate multiple scalar Refs (e.g., via tree of refs), not one shaped Ref
3. Verify no transform chain leaves a non-empty shape on the ref

Example fix

// before
sem = alloc_buffer(shape=(1,), dtype=semaphore_dtype)
// after
sem = alloc_buffer(shape=(), dtype=semaphore_dtype)

assert sem_aval.shape == (), f"semaphore must be scalar, got {sem_aval.shape}"

def is_scalar_ref(ref) -> bool:
    return getattr(getattr(ref, 'aval', ref), 'shape', None) == ()

When it happens

Trigger: Allocating a semaphore Ref with shape (n,) instead of (), or applying an indexer/transform that yields a non-scalar shape, then calling semaphore_signal/wait/read.

Common situations: Trying to create an array of semaphores; applying slicing that leaves dimensionality on the semaphore ref.

Related errors

• Cannot signal on a non-()-shaped semaphore
• Cannot signal on a non-()-shaped semaphore
• axis is out of range.
• Cannot on a non-semaphore Ref
• `device_id`s must be an int32 value, but got

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

Appendix: source

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

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."""
  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:

View on GitHub (pinned to 1e1c6a8fc0)