{"record":{"id":"3039a8ef9e7fac6f","repo":"jax-ml/jax","slug":"unsupported-reduction-for-f32-only-min-max-absm","errorCode":null,"errorMessage":"Unsupported reduction for f32. Only min, max, absmin, and absmax are supported, got: {reduce}","messagePattern":"Unsupported reduction for f32\\. Only min, max, absmin, and absmax are supported, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1347,"sourceCode":"        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`.\n    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to\n    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.\n    # This restriction can be lifted if `TiledLayout` supports multiple\n    # vector dims.\n    if self.layout == sparse_meta_layout():\n      raise NotImplementedError(\"Sparse meta layout loads unsupported.\")\n    if regs_shape[0] != 1:  # We'll need to issue multiple loads below.\n      raise NotImplementedError(\"Loading multiple row tiles\")","sourceCodeStart":1329,"sourceCodeEnd":1365,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1329-L1365","documentation":"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.","triggerScenarios":"Calling tmem.load(layout, reduce='sum') (or any reduce other than min/max/absmin/absmax) on a TensorMem whose dtype is f32.","commonSituations":"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.","solutions":["Use only 'min', 'max', 'absmin', or 'absmax' for f32 reductions","If you need a sum, load without reduce and accumulate in registers explicitly","Check self.dtype before choosing the reduce string"],"exampleFix":"// before\narr = tmem.load(layout, reduce='sum')\n// after\narr = tmem.load(layout)  # accumulate manually, or use reduce='max'","handlingStrategy":"validation","validationCode":"allowed = (\"min\", \"max\", \"absmin\", \"absmax\")\nif isinstance(tmem.dtype, ir.F32Type) and reduce not in allowed:\n    raise ValueError(f'reduce must be one of {allowed} for f32')","typeGuard":"def is_supported_f32_reduce(reduce: str) -> bool:\n    return reduce in ('min', 'max', 'absmin', 'absmax')","tryCatchPattern":"try:\n    arr, red = tmem.load(layout, reduce=reduce)\nexcept ValueError as e:\n    if 'Unsupported reduction' in str(e):\n        arr, _ = tmem.load(layout); red = manual_reduce(arr)\n    else: raise","preventionTips":["Keep a module-level tuple of supported reduces and validate kernel params against it","Unit-test kernel wrappers with each reduce kind you expose"],"tags":["mosaic","gpu","tcgen05","reduction","tensor-core","jax"],"backgroundTag":"unsupported-operation-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}