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

}

// -----

// Tests collecting compilation and execution devices results in an error.

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
  func.func @bad_devices() {
    // expected-error@+1 {{error in fetching TPU compilation/execution devices: no TPU_SYSTEM devices found}}

  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;

// configuring the persistor used in the DeviceCompiler. Please note that
// non-XLA devices aren't supported yet. This is because:
// 1. PjRtClient doesn't support data transfer for non-XLA devices yet
// 2. Fetching the PjRtClient for non-XLA devices is also not supported yet
Status GetOrCreatePjRtDeviceCompilerAndProfiler(
    const OpKernelContext& ctx, const XlaPlatformInfo& platform_info,

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

  auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
  // Get runtime devices information from the closest parent module.
  auto module = getOperation();
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices)))
    return signalPassFailure();

  OpBuilder builder(&getContext());
  llvm::SmallVector<mlir::tf_device::ClusterOp, 4> clusters;

  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));
  device_mgr_ = std::make_unique<StaticDeviceMgr>(std::move(devices));

  OptimizerOptions opts;
  lib_def_ = std::make_unique<FunctionLibraryDefinition>(OpRegistry::Global(),

    Builder* builder, OperationState* state, int n,
    std::optional<DictionaryAttr> devices,
    llvm::ArrayRef<std::pair<ValueRange, Type>> replicated_inputs,
    ValueRange packed_inputs, TypeRange replica_output_types) {
  DCHECK_GE(n, 2);
  state->addAttribute("n", builder->getI32IntegerAttr(n));

  if (devices.has_value()) state->addAttribute("devices", devices.value());

  Region* region = state->addRegion();

// -----

// CHECK-LABEL: func @replicate_with_devices
func.func @replicate_with_devices() {
  tf_device.replicate() {n = 2 : i32, devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}} {
    tf_device.return
  }
  func.return

// CHECK:      tf_device.replicate
// CHECK-SAME: devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}
// CHECK-SAME: n = 2
// CHECK-NEXT:   tf_device.return
}

// -----