jax-ml/jax · error · ValueError

Unsupported reduction for f32. Only min, max, absmin, and…

Error message

Unsupported reduction for f32. Only min, max, absmin, and absmax are supported, got: {reduce}

What it means

Raised by TensorMem.load in Mosaic GPU's tcgen05 module when a fused load-reduce is requested with an f32 tensor but the reduction kind is not one of min, max, absmin, absmax. Only those four reductions are implemented for 32-bit floats. Any other string (e.g. 'sum', 'add') reaches this ValueError.

Solutions

1. Use only 'min', 'max', 'absmin', or 'absmax' for f32 reductions
2. If you need a sum, load without reduce and accumulate in registers explicitly
3. Check self.dtype before choosing the reduce string

Example fix

// before
arr = tmem.load(layout, reduce='sum')
// after
arr = tmem.load(layout)  # accumulate manually, or use reduce='max'

allowed = ("min", "max", "absmin", "absmax")
if isinstance(tmem.dtype, ir.F32Type) and reduce not in allowed:
    raise ValueError(f'reduce must be one of {allowed} for f32')

def is_supported_f32_reduce(reduce: str) -> bool:
    return reduce in ('min', 'max', 'absmin', 'absmax')

try:
    arr, red = tmem.load(layout, reduce=reduce)
except ValueError as e:
    if 'Unsupported reduction' in str(e):
        arr, _ = tmem.load(layout); red = manual_reduce(arr)
    else: raise

When it happens

Trigger: Calling tmem.load(layout, reduce='sum') (or any reduce other than min/max/absmin/absmax) on a TensorMem whose dtype is f32.

Common situations: Porting a kernel from Triton/wgmma-style accumulate-on-load code where 'sum' reductions are common; assuming all reduction kinds supported for integer dtypes also apply to f32.

Related errors

• Fused load-reduce is not supported for this layout
• Sparse MMA not supported for M=64
• Unsupported dtype for reduction
• Unsupported reduction kind
• Loading multiple row tiles

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

Appendix: source

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

        layout = fa.WGMMA_LAYOUT
      elif is_at_least_16b and columns % 16 == 0 and self.layout == tmem_m64_collective_layout(columns, packing):
        layout = fa_m64_collective_layout(columns)
      elif packing * bitwidth == 32:
        layout = self.layout.as_tiled_layout()
      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")

View on GitHub (pinned to 1e1c6a8fc0)