| 668 | } |
| 669 | |
| 670 | Literal LiteralBase::Transpose(absl::Span<const int64> permutation) const { |
| 671 | CHECK(shape().IsArray()) << "Tuple is not supported for transpose"; |
| 672 | CHECK(IsPermutation(permutation, shape().rank())) |
| 673 | << "Given permutation is not a permutation of dimension numbers"; |
| 674 | // To transpose the array, we just permute the dimensions and layout, and |
| 675 | // do a straight memory copy of the raw data set. |
| 676 | // This is considerably faster than iterating over every array element using |
| 677 | // the EachCell<>() and Set<>() APIs. |
| 678 | std::vector<int64> inverse_permutation = InversePermutation(permutation); |
| 679 | Shape permuted_shape = |
| 680 | ShapeUtil::PermuteDimensions(inverse_permutation, shape()); |
| 681 | // Replace the layout with one affine to this shape, such that a |
| 682 | // transpose operation can be performed by leaving the flat values |
| 683 | // representation intact. |
| 684 | // For example, consider the shape F32[11,8]{1,0} under a {1,0} permutation. |
| 685 | // The shape with affine layout resulting from that operation will be |
| 686 | // F32[8,11]{0,1}, since it leaves the original most minor (the 8 sized), the |
| 687 | // most minor. |
| 688 | // |
| 689 | // Essentially, given MinMaj(Di) the position of the Di dimension within the |
| 690 | // minor to major vector, and given T(Di) the index that the original Di |
| 691 | // dimension has within the transposed array, a layout is affine if |
| 692 | // MinMaj(Di) == TMinMaj(T(Di)), with TMinMaj() being the minor to major |
| 693 | // vector of the affine layout. |
| 694 | CHECK(LayoutUtil::IsDenseArray(permuted_shape)); |
| 695 | Layout* layout = permuted_shape.mutable_layout(); |
| 696 | layout->clear_minor_to_major(); |
| 697 | for (auto index : LayoutUtil::MinorToMajor(shape())) { |
| 698 | layout->add_minor_to_major(inverse_permutation[index]); |
| 699 | } |
| 700 | Literal new_literal(permuted_shape); |
| 701 | DCHECK_EQ(ShapeUtil::ByteSizeOf(new_literal.shape()), |
| 702 | ShapeUtil::ByteSizeOf(shape())); |
| 703 | std::memcpy(new_literal.untyped_data(), untyped_data(), size_bytes()); |
| 704 | return new_literal; |
| 705 | } |
| 706 | |
| 707 | template <typename NativeT> |
| 708 | Literal LiteralBase::SliceInternal( |