docs: clarify shape of values and order of counts (#752)

src/array_api_stubs/_draft/set_functions.py

@@ -34,10 +34,10 @@ def unique_all(x: array, /) -> Tuple[array, array, array, array]:
     out: Tuple[array, array, array, array]
         a namedtuple ``(values, indices, inverse_indices, counts)`` whose
 
-        - first element must have the field name ``values`` and must be an array containing the unique elements of ``x``. The array must have the same data type as ``x``.
-        - second element must have the field name ``indices`` and must be an array containing the indices (first occurrences) of ``x`` that result in ``values``. The array must have the same shape as ``values`` and must have the default array index data type.
+        - first element must have the field name ``values`` and must be a one-dimensional array containing the unique elements of ``x``. The array must have the same data type as ``x``.
+        - second element must have the field name ``indices`` and must be an array containing the indices (first occurrences) of a flattened ``x`` that result in ``values``. The array must have the same shape as ``values`` and must have the default array index data type.
         - third element must have the field name ``inverse_indices`` and must be an array containing the indices of ``values`` that reconstruct ``x``. The array must have the same shape as ``x`` and must have the default array index data type.
-        - fourth element must have the field name ``counts`` and must be an array containing the number of times each unique element occurs in ``x``. The returned array must have same shape as ``values`` and must have the default array index data type.
+        - fourth element must have the field name ``counts`` and must be an array containing the number of times each unique element occurs in ``x``. The order of the returned counts must match the order of ``values``, such that a specific element in ``counts`` corresponds to the respective unique element in ``values``. The returned array must have same shape as ``values`` and must have the default array index data type.
 
         .. note::
            The order of unique elements is not specified and may vary between implementations.
@@ -78,8 +78,8 @@ def unique_counts(x: array, /) -> Tuple[array, array]:
     out: Tuple[array, array]
         a namedtuple `(values, counts)` whose
 
-        -   first element must have the field name ``values`` and must be an array containing the unique elements of ``x``. The array must have the same data type as ``x``.
-        -   second element must have the field name `counts` and must be an array containing the number of times each unique element occurs in ``x``. The returned array must have same shape as ``values`` and must have the default array index data type.
+        -   first element must have the field name ``values`` and must be a one-dimensional array containing the unique elements of ``x``. The array must have the same data type as ``x``.
+        -   second element must have the field name `counts` and must be an array containing the number of times each unique element occurs in ``x``. The order of the returned counts must match the order of ``values``, such that a specific element in ``counts`` corresponds to the respective unique element in ``values``. The returned array must have same shape as ``values`` and must have the default array index data type.
 
         .. note::
            The order of unique elements is not specified and may vary between implementations.
@@ -120,7 +120,7 @@ def unique_inverse(x: array, /) -> Tuple[array, array]:
     out: Tuple[array, array]
         a namedtuple ``(values, inverse_indices)`` whose
 
-        -   first element must have the field name ``values`` and must be an array containing the unique elements of ``x``. The array must have the same data type as ``x``.
+        -   first element must have the field name ``values`` and must be a one-dimensional array containing the unique elements of ``x``. The array must have the same data type as ``x``.
         -   second element must have the field name ``inverse_indices`` and must be an array containing the indices of ``values`` that reconstruct ``x``. The array must have the same shape as ``x`` and have the default array index data type.
 
         .. note::
@@ -158,7 +158,7 @@ def unique_values(x: array, /) -> array:
     Returns
     -------
     out: array
-        an array containing the set of unique elements in ``x``. The returned array must have the same data type as ``x``.
+        a one-dimensional array containing the set of unique elements in ``x``. The returned array must have the same data type as ``x``.
 
         .. note::
            The order of unique elements is not specified and may vary between implementations.