if (!MatchTopKComparator<mhlo::ReturnOp>(op.getComparator()))
      return rewriter.notifyMatchFailure(op, "only match for GT comparator");
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    Value k_cst = BuildIntConstOp(builder, rewriter, k, rewriter.getI32Type());
    rewriter.replaceOpWithNewOp<TF::TopKV2Op>(op, keys.getType(),
                                              indices.getType(), keys, k_cst);
    return success();
  };