// Get the number of devices per host.
  int device_num = 0;
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(
          getOperation()->getParentOfType<ModuleOp>(), &devices)))
    return signalPassFailure();
  for (const auto& device_name : devices.device_names()) {
    if (device_name.has_type && device_name.type == "TPU") ++device_num;
  }

  if (device_num == 0) return;

                                                    void* device_info,
                                                    TF_Status* status);

// Returns whether `device_name` maps to a registered custom device.
TF_CAPI_EXPORT extern bool TFE_IsCustomDevice(TFE_Context* ctx,
                                              const char* device_name);

                                               const string& device_name,
                                               int device_ordinal) {
  return Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
      DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
      absl::StrCat("device: ", device_name, " device"));
}

}  // namespace

    const std::optional<FunctionDef>& fdef) {
  SessionOptions options;
  auto* device_count = options.config.mutable_device_count();
  for (const auto& device_name : device_names) {
    device_count->insert({device_name, 1});
  }

  std::vector<std::unique_ptr<Device>> devices;
  TF_CHECK_OK(DeviceFactory::AddDevices(
      options, "/job:localhost/replica:0/task:0", &devices));

  for (auto& func : *funcs) {
    func.walk([&hardware_names, &index](mlir::Operation* op) {
      auto device_name = GetDeviceName(op);
      if (!device_name.has_value()) return;

      auto iter = hardware_names->find(*device_name);
      if (iter == hardware_names->end()) {
        hardware_names->insert({*device_name, index++});
      }
    });
  }

  // Build the flatbuffer.

  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
  TFE_OpSetAttrString(op, "container", "localhost", 0);
  TFE_OpSetAttrString(op, "shared_name", "", 0);
  if (!device_name.empty()) {
    TFE_OpSetDevice(op, device_name.c_str(), status);
  }
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_TensorHandle* var_handle = nullptr;
  int num_retvals = 1;
  TFE_Execute(op, &var_handle, &num_retvals, status);

static std::unique_ptr<OpKernel> GetFakeKernel(const char* device_name,
                                               const char* op_name,
                                               const char* node_name,
                                               Status* status) {
  NodeDef def;
  def.set_op(op_name);
  def.set_name(node_name);
  def.set_device(device_name);
  def.add_input("input1");
  def.add_input("input2");

// `device_name`.
TFE_TensorHandle* CreateVarHandle(TFE_Context* ctx,
                                  const tensorflow::string& device_name,
                                  const tensorflow::string& variable_name);

// Create a variable with value `value` and name `variable_name` on a device
// with name `device_name`.
TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value,

        shape_determination_fns;

    Options(std::string device_name_prefix, std::string device_name,
            int device_ordinal, std::string compilation_device_name,
            std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                shape_determination_fns)
        : device_name_prefix(device_name_prefix),
          device_name(device_name),
          device_ordinal(device_ordinal),

  }
  Status s;
  const char* device_name = DeviceName(&s);
  if (!s.ok()) {
    device_name = "<error fetching device name>";
  }
  return absl::StrCat("TensorHandle(", value_string, ", shape=", shape_string,
                      ", dtype=", DataType_Name(DataType()), ", device=\"",
                      device_name, "\")");
}

Status ImmediateExecutionTensorHandle::SummarizeValue(