TPUColocateCompositeResourceOps> {
  void runOnOperation() override;
};

// Wraps single op in `tf_device.launch` for explicit device assignment.
void WrapOpInLaunch(OpBuilder* builder, Location loc, Operation* op,
                    llvm::StringRef device) {
  builder->setInsertionPoint(op);
  auto launch = builder->create<tf_device::LaunchOp>(

struct DeviceAttributeToLaunch
    : public impl::DeviceAttributeToLaunchPassBase<DeviceAttributeToLaunch> {
  void runOnOperation() override;
};

void WrapOpInLaunch(Operation* op, llvm::StringRef device) {
  OpBuilder builder(op);

  auto launch_op = builder.create<tf_device::LaunchOp>(
      op->getLoc(), builder.getStringAttr(device),
      /*result_types=*/op->getResultTypes());

  mlir::OpBuilder builder(&context);
  auto loc = cluster->getLoc();

  // Wrap the cluster op into a Launch op
  auto launch_op = tensorflow::WrapOpInLaunch(&builder, loc, cluster, device);

  EXPECT_TRUE(llvm::isa<mlir::tf_device::LaunchOp>(launch_op));
  launch_op->erase();
}

}  // namespace

    mlir::tf_device::ParallelExecuteOp old_parallel_execute,
    mlir::tf_device::ParallelExecuteOp* new_parallel_execute);

// Wraps single op in `tf_device.launch` for explicit device assignment.
mlir::tf_device::LaunchOp WrapOpInLaunch(mlir::OpBuilder* builder,
                                         mlir::Location loc,
                                         mlir::Operation* op,
                                         llvm::StringRef device);

    new_parallel_execute->getRegions()[new_idx].takeBody(
        old_parallel_execute.getRegions()[old_idx]);
  }

  return cluster_idx;
}

mlir::tf_device::LaunchOp WrapOpInLaunch(mlir::OpBuilder* builder,
                                         mlir::Location loc,
                                         mlir::Operation* op,
                                         llvm::StringRef device) {

                builder->getStringAttr(GetRandomStateVariableName()))}));
  }
  return state_vars;
}

// Wraps single op in `tf_device.launch` for explicit device assignment.
void WrapOpInLaunch(OpBuilder* builder, Location loc, Operation* op,
                    llvm::StringRef device) {
  OpBuilder::InsertPoint insert_point = builder->saveInsertionPoint();

  auto launch = builder->create<tf_device::LaunchOp>(

      /*compilation_status=*/compilation_status_type, /*program=*/
      llvm::SmallVector<Type, 8>(num_cores_per_replica, program_type),
      compile_op_operands, txt_module, txt_metadata);

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

// Assigns explicit devices to replicate op. An aliased device is created per

                       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 =
      CreateAllReduce(replicate, inner_builder, launch.getResult(0));

  orig_to_new[val_bcast] = all_reduce;