| 164 | } |
| 165 | |
| 166 | XlaOp ArgMinMaxTwoPass(XlaOp input, PrimitiveType output_type, int axis, |
| 167 | bool is_min) { |
| 168 | XlaBuilder* builder = input.builder(); |
| 169 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 170 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 171 | XlaOp init_value; |
| 172 | XlaComputation reducer; |
| 173 | if (is_min) { |
| 174 | init_value = MaxValue(builder, input_shape.element_type()); |
| 175 | reducer = CreateScalarMinComputation(input_shape.element_type(), builder); |
| 176 | } else { |
| 177 | init_value = MinValue(builder, input_shape.element_type()); |
| 178 | reducer = CreateScalarMaxComputation(input_shape.element_type(), builder); |
| 179 | } |
| 180 | |
| 181 | XlaOp input_max = Reduce(input, init_value, reducer, |
| 182 | /*dimensions_to_reduce=*/{axis}); |
| 183 | std::vector<int64> broadcast_dims(input_shape.rank() - 1); |
| 184 | std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); |
| 185 | std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); |
| 186 | // Compute a mask that has 1s for elements equal to the maximum. |
| 187 | XlaOp partial_mask = |
| 188 | ConvertElementType(Eq(input, input_max, broadcast_dims), output_type); |
| 189 | |
| 190 | // In order to make identity elements for a bitwise And, we: |
| 191 | // Left shift the 1 to the leftmost bit, yielding 0x10...0 |
| 192 | // Arithmetic right shift the 1 back to the rightmost bit, yielding |
| 193 | // 0xFF...F |
| 194 | int32 bits_in_type = |
| 195 | ShapeUtil::ByteSizeOfPrimitiveType(output_type) * 8 - 1; |
| 196 | XlaOp shift_amount = ConstantR0WithType(builder, output_type, bits_in_type); |
| 197 | XlaOp full_mask = ShiftRightArithmetic( |
| 198 | ShiftLeft(partial_mask, shift_amount), shift_amount); |
| 199 | |
| 200 | // And with the vector [0, 1, 2, ...] to convert each 0xFF...F into its |
| 201 | // index. |
| 202 | |
| 203 | const int64 axis_size = ShapeUtil::GetDimension(input_shape, axis); |
| 204 | XlaOp iota = Iota(builder, output_type, axis_size); |
| 205 | XlaOp product = And(full_mask, iota, /*broadcast_dimensions=*/{axis}); |
| 206 | |
| 207 | // If there are multiple maximum elements, choose the one with the highest |
| 208 | // index. |
| 209 | return Reduce(product, MinValue(builder, output_type), |
| 210 | CreateScalarMaxComputation(output_type, builder), |
| 211 | /*dimensions_to_reduce=*/{axis}); |
| 212 | }); |
| 213 | } |
| 214 | } // namespace |
| 215 | |
| 216 | XlaOp ArgMax(XlaOp input, PrimitiveType output_type, int axis) { |
no test coverage detected