MatX icon indicating copy to clipboard operation
MatX copied to clipboard

[FEA] add argminmax function

Open tylera-nvidia opened this issue 1 year ago • 0 comments

Is your feature request related to a problem? Please describe. MatX currently only has a argmin and argmax function, requiring separate kernel launches for each. If a user needs both the min and max, we should be able to submit a single kernel to calculate both (really all 4 outputs).

Describe the solution you'd like

A single call returns minValue, minIdx, maxValue, and maxIdx, such as the call below:

(matx::mtie(minVal, minIdx, maxVal, maxIdx) = matx::argminmax(inFlattened)).run();

This should work for batched reductions as well( outputs are all some higher dimension tensor, matching the sizes of inFlattened).

Describe alternatives you've considered We can currently compute these results with separate calls, but this is less performant.

tylera-nvidia avatar Aug 05 '24 17:08 tylera-nvidia