{"record":{"id":"595d48cceafa5b6b","repo":"jax-ml/jax","slug":"reduction-op-not-supported","errorCode":null,"errorMessage":"reduction_op not supported","messagePattern":"reduction_op not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":1412,"sourceCode":"    src_transforms: tuple[Any, ...],\n    dst_allocation_key_as_array: jax.Array,\n    dst_transforms: tuple[Any, ...],\n    predicate: jax.Array | None,\n    source_info: source_info_util.SourceInfo,\n    commit_group: bool,\n    reduction_op: mgpu.TMAReductionOp,\n):\n  # TODO(jburnim,paulbib): Implement commit_group.\n  del commit_group\n  src_allocation_key = HostAllocationKey.from_array(src_allocation_key_as_array)\n  src_transforms = jax.tree.map(int, _remove_noop_transforms(src_transforms))\n  dst_allocation_key = HostAllocationKey.from_array(dst_allocation_key_as_array)\n  dst_transforms = jax.tree.map(int, _remove_noop_transforms(dst_transforms))\n\n  if predicate is not None:\n    raise NotImplementedError(\"predicate not supported\")\n  if reduction_op is not None:\n    raise NotImplementedError(\"reduction_op not supported\")\n\n  clock = None\n\n  shared_memory = _get_shared_memory()\n  if shared_memory.detect_races:\n    clock = shared_memory.incr_clock(thread)\n\n  task = AsyncCopySmemToGmemTask(\n      mesh_location=mesh_location,\n      thread=thread,\n      src_allocation_key=src_allocation_key,\n      src_transforms=src_transforms,\n      dst_allocation_key=dst_allocation_key,\n      dst_transforms=dst_transforms,\n      source_info=source_info,\n      clock=clock,\n  )\n","sourceCodeStart":1394,"sourceCodeEnd":1430,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L1394-L1430","documentation":"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.","triggerScenarios":"Calling copy_smem_to_gmem with reduction_op set (e.g. ReductionOp.ADD) while running under GPU interpret mode.","commonSituations":"Kernels using reducing copies to accumulate partial results from shared memory; debugging such kernels in interpret mode on CPU.","solutions":["Replace the reducing copy with an explicit load + atomic/store sequence the interpreter supports","Run that path on device only, and exclude it from interpret-mode tests","Update jax in case reduction support was added to the interpreter","Guard interpret-mode runs with a flag that uses a simpler code path"],"exampleFix":"# before\ncopy_smem_to_gmem(src, dst, reduction_op=ReductionOp.ADD)\n# after\ntmp = smem_block[...]\ndst_ref.atomic_add(tmp)  # interpreter-supported path","handlingStrategy":"fallback","validationCode":"if interpret_mode and reduction_op is not None:\n    use_explicit_atomic_path()","typeGuard":null,"tryCatchPattern":"try:\n    copy_smem_to_gmem(src, dst, reduction_op=op)\nexcept NotImplementedError:\n    tmp = smem[...]; dst_ref.atomic_add(tmp)","preventionTips":["Wrap reducing copies behind a helper with an interpret-safe fallback"],"tags":["pallas","mosaic-gpu","interpret-mode","reduction","not-implemented"],"backgroundTag":"unsupported-operation-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}