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 =