func::FuncOp func = cluster_func.getFuncOp();
    Block& block = func.front();
    for (Value read_operand : read_operands)
      block.addArgument(read_operand.getType(), loc);

    func.setType(FunctionType::get(&getContext(), block.getArgumentTypes(),
                                   func.getResultTypes()));
    cluster_func.erase();
  }
}

}  // namespace

      arg.setType(resultType);
    }
    llvm::SmallVector<Type, 8> arg_types;
    for (auto arg : func.getArguments()) arg_types.push_back(arg.getType());
    func.setType(
        FunctionType::get(func.getContext(), arg_types,
                          func.front().getTerminator()->getOperandTypes()));
    return WalkResult::advance();
  });

}

func::FuncOp CreateSessionInitFunc(ModuleOp module) {
  constexpr char kSessionInitFuncName[] = "SessionInitializerFunction";

  mlir::OpBuilder builder(module.getBodyRegion());
  auto func_type =
      FunctionType::get(module.getContext(), /*inputs=*/{}, /*results=*/{});
  auto func = builder.create<func::FuncOp>(module->getLoc(),
                                           kSessionInitFuncName, func_type);

  auto tensor_string_type =
      RankedTensorType::get({}, TF::StringType::get(builder.getContext()));
  main_func.getArguments().front().setType(tensor_string_type);
  main_func.setType(
      FunctionType::get(builder.getContext(), {tensor_string_type},
                        main_func.getFunctionType().getResults()));
  // Name of the main function for the eventual executable needs to be set.

        m.print(os);
        return os.str();
      });

  py::class_<mlir::func::FuncOp>(m, "FuncOp")
      .def("create",
           [](mlir::Location location, std::string name,
              mlir::FunctionType type) {
             auto func = mlir::func::FuncOp::create(location, name, type);
             func.addEntryBlock();
             return func;
           })
      .def(
          "getBody",

              module.getContext()));

      arg.setType(new_arg_type);
    }

    // Update the function type.
    func.setType(mlir::FunctionType::get(module.getContext(),
                                         func.getBody().getArgumentTypes(),
                                         func.getFunctionType().getResults()));
  }
  return success();
}

  // Handle the entry functions only.
  if (func.getName() != "main" && (!attrs || attrs.empty())) {
    return;
  }

  OpBuilder builder(func);
  FunctionType func_type = func.getFunctionType();
  llvm::SmallVector<Type, 4> new_input_types(func_type.getInputs().begin(),
                                             func_type.getInputs().end());

      op.erase();
    }
  }
  if (lifted_op_and_arg_idx.empty()) return success();

  // Update the function signature as well as its uses.
  target_func.setType(FunctionType::get(target_func.getContext(),
                                        block.getArgumentTypes(),
                                        func_type.getResults()));

  IRMapping mapping;
  OpBuilder builder(module_op);

                                    PatternRewriter& rewriter,
                                    Block* insertion_point) {
  rewriter.setInsertionPointToStart(insertion_point);

  auto ftype = FunctionType::get(rewriter.getContext(), {type}, {type});
  auto name = rewriter.getStringAttr("gelu_decomp");
  func::FuncOp new_func = rewriter.create<func::FuncOp>(
      insertion_point->front().getLoc(), name, ftype);

  builder.setInsertionPointToEnd(&module_op.getBodyRegion().front());

  auto filename_input_type = RankedTensorType::get(
      /*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());

  FunctionType func_type = builder.getFunctionType(
      /*inputs=*/{filename_input_type}, /*results=*/{});
  auto save_func = builder.create<func::FuncOp>(
      NameLoc::get(builder.getStringAttr(kTfQuantSaveFuncName)),