compilation_device = nullptr;
  }
  id_to_compilation_device_.push_back(compilation_device);

  return DeviceId(new_id);
}

string DeviceInfoCache::DebugString(const DeviceSet& device_set) const {
  std::vector<string> names;
  device_set.ForEach([&](DeviceId device_id) {
    names.push_back(string(GetNameFor(device_id)));
    return true;
  });

  return absl::StrCat("[", absl::StrJoin(names, ","), "]");
}

#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "tensorflow/compiler/jit/shape_inference.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/framework/function.h"

  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;
  session_options.config.mutable_graph_options()

  // Whether to run this pass. If this is enabled, the GraphDef will be imported
  // to MLIR even no tf composition file is found.
  ::tensorflow::MlirOptimizationPassState GetPassState(
      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const override;

  // This should be used as a thin mapper around mlir::ModulePass::runOnModule

    device_set.Insert(device_id);
  }

  TF_ASSIGN_OR_RETURN(jit::DeviceId result,
                      PickDeviceForXla(*device_info_cache, device_set,
                                       /*allow_mixing_unknown_and_cpu=*/true));
  VLOG(2) << "For " << function_name << " PickDeviceForXla("
          << device_info_cache->DebugString(device_set) << ") -> "

#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_FAKE_SESSION_H_

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/public/session.h"

namespace mlir {
namespace TF {

#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/utils/string_container_utils.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/resource_var.h"

namespace mlir {
namespace tf_saved_model {

std::string GetVariableName(TF::VarHandleOp var_handle_op) {

    const std::string& function_name, const ConfigProto& config_proto,
    ModuleOp module, const Graph& graph,
    const tensorflow::FunctionLibraryDefinition& function_library) {
  if (GetPassState(/*device_set=*/nullptr, config_proto, graph,
                   function_library) ==
      ::tensorflow::MlirOptimizationPassState::Disabled) {
    VLOG(1) << "Skipping MLIR Graph Optimization Pass"

#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"

#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"