Fix argmax. Higher index should also be taken into account (#3179)

ivarflakstad · web-flow · commit 87653ca021b0 · 2025-11-11T22:13:36.000+01:00
diff --git a/candle-metal-kernels/src/metal_src/reduce.metal b/candle-metal-kernels/src/metal_src/reduce.metal
@@ -100,7 +100,7 @@ constexpr METAL_FUNC bool operator<(indexed<T> lhs, indexed<T> rhs) {
 
 template<typename T>
 constexpr METAL_FUNC bool operator>(indexed<T> lhs, indexed<T> rhs) {
-    return lhs.val > rhs.val || (lhs.val == rhs.val && lhs.i < rhs.i);
+    return lhs.val > rhs.val || (lhs.val == rhs.val && lhs.i > rhs.i);
 }
 
 template<typename T>

Original file line number	Diff line number	Diff line change
`@@ -100,7 +100,7 @@ constexpr METAL_FUNC bool operator<(indexed<T> lhs, indexed<T> rhs) {`
`100`	`100`
`101`	`101`	`template<typename T>`
`102`	`102`	`constexpr METAL_FUNC bool operator>(indexed<T> lhs, indexed<T> rhs) {`
`103`		`- return lhs.val > rhs.val \|\| (lhs.val == rhs.val && lhs.i < rhs.i);`
	`103`	`+ return lhs.val > rhs.val \|\| (lhs.val == rhs.val && lhs.i > rhs.i);`
`104`	`104`	`}`
`105`	`105`
`106`	`106`	`template<typename T>`