| 1787 | } |
| 1788 | |
| 1789 | XlaOp XlaBuilder::Gather(XlaOp input, XlaOp start_indices, |
| 1790 | const GatherDimensionNumbers& dimension_numbers, |
| 1791 | absl::Span<const int64> slice_sizes, |
| 1792 | bool indices_are_sorted) { |
| 1793 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1794 | HloInstructionProto instr; |
| 1795 | instr.set_indices_are_sorted(indices_are_sorted); |
| 1796 | |
| 1797 | TF_ASSIGN_OR_RETURN(const Shape* input_shape, GetShapePtr(input)); |
| 1798 | TF_ASSIGN_OR_RETURN(const Shape* start_indices_shape, |
| 1799 | GetShapePtr(start_indices)); |
| 1800 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGatherShape( |
| 1801 | *input_shape, *start_indices_shape, |
| 1802 | dimension_numbers, slice_sizes)); |
| 1803 | *instr.mutable_shape() = shape.ToProto(); |
| 1804 | |
| 1805 | *instr.mutable_gather_dimension_numbers() = dimension_numbers; |
| 1806 | for (int64 bound : slice_sizes) { |
| 1807 | instr.add_gather_slice_sizes(bound); |
| 1808 | } |
| 1809 | |
| 1810 | return AddInstruction(std::move(instr), HloOpcode::kGather, |
| 1811 | {input, start_indices}); |
| 1812 | }); |
| 1813 | } |
| 1814 | |
| 1815 | XlaOp XlaBuilder::Scatter(XlaOp input, XlaOp scatter_indices, XlaOp updates, |
| 1816 | const XlaComputation& update_computation, |
no test coverage detected