jax-ml/jax · error · ValueError

DMA start priority must be non-negative: {priority}

Error message

DMA start priority must be non-negative: {priority}

What it means

dma_start accepts a priority used to order concurrent DMA operations on TPU hardware; the abstract eval validates that priority is a non-negative integer since negative priorities have no hardware meaning and would corrupt queue ordering.

Source

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

    if jax_core.typeof(_get_ref(src_sem)).is_high:
      raise NotImplementedError("dma_start not implemented in LoJAX yet.")
  src_ref_aval.inner_aval.dma_start(
      src_ref,
      dst_ref,
      src_sem,
      dst_sem,
      device_id=device_id,
      priority=priority,
      device_id_type=device_id_type,
      add=add,
  )
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use non-negative priorities, remapping 'highest priority' to 0 and increasing numbers for lower priority
2. Clamp/validate priority >= 0 before the call
3. Replace -1 sentinel defaults with 0 or None

Example fix

# before
dma_start(src, dst, sem, priority=-1)

# after
# invert convention: most urgent = 0
dma_start(src, dst, sem, priority=max(0, -priority))

priority = max(0, priority)
assert priority >= 0, f"DMA priority must be >= 0, got {priority}"

When it happens

Trigger: dma_start(..., priority=k) with k < 0 — e.g. priority derived from a subtraction, a config default of -1, or an inverted priority convention.

Common situations: Porting code where lower number meant higher priority (so 0, -1, -2 were used); exposing priority as a user knob without clamping; default sentinel -1 leaking into the call.

Related errors

• Either both or neither `src_sem` and `device_id` can be set.
• shift must be non-negative.
• stride and stride_axis must be both specified or not.
• stride must be non-negative.
• expected axis and stride_axis are different.

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