jax-ml/jax · error · ValueError

Cannot `wait_send` on a local copy.

Error message

Cannot `wait_send` on a local copy.

What it means

wait_send blocks until a REMOTE DMA's source-side semaphore is signaled (it waits on the send completing on the other device). On a purely local copy there is no send side to wait on, so calling wait_send on a non-remote DMA is a user error and raises ValueError.

Source

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

    )

  def wait(self):
    if self.is_remote:
      self.wait_send()
    self.wait_recv()

  def wait_recv(self):
    self._used = True
    flat_args, tree = self._get_args_and_tree()
    dma_wait_p.bind(
        *flat_args, tree=tree, device_id_type=self.device_id_type,
        insert_dummy_device=False, is_wait_send=False
    )

  def wait_send(self):
    self._used = True
    if not self.is_remote:
      raise ValueError("Cannot `wait_send` on a local copy.")
    # We swap src and dst since by default dma_wait_p waits on the dst_sem
    # TODO(rdyro): Update the lowering to use `is_wait_send` instead of
    # swapping src and dst.
    flat_args, tree = self._get_args_and_tree(
        swap_src_and_dst=True,
    )
    dma_wait_p.bind(
        *flat_args, tree=tree, device_id_type=self.device_id_type,
        insert_dummy_device=self.is_remote,
        is_wait_send=True,
    )


def _dma_flatten(*args):
  flat_tree = ft.flatten(args)
  return flat_tree.vals, flat_tree.tree

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Branch on is_remote: only call wait_send when dma.is_remote is true
2. Use the local equivalent (e.g. wait on the destination semaphore) for local copies
3. Gate send-waiting behind whether src_sem/device_id were provided at DMA construction

Example fix

# before
for dma in dmas:
  dma.wait_send()

# after
for dma in dmas:
  if dma.is_remote:
    dma.wait_send()

def safe_wait_send(dma):
  if getattr(dma, 'is_remote', False):
    dma.wait_send()
  # local copies: no send side to wait on

try:
  dma.wait_send()
except ValueError:
  pass  # local copy; nothing to wait on

When it happens

Trigger: Calling descriptor.wait_send() on a DMA started without a device_id (i.e., is_remote is False). Common when generic pipelining code (scalar_subcore_fn body) unconditionally waits on sends for all DMAs including local ones.

Common situations: Writing a pipeline that issues many DMAs and calls wait_send on each; enabling send-wait only needed for cross-die copies but applying it uniformly; forgetting that local copies only support wait_receive-style waiting.

Related errors

• Either both or neither `src_sem` and `device_id` can be set.
• dma_start not implemented in LoJAX yet.
• DMA start priority must be non-negative: {priority}
• DMA source/destination/semaphore arguments must be Refs.
• Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}

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