class CustomThreadPoolImpl : public thread::ThreadPoolInterface {
 public:
  explicit CustomThreadPoolImpl(int numThreads) {
    underlying_threadpool_.reset(new thread::ThreadPool(
        tensorflow::Env::Default(), "custom_threadpool", numThreads));
    num_schedule_called_ = 0;
  }

  void Schedule(std::function<void()> fn) override {
    num_schedule_called_ += 1;
    underlying_threadpool_->Schedule(std::move(fn));
  }

#     # copybara:uncomment driver = "//tensorflow/compiler/mlir:run_lit.sh",
#     exclude = [
#         "compile.benchmark.large.mlir",
#         "batch_function_fallback.mlir",
#         "create_op.mlir",
#         "custom_thread_pool.mlir",
#     ],
#     # copybara:uncomment flaky = ["compile.error.mlir"],
#     size_override = {
#         "compile.benchmark.small.mlir": "medium",
#         "batching_fallback.mlir": "medium",
#     },