using OpRewritePattern<CallOp>::OpRewritePattern;

 public:
  explicit RewriteTFRCallOp(MLIRContext* context, const SymbolTable& table,
                            bool materialize_derived_attrs)
      : OpRewritePattern<CallOp>(context),
        symbol_table_(table),
        materialize_derived_attrs_(materialize_derived_attrs) {}

  LogicalResult matchAndRewrite(CallOp call_op,

      const SymbolTable& symbol_table, FuncOp function,
      FunctionLibraryDefinition* flib_def,
      llvm::SmallDenseSet<FuncOp>& visited_functions,
      absl::flat_hash_set<Node*>* control_ret_nodes);

 private:
  explicit Exporter(const GraphExportConfig* configs, Graph* graph,
                    const Dialect* tf_dialect, const SymbolTable* symbol_table)
      : configs_(*configs),

      const SymbolTable& symbol_table, FuncOp function,
      FunctionLibraryDefinition* flib_def,
      llvm::SmallDenseSet<FuncOp>& visited_functions,
      absl::flat_hash_set<Node*>* control_ret_nodes);

 private:
  explicit Exporter(const GraphExportConfig* configs, Graph* graph,
                    const Dialect* tf_dialect, const SymbolTable* symbol_table)
      : configs_(*configs),

FailureOr<StringAttr> RenameStablehloFunctions(
    MLIRContext *context, SymbolTableCollection &symbol_tables,
    ModuleOp tf_module, ModuleOp stablehlo_module) {
  SymbolTable &tf_symbol_table = symbol_tables.getSymbolTable(tf_module);
  // `stablehlo_module` is deleted right after the deserialization, so no need
  // to store its `SymbolTable` to `SymbolTableCollection`.
  SymbolTable stablehlo_symbol_table(stablehlo_module);

  Builder builder(context);

  return success();
}

LogicalResult SessionInitializerOp::verify() {
  SessionInitializerOp session_initializer = *this;
  mlir::SymbolTable symbol_table(
      session_initializer->getParentOfType<ModuleOp>());

  for (auto sym_ref : session_initializer.getInitializers()) {
    auto init_func_op = symbol_table.lookup<mlir::func::FuncOp>(
        mlir::cast<FlatSymbolRefAttr>(sym_ref).getValue());

    if (!init_func_op)

  // Adds the new function to symbol table.
  SymbolTable symbol_table(module_op);
  symbol_table.insert(main_func);
  return true;
}

// Creates a new function name by attaching a number suffix
// (`main_func_name_{i}`) and incrementing it until there are no conflicts.
std::string CreateNewFuncName(const StringRef main_func_name,
                              SymbolTable& symbol_table) {
  int suffix_id = 0;

    Type quantize_result_type) {
  rewriter.create<func::ReturnOp>(input.getLoc(), ArrayRef<Value>({output}));

  quantization_func.setVisibility(func::FuncOp::Visibility::Private);
  SymbolTable symbol_table(quantized_op->getParentOfType<ModuleOp>());

  symbol_table.insert(quantization_func);

  FlatSymbolRefAttr func_name_attr =
      FlatSymbolRefAttr::get(rewriter.getStringAttr(func_name));

  rewriter.restoreInsertionPoint(original_point);

  // counter.
  std::string CreateNewFuncName(const StringRef func_name,
                                SymbolTable &symbol_table) const {
    int suffix_id = 0;
    std::string new_func_name = absl::StrCat(kStablehloFuncNamePrefix, "_",
                                             func_name.str(), "_", suffix_id);
    while (symbol_table.lookup(new_func_name)) {
      suffix_id++;
      new_func_name = absl::StrCat(kStablehloFuncNamePrefix, "_",

    op->setOperand(weight_operand_idx, reshape_op);

    // Create a new function with preprocessed types.
    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    OperandRange func_args = op.getArgs();
    func::FuncOp new_float_func = float_func.clone();

    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);

    const func::FuncOp ref_func =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(f_attr.getValue()));
    func::FuncOp new_ref_func = dyn_cast<func::FuncOp>(ref_func->clone());
    return symbol_table.insert(new_ref_func);
  }

  LogicalResult match(LiftedOpT op) const override {