jax-ml/jax · error · NotImplementedError

DMA partial discharge add=True not yet implemented.

Error message

DMA partial discharge add=True not yet implemented.

What it means

The discharge rule for dma_start lowers queued DMAs into explicit memory effects when discharging state; partial-discharge accumulation (add=True, which would accumulate the DMA'd data additively) is not implemented, so it raises NotImplementedError.

Source

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

      pp.text(" "),
      sp.pp_ref_transforms(context, src_ref),
      pp.text(" -> "),
      sp.pp_ref_transforms(context, dst_ref),
      pp.text(" "),
      sp.pp_ref_transforms(context, dst_sem),
  ])

jax_core.pp_eqn_rules[dma_start_p] = _dma_start_pp_eqn


def dma_start_discharge_rule(
    ctx, *args, tree, device_id_type,
    priority, add
):
  # Note: we ignore the DMA priority in discharge rules.
  del priority
  if add:
    raise NotImplementedError(
        "DMA partial discharge add=True not yet implemented.")
  src_ref, dst_ref, dst_sem, src_sem, device_id = _dma_unflatten(tree, args)
  src_ref, src_transforms = _get_ref_and_transforms(src_ref)
  dst_ref, dst_transforms = _get_ref_and_transforms(dst_ref)
  dst_sem, dst_sem_transforms = _get_ref_and_transforms(dst_sem)
  src_sem, src_sem_transforms = _get_ref_and_transforms(src_sem)

  src_ref_aval, dst_ref_aval, dst_sem_aval, src_sem_aval, _ = _dma_unflatten(
      tree, ctx.in_avals
  )

  _, dst_discharge, dst_sem_discharge, *maybe_src_sem_discharge = (
      _dma_unflatten(tree, ctx.should_discharge)
  )
  dst_discharge = _get_ref(dst_discharge)
  dst_sem_discharge = _get_ref(dst_sem_discharge)
  is_remote = device_id is not None
  src_sem_discharge = None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use add=False for the DMA (full discharge instead of partial accumulation)
2. Restructure the kernel so accumulation happens via explicit loads/adds/stores after the DMA rather than during discharge
3. Update JAX — this is a known unimplemented corner and may land in a newer release
4. Report/check the JAX issue tracker for dma partial discharge support

Example fix

# before
dma_start(src, dst, sem, add=True)  # later hits discharge rule

# after
dma_start(src, dst, sem, add=False)
# then accumulate manually:
# dst_ref[...] = dst_ref[...] + value

try:
  run_with_discharge(...)
except NotImplementedError as e:
  if 'partial discharge' in str(e):
    # retry with add=False and accumulate manually
    run_with_discharge(..., dma_add=False)

When it happens

Trigger: Running a transformation that triggers the dma_start discharge rule with add=True (partial discharge path), typically via certain pipelining or state-discharge passes in the Mosaic pipeline.

Common situations: Using advanced pipelining features or custom pipeline transformations that discharge DMA state with accumulation enabled; cutting-edge kernel code that combines DMA with partial state discharge; version changes where discharge logic became stricter.

Related errors

• dma_start not implemented in LoJAX yet.
• {uninitialized_memory} + {str(dtype)}
• Batching over dynamic grid values is not supported yet.
• Either both or neither `src_sem` and `device_id` can be set.
• Cannot `wait_send` on a local copy.

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