{"record":{"id":"7f68322af6ba72fd","repo":"jax-ml/jax","slug":"unsupported-reduction-kind-reduce","errorCode":null,"errorMessage":"Unsupported reduction kind: {reduce}","messagePattern":"Unsupported reduction kind: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1771,"sourceCode":"  regs = [None] * (cols // reg_packing)\n  red_reg = None\n  for addr_row_col, instr_num, lane_step, num_slice in it:\n    assert lane_step == 0, lane_step\n    instr_regs = _tmem_load(addr_row_col, load_shape, instr_num, pack, reduce, dtype)\n    if reduce:\n      *instr_regs, instr_red_reg = instr_regs\n      instr_red_reg = utils.bitcast(instr_red_reg, dtype)\n      if red_reg is None:\n        red_reg = instr_red_reg\n      elif isinstance(dtype, ir.F32Type):\n        # abs is applied by the TMEM load, so both red_regs are non-negative.\n        match reduce:\n          case \"min\" | \"absmin\":\n            red_reg = arith.minimumf(red_reg, instr_red_reg)\n          case \"max\" | \"absmax\":\n            red_reg = arith.maximumf(red_reg, instr_red_reg)\n          case _:\n            raise ValueError(f\"Unsupported reduction kind: {reduce}\")\n      elif isinstance(dtype, ir.IntegerType):\n        match reduce:\n          case \"min\":\n            red_reg = arith.minsi(red_reg, instr_red_reg)\n          case \"max\":\n            red_reg = arith.maxsi(red_reg, instr_red_reg)\n          case \"absmin\":\n            red_reg = arith.minui(red_reg, instr_red_reg)\n          case \"absmax\":\n            red_reg = arith.maxui(red_reg, instr_red_reg)\n          case _:\n            raise ValueError(f\"Unsupported reduction kind: {reduce}\")\n      else:\n        raise ValueError(f\"Unsupported reduction dtype: {dtype}\")\n    if reg_packing == 1 and vector_length == 2:\n      regs[num_slice] = [llvm.bitcast(dtype, r) for r in instr_regs]\n    else:\n      regs[num_slice] = [utils.bitcast(r, vec_ty) for r in instr_regs]","sourceCodeStart":1753,"sourceCodeEnd":1789,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1753-L1789","documentation":"In _load_32xcols_native's fused reduction loop, floating-point reductions only implement min/absmin (minimumf) and max/absmax (maximumf). Any other reduce string that survived earlier validation reaches the match's default branch and raises ValueError.","triggerScenarios":"Internal: calling _load_32xcols_native with a reduce kind other than the float-supported set for an F32 dtype — normally guarded by TensorMem.load's earlier checks, so hitting this means a new/unvalidated reduce kind was passed through.","commonSituations":"Adding a new reduction kind (e.g. 'sum') to Mosaic and wiring it past the outer validation; calling the private helper directly in custom lowering code.","solutions":["Use only min, max, absmin, absmax for float load-reduces","If extending Mosaic, add the case to the match (e.g. arith.addf for 'sum') and update the outer validation","Call the public TensorMem.load rather than the private helper so validation runs"],"exampleFix":null,"handlingStrategy":"try-catch","validationCode":null,"typeGuard":"def is_native_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 kind' in str(e):\n        arr, _ = tmem.load(layout); red = manual_reduce(arr)\n    else: raise","preventionTips":["Don't call private helpers (_load_32xcols_native) directly","Validate reduce strings at your kernel API boundary"],"tags":["mosaic","gpu","tcgen05","reduction","internal","jax"],"backgroundTag":"unsupported-operation-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}