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.
LogicalResult CreateNewDeviceLaunchOp(

    // 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(),
                                            yield_op.getOperands());

  // arguments.
  Region& op_region = op.getBody();
  for (auto p : llvm::zip(live_ins, outlined_func_block->getArguments())) {
    replaceAllUsesInRegionWith(std::get<0>(p), std::get<1>(p), op_region);
  }

  // Move all instructions in cluster_op into outlined_function's only block.
  auto& op_body = op.GetBody().getOperations();
  outlined_func_block->getOperations().splice(

      removed_params.set(i);
      callee.getBody().addArgument(param.getType(), param.getLoc());
      param.replaceAllUsesWith(callee.getArguments().back());
      removed_params.push_back(false);
    }
  }
  // Remove old resource-type parameters.
  callee.getBody().front().eraseArguments(removed_params);
  // Update function type.

  launch.replaceAllUsesWith(launch.GetBody().getTerminator()->getOperands());

  // For all inner ops, assign the launch device as a `device` attribute.
  if (failed(AssignDevicesInRegion(tf_dialect, launch, launch.getBody())))
    return failure();

  // Move all inner ops of the launch to the block containing the launch.
  auto body = launch.GetBody().without_terminator();

    }

    public String getTitle() {
        return title;
    }

    public void setTitle(final String title) {
        this.title = title;
    }

    public String getBody() {
        return body;
    }

    public void setBody(final String body) {
        this.body = body;
    }

    MakeShapeOpInvariant(replicate_op, num_replicas, replicate_block, shape_op);
  });

  Region* replicate_region = &replicate_op.getBody();
  std::optional<DictionaryAttr> virtual_device_list = replicate_op.getDevices();
  for (Operation& inner_op :
       llvm::make_early_inc_range(replicate_op.GetBody())) {
    if (llvm::isa<tf_device::ReturnOp>(inner_op)) continue;
    // Skip hoisting if the inner op device attribute is a virtual device

SmallVector<TF::VarHandleOp> CloneVarHandleOpsIntoSaveFunc(
    func::FuncOp save_func, const ArrayRef<TF::VarHandleOp> var_handle_ops) {
  Block& save_op_block = save_func.getBody().front();

  IRMapping mapper{};
  SmallVector<TF::VarHandleOp> cloned_var_handle_ops = {};
  for (auto var_handle_op : var_handle_ops) {
    Operation* cloned_var_handle_op = var_handle_op->clone(mapper);

           func_op.getArguments().take_back(call_op.getNumOperands()),
           call_op.getOperands())) {
    arg_mapper.map(func_arg, operand);
  }

  Region& function_body = func_op.getBody();
  IRMapping new_op_mapper;
  for (Operation& op : function_body.getOps()) {
    if (llvm::isa<func::ReturnOp>(op)) {
      for (auto [call_result, return_value] :

      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {
        // Only mirrored args of ReplicateOp can be annotated.
        if (mirrored_replicate_args.count(block_arg.getArgNumber()) == 0) {
          continue;
        }
      } else if (!operand.getParentRegion()->isProperAncestor(
                     &replicate.getBody())) {
        // Not a replication-invariant operand.
        continue;
      }