ArgMax/Min `selectLastIndex` is not supported on CoreML

[philloooo]

For argmax/min on CoreML: In case of ties, the identity of the return value is not guaranteed.

Current spec seems to indicate if selectLastIndex is not passed, the first index will be chosen. This can't be satisfied by the CoreML backend.

How much does models actually depend on this? Do we have examples where selectLastIndex is useful?

If not, I'd suggest removing this parameter.

[fdwr]

For element value ties, PyTorch chooses the last index whereas Tensorflow chooses the first index, and because WebNN can be called by any front end, WebNN being able to accommodate either is valuable. It sounds like for CoreML, it either chooses:

• (a) semiarbitrarily between ties (which would be most unfortunate 😐)
• (b) chooses one over the other, but just doesn't state which it is.

Can you experiment to see which happens? e.g. Given input [0, 1, -5, 3, -5, 4] and axis = 0, does the reduceMin operation pick -5 from index 2 or 4?

With that info (assuming it is deterministic (b)), it should be possible to (not as efficiently, but possible) to use a transpose along the active dimension followed by an inversion of the indices. So CoreML selected index index 2 with selectLastIndex = false, then when selectLastIndex = true, it's the same as passing the transpose [4, -5, 3, -5, 1, 0] to get index 1 followed by a sub(5, indices) to adjust the index to 4.

How much does models actually depend on this?

🤔 That I'm unsure about, but it could make the difference in a list of equal probabilities between presenting one label to the user vs another. Granted, floating point math is imprecise anyway, and you can't expect the exact same labels in the same version of the same library even between different computers, but it would make a bigger difference on integer inputs.

[philloooo]

During my local testing, argmax always returns the smaller index. Although I don't think we should rely on that behavior since the documentation clearly states it's not guaranteed. @mwyrzykowski do you know if argmax is guaranteed to return smaller index?

[mwyrzykowski]

@philloooo It is not guaranteed and may change

[fdwr]

@philloooo It is not guaranteed and may change

Shoot, so sounds like CoreML's tie breaker behavior is undocumented and nondeterministic. So some options include:

• Request that the CoreML behavior be codified one way or another. 😉 (or request some future version of CoreML add this tie breaker flag 📅)
• See if any of the other Apple ML libraries (BNNS, MPS, MLX) support this operator? (I'm really unsure how interoperable each of these libraries is, but I presume CoreML is implemented in terms of these lower-level layers and there is some way to interop between them?)
• Defeat the nondeterminism of lowest index vs highest index by using more operators. For example, one could use min(argMin(x), sub(axisLength, transpose(argMin(x), ...)) or max(argMin(x), sub(axisLength, transpose(argMin(x), ...)) to respectively select the lowest or the highest index. Though, this degree of overhead for an uncommon case warrants a potential 3rd enum of "don't care". So rather than the current boolean select-first-index-upon-ties and select-last-index-upon-ties, you'd also have select-indeterminate-index-on-ties, which should probably just be the default value unless the caller needs one over the other. The quandary here is what happens for a three-way tie? If you have argMax([2,1,1,1,3]), does it return index 1, 2, or 3? If 1 or 3, the above code would work, but if ever 2 (the middle index), then I see no way to even emulate it.
• Keep the tiebreaker flag so backends which do support it will return desired results, with the caveat that different backends may diverge. :/
• Remove the tiebreaker flag and say in WebNN all ties are undocumented behavior. 😶

Any others come to mind?