| 134 | } |
| 135 | |
| 136 | XlaOp ArgMinMax(XlaOp input, PrimitiveType output_type, int axis, bool is_min) { |
| 137 | XlaBuilder* builder = input.builder(); |
| 138 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 139 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 140 | XlaOp value_init_value; |
| 141 | if (is_min) { |
| 142 | value_init_value = MaxValue(builder, input_shape.element_type()); |
| 143 | } else { |
| 144 | value_init_value = MinValue(builder, input_shape.element_type()); |
| 145 | } |
| 146 | int64 dimension_size = input_shape.dimensions(axis); |
| 147 | auto index_type = dimension_size <= INT32_MAX ? S32 : output_type; |
| 148 | XlaOp index_init_value = Zero(builder, index_type); |
| 149 | auto iota_shape = input_shape; |
| 150 | iota_shape.set_element_type(index_type); |
| 151 | XlaOp iota = Iota(builder, iota_shape, axis); |
| 152 | |
| 153 | XlaComputation reducer = CreateMinMaxComputation( |
| 154 | builder, input_shape.element_type(), index_type, is_min); |
| 155 | XlaOp max_argmax = Reduce(builder, {input, iota}, |
| 156 | {value_init_value, index_init_value}, reducer, |
| 157 | /*dimensions_to_reduce=*/{axis}); |
| 158 | XlaOp argmax = GetTupleElement(max_argmax, 1); |
| 159 | if (index_type != output_type) { |
| 160 | argmax = ConvertElementType(argmax, output_type); |
| 161 | } |
| 162 | return argmax; |
| 163 | }); |
| 164 | } |
| 165 | |
| 166 | XlaOp ArgMinMaxTwoPass(XlaOp input, PrimitiveType output_type, int axis, |
| 167 | bool is_min) { |
no test coverage detected