{"record":{"id":"a3b0a6338f6ee12d","repo":"jax-ml/jax","slug":"reduction-op-reduction-op-not-supported-by-the-t","errorCode":null,"errorMessage":"Reduction op {reduction_op} not supported by the TMA implementation for element type {element_type}","messagePattern":"Reduction op (.+?) not supported by the TMA implementation for element type (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1281,"sourceCode":"      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\n    else:\n      raise ValueError(\"Only SMEM <-> GMEM copies supported\")","sourceCodeStart":1263,"sourceCodeEnd":1299,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1263-L1299","documentation":"Even with the TMA implementation, only certain (reduction_op, element_type) pairs are supported by TMA reduction semantics (checked by _is_tma_reduction_op_supported). Passing an unsupported combination, such as a min/max or non-add reduction on an integer type, raises this ValueError.","triggerScenarios":"Calling async_copy(reduction_op=..., implementation=TMA) where _is_tma_reduction_op_supported(reduction_op, element_type) is False, e.g. an add reduction on f8 types or logical ops on floats.","commonSituations":"Experimenting with newer dtypes (fp8, int4) in Mosaic kernels with accumulate-on-load; assuming all ReductionOp values work for every element type.","solutions":["Use a supported (op, dtype) pair — typically add on floating-point types per TMA hardware semantics.","Remove reduction_op and implement the reduction manually after the load.","Check the _is_tma_reduction_op_supported table in launch_context.py for the exact supported set on your JAX version."],"exampleFix":"// before\nctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.MAX, implementation=mgpu.AsyncCopyImplementation.TMA)  # unsupported for this dtype\n// after\nctx.async_copy(src, dst, ..., implementation=mgpu.AsyncCopyImplementation.TMA)\n# then compute max manually over the loaded slice","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import launch_context as lc\nif reduction_op is not None:\n    assert lc._is_tma_reduction_op_supported(reduction_op, element_type), \\\n        f'unsupported (op, dtype) pair: {reduction_op}, {element_type}'","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(..., reduction_op=op, implementation=mgpu.AsyncCopyImplementation.TMA)\nexcept ValueError as e:\n    if 'not supported by the TMA implementation' in str(e):\n        # fall back: plain copy + manual reduction\n        ctx.async_copy(..., implementation=mgpu.AsyncCopyImplementation.TMA)\n    else:\n        raise","preventionTips":["Stick to ADD reductions on standard float dtypes.","Check the supported-op table for your JAX version before using exotic ops.","Test fp8/int kernels with reductions separately."],"tags":["jax","mosaic-gpu","tma","reduction","dtype-mismatch","unsupported-operation"],"backgroundTag":"unsupported-operation-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}