// CHECK-SAME: devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]
          %rep:2 = tf_device.replicate([%arg0, %arg1] as %arg30: tensor<*x!tf_type.resource<tensor<f32>>>,
                              [%arg2, %arg3] as %arg31: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>)

      TF_NewStatus(), TF_DeleteStatus);
  TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get());
  CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  const int num_devices = TF_DeviceListCount(devices);
  for (int i = 0; i < num_devices; ++i) {
    const string dev_type(TF_DeviceListType(devices, i, status.get()));

                                                   size_t proto_len,
                                                   TF_Status* status);

// TODO(b/166642410): It would be nice, for custom devices and for other users,
// to have a non-string representation of devices (TF_Device) extracted from
// tensors/ops/etc. and usable in APIs like OpSetDevice/ResetOp/etc.

#define TFE_CUSTOM_DEVICE_VERSION 4

		expectedResult []kubecontainer.CDIDevice
	}{
		{
			description: "empty CDI devices",
			claimInfo: &ClaimInfo{
				ClaimInfoState: state.ClaimInfoState{
					CDIDevices: map[string][]string{},
				},
			},
		},
		{
			description: "nil CDI devices",
			claimInfo:   &ClaimInfo{},
		},
		{
			description: "valid CDI devices",
			claimInfo: &ClaimInfo{
				ClaimInfoState: state.ClaimInfoState{

    // Once the rollout to GPU is complete, this flag can be deprecated.
    bool enabled_for_gpu_;

   private:
    // Devices for which using Device API (PjRt) is allowed in the XlaLaunch op.
    // This can only be modified programmatically.
    absl::flat_hash_set<std::string> xla_launch_allowed_devices_;
    // Devices for which using Device API (PjRt) is allowed in the
    // XlaCompileOnDemand op. This can only be modified programmatically.

				annotations = append(annotations, claimAnnotations...)

				devices := claimInfo.cdiDevicesAsList()
				klog.V(3).InfoS("Add CDI devices", "claim", *claimName, "CDI devices", devices)
				cdiDevices = append(cdiDevices, devices...)

				return nil
			})
			if err != nil {
				return nil, fmt.Errorf("locked cache operation: %w", err)
			}
		}
	}

    int64_t device_id;
    if (!device_attr) {
      device_id = kAnyDevice;
    } else if (devices.find(device_attr) != devices.end()) {
      device_id = devices[device_attr];
    } else {
      device_id = 1 + devices.size();
      devices[device_attr] = device_id;
    }

    auto& alias_analysis = side_effect_analysis.GetAliasAnalysis();

  "tf_device.replicate" () ({
// expected-error@-1 {{'tf_device.replicate' op expects number of devices (2) to be equal to 'n' (3)}}
  ^entry:
    tf_device.return
  }) {devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}, n = 3 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
}

// -----

// Check that replicate op's 'devices' attribute must consist of dictionary
// with values as list with size equal to 'n' attribute.

LogicalResult GetTpuDeviceAssignment(
    ClusterOp cluster, ReplicateOp replicate, mlir::ModuleOp module,
    absl::StatusOr<TPUDeviceAssignment>& status_or_tpu_device_assignment) {
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();

  uint32_t num_replicas = replicate.getN();

  auto num_cores_per_replica_attr = cluster->getAttrOfType<mlir::IntegerAttr>(

            "operate on resources that live on composite devices."),
       Flag("tf_mlir_enable_strict_clusters", &enable_mlir_strict_clusters,
            "Do not allow clusters that have cyclic control dependencies."),
       Flag("tf_mlir_enable_multiple_local_cpu_devices",
            &enable_mlir_multiple_local_cpu_devices,
            "Enable multiple local CPU devices. CPU ops which are outside "