Document sort stability and NaN handling (#3400)

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.