jax-ml/jax · error · ValueError

DMA source semaphore must be a Ref.

Error message

DMA source semaphore must be a Ref.

What it means

The optional source semaphore of dma_start must, like the destination, be a state ref — the abstract eval checks src_sem_aval (when not None) is a state.AbstractRef or TransformedRef. Passing a plain array or other object raises this ValueError.

Source

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

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

def _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,
                      context: jax_core.JaxprPpContext,
                      settings: jax_core.JaxprPpSettings):
  invars = eqn.invars

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the source semaphore as a Ref obtained from the kernel's ref arguments
2. Verify nothing unwrapped the ref (e.g. sem[()] reads) before dma_start
3. If you did not mean to signal the source, pass src_sem=None (and device_id=None)

Example fix

# before
dma_start(src, dst, dst_sem, src_sem=sem_val, device_id=dev)  # sem_val is an Array

# after
dma_start(src, dst, dst_sem, src_sem=sem_ref, device_id=dev)

if src_sem is not None:
  from jax._src import state
  assert isinstance(getattr(src_sem, 'aval', src_sem), (state.AbstractRef, state.TransformedRef))

def is_ref(v) -> bool:
  from jax._src import state
  return isinstance(getattr(v, 'aval', v), (state.AbstractRef, state.TransformedRef))

When it happens

Trigger: dma_start(..., src_sem=<jax.Array>, device_id=dev) where src_sem is not a Ref.

Common situations: Adding remote-copy signaling later and passing a semaphore value (e.g. an int read from a sem) instead of the sem ref; inconsistent plumbing where some call sites hold refs and others hold arrays.

Related errors

• DMA source/destination/semaphore arguments must be Refs.
• Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}
• Cannot signal on a non-()-shaped semaphore: {src_sem_shape}
• Semaphore {sem_id} occurs as both fixed-id and internal.
• Either both or neither `src_sem` and `device_id` can be set.

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