Document sort stability and NaN handling

mlx/ops.cpp

@@ -2546,13 +2546,15 @@ array blackman(int M, StreamOrDevice s /* = {} */) {
   return add(subtract(alpha, term1, s), term2, s);
 }
 
-/** Returns a sorted copy of the flattened array. */
+/** Returns a sorted copy of the flattened array.
+ * The sort is stable and NaN values are placed at the end. */
 array sort(const array& a, StreamOrDevice s /* = {} */) {
   int size = a.size();
   return sort(reshape(a, {size}, s), 0, s);
 }
 
-/** Returns a sorted copy of the array along a given axis. */
+/** Returns a sorted copy of the array along a given axis.
+ * The sort is stable and NaN values are placed at the end. */
 array sort(const array& a, int axis, StreamOrDevice s /* = {} */) {
   // Check for valid axis
   if (axis + static_cast<int>(a.ndim()) < 0 ||
@@ -2567,13 +2569,15 @@ array sort(const array& a, int axis, StreamOrDevice s /* = {} */) {
       a.shape(), a.dtype(), std::make_shared<Sort>(to_stream(s), axis), {a});
 }
 
-/** Returns indices that sort the flattened array. */
+/** Returns indices that sort the flattened array.
+ * The sort is stable and NaN values are placed at the end. */
 array argsort(const array& a, StreamOrDevice s /* = {} */) {
   int size = a.size();
   return argsort(reshape(a, {size}, s), 0, s);
 }
 
-/** Returns indices that sort the array along a given axis. */
+/** Returns indices that sort the array along a given axis.
+ * The sort is stable and NaN values are placed at the end. */
 array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) {
   // Check for valid axis
   if (axis + static_cast<int>(a.ndim()) < 0 ||

mlx/ops.h

@@ -770,16 +770,28 @@ inline array argmax(const array& a, StreamOrDevice s = {}) {
 MLX_API array
 argmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
 
-/** Returns a sorted copy of the flattened array. */
+/**
+ * Returns a sorted copy of the flattened array.
+ * The sort is stable and NaN values are placed at the end.
+ */
 MLX_API array sort(const array& a, StreamOrDevice s = {});
 
-/** Returns a sorted copy of the array along a given axis. */
+/**
+ * Returns a sorted copy of the array along a given axis.
+ * The sort is stable and NaN values are placed at the end.
+ */
 MLX_API array sort(const array& a, int axis, StreamOrDevice s = {});
 
-/** Returns indices that sort the flattened array. */
+/**
+ * Returns indices that sort the flattened array.
+ * The sort is stable and NaN values are placed at the end.
+ */
 MLX_API array argsort(const array& a, StreamOrDevice s = {});
 
-/** Returns indices that sort the array along a given axis. */
+/**
+ * Returns indices that sort the array along a given axis.
+ * The sort is stable and NaN values are placed at the end.
+ */
 MLX_API array argsort(const array& a, int axis, StreamOrDevice s = {});
 
 /**

python/src/ops.cpp

@@ -2821,6 +2821,9 @@ void init_ops(nb::module_& m) {
       R"pbdoc(
         Returns a sorted copy of the array.
 
+        The sort is stable, meaning equal elements preserve their relative
+        order. ``NaN`` values are placed at the end.
+
         Args:
             a (array): Input array.
             axis (int or None, optional): Optional axis to sort over.
@@ -2848,6 +2851,9 @@ void init_ops(nb::module_& m) {
       R"pbdoc(
         Returns the indices that sort the array.
 
+        The sort is stable, meaning equal elements preserve their relative
+        order. ``NaN`` values are placed at the end.
+
         Args:
             a (array): Input array.
             axis (int or None, optional): Optional axis to sort over.