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