mlir::SymbolTableCollection* symbol_table = nullptr);

// Creates a new MLIR module that contains only the given functions and all
// reachable functions from them.
mlir::FailureOr<mlir::OwningOpRef<mlir::ModuleOp>> CreatePrunedModule(
    mlir::ModuleOp module, llvm::ArrayRef<llvm::StringRef> function_names);

}  // namespace TF
}  // namespace mlir

      }

      return WalkResult::advance();
    });
    if (result.wasInterrupted()) {
      return result;
    }
  }

  return WalkResult::advance();
}

FailureOr<OwningOpRef<ModuleOp>> CreatePrunedModule(
    ModuleOp module, llvm::ArrayRef<llvm::StringRef> function_names) {
  SymbolTableCollection symbol_table;
  OpBuilder builder(module.getContext());

  OwningOpRef<ModuleOp> pruned =