{"record":{"id":"957450bd7146ddd6","repo":"jax-ml/jax","slug":"unsupported-reduction-kind-op-kind","errorCode":null,"errorMessage":"Unsupported reduction kind: {op.kind}","messagePattern":"Unsupported reduction kind: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":871,"sourceCode":"  element_type = op.vector.type.element_type\n  scratch = _slice_smem(\n      ir.MemRefType.get([4], element_type, memory_space=utils.smem()),\n      ir.IntegerAttr(op.attributes[\"offset\"]).value,\n      ctx.smem_requested_bytes,\n  )\n  axes = range(op.vector.type.rank)\n  op_kind = _combining_kind(op.kind)\n  is_signed = _is_reduction_signed(op_kind)\n  a = _fragmented_array_from_ir(op.vector, layout, is_signed)\n  match op_kind:\n    case vector.CombiningKind.ADD:\n      result = a.reduce(\"add\", axes, scratch)\n    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:\n      result = a.reduce(\"max\", axes, scratch)\n    case vector.CombiningKind.MINUI | vector.CombiningKind.MINSI | vector.CombiningKind.MINIMUMF:\n      result = a.reduce(\"min\", axes, scratch)\n    case _:\n      raise NotImplementedError(f\"Unsupported reduction kind: {op.kind}\")\n  assert isinstance(result.layout, fa.WGSplatFragLayout)\n  return [result.registers.item()]\n\n\n@_register_lowering(vector.MultiDimReductionOp)\ndef _vector_multi_dim_reduction_op_lowering_rule(\n    ctx: LoweringContext, op: vector.MultiDimReductionOp\n) -> Sequence[ir.Value]:\n  [in_layout, acc_layout] = inference_utils.in_layouts(op)\n  [out_layout] = inference_utils.out_layouts(op)\n  if out_layout != acc_layout:\n    raise ValueError(\n        f\"Output layout {out_layout} must match the accumulator layout\"\n        f\" {acc_layout}\"\n    )\n\n  if len(op.reduction_dims) != 1:\n    raise NotImplementedError(\"Only 1 reduction dimension is supported.\")","sourceCodeStart":853,"sourceCodeEnd":889,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L853-L889","documentation":"vector.reduction lowering supports only add/max/min combining kinds; anything else (e.g. AND, OR, XOR) falls through to NotImplementedError.","triggerScenarios":"Emitting vector.reduction with kind like AND/OR/XOR in a Mosaic kernel.","commonSituations":"Porting code that uses bitwise vector reductions, which Mosaic GPU does not lower.","solutions":["Implement bitwise reduction manually via elementwise ops + add reduction on integer representation","Restrict to add/max/min reductions"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":"SUPPORTED = {'add', 'max', 'min'}\nassert kind in SUPPORTED, f'reduction kind {kind} unsupported'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Stick to add/min/max vector reductions; emulate bitwise ops manually"],"tags":["jax","mosaic","gpu","reduction","vector"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}