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

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

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

        std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
      }

      rewriter.setInsertionPoint(cloned_func.getBody().front().getTerminator());
      rewriter.create<TF::_XlaSendFromHostOp>(
          func.getLoc(),
          cloned_func.getBody().front().getTerminator()->getOperands(),
          /*dynamic_key=*/dynamic_key, op.getRecvKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),

  llvm::SmallVector<Type, 8> arg_types;
  for (auto arg : func.getArguments()) arg_types.push_back(arg.getType());
  func.setType(
      FunctionType::get(func.getContext(), arg_types,
                        func.front().getTerminator()->getOperandTypes()));
}

// Holds the size value of a tensor list and whether the size is statically
// known (fixed).
struct SizeInfo {
  Value size;
  bool fixed;
};

  const int size_before = output_types->size();
  for (int index = start; index < end; ++index) {
    Block& block = parallel_execute.GetRegionBlockWithIndex(index);
    auto terminator_operand_types = block.getTerminator()->getOperandTypes();
    output_types->append(terminator_operand_types.begin(),
                         terminator_operand_types.end());
  }
  return output_types->size() - size_before;
}