class MockMlirOptimizationPass : public MlirOptimizationPass {
 public:
  MOCK_METHOD(llvm::StringRef, name, (), (const, override));
  MOCK_METHOD(MlirOptimizationPassState, GetPassState,
              (const DeviceSet* device_set, const ConfigProto& config_proto,
               const Graph& graph,
               const FunctionLibraryDefinition& function_library),
              (const, override));
  MOCK_METHOD(Status, Run,

  std::vector<std::unique_ptr<Device>> devices;
  devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU));
  devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU));

  std::unique_ptr<DeviceSet> device_set(new DeviceSet());
  for (auto& device : devices) {
    device_set->AddDevice(device.get());
  }

  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  SessionOptions session_options;

  mlir::func::registerAllExtensions(registry);
  // clang-format on
}

Status MlirFunctionOptimizationPass::Run(
    const std::string& function_name, const DeviceSet& device_set,
    const ConfigProto& config_proto,
    const FunctionOptimizationPass::FunctionOptions& function_options,
    std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,

      : error_node_(error_node), error_device_(error_device) {}
  tensorflow::Status Run(const std::string& function_name,
                         const tensorflow::DeviceSet& device_set,
                         const tensorflow::ConfigProto& config_proto,
                         const FunctionOptions& function_options,
                         std::unique_ptr<tensorflow::Graph>* graph,

    const string& function_name, const FunctionLibraryDefinition& flib_def) {
  const FunctionDef* func_def = flib_def.Find(function_name);
  TF_RET_CHECK(func_def) << "Could not find " << function_name;

  jit::DeviceSet device_set;

  for (const NodeDef& ndef : func_def->node_def()) {
    VLOG(3) << ndef.DebugString();
    if (!ndef.device().empty()) {
      TF_ASSIGN_OR_RETURN(jit::DeviceId device_id,