jax-ml/jax · error · ValueError

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

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.

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)

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

• Sparse MMA not supported for M=64
• Unsupported dtype for reduction: {self.dtype}
• Fused load-reduce is not supported for this layout
• Unsupported reduction kind: {reduce}
• Unsupported reduction kind: {op.kind}

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