Operation* tpu_copy_with_dynamic_shape_op) {
  llvm::SmallSetVector<Value, 4> new_launch_op_results;

  new_launch_op_results.insert(
      old_launch_op->GetBody().getTerminator()->getOperands().begin(),
      old_launch_op->GetBody().getTerminator()->getOperands().end());

  for (Value operand : tpu_copy_with_dynamic_shape_op->getOperands()) {
    if (GetOpOfValue(operand)->getParentRegion() ==

    cloned_cond.eraseArgument(idx);
    cloned_body.front().getTerminator()->eraseOperand(idx);
    cloned_body.eraseArgument(idx);
  }

  // Patch up branch function types.
  for (func::FuncOp func : {cloned_cond, cloned_body}) {
    func.setType(
        FunctionType::get(func.getContext(), func.front().getArgumentTypes(),
                          func.front().getTerminator()->getOperandTypes()));
  }

  }

  mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
      builder, cluster, /*before=*/true, launch_block, host_device);

  for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
                               launch.getResults()))
    replaceAllUsesInRegionWith(std::get<0>(result), std::get<1>(result),
                               cluster.getBody());
}

  for (OpOperand& operand :
       reduced_launch_op.GetBody().getTerminator()->getOpOperands()) {
    reduced_return_operand_map.insert(
        {operand.get(), operand.getOperandNumber()});
  }

  for (int i = 0; i < original_launch_op->getNumResults(); i++) {
    Value operand = original_launch_op.GetBody().getTerminator()->getOperand(i);
    auto defining_op = operand.getDefiningOp();

  // If `destination_block` isn't empty, erase its terminator to ensure that it
  // never contains two terminator-like ops after merging.
  if (!destination_block->empty())
    rewriter.eraseOp(destination_block->getTerminator());

  destination_block->addArguments(
      block_arguments_type,
      SmallVector<Location>(block_arguments_type.size(),
                            source_block->getParent()->getLoc()));

    auto results = parallel_execute.GetRegionOutputs(i);
    auto& results_to_remove = results_to_remove_by_region[i];
    results_to_remove.reserve(results.size());
    Operation* terminator = block.getTerminator();
    for (auto result : llvm::enumerate(results)) {
      TF::AssignVariableOp assign_var =
          GetSingleUseResourceWrite(parallel_execute, result.value());
      if (!assign_var) {

        replaceAllUsesInRegionWith(
            result,
            it->getThenBranch().front().getTerminator()->getOperand(
                result.getResultNumber()),
            it2->getThenBranch());
        replaceAllUsesInRegionWith(
            result,
            it->getElseBranch().front().getTerminator()->getOperand(
                result.getResultNumber()),
            it2->getElseBranch());
      }
    }

          GetOpOfValue(cluster_func_operand.value()));
      if (!device_launch_op) continue;
      for (auto result : llvm::zip(
               device_launch_op.getResults(),
               device_launch_op.GetBody().getTerminator()->getOperands())) {
        if (std::get<0>(result) == cluster_func_operand.value() &&
            llvm::isa<TF::TPUAnnotateTensorsWithDynamicShapeOp>(
                std::get<1>(result).getDefiningOp())) {

              {host_if.getThenBranch().front().getTerminator()},
              core_to_compilation_key, core_to_device_ordinal,
              /*control_above=*/true, is_map_oc, communication_key_index)))
        return WalkResult::interrupt();
      if (failed(MoveToHostMultiCluster(
              device_cluster, &if_op.getElseBranch().front(),
              {host_if.getElseBranch().front().getTerminator()},

    RegionRange regions) {
  for (Region* region : regions) {
    assert(llvm::hasSingleElement(*region) && "Expected single block region");
    Block& front = region->front();
    auto old_return = front.getTerminator();
    assert(old_return->getNumOperands() == op_->getNumResults());
    auto new_return_operands = llvm::to_vector<4>(old_return->getOperands());
    new_return_operands.resize(num_new_results_);