| 734 | } |
| 735 | |
| 736 | StatusOr<HloInstruction*> HloComputation::DeepCopyHelper( |
| 737 | HloInstruction* instruction, ShapeIndex* index, |
| 738 | const std::function< |
| 739 | HloInstruction*(HloInstruction* leaf, const ShapeIndex& leaf_index, |
| 740 | HloComputation* computation)>& copy_leaf) { |
| 741 | if (instruction->shape().IsTuple()) { |
| 742 | std::vector<HloInstruction*> elements; |
| 743 | for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); |
| 744 | i++) { |
| 745 | HloInstruction* gte = |
| 746 | AddInstruction(HloInstruction::CreateGetTupleElement( |
| 747 | ShapeUtil::GetTupleElementShape(instruction->shape(), i), |
| 748 | instruction, i)); |
| 749 | |
| 750 | index->push_back(i); |
| 751 | TF_ASSIGN_OR_RETURN(HloInstruction * element, |
| 752 | DeepCopyHelper(gte, index, copy_leaf)); |
| 753 | elements.push_back(element); |
| 754 | index->pop_back(); |
| 755 | } |
| 756 | return AddInstruction(HloInstruction::CreateTuple(elements)); |
| 757 | } |
| 758 | if (instruction->shape().IsToken()) { |
| 759 | // Tokens have no on-device representation and cannot be copied. Pass |
| 760 | // through transparently. |
| 761 | return instruction; |
| 762 | } |
| 763 | |
| 764 | // Array shape. |
| 765 | TF_RET_CHECK(instruction->shape().IsArray()); |
| 766 | return copy_leaf(instruction, *index, this); |
| 767 | } |
| 768 | |
| 769 | StatusOr<HloInstruction*> HloComputation::DeepCopyInstruction( |
| 770 | HloInstruction* instruction, const ShapeTree<bool>* indices_to_copy, |