jax-ml/jax · error · ValueError

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

Error message

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

What it means

Like the destination semaphore, the optional source semaphore of dma_start must be scalar (() - shaped); a non-empty shape means per-element signaling which the hardware does not support for source-side completion.

Source

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

      _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,
  )

def _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,
                      context: jax_core.JaxprPpContext,
                      settings: jax_core.JaxprPpSettings):
  invars = eqn.invars
  tree = eqn.params["tree"]
  priority = eqn.params["priority"]
  add = eqn.params["add"]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a scalar () semaphore for src_sem
Select one element from a semaphore array if you manage several
Reallocate the semaphore with shape ()

Example fix

# before
dma_start(src, dst, dst_sem, src_sem=sem_vector, device_id=dev)

# after
sem_scalar = sem_vector[i]  # scalar ref
dma_start(src, dst, dst_sem, src_sem=sem_scalar, device_id=dev)

Defensive patterns

Strategy: validation

Validate before calling

if src_sem is not None:
  assert src_sem.shape == (), f"src semaphore must be scalar, got {src_sem.shape}"

Type guard

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

Prevention

One scalar semaphore per signal path
Mirror the () shape check for both src and dst semaphores

When it happens

Trigger: dma_start with src_sem whose ref shape is non-empty, e.g. shape=(num_signals,) or (1,).

Common situations: Allocating a vector of semaphores to track multiple sends and passing the whole vector as src_sem; symmetric copy-paste from the destination semaphore that also had a wrong shape.

Related errors

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