jax-ml/jax · error · NotImplementedError

reduction_op not supported

Error message

reduction_op not supported

What it means

copy_smem_to_gmem in the GPU interpreter does not implement the optional reduction_op argument (e.g. atomic add/min reductions on copy). Passing reduction_op raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:1412

    src_transforms: tuple[Any, ...],
    dst_allocation_key_as_array: jax.Array,
    dst_transforms: tuple[Any, ...],
    predicate: jax.Array | None,
    source_info: source_info_util.SourceInfo,
    commit_group: bool,
    reduction_op: mgpu.TMAReductionOp,
):
  # TODO(jburnim,paulbib): Implement commit_group.
  del commit_group
  src_allocation_key = HostAllocationKey.from_array(src_allocation_key_as_array)
  src_transforms = jax.tree.map(int, _remove_noop_transforms(src_transforms))
  dst_allocation_key = HostAllocationKey.from_array(dst_allocation_key_as_array)
  dst_transforms = jax.tree.map(int, _remove_noop_transforms(dst_transforms))

  if predicate is not None:
    raise NotImplementedError("predicate not supported")
  if reduction_op is not None:
    raise NotImplementedError("reduction_op not supported")

  clock = None

  shared_memory = _get_shared_memory()
  if shared_memory.detect_races:
    clock = shared_memory.incr_clock(thread)

  task = AsyncCopySmemToGmemTask(
      mesh_location=mesh_location,
      thread=thread,
      src_allocation_key=src_allocation_key,
      src_transforms=src_transforms,
      dst_allocation_key=dst_allocation_key,
      dst_transforms=dst_transforms,
      source_info=source_info,
      clock=clock,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace the reducing copy with an explicit load + atomic/store sequence the interpreter supports
2. Run that path on device only, and exclude it from interpret-mode tests
3. Update jax in case reduction support was added to the interpreter
4. Guard interpret-mode runs with a flag that uses a simpler code path

Example fix

# before
copy_smem_to_gmem(src, dst, reduction_op=ReductionOp.ADD)
# after
tmp = smem_block[...]
dst_ref.atomic_add(tmp)  # interpreter-supported path

if interpret_mode and reduction_op is not None:
    use_explicit_atomic_path()

try:
    copy_smem_to_gmem(src, dst, reduction_op=op)
except NotImplementedError:
    tmp = smem[...]; dst_ref.atomic_add(tmp)

When it happens

Trigger: Calling copy_smem_to_gmem with reduction_op set (e.g. ReductionOp.ADD) while running under GPU interpret mode.

Common situations: Kernels using reducing copies to accumulate partial results from shared memory; debugging such kernels in interpret mode on CPU.

Related errors

• predicate not supported
• Dynamic grid bounds not (yet) supported in GPU interpret mod
• Dynamic grid bounds not (yet) supported on GPU
• Scalar arguments not (yet) supported on GPU
• Unsupported memory space: {space}

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