jax-ml/jax · error · ValueError

Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}

Error message

Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}

What it means

DMA completion semaphores are single boolean flags in hardware; signaling requires a scalar ()-shaped semaphore ref. If dst_sem has any non-empty shape the abstract eval rejects it, since there is no per-element signaling.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:360

  return []
dma_start_p.to_lojax = _dma_start_to_lojax

@dma_start_p.def_effectful_abstract_eval
def _dma_start_abstract_eval(*args, tree, device_id_type, priority, add):
  if priority < 0:
    raise ValueError(f"DMA start priority must be non-negative: {priority}")
  src_ref_aval, dst_ref_aval, dst_sem_aval, src_sem_aval, device_id_aval = (
      _dma_unflatten(tree, args)
  )
  if not all(
      isinstance(x, (state.AbstractRef, state.TransformedRef))
      for x in [src_ref_aval, dst_ref_aval, dst_sem_aval]
  ):
    raise ValueError(
        "DMA source/destination/semaphore arguments must be Refs.")
  dst_sem_shape = dst_sem_aval.shape
  if dst_sem_shape:
    raise ValueError(
        f"Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}"
    )
  if src_sem_aval is not None:
    if not isinstance(src_sem_aval, (state.AbstractRef, state.TransformedRef)):
      raise ValueError("DMA source semaphore must be a Ref.")
    src_sem_shape = src_sem_aval.shape
    if src_sem_shape:
      raise ValueError(
          f"Cannot signal on a non-()-shaped semaphore: {src_sem_shape}"
      )
  return [], _get_dma_effects(
      src_ref_aval,
      dst_ref_aval,
      dst_sem_aval,
      src_sem_aval,
      device_id_aval,
      device_id_type,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Allocate the destination semaphore with shape () and pass that scalar ref
If you have an array of semaphores, pass a single element (e.g. sem_ref[i] via slicing that yields a scalar ref)
Use one scalar semaphore per DMA or share one scalar semaphore across DMAs

Example fix

# before
sem = alloc_semaphore(shape=(1,))
dma_start(src, dst, sem)

# after
sem = alloc_semaphore(shape=())
dma_start(src, dst, sem)

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

def is_scalar_sem(sem_ref) -> bool:
  return sem_ref.shape == ()

Prevention

Allocate semaphores with shape ()
Slice one element when using semaphore arrays

When it happens

Trigger: dma_start where the destination semaphore ref was allocated with a shape, e.g. shape=(1,) or shape=(num_blocks,).

Common situations: Allocating a semaphore array for multiple DMAs and passing the whole array instead of one element; mirroring the shape of the buffer being copied onto the semaphore by copy-paste; shape defaults like (1,) instead of ().

Related errors

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