EXPECT_EQ(attr_names.size(), 2);
  EXPECT_EQ(attr_names[0], "tf.entry_function");
  EXPECT_EQ(attr_names[1],
            mlir::tf_saved_model::kTfSavedModelInitializerTypeAttr);
}

TEST(CallGraphUtilTest, GetEntryFunctions) {
  const char *const code = R"mlir(
func.func @entry_func_1(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {

    }
  }
  return success();
}

void XlaValidateInputsPass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symtab(module);
  llvm::SmallVector<func::FuncOp> entry_funcs = GetEntryFunctions(module);
  if (entry_funcs.empty()) {
    LOG(WARNING) << "missing entry functions";
  }

  if (HasNoNestedEntryFunctions(entry_funcs, symtab).failed()) {
    return signalPassFailure();
  }

  for (const auto &attr : GetEntryFunctionAttributeNames()) {
    if (func->hasAttr(attr)) {
      return true;
    }
  }
  return false;
}

llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module) {
  llvm::SmallVector<func::FuncOp> entry_funcs;
  module.walk([&](func::FuncOp func) {
    // A model may have multiple graphs, with each graph having its own entry.

// Check if a function is an entry in an MLIR module.
bool IsEntryFunction(func::FuncOp func);

// Get all the entry functions in an MLIR module.
llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module);

// Get all the functions referenced in a symber user op and save them in
// `callees`.
LogicalResult GetCallees(SymbolUserOpInterface op, SymbolTable &symtab,

  ModuleOp module = getOperation();
  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()) {