jax-ml/jax · error · ValueError

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

Error message

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

What it means

Pallas semaphores must be scalar-shaped (). This error fires when the semaphore Ref (after applying any ref transforms/indexers) has a non-empty shape.

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)

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: {dst_sem_shape}
• Cannot signal on a non-()-shaped semaphore: {src_sem_shape}
• Non-decrementing wait is not supported.
• Semaphore {sem_id} occurs as both fixed-id and internal.
• Seed key_data must be 1D.

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