jax-ml/jax · error · ValueError

Semaphore value shape {ref_value.shape} does not match aval

Error message

Semaphore value shape {ref_value.shape} does not match aval shape {ref_aval.shape}

What it means

During state discharge of a semaphore, the runtime value's shape must match the ref aval's shape (or be scalar). A mismatch means the stored semaphore value was shaped differently from what the ref abstraction declares.

Source

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

  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}"
    )


semaphore_read_p = jax_core.Primitive("semaphore_read")
semaphore_read_p.multiple_results = False


def semaphore_read(sem_or_view) -> jax_typing.Array:
  """Reads the value of a semaphore.

  Args:
    sem_or_view: A Ref (or view) representing a semaphore.

  Returns:
    A scalar Array containing the value of the semaphore.
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure semaphores are scalar () Refs with the correct semaphore dtype
2. Avoid custom ref transforms on semaphores
3. Update JAX / report a bug with a reproducer if shapes are correct

assert ref_value.shape == ref_aval.shape or ref_value.shape == ()

When it happens

Trigger: Interpret/discharge paths where the semaphore value stored has a shape differing from the Ref's declared shape, e.g. after improper allocation or transform of the semaphore.

Common situations: Mostly an internal invariant violation seen in interpret mode with incorrectly shaped semaphores; rare user-facing unless custom ref transforms are used.

Related errors

• Block shape for {origin} (= {block_shape}) must have the sam
• Stacking only supported when the block size along the stack
• Sum of sizes {n} must be equal to dimension {axis} of the op
• Block size must be a multiple of the input size. Got block {
• Every block dimension must be either a multiple or factor of

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