}
}

ParallelDevice::ParallelDevice(const std::vector<std::string>& devices,
                               bool is_async, int in_flight_nodes_limit)
    : underlying_devices_(devices),
      default_cancellation_manager_(absl::make_unique<CancellationManager>()) {
  device_threads_.reserve(devices.size());
  for (int device_index = 0; device_index < devices.size(); ++device_index) {
    device_threads_.emplace_back(new DeviceThread(

void LayoutAssignmentPass::runOnOperation() {
  func::FuncOp func = getOperation();

  // Get runtime devices information from the closest parent module.
  RuntimeDevices devices;
  if (failed(::tensorflow::GetDevicesFromOp(func->getParentOfType<ModuleOp>(),
                                            &devices)))
    return signalPassFailure();

  // If there is no runtime device information and data format is not explicitly

      location, concat_dim_type, concat_dimension_attr);

  // Correctly set output shapes of concat op output if output shape is
  // statically known. Since the shape of TPUExecute op must be the same
  // across logical devices, we refer to the shape of 0th logical device
  // computation output.
  mlir::Type output_type;
  auto input_type = mlir::cast<mlir::TensorType>(inputs[0].getType());

  if (input_type.hasRank()) {

func.func @do_not_hoist_ops_with_virtual_device(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) {
  %0:8 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>) {devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]}, n = 2: i32} {
    %1 = "tf.Shape"(%ri) {device = "", T = "tfdtype$DT_FLOAT", out_type = "tfdtype$DT_INT32"} : (tensor<*xf32>) -> tensor<?xi32>

    rollout_config.enabled_for_xla_launch_ = true;
    rollout_config.enabled_for_compile_on_demand_ = true;
    rollout_config.enabled_for_gpu_ = true;

    // Set flag to enable using XLA devices. PJRT currently is only supported
    // for XLA devices.
    GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;

    // Add and setup the GPU device.
    auto device_type = DeviceType(DEVICE_GPU);

  // in an initialized context.
  for (auto d = devices.begin(); d != devices.end();) {
    if (absl::StrContains(d->get()->name(), "CPU:0")) {
      d = devices.erase(d);
    } else {
      ++d;
    }
  }

  status->status = tensorflow::unwrap(ctx)->AddDevices(std::move(devices));
}

void TFE_InsertConfigKeyValue(TFE_Context* ctx, const char* key,

  llvm::SmallVector<Value, 4> inputs;
  // The result values of the function.
  llvm::SmallVector<Value, 4> results;
  // The devices of the input values. It should have the same size as inputs.
  llvm::SmallVector<std::string, 4> input_devices;
  // The devices of the result values. It should have the same size as results.
  llvm::SmallVector<std::string, 4> result_devices;

                                DenseIntElementsAttr replica_groups,
                                Operation* op) {
  // Use special group_key 0 to represent "all available devices". This
  // shall resolve to a DeviceAssignment that includes all devices intended for
  // replica_groups.
  IntegerAttr group_size = builder.getI32IntegerAttr(replica_groups.size());
  IntegerAttr group_key = builder.getI32IntegerAttr(0);

      TFE_ContextListDevices(context.get(), status.get()), TF_DeleteDeviceList);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
  bool has_tpu = false;
  for (int device_index = 0; device_index < TF_DeviceListCount(devices.get());
       ++device_index) {
    std::string device_type =
        TF_DeviceListType(devices.get(), device_index, status.get());

  TF_ASSERT_OK(s.ToGraph(graph.get()));

  // This is needed to register the XLA_GPU device.
  std::vector<std::unique_ptr<Device>> devices;
  TF_ASSERT_OK(DeviceFactory::AddDevices(
      SessionOptions(), "/job:localhost/replica:0/task:0", &devices));

  // Scope::ToGraph loses the assigned device name since it goes through
  // GraphDef/NodeDef which does not have a field for the assigned device name.