op->erase();
      return WalkResult::advance();
    }

    if (!devices.has_value()) return WalkResult::advance();

    // Map aliased devices to explicit devices based on replica.
    if (auto launch = dyn_cast<tf_device::LaunchOp>(op))
      if (auto device_by_replica = devices.value().get(launch.getDevice()))
        launch->setAttr(
            kDeviceAttr,

// CHECK: "tf.opA"
// device = "CORE_0"
// CHECK: _parallel_execution_ids = "r0:1"


// Tests devices are not remapped if device is not in replicate devices.
// CHECK-LABEL: func @no_override_device
func.func @no_override_device() {
  tf_executor.graph {
    %0 = tf_executor.island {
      tf_device.replicate {n = 2 : i32, devices = {CORE_0 = ["/CPU:0", "/GPU:1"]}} {
        "tf_device.launch"() ({
          "tf.opA"() : () -> ()

The region held by this operation represents a computation that is replicated
across multiple devices. The number of replications is based on the `n`
attribute. Explicit devices can be populated in the `devices` attribute, and it
must be a mapping of device alias to list of explicit or aliased device names
from the outer scope. The device name map specifies devices on which replicated
ops inside tf_device.replicate will be executed.

			return err
		}
	} else {
		// Add single devicepath to devices
		devices = append(devices, dstPath)
	}
	klog.V(4).Infof("fc: DetachDisk devicePath: %v, dstPath: %v, devices: %v", devicePath, dstPath, devices)
	var lastErr error
	for _, device := range devices {
		err := util.detachFCDisk(c.io, c.exec, device)
		if err != nil {
			klog.Errorf("fc: detachFCDisk failed. device: %v err: %v", device, err)

 private:
  void runOnOperation() override;

  // Given a func and targeted devices, we will try to clonse the func &
  // transform/optimize for those devices.
  // This will only happen if the whole subgraph can be supported by the target
  // or can be supported after some transformations.
  void GetAlternativeGraphForFunc(ArrayRef<std::string> devices,
                                  func::FuncOp func, ModuleOp module,

            lhs_device_attr.getValue() == rhs_device_attr.getValue());
  };

  // Check if tf_executor.NextIteration.Source/tf_executor.NextIteration.Sink
  // pair has matching devices or no devices.
  if (auto source = llvm::dyn_cast<tf_executor::NextIterationSourceOp>(op)) {
    return ops_have_same_device(source, source.GetSink());
  } else if (auto sink = llvm::dyn_cast<tf_executor::NextIterationSinkOp>(op)) {

  if (execute_arg_to_outer_args.empty()) return false;

  // Extract the replicated devices.
  auto devices_attr = replicate.getDevices();
  if (!devices_attr) return false;

  auto device_map = devices_attr.value();
  llvm::SmallDenseMap<llvm::StringRef, llvm::SmallVector<StringRef, 4>> devices;
  devices.reserve(device_map.size());

  for (auto it : device_map) {
    auto device_alias = it.getName().strref();

    // If padded_shape_fn is empty, a default implementation that returns
    // the logical on-device shape without padding is used.
    PaddedShapeFn padded_shape_fn;

    // Set of devices to use. This controls which of the devices on the given
    // platform will have resources allocated. For GPUs this will be
    // filled from visible_gpu_devices list from session configuration.
    std::optional<std::set<int>> allowed_devices;

  // A parallel tensor with scalar integers numbering component devices.
  std::unique_ptr<ParallelTensor> DeviceIDs(TFE_Context* context,
                                            TF_Status* status) const;

  // The number of devices operations run on.
  size_t num_underlying_devices() const { return underlying_devices_.size(); }

  // The devices operations run on.
  const std::vector<std::string>& underlying_devices() const {

  return tensorflow::WrapOpInLaunch(builder, compile_op.getLoc(), compile_op,
                                    compilation_device);
}

// Assigns explicit devices to replicate op. An aliased device is created per
// core, and all replica devices per core are grouped together.
void AssignDevicesToReplicate(
    tf_device::ReplicateOp replicate,
    llvm::ArrayRef<llvm::SmallVector<tensorflow::TPUDeviceAndHost, 8>>