| 218 | } |
| 219 | |
| 220 | XlaOp ArgMin(XlaOp input, PrimitiveType output_type, int axis) { |
| 221 | return ArgMinMax(input, output_type, axis, /*is_min=*/true); |
| 222 | } |
| 223 | |
| 224 | XlaOp ArgMaxTwoPass(XlaOp input, PrimitiveType output_type, int axis) { |
| 225 | return ArgMinMaxTwoPass(input, output_type, axis, /*is_min=*/false); |