| 816 | } |
| 817 | |
| 818 | XlaOp XlaBuilder::Slice(XlaOp operand, absl::Span<const int64> start_indices, |
| 819 | absl::Span<const int64> limit_indices, |
| 820 | absl::Span<const int64> strides) { |
| 821 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 822 | HloInstructionProto instr; |
| 823 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 824 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferSliceShape( |
| 825 | *operand_shape, start_indices, |
| 826 | limit_indices, strides)); |
| 827 | *instr.mutable_shape() = shape.ToProto(); |
| 828 | for (int i = 0; i < start_indices.size(); i++) { |
| 829 | auto* slice_config = instr.add_slice_dimensions(); |
| 830 | slice_config->set_start(start_indices[i]); |
| 831 | slice_config->set_limit(limit_indices[i]); |
| 832 | slice_config->set_stride(strides[i]); |
| 833 | } |
| 834 | |
| 835 | return AddInstruction(std::move(instr), HloOpcode::kSlice, {operand}); |
| 836 | }); |
| 837 | } |
| 838 | |
| 839 | XlaOp XlaBuilder::SliceInDim(XlaOp operand, int64 start_index, |
| 840 | int64 limit_index, int64 stride, int64 dimno) { |
no test coverage detected