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";

      return failure();
    }
  }
  return success();
}

LogicalResult HasTopLevelCompilationMarker(
    llvm::SmallVector<func::FuncOp> &entry_funcs) {
  for (auto &entry_func : entry_funcs) {
    if (entry_func->hasAttr(mlir::TF::kCompileDeviceTypeAttr)) {
      entry_func->emitError() << "TF2XLA MLIR Non-replicated Phase 1 Bridge "
                                 "does not support top-level compilation "

  SymbolTableCollection symbol_table_collection;
  SymbolTable symtab = symbol_table_collection.getSymbolTable(module);
  llvm::SmallVector<FuncOp> entry_funcs = GetEntryFunctions(module);
  for (auto &entry_func : entry_funcs) {
    if (EncapsulateFirstXlaCompilablePartitionedCalls(
            entry_func, symbol_table_collection, symtab)
            .failed()) {
      return signalPassFailure();
    }
  }
}

  // 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,

    SymbolTable symbol_table(module_op);

    auto entry_func_op = dyn_cast_or_null<func::FuncOp>(
        symbol_table.lookup(GetEntryFunctionName(op)));
    if (!entry_func_op) return failure();

    // Replace the XlaCallModuleOp with a new CallOp.
    rewriter.replaceOpWithNewOp<func::CallOp>(op, entry_func_op, op.getArgs());
    return success();
  }
};

    // have `tf_saved_model.initializer_type` instead.
    if (IsEntryFunction(func)) {
      entry_funcs.push_back(func);
    }
  });
  return entry_funcs;
}

LogicalResult GetCallees(SymbolUserOpInterface op, SymbolTable &symtab,
                         llvm::SmallVector<func::FuncOp> &callees) {
  for (auto attr : op->getAttrs()) {
    auto sym = mlir::dyn_cast<SymbolRefAttr>(attr.getValue());

                         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());

  if (!entry_func_symbol) {
    return op.emitOpError() << "does not have "
                            << kStablehloEntryFunctionAttrName << " attribute";
  }
  auto entry_func =
      symbol_table.lookupSymbolIn<func::FuncOp>(module, entry_func_symbol);
  if (!entry_func) {
    return op.emitOpError()
           << "references an unknown entry function " << entry_func_symbol;
  }

  OpBuilder builder(module.getContext());