jax-ml/jax · error · ValueError

Either both or neither `src_sem` and `device_id` can be set.

Error message

Either both or neither `src_sem` and `device_id` can be set.

What it means

A remote DMA descriptor couples the source semaphore with a device id: signaling a remote source requires knowing which device to signal. The dataclass __post_init__ enforces that src_sem and device_id are either both set or both None, and XOR of the None-ness triggers this error.

Source

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

mlir.register_lowering(roll_p, _roll_lowering_rule)


@dataclasses.dataclass
class AsyncCopyDescriptor:
  src_ref: Any
  dst_ref: Any
  dst_sem: Any
  src_sem: Any | None
  device_id: MultiDimDeviceId | IntDeviceId | None
  device_id_type: primitives.DeviceIdType = primitives.DeviceIdType.MESH
  _used: bool = dataclasses.field(
      default=False, init=False, compare=False, hash=False
  )

  def __post_init__(self):
    if (self.src_sem is None) ^ (self.device_id is None):
      raise ValueError("Either both or neither `src_sem` and `device_id` "
                       "can be set.")

  def __del__(self):
    if not self._used:
      # Exceptions in ``__del__`` are ignored, so logging is our only option.
      logging.error(
          "AsyncCopyDescriptor was not used."
          " Did you mean to call `start` or `wait` on it?"
      )

  @property
  def is_remote(self):
    return self.src_sem is not None

  def _get_args_and_tree(
      self,
      swap_src_and_dst: bool = False,
      device_id: MultiDimDeviceId | IntDeviceId | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide both src_sem and device_id together
2. If you do not need source-side signaling, pass src_sem=None and device_id=None and rely on the destination semaphore only
3. Add a unit assertion that the two fields' presence matches in your kernel launch path

Example fix

# before
dma = Descriptor(src_sem=src_sem, device_id=None, ...)

# after
dma = Descriptor(src_sem=src_sem, device_id=peer_device_id, ...)
# or
dma = Descriptor(src_sem=None, device_id=None, ...)

assert (src_sem is None) == (device_id is None), "src_sem and device_id are coupled"
kwargs = {} if src_sem is None else dict(src_sem=src_sem, device_id=device_id)

When it happens

Trigger: Constructing the DMA descriptor (the object returned by dma_start-related setup in mosaic.primitives) with src_sem set but device_id=None, or vice versa.

Common situations: Enabling source-completion signaling for cross-chip copies but forgetting the peer device id; refactoring where device_id plumbing is dropped behind a config flag while src_sem remains.

Related errors

• Cannot `wait_send` on a local copy.
• DMA start priority must be non-negative: {priority}
• shift must be non-negative.
• stride and stride_axis must be both specified or not.
• stride must be non-negative.

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