{"record":{"id":"436e11e7a44cb81a","repo":"jax-ml/jax","slug":"dma-source-semaphore-must-be-a-ref","errorCode":null,"errorMessage":"DMA source semaphore must be a Ref.","messagePattern":"DMA source semaphore must be a Ref\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":365,"sourceCode":"  if priority < 0:\n    raise ValueError(f\"DMA start priority must be non-negative: {priority}\")\n  src_ref_aval, dst_ref_aval, dst_sem_aval, src_sem_aval, device_id_aval = (\n      _dma_unflatten(tree, args)\n  )\n  if not all(\n      isinstance(x, (state.AbstractRef, state.TransformedRef))\n      for x in [src_ref_aval, dst_ref_aval, dst_sem_aval]\n  ):\n    raise ValueError(\n        \"DMA source/destination/semaphore arguments must be Refs.\")\n  dst_sem_shape = dst_sem_aval.shape\n  if dst_sem_shape:\n    raise ValueError(\n        f\"Cannot signal on a non-()-shaped semaphore: {dst_sem_shape}\"\n    )\n  if src_sem_aval is not None:\n    if not isinstance(src_sem_aval, (state.AbstractRef, state.TransformedRef)):\n      raise ValueError(\"DMA source semaphore must be a Ref.\")\n    src_sem_shape = src_sem_aval.shape\n    if src_sem_shape:\n      raise ValueError(\n          f\"Cannot signal on a non-()-shaped semaphore: {src_sem_shape}\"\n      )\n  return [], _get_dma_effects(\n      src_ref_aval,\n      dst_ref_aval,\n      dst_sem_aval,\n      src_sem_aval,\n      device_id_aval,\n      device_id_type,\n  )\n\ndef _dma_start_pp_eqn(eqn: jax_core.JaxprEqn,\n                      context: jax_core.JaxprPpContext,\n                      settings: jax_core.JaxprPpSettings):\n  invars = eqn.invars","sourceCodeStart":347,"sourceCodeEnd":383,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L347-L383","documentation":"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.","triggerScenarios":"dma_start(..., src_sem=<jax.Array>, device_id=dev) where src_sem is not a Ref.","commonSituations":"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.","solutions":["Pass the source semaphore as a Ref obtained from the kernel's ref arguments","Verify nothing unwrapped the ref (e.g. sem[()] reads) before dma_start","If you did not mean to signal the source, pass src_sem=None (and device_id=None)"],"exampleFix":"# before\ndma_start(src, dst, dst_sem, src_sem=sem_val, device_id=dev)  # sem_val is an Array\n\n# after\ndma_start(src, dst, dst_sem, src_sem=sem_ref, device_id=dev)","handlingStrategy":"type-guard","validationCode":"if src_sem is not None:\n  from jax._src import state\n  assert isinstance(getattr(src_sem, 'aval', src_sem), (state.AbstractRef, state.TransformedRef))","typeGuard":"def is_ref(v) -> bool:\n  from jax._src import state\n  return isinstance(getattr(v, 'aval', v), (state.AbstractRef, state.TransformedRef))","tryCatchPattern":null,"preventionTips":["Pass the semaphore ref, never a value read from it","Keep src_sem and device_id plumbing symmetrical"],"tags":["jax","pallas","mosaic","dma","semaphore","ref-vs-array","type-validation"],"backgroundTag":"expected-ref-got-array","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}