std::string cpu0_device;
  if (failed(tensorflow::GetNonReplicatedCPU0(op, &cpu0_device)))
    return failure();
  if (device_str != tensorflow::GetDeviceAliasForHostOfLogicalCore(0) &&
      device_str != cpu0_device) {
    op->emitError()
        << "TPUCopyWithDynamicShapeOp's device is not a recognized host 0: "
        << device_str;
    return failure();
  }

std::string GetDeviceAliasForLogicalCore(int core_index);

// Virtual device name of the host that is associated with the passed logical
// core. The logical core is the index of a core within a replica.
std::string GetDeviceAliasForHostOfLogicalCore(int core_index);

// Returns true if cluster contains model parallelism based on
// `num_cores_per_replica_attribute`. Otherwise returns false.
bool HasModelParallelism(mlir::tf_device::ClusterOp cluster);

    std::string& host_device) {
  mlir::tf_device::ReplicateOp replicate =
      cluster->getParentOfType<mlir::tf_device::ReplicateOp>();
  if (replicate) {
    host_device = GetDeviceAliasForHostOfLogicalCore(0);
    return mlir::success();
  }

  tensorflow::TPUDevicesAndHosts devices_and_hosts;
  if (failed(GetTPUDevicesAndHostsNotReplicated(devices, cluster,

  assigned_id->setAttr(kICIWeightDistributionMlirBridgeMarker,
                       before_cluster_builder.getBoolAttr(true));
  std::string device = tensorflow::GetDeviceAliasForHostOfLogicalCore(0);
  LaunchOp launch = tensorflow::WrapOpInLaunch(
      &before_cluster_builder, val_bcast.getLoc(), assigned_id, device);

  Value all_reduce =

    // For data parallelism, also add replicated host devices, as these are
    // necessary for outside compilation.
    device_attrs.push_back(builder->getNamedAttr(
        tensorflow::GetDeviceAliasForHostOfLogicalCore(core),
        builder->getStrArrayAttr(hosts_by_core)));
  }

  replicate->setAttr(kDevicesAttr, builder->getDictionaryAttr(device_attrs));
}

  mlir::TF::RuntimeDevices runtime_devices;
  std::string host_device;
  EXPECT_TRUE(mlir::succeeded(
      GetHostDeviceOutsideComputation(runtime_devices, cluster, &host_device)));
  EXPECT_EQ(host_device, GetDeviceAliasForHostOfLogicalCore(0));
}

TEST(TPURewriteDeviceUtilTest, TestGetHostDeviceNotReplicated) {
  mlir::MLIRContext context;
  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();