Normalize negative axes in sort and argsort

Author: deBrian07

mlx/ops.cpp

@@ -2784,17 +2784,9 @@ array sort(const array& a, StreamOrDevice s /* = {} */) {
 
 /** Returns a sorted copy of the array along a given axis. */
 array sort(const array& a, int axis, StreamOrDevice s /* = {} */) {
-  // Check for valid axis
-  if (axis + static_cast<int>(a.ndim()) < 0 ||
-      axis >= static_cast<int>(a.ndim())) {
-    std::ostringstream msg;
-    msg << "[sort] Received invalid axis " << axis << " for array with "
-        << a.ndim() << " dimensions.";
-    throw std::invalid_argument(msg.str());
-  }
-
+  auto ax = normalize_axis_index(axis, a.ndim(), "[sort] ");
   return array(
-      a.shape(), a.dtype(), std::make_shared<Sort>(to_stream(s), axis), {a});
+      a.shape(), a.dtype(), std::make_shared<Sort>(to_stream(s), ax), {a});
 }
 
 /** Returns indices that sort the flattened array. */
@@ -2805,17 +2797,9 @@ array argsort(const array& a, StreamOrDevice s /* = {} */) {
 
 /** Returns indices that sort the array along a given axis. */
 array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) {
-  // Check for valid axis
-  if (axis + static_cast<int>(a.ndim()) < 0 ||
-      axis >= static_cast<int>(a.ndim())) {
-    std::ostringstream msg;
-    msg << "[argsort] Received invalid axis " << axis << " for array with "
-        << a.ndim() << " dimensions.";
-    throw std::invalid_argument(msg.str());
-  }
-
+  auto ax = normalize_axis_index(axis, a.ndim(), "[argsort] ");
   return array(
-      a.shape(), uint32, std::make_shared<ArgSort>(to_stream(s), axis), {a});
+      a.shape(), uint32, std::make_shared<ArgSort>(to_stream(s), ax), {a});
 }
 
 /**

python/tests/test_vmap.py

@@ -1023,6 +1023,20 @@ def fn(x, y):
         self.assertTrue(mx.array_equal(expected, out))
         self.assertEqual(6, counter[0])
 
+    def test_vmap_sort(self):
+        a = mx.random.uniform(shape=(3, 5))
+        expected = mx.stack([mx.sort(a[:, i]) for i in range(a.shape[1])], axis=1)
+        for axis in (0, -1):
+            out = mx.vmap(lambda x: mx.sort(x, axis=axis), in_axes=1, out_axes=1)(a)
+            self.assertTrue(mx.array_equal(out, expected))
+
+    def test_vmap_argsort(self):
+        a = mx.random.uniform(shape=(3, 5))
+        expected = mx.stack([mx.argsort(a[:, i]) for i in range(a.shape[1])], axis=1)
+        for axis in (0, -1):
+            out = mx.vmap(lambda x: mx.argsort(x, axis=axis), in_axes=1, out_axes=1)(a)
+            self.assertTrue(mx.array_equal(out, expected))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()