}
  return false;
}

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action) {
  for (int arg_index = 0; arg_index < func.getNumArguments(); ++arg_index) {
    if (IsArgConsumedByFallback(func, arg_index)) action(arg_index);
  }
}

void ForEachArgConsumedByFallback(
    mlir::ModuleOp module,

namespace tfrt {
namespace fallback_async {

bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index);

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action);

void ForEachArgConsumedByFallback(
    mlir::ModuleOp module,
    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action);

}  // namespace fallback_async

  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

  std::vector<std::pair<std::string, int>> func_and_index;
  ForEachArgConsumedByFallback(
      module.get(),
      [&func_and_index](llvm::StringRef func_name, int arg_index) {
        func_and_index.push_back({func_name.str(), arg_index});
      });

  ASSERT_EQ(func_and_index.size(), 1);