{"record":{"id":"5437ca7aeb76949d","repo":"jax-ml/jax","slug":"only-the-tma-implementation-supports-reductions","errorCode":null,"errorMessage":"Only the TMA implementation supports reductions","messagePattern":"Only the TMA implementation supports reductions","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1279,"sourceCode":"    element_bitwidth = utils.bitwidth(element_type)\n    if element_type != dst_ref_ty.element_type:\n      raise ValueError(\n          f\"Expected same element type, got {element_type} and\"\n          f\" {dst_ref_ty.element_type}\"\n      )\n\n    if isinstance(collective, gpu.Dimension):\n      collective = (collective,)\n    elif collective is None:\n      collective = ()\n    if not isinstance(gmem_transform, tuple):\n      gmem_transform = (gmem_transform,)\n    if not isinstance(gmem_slice, tuple):\n      gmem_slice = (gmem_slice,)\n\n    if reduction_op is not None:\n      if implementation != AsyncCopyImplementation.TMA:\n        raise ValueError(\"Only the TMA implementation supports reductions\")\n      if not _is_tma_reduction_op_supported(reduction_op, element_type):\n        raise ValueError(\n            f\"Reduction op {reduction_op} not supported by the TMA\"\n            f\" implementation for element type {element_type}\"\n        )\n\n    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):\n      gmem_ref, smem_ref = src_ref, dst_ref\n      if implementation == AsyncCopyImplementation.TMA and barrier is None:\n        raise ValueError(\"Barriers are required for TMA GMEM -> SMEM copies\")\n      if arrive is None:\n        arrive = True  # Arrive by default\n    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:\n      gmem_ref, smem_ref = dst_ref, src_ref\n      if barrier is not None:\n        raise ValueError(\"Barriers are unsupported for SMEM -> GMEM copies\")\n      if arrive is None:\n        arrive = True  # Commit this copy to the async group by default","sourceCodeStart":1261,"sourceCodeEnd":1297,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1261-L1297","documentation":"async_copy supports reduction operations (e.g. accumulate-on-load via TMA reduction semantics) only through the TMA implementation. Passing reduction_op with implementation != AsyncCopyImplementation.TMA raises this ValueError.","triggerScenarios":"Calling async_copy(reduction_op=..., implementation=AsyncCopyImplementation.LDGSTS) (or any non-TMA implementation), which has no hardware support for fused reductions.","commonSituations":"Porting kernels between GPUs: using reductions on loads that relied on TMA on Hopper but running/configuring a CUDA-core or LDGSTS path; leaving a legacy implementation argument set when adding reduction_op.","solutions":["Set implementation=AsyncCopyImplementation.TMA when reduction_op is not None.","Or drop reduction_op and perform the reduction manually after the copy (e.g. accumulate in SMEM/registers).","Confirm the target GPU supports TMA reductions for the element type."],"exampleFix":"// before\nctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.ADD, implementation=mgpu.AsyncCopyImplementation.LDGSTS)\n// after\nctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.ADD, implementation=mgpu.AsyncCopyImplementation.TMA)","handlingStrategy":"validation","validationCode":"if reduction_op is not None:\n    assert implementation == mgpu.AsyncCopyImplementation.TMA, \\\n        'reductions require the TMA implementation'","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(..., reduction_op=op)\nexcept ValueError as e:\n    if 'Only the TMA implementation supports reductions' in str(e):\n        ctx.async_copy(...)  # retry without reduction_op, reduce manually\n    else:\n        raise","preventionTips":["Gate reduction usage on TMA availability of the target GPU.","Keep implementation explicit in all async_copy calls.","Reduce manually when portability matters more than the fused reduction."],"tags":["jax","mosaic-gpu","tma","reduction","unsupported-operation"],"backgroundTag":"unsupported-operation-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}