ASSERT_TRUE(module);
  auto entry_funcs = GetEntryFunctions(*module);
  EXPECT_EQ(entry_funcs.size(), 2);
  EXPECT_EQ(entry_funcs[0].getSymName(), "entry_func_1");
  EXPECT_EQ(entry_funcs[1].getSymName(), "entry_func_2");
}

TEST(CallGraphUtilTest, GetCallees) {
  const char *const code = R"mlir(
func.func @entry_func(%arg0: tensor<i32>) -> tensor<i32> attributes {tf_saved_model.initializer_type = ""} {

  EXPECT_TRUE(
      IsSupportedByReplicatedBridge(graph, /*function_library=*/nullptr));
}

TEST(IsSupportedByReplicatedBridge, ReplicatedModule) {
  const char* const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {
  %0 = "tf.Identity"(%arg0) {_tpu_replicate = "cluster"} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}
)mlir";

  // Returns `success()` if `entry_func_op`'s body is eligible for rewriting. At
  // this point `entry_func_op`'s signature has not been reset with quantized
  // types.
  virtual LogicalResult match(func::FuncOp entry_func_op,
                              const Method& quantization_method) const = 0;

  // Rewrites the `entry_func_op`'s body.
  virtual void rewrite(func::FuncOp entry_func_op,

                         FunctionLibraryDefinition* flib_def,
                         absl::flat_hash_set<Node*>* control_ret_nodes) {
  mlir::StringAttr entry_func_id =
      mlir::StringAttr::get(module.getContext(), "main");
  std::optional<FuncOp> entry_func;
  FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
                                          FunctionDefLibrary());

                         FunctionLibraryDefinition* flib_def,
                         absl::flat_hash_set<Node*>* control_ret_nodes) {
  mlir::StringAttr entry_func_id =
      mlir::StringAttr::get(module.getContext(), "main");
  std::optional<FuncOp> entry_func;
  FunctionLibraryDefinition temp_flib_def(OpRegistry::Global(),
                                          FunctionDefLibrary());

                                          func::FuncOp entry_func,
                                          std::string* serialized_func_module) {
  ModuleOp module = entry_func->getParentOfType<ModuleOp>();
  SymbolTable entry_module_table(module);
  llvm::SmallVector<func::FuncOp, 4> referenced({entry_func});

  // Create a new module to hold func and all referenced functions.

      lifted_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
  SymbolTable symbol_table(*module_op);
  auto entry_func = dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));
  auto lifted_dot_general_op =
      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(entry_func);

  EXPECT_TRUE(isa<TF::XlaCallModuleOp>(lifted_op));
  EXPECT_EQ(