jax-ml/jax · error · ValueError

Reduction op {reduction_op} not supported by the TMA impleme

Error message

Reduction op {reduction_op} not supported by the TMA implementation for element type {element_type}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1281

      raise ValueError(
          f"Expected same element type, got {element_type} and"
          f" {dst_ref_ty.element_type}"
      )

    if isinstance(collective, gpu.Dimension):
      collective = (collective,)
    elif collective is None:
      collective = ()
    if not isinstance(gmem_transform, tuple):
      gmem_transform = (gmem_transform,)
    if not isinstance(gmem_slice, tuple):
      gmem_slice = (gmem_slice,)

    if reduction_op is not None:
      if implementation != AsyncCopyImplementation.TMA:
        raise ValueError("Only the TMA implementation supports reductions")
      if not _is_tma_reduction_op_supported(reduction_op, element_type):
        raise ValueError(
            f"Reduction op {reduction_op} not supported by the TMA"
            f" implementation for element type {element_type}"
        )

    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):
      gmem_ref, smem_ref = src_ref, dst_ref
      if implementation == AsyncCopyImplementation.TMA and barrier is None:
        raise ValueError("Barriers are required for TMA GMEM -> SMEM copies")
      if arrive is None:
        arrive = True  # Arrive by default
    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:
      gmem_ref, smem_ref = dst_ref, src_ref
      if barrier is not None:
        raise ValueError("Barriers are unsupported for SMEM -> GMEM copies")
      if arrive is None:
        arrive = True  # Commit this copy to the async group by default
    else:
      raise ValueError("Only SMEM <-> GMEM copies supported")

View on GitHub (pinned to 1e1c6a8fc0)

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.

Example fix

// before
ctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.MAX, implementation=mgpu.AsyncCopyImplementation.TMA)  # unsupported for this dtype
// after
ctx.async_copy(src, dst, ..., implementation=mgpu.AsyncCopyImplementation.TMA)
# then compute max manually over the loaded slice

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mosaic.gpu import launch_context as lc
if reduction_op is not None:
    assert lc._is_tma_reduction_op_supported(reduction_op, element_type), \
        f'unsupported (op, dtype) pair: {reduction_op}, {element_type}'

Try / catch

try:
    ctx.async_copy(..., reduction_op=op, implementation=mgpu.AsyncCopyImplementation.TMA)
except ValueError as e:
    if 'not supported by the TMA implementation' in str(e):
        # fall back: plain copy + manual reduction
        ctx.async_copy(..., implementation=mgpu.AsyncCopyImplementation.TMA)
    else:
        raise

Prevention

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.

When it happens

Trigger: 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.

Common situations: Experimenting with newer dtypes (fp8, int4) in Mosaic kernels with accumulate-on-load; assuming all ReductionOp values work for every element type.

Related errors

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