jax-ml/jax · error · ValueError

Unsupported reduction for i32. Only min and max are supporte

Error message

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

What it means

When tcgen05.load(reduce=...) is used on an i32 TMEM ref, the hardware TMEM reduction (tcgen05.reduce) only supports min and max for 32-bit integers. Other reduction kinds ('sum', 'add', 'and', …) raise this error. For unsigned i32 the code rewrites to umin/umax internally.

Source

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

    packing = self.packing
    bitwidth = utils.bitwidth(self.dtype)
    is_at_least_16b = bitwidth in {16, 32}
    columns = self.shape[1]
    if layout is None:
      if is_at_least_16b and self.layout == tmem_default_layout(packing):
        layout = LAYOUT
      elif is_at_least_16b and packing <= columns // 2 and self.layout == tmem_half_lane_layout(columns, packing):
        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`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use reduce='min' or reduce='max' for i32
2. For sums, load to registers and reduce there (or use f32 accumulator)
3. Keep the accumulator in f32/bf16 if TMEM-side 'sum' reduction is required

Example fix

# before
acc = tcgen05.load(ref_i32, reduce='sum')
# after
acc = tcgen05.load(ref_i32, reduce='max')  # or load and sum in registers / use f32 accumulator

if ref.dtype in (i32,) and reduce_ is not None:
    assert reduce_ in ('min', 'max'), f'i32 TMEM reduce supports min/max only, got {reduce_}'

def tmem_reduce_ok(dtype, reduce_kind: str) -> bool:
    from jaxlib.mlir import ir
    if isinstance(dtype, ir.IntegerType) and dtype.width == 32:
        return reduce_kind in ('min', 'max')
    return True

When it happens

Trigger: tcgen05.load(tmem_i32_ref, reduce='sum'); any reduce other than 'min'/'max' on an i32 TMEM ref.

Common situations: Trying to accumulate integer sums directly from TMEM; reusing f32 reduction code ('sum') after switching the accumulator dtype to i32 (e.g. for integer MMA or i8 experiments).

Related errors

• Unsupported dtype for reduction: {val_aval.dtype}. Only floa
• tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}
• TMEM can only be sliced, not indexed
• TMEM cannot be sliced along rows
• Cannot slice TMEM with multiple tiles along rows.

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