auto new_launch_op = builder.create<tf_device::LaunchOp>(
      old_launch_op->getLoc(), old_launch_op->getDeviceAttr(),
      /*result_types=*/new_launch_op_results_types);

  new_launch_op.getBody().takeBody(old_launch_op->getBody());
  new_launch_op.GetBody().getTerminator()->setOperands(new_launch_op_results);

  return new_launch_op;
}

// Create the new device launch op which wraps the copy op.

    region.push_back(new_block);
    builder.setInsertionPointToEnd(&region.front());
    Operation *inner_op = builder.clone(*op);
    builder.create<YieldOp>(loc, inner_op->getResults());
    outer_op.getBody().takeBody(region);
    // Careful: We can't use outer_op.getResults(), because that also includes
    // the control token.
    op->replaceAllUsesWith(outer_op.getOutputs());
    op->erase();

    // We will "steal" the body of the island and replace it with a call to the
    // new function later.
    {
      YieldOp yield_op = island_op.GetYield();
      outlined_func.getBody().takeBody(island_op.getBody());

      // Replace the yield with a return
      OpBuilder replacer(yield_op);
      island_op.getBody().push_back(new Block);
      replacer.create<mlir::func::ReturnOp>(yield_op.getLoc(),