SmallVector<Type> result_types{op.getResult().getType()};
    ReplaceTfOpWithNewOp<TF::XlaVariadicReduceV2Op>(
        rewriter, op, result_types, inputs, init_values,
        op.getDimensionsToReduce(), op.getReducer());
    return ::mlir::success();
  };
};

void XlaReduceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {

    } else if (auto xla_variadic_reduce_v2_op =
                   dyn_cast<TF::XlaVariadicReduceV2Op>(op)) {
      return propagate_shape_to(xla_variadic_reduce_v2_op.getReducer());
    } else if (auto xla_variadic_sort_op =
                   dyn_cast<TF::XlaVariadicSortOp>(op)) {
      return propagate_shape_to(xla_variadic_sort_op.getComparator());
    }
  }

  LogicalResult matchAndRewrite(TF::XlaVariadicReduceV2Op op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();

    mlir::SymbolRefAttr func = op.getReducer();
    auto func_op = cast<mlir::func::FuncOp>(SymbolTable::lookupSymbolIn(
        op->getParentOfType<mlir::ModuleOp>(), func));
    auto func_ty = func_op.getFunctionType();