// 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;

// `tf.devices` dictionary attribute with a full device name as a key, and
// device metadata as a value.
//
// Device names added in full parsed device form:
//   /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num>
//
// Supported device metadata types:
// (1) GpuDeviceMetadata: GPU device compute capability.
void AddDevicesToOp(mlir::Operation* op, const DeviceSet* device_set);

                                                    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

      cpu_value.get(), context.get(), device_name, status.get()));
  EXPECT_EQ(TF_GetCode(status.get()), TF_UNIMPLEMENTED);

  std::array<TFE_TensorHandle*, 2> components{cpu_value.get(), cpu_value.get()};
  TensorHandlePtr device_value = CreatePerDeviceValues(
      context.get(), components, device_name, status.get());
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

    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.

namespace tensorflow {
namespace parallel_device {

// Allocate a parallel device named `device_name` which forwards operations to
// `underlying_devices`, maintaining "parallel tensors" with components placed
// on each underlying device.
//
// For example if `device_name` is
//   "/job:localhost/replica:0/task:0/device:CUSTOM:0"
// and `underlying_devices` is

  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);

#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/platform/test.h"

namespace {

struct LoggingDevice {
  tensorflow::string device_name;
  tensorflow::string underlying_device;
  // Set to true whenever a TensorHandle is copied onto the device
  bool* arrived_flag;
  // Set to true whenever an operation is executed
  bool* executed_flag;