Fixes unstable argsort() test logic

Author: honno

array_api_tests/test_sorting_functions.py

@@ -5,19 +5,17 @@
 from hypothesis import strategies as st
 from hypothesis.control import assume
 
-from .typing import Scalar, ScalarType, Shape
-
 from . import _array_module as xp
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
+from .typing import Scalar, Shape
 
 
 def assert_scalar_in_set(
     func_name: str,
-    type_: ScalarType,
     idx: Shape,
     out: Scalar,
     set_: Set[Scalar],
@@ -62,13 +60,12 @@ def test_argsort(x, data):
     scalar_type = dh.get_scalar_type(x.dtype)
     for indices in sh.axes_ndindex(x.shape, axes):
         elements = [scalar_type(x[idx]) for idx in indices]
-        orders = sorted(
-            range(len(elements)),
-            key=elements.__getitem__,
-            reverse=kw.get("descending", False),
+        orders = list(range(len(elements)))
+        sorders = sorted(
+            orders, key=elements.__getitem__, reverse=kw.get("descending", False)
         )
         if kw.get("stable", True):
-            for idx, o in zip(indices, orders):
+            for idx, o in zip(indices, sorders):
                 ph.assert_scalar_equals("argsort", int, idx, int(out[idx]), o, **kw)
         else:
             idx_elements = dict(zip(indices, elements))
@@ -78,17 +75,17 @@ def test_argsort(x, data):
                 element_orders[e] = [
                     idx_orders[idx] for idx in indices if idx_elements[idx] == e
                 ]
-            for idx, e in zip(indices, elements):
-                o = int(out[idx])
+            selements = [elements[o] for o in sorders]
+            for idx, e in zip(indices, selements):
                 expected_orders = element_orders[e]
+                out_o = int(out[idx])
                 if len(expected_orders) == 1:
-                    expected_order = expected_orders[0]
                     ph.assert_scalar_equals(
-                        "argsort", int, idx, o, expected_order, **kw
+                        "argsort", int, idx, out_o, expected_orders[0], **kw
                     )
                 else:
                     assert_scalar_in_set(
-                        "argsort", int, idx, o, set(expected_orders), **kw
+                        "argsort", idx, out_o, set(expected_orders), **kw
                     )
 
 
@@ -129,6 +126,7 @@ def test_sort(x, data):
         )
         for out_idx, o in zip(indices, orders):
             x_idx = indices[o]
+            # TODO: error message when unstable should not imply just one idx
             ph.assert_0d_equals(
                 "sort",
                 f"x[{x_idx}]",