| 852 | } |
| 853 | |
| 854 | XlaOp XlaBuilder::DynamicSlice(XlaOp operand, XlaOp start_indices, |
| 855 | absl::Span<const int64> slice_sizes) { |
| 856 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 857 | HloInstructionProto instr; |
| 858 | |
| 859 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 860 | TF_ASSIGN_OR_RETURN(const Shape* start_indices_shape, |
| 861 | GetShapePtr(start_indices)); |
| 862 | TF_ASSIGN_OR_RETURN( |
| 863 | Shape shape, ShapeInference::InferDynamicSliceShape( |
| 864 | *operand_shape, {*start_indices_shape}, slice_sizes)); |
| 865 | *instr.mutable_shape() = shape.ToProto(); |
| 866 | |
| 867 | for (int64 size : slice_sizes) { |
| 868 | instr.add_dynamic_slice_sizes(size); |
| 869 | } |
| 870 | |
| 871 | return AddInstruction(std::move(instr), HloOpcode::kDynamicSlice, |
| 872 | {operand, start_indices}); |
| 873 | }); |
| 874 | } |
| 875 | |
| 876 | XlaOp XlaBuilder::DynamicSlice(XlaOp operand, |
| 877 | absl::Span<const XlaOp> start_indices, |
no test coverage detected