{"record":{"id":"197ca76262b57b7f","repo":"jax-ml/jax","slug":"unsupported-reduction-for-i32-only-min-and-max-ar","errorCode":null,"errorMessage":"Unsupported reduction for i32. Only min and max are supported, got: {reduce}","messagePattern":"Unsupported reduction for i32\\. Only min and max are supported, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1339,"sourceCode":"    packing = self.packing\n    bitwidth = utils.bitwidth(self.dtype)\n    is_at_least_16b = bitwidth in {16, 32}\n    columns = self.shape[1]\n    if layout is None:\n      if is_at_least_16b and self.layout == tmem_default_layout(packing):\n        layout = LAYOUT\n      elif is_at_least_16b and packing <= columns // 2 and self.layout == tmem_half_lane_layout(columns, packing):\n        layout = fa.WGMMA_LAYOUT\n      elif is_at_least_16b and columns % 16 == 0 and self.layout == tmem_m64_collective_layout(columns, packing):\n        layout = fa_m64_collective_layout(columns)\n      elif packing * bitwidth == 32:\n        layout = self.layout.as_tiled_layout()\n      else:\n        raise ValueError(f\"TMEM layout {self.layout} is not supported\")\n    if reduce is not None:\n      if isinstance(self.dtype, ir.IntegerType) and bitwidth == 32:\n        if reduce not in (\"min\", \"max\"):\n          raise ValueError(\n              \"Unsupported reduction for i32. Only min and max are supported,\"\n              f\" got: {reduce}\"\n          )\n        if not is_signed:\n          reduce = \"abs\" + reduce  # type: ignore\n      elif isinstance(self.dtype, ir.F32Type):\n        if reduce not in (\"min\", \"max\", \"absmin\", \"absmax\"):\n          raise ValueError(\n              \"Unsupported reduction for f32. Only min, max, absmin, and\"\n              f\" absmax are supported, got: {reduce}\"\n          )\n      else:\n        raise ValueError(f\"Unsupported dtype for reduction: {self.dtype}\")\n\n    has_default_layout = self.layout == tmem_default_layout(packing)\n    regs_shape = layout.registers_shape(self.shape)\n    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the\n    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.","sourceCodeStart":1321,"sourceCodeEnd":1357,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1321-L1357","documentation":"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.","triggerScenarios":"tcgen05.load(tmem_i32_ref, reduce='sum'); any reduce other than 'min'/'max' on an i32 TMEM ref.","commonSituations":"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).","solutions":["Use reduce='min' or reduce='max' for i32","For sums, load to registers and reduce there (or use f32 accumulator)","Keep the accumulator in f32/bf16 if TMEM-side 'sum' reduction is required"],"exampleFix":"# before\nacc = tcgen05.load(ref_i32, reduce='sum')\n# after\nacc = tcgen05.load(ref_i32, reduce='max')  # or load and sum in registers / use f32 accumulator","handlingStrategy":"validation","validationCode":"if ref.dtype in (i32,) and reduce_ is not None:\n    assert reduce_ in ('min', 'max'), f'i32 TMEM reduce supports min/max only, got {reduce_}'","typeGuard":"def tmem_reduce_ok(dtype, reduce_kind: str) -> bool:\n    from jaxlib.mlir import ir\n    if isinstance(dtype, ir.IntegerType) and dtype.width == 32:\n        return reduce_kind in ('min', 'max')\n    return True","tryCatchPattern":null,"preventionTips":["Reserve TMEM-side 'sum' reductions for floating-point accumulators","For integer sums: load to registers and reduce there, or keep f32 accumulators"],"tags":["jax","mosaic","tmem","reduction","dtype","unsupported-operation"],"backgroundTag":"unsupported-operation-for-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}