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();

          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_);

    BlockArgument arg = std::get<1>(sharding_and_arg);
    if (failed(VerifySharding(arg.getType(), sharding))) return mlir::failure();
  }
  Operation* terminator = function_block.getTerminator();
  for (auto sharding_and_retval :
       llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
    const auto& sharding = std::get<0>(sharding_and_retval);

  FunctionType type;
  if (passthru_extra_args) {
    type = FunctionType::get(context, types, types);
  } else {
    SmallVector<Type, 4> result_types;
    auto operands = region.front().getTerminator()->getOperandTypes();
    result_types.append(operands.begin(), operands.end());
    type = FunctionType::get(context, types, result_types);
  }

  auto outlined_func = builder.create<func::FuncOp>(loc, name, type);

    unsigned region_index) {
  int num_region_results =
      GetRegionBlockWithIndex(region_index).getTerminator()->getNumOperands();

  int return_value_offset = 0;
  for (int region_id = 0; region_id < region_index; ++region_id)
    return_value_offset +=
        GetRegionBlockWithIndex(region_id).getTerminator()->getNumOperands();

  return getResults().slice(return_value_offset, num_region_results);
}