| 721 | } |
| 722 | |
| 723 | Literal LiteralBase::Slice(absl::Span<const int64> start_indices, |
| 724 | absl::Span<const int64> limit_indices) const { |
| 725 | CHECK(shape().IsArray()) << "tuple is not supported for slice"; |
| 726 | |
| 727 | DimensionVector result_dimensions; |
| 728 | for (int64 dnum = 0; dnum < shape().rank(); ++dnum) { |
| 729 | CHECK_GE(start_indices[dnum], 0); |
| 730 | CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)) |
| 731 | << "dnum = " << dnum; |
| 732 | int64 dimension = limit_indices[dnum] - start_indices[dnum]; |
| 733 | CHECK_GE(dimension, 0) << "dnum = " << dnum; |
| 734 | result_dimensions.push_back(dimension); |
| 735 | } |
| 736 | const auto result_shape = |
| 737 | ShapeUtil::MakeShapeWithLayout(shape().element_type(), result_dimensions, |
| 738 | LayoutUtil::MinorToMajor(shape())); |
| 739 | switch (result_shape.element_type()) { |
| 740 | case PRED: |
| 741 | return SliceInternal<bool>(result_shape, start_indices); |
| 742 | case U8: |
| 743 | return SliceInternal<uint8>(result_shape, start_indices); |
| 744 | case U16: |
| 745 | return SliceInternal<uint16>(result_shape, start_indices); |
| 746 | case U32: |
| 747 | return SliceInternal<uint32>(result_shape, start_indices); |
| 748 | case U64: |
| 749 | return SliceInternal<uint64>(result_shape, start_indices); |
| 750 | case S8: |
| 751 | return SliceInternal<int8>(result_shape, start_indices); |
| 752 | case S16: |
| 753 | return SliceInternal<int16>(result_shape, start_indices); |
| 754 | case S32: |
| 755 | return SliceInternal<int32>(result_shape, start_indices); |
| 756 | case S64: |
| 757 | return SliceInternal<int64>(result_shape, start_indices); |
| 758 | case F16: |
| 759 | return SliceInternal<half>(result_shape, start_indices); |
| 760 | case BF16: |
| 761 | return SliceInternal<bfloat16>(result_shape, start_indices); |
| 762 | case F32: |
| 763 | return SliceInternal<float>(result_shape, start_indices); |
| 764 | case F64: |
| 765 | return SliceInternal<double>(result_shape, start_indices); |
| 766 | case C64: |
| 767 | return SliceInternal<complex64>(result_shape, start_indices); |
| 768 | case C128: |
| 769 | return SliceInternal<complex128>(result_shape, start_indices); |
| 770 | default: |
| 771 | LOG(FATAL) << "not yet implemented: " |
| 772 | << PrimitiveType_Name(result_shape.element_type()); |
| 773 | } |
| 774 | } |
| 775 | |
| 776 | Literal LiteralBase::Clone() const { |
| 777 | Literal result(shape()); |