add argmax and argmin to stablehlo

[yuanfz98]

Hello,

Could community evaluate the possibiliy of adding argmax and argmin to stablehlo op set ?

Ref to llvm/llvm-project#166119

[dimitar-asenov]

Adding more ops does come with a cost. Could you please explain in more detail why you need these ops, why the current ops are not a good substitute and why having argmax/argmix will lead to significant improvements for a broad set of problems.

[ai-mannamalai]

Alternatively @dimitar-asenov could you layout how to compute argmax as a sequence of existing stablhlo ops ?

[dimitar-asenov]

E.g. for a 1D array you could:

• use a reduction to find out the maximum value.
• Compare the original array to a broadcast of the maximum value. This will give you a bool array with 1 where the maximum is, and 0 in other places. You can use this as a mask.
• Use iota to make an array 0, 1, 2, ... with the same number of elements as the original array.
• Multiply the indices array with the mask, then use a reduction with max to get the index of the max element = argmax.

Depending on the use-case this might be fast enough.