void AddGpuDevice(const ParsedName& device,
                    const GpuDeviceMetadata& metadata);

  llvm::ArrayRef<ParsedName> device_names() const { return device_names_; }
  size_t NumDevices() const { return device_names_.size(); }

  // Returns GPU device metadata if it is available, otherwise returns None.
  std::optional<GpuDeviceMetadata> GetGpuDeviceMetadata(
      const ParsedName& device) const;

  (*module_ref)->setAttr("tf.devices", device_dict);

  mlir::TF::RuntimeDevices devices;
  EXPECT_TRUE(mlir::succeeded(GetDevicesFromOp(*module_ref, &devices)));

  ASSERT_EQ(devices.NumDevices(), 1);
  ASSERT_EQ(devices.device_names().size(), 1);
  ASSERT_EQ(DeviceNameUtils::ParsedNameToString(devices.device_names()[0]),
            "/job:worker/replica:0/task:0/device:CPU:0");
}

    return signalPassFailure();

  // If there is no runtime device information and data format is not explicitly
  // forced, there is nothing to do.
  if (devices.NumDevices() == 0 && force_data_format_.empty()) return;

  func.walk([&](LayoutSensitiveInterface layout_sensitive_interface) {
    // Get desired op data format.
    StringRef target_data_format = force_data_format_;