| 1813 | } |
| 1814 | |
| 1815 | XlaOp XlaBuilder::Scatter(XlaOp input, XlaOp scatter_indices, XlaOp updates, |
| 1816 | const XlaComputation& update_computation, |
| 1817 | const ScatterDimensionNumbers& dimension_numbers, |
| 1818 | bool indices_are_sorted, bool unique_indices) { |
| 1819 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1820 | HloInstructionProto instr; |
| 1821 | instr.set_indices_are_sorted(indices_are_sorted); |
| 1822 | |
| 1823 | instr.set_unique_indices(unique_indices); |
| 1824 | |
| 1825 | TF_ASSIGN_OR_RETURN(const Shape* input_shape, GetShapePtr(input)); |
| 1826 | TF_ASSIGN_OR_RETURN(const Shape* scatter_indices_shape, |
| 1827 | GetShapePtr(scatter_indices)); |
| 1828 | TF_ASSIGN_OR_RETURN(const Shape* updates_shape, GetShapePtr(updates)); |
| 1829 | TF_ASSIGN_OR_RETURN(const ProgramShape& to_apply_shape, |
| 1830 | update_computation.GetProgramShape()); |
| 1831 | TF_ASSIGN_OR_RETURN( |
| 1832 | Shape shape, ShapeInference::InferScatterShape( |
| 1833 | *input_shape, *scatter_indices_shape, *updates_shape, |
| 1834 | to_apply_shape, dimension_numbers)); |
| 1835 | *instr.mutable_shape() = shape.ToProto(); |
| 1836 | |
| 1837 | *instr.mutable_scatter_dimension_numbers() = dimension_numbers; |
| 1838 | |
| 1839 | AddCalledComputation(update_computation, &instr); |
| 1840 | return AddInstruction(std::move(instr), HloOpcode::kScatter, |
| 1841 | {input, scatter_indices, updates}); |
| 1842 | }); |
| 1843 | } |
| 1844 | |
| 1845 | XlaOp XlaBuilder::Conditional(XlaOp predicate, XlaOp true_operand, |
| 1846 | const XlaComputation& true_computation, |
no test coverage detected