jax-ml/jax · error · ValueError

Unsupported dtype for reduction: {self.dtype}

Error message

Unsupported dtype for reduction: {self.dtype}

What it means

Raised by TensorMem.load when a fused load-reduce is requested but the tensor memory dtype is neither an integer type nor f32. The reduction path only knows how to lower reductions for integer and f32 element types.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1352

      else:
        raise ValueError(f"TMEM layout {self.layout} is not supported")
    if reduce is not None:
      if isinstance(self.dtype, ir.IntegerType) and bitwidth == 32:
        if reduce not in ("min", "max"):
          raise ValueError(
              "Unsupported reduction for i32. Only min and max are supported,"
              f" got: {reduce}"
          )
        if not is_signed:
          reduce = "abs" + reduce  # type: ignore
      elif isinstance(self.dtype, ir.F32Type):
        if reduce not in ("min", "max", "absmin", "absmax"):
          raise ValueError(
              "Unsupported reduction for f32. Only min, max, absmin, and"
              f" absmax are supported, got: {reduce}"
          )
      else:
        raise ValueError(f"Unsupported dtype for reduction: {self.dtype}")

    has_default_layout = self.layout == tmem_default_layout(packing)
    regs_shape = layout.registers_shape(self.shape)
    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the
    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.
    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to
    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.
    # This restriction can be lifted if `TiledLayout` supports multiple
    # vector dims.
    if self.layout == sparse_meta_layout():
      raise NotImplementedError("Sparse meta layout loads unsupported.")
    if regs_shape[0] != 1:  # We'll need to issue multiple loads below.
      raise NotImplementedError("Loading multiple row tiles")
    if (
        layout == LAYOUT
        and self.layout == tmem_default_layout(packing)
        and is_at_least_16b
    ):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the reduce argument and perform the reduction manually in registers
2. Redo arithmetic so the reduced tensor lives in f32 or an integer dtype
3. Check isinstance(tmem.dtype, (ir.IntegerType, ir.F32Type)) before requesting a reduction

Example fix

// before
arr = tmem.load(layout, reduce='max')  # tmem dtype is f16
// after
arr = tmem.load(layout)
red = mx.maximum(arr, axis=...)  # explicit reduction

from mlir import ir
if reduce is not None and not isinstance(tmem.dtype, (ir.IntegerType, ir.F32Type)):
    raise ValueError('fused load-reduce needs integer or f32 dtype')

def can_fused_reduce(tmem) -> bool:
    return isinstance(tmem.dtype, (ir.IntegerType, ir.F32Type))

try:
    arr, red = tmem.load(layout, reduce=reduce)
except ValueError:
    arr, _ = tmem.load(layout)
    red = reduce_in_registers(arr, reduce)

When it happens

Trigger: Calling tmem.load(..., reduce=...) where the TensorMem was allocated with a dtype like f16, bf16, or f64 — anything other than an ir.IntegerType or ir.F32Type.

Common situations: Allocating TMEM with half-precision dtypes (common in attention kernels) and then attempting a fused load-reduce; recent dtype additions to Mosaic that lack reduction support.

Related errors

• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block
• MMA with element type {elem_type_str} only supports accumula
• MMA with element type {elem_type_str} only supports accumula
• Unsupported reduction for f32. Only min, max, absmin, and ab

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