From aab446988875d0f50347ff620547c16e7ed19195 Mon Sep 17 00:00:00 2001
From: Tomer Michaeli <tomer.michaeli2@cs.huji.ac.il>
Date: Sun, 1 Oct 2023 17:01:11 +0300
Subject: [PATCH 01/23] added argsort function and tests

---
 heat/core/manipulations.py            |  41 ++++++++++
 heat/core/tests/test_manipulations.py | 112 ++++++++++++++++++++++++++
 2 files changed, 153 insertions(+)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 0d986a8f34..807e516aa3 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -23,6 +23,7 @@
 from . import _operations
 
 __all__ = [
+    "argsort",
     "balance",
     "broadcast_arrays",
     "broadcast_to",
@@ -63,6 +64,46 @@
 ]
 
 
+def argsort(a: DNDarray, axis: int = -1, descending: bool = False) -> DNDarray:
+    """
+    Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
+    The sorting is not stable which means that equal elements in the result may have a different ordering than in the
+    original array.
+    Sorting where `axis==a.split` needs a lot of communication between the processes of MPI.
+
+    Parameters
+    ----------
+    a : DNDarray
+        Input array to be sorted.
+    axis : int, optional
+        The dimension to sort along.
+        Default is the last axis.
+    descending : bool, optional
+        If set to `True`, indices are sorted in descending order.
+
+    Raises
+    ------
+    ValueError
+        If `axis` is not consistent with the available dimensions.
+
+    Examples
+    --------
+    >>> x = ht.array([[4, 1], [2, 3]], split=0)
+    >>> x.shape
+    (1, 2)
+    (1, 2)
+    >>> y = ht.argsort(x, axis=0)
+    >>> y
+    (array([[1, 0],
+            [0, 1]]))
+    >>> ht.argsort(x, descending=True)
+    (array([[0, 1],
+            [1, 0]]))
+    """
+    _, indices = sort(a=a, axis=axis, descending=descending, out=None)
+    return indices
+
+
 def balance(array: DNDarray, copy=False) -> DNDarray:
     """
     Out of place balance function. More information on the meaning of balance can be found in
diff --git a/heat/core/tests/test_manipulations.py b/heat/core/tests/test_manipulations.py
index 9825d333e9..89aec84cce 100644
--- a/heat/core/tests/test_manipulations.py
+++ b/heat/core/tests/test_manipulations.py
@@ -6,6 +6,118 @@
 
 
 class TestManipulations(TestCase):
+    def test_argsort(self):
+        size = ht.MPI_WORLD.size
+        rank = ht.MPI_WORLD.rank
+        tensor = (
+            torch.arange(size, device=self.device.torch_device).repeat(size).reshape(size, size)
+        )
+
+        data = ht.array(tensor, split=None)
+        result_indices = ht.argsort(data, axis=0, descending=True)
+        exp_indices = torch.argsort(tensor, dim=0, descending=True)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        result_indices = ht.argsort(data, axis=1, descending=True)
+        exp_indices = torch.argsort(tensor, dim=1, descending=True)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        data = ht.array(tensor, split=0)
+
+        exp_indices = torch.tensor([[rank] * size], device=self.device.torch_device)
+        result_indices = ht.argsort(data, descending=True, axis=0)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        exp_indices = (
+            torch.arange(size, device=self.device.torch_device)
+            .reshape(1, size)
+            .argsort(dim=1, descending=True)
+        )
+        result_indices = ht.argsort(data, descending=True, axis=1)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        indices1 = ht.argsort(data, axis=1, descending=True)
+        indices2 = ht.argsort(data, descending=True)
+        self.assertTrue(ht.equal(indices1, indices2))
+
+        data = ht.array(tensor, split=1)
+
+        indices_axis_zero = torch.arange(
+            size, dtype=torch.int64, device=self.device.torch_device
+        ).reshape(size, 1)
+        result_indices = ht.argsort(data, axis=0, descending=True)
+        # comparison value is only true on CPU
+        if result_indices.larray.is_cuda is False:
+            self.assertTrue(torch.equal(result_indices.larray, indices_axis_zero.int()))
+
+        exp_axis_one = (
+            torch.tensor(size - rank - 1, device=self.device.torch_device)
+            .repeat(size)
+            .reshape(size, 1)
+        )
+        result_indices = ht.argsort(data, descending=True, axis=1)
+        self.assertTrue(torch.equal(result_indices.larray, exp_axis_one.int()))
+
+        tensor = torch.tensor(
+            [
+                [[2, 8, 5], [7, 2, 3]],
+                [[6, 5, 2], [1, 8, 7]],
+                [[9, 3, 0], [1, 2, 4]],
+                [[8, 4, 7], [0, 8, 9]],
+            ],
+            dtype=torch.int32,
+            device=self.device.torch_device,
+        )
+
+        data = ht.array(tensor, split=0)
+        if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=self.device.torch_device
+            )
+        else:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device.torch_device
+            )
+        result_indices = ht.argsort(data, axis=0)
+        first_indices = result_indices[0].larray
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_zero))
+
+        data = ht.array(tensor, split=1)
+        indices_axis_one = torch.tensor(
+            [[0, 1, 1]], dtype=torch.int32, device=self.device.torch_device
+        )
+        result_indices = ht.argsort(data, axis=1)
+        first_indices = result_indices[0].larray[:1]
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_one))
+
+        data = ht.array(tensor, split=2)
+        indices_axis_two = torch.tensor(
+            [[0], [1]], dtype=torch.int32, device=self.device.torch_device
+        )
+        result_indices = ht.argsort(data, axis=2)
+        first_indices = result_indices[0].larray[:, :1]
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_two))
+
+        # test exceptions
+        with self.assertRaises(ValueError):
+            ht.argsort(data, axis=3)
+        with self.assertRaises(TypeError):
+            ht.argsort(data, axis="1")
+
+        rank = ht.MPI_WORLD.rank
+        ht.random.seed(1)
+        data = ht.random.randn(100, 1, split=0)
+        indices = ht.argsort(data, axis=0)
+        result = data[indices]
+        counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
+        for i, c in enumerate(counts):
+            for idx in range(c - 1):
+                if rank == i:
+                    self.assertTrue(torch.lt(result.larray[idx], result.larray[idx + 1]).all())
+
     def test_broadcast_arrays(self):
         a = ht.array([[1], [2]])
         b = ht.array([[0, 1]])

From 95ce904be02b8b0ce76f847f5120fa0fca08103b Mon Sep 17 00:00:00 2001
From: Tomer Michaeli <tomer.michaeli2@cs.huji.ac.il>
Date: Sun, 1 Oct 2023 17:01:11 +0300
Subject: [PATCH 02/23] added argsort function and tests

---
 heat/core/manipulations.py            |  41 ++++++++++
 heat/core/tests/test_manipulations.py | 112 ++++++++++++++++++++++++++
 2 files changed, 153 insertions(+)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 0d986a8f34..807e516aa3 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -23,6 +23,7 @@
 from . import _operations
 
 __all__ = [
+    "argsort",
     "balance",
     "broadcast_arrays",
     "broadcast_to",
@@ -63,6 +64,46 @@
 ]
 
 
+def argsort(a: DNDarray, axis: int = -1, descending: bool = False) -> DNDarray:
+    """
+    Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
+    The sorting is not stable which means that equal elements in the result may have a different ordering than in the
+    original array.
+    Sorting where `axis==a.split` needs a lot of communication between the processes of MPI.
+
+    Parameters
+    ----------
+    a : DNDarray
+        Input array to be sorted.
+    axis : int, optional
+        The dimension to sort along.
+        Default is the last axis.
+    descending : bool, optional
+        If set to `True`, indices are sorted in descending order.
+
+    Raises
+    ------
+    ValueError
+        If `axis` is not consistent with the available dimensions.
+
+    Examples
+    --------
+    >>> x = ht.array([[4, 1], [2, 3]], split=0)
+    >>> x.shape
+    (1, 2)
+    (1, 2)
+    >>> y = ht.argsort(x, axis=0)
+    >>> y
+    (array([[1, 0],
+            [0, 1]]))
+    >>> ht.argsort(x, descending=True)
+    (array([[0, 1],
+            [1, 0]]))
+    """
+    _, indices = sort(a=a, axis=axis, descending=descending, out=None)
+    return indices
+
+
 def balance(array: DNDarray, copy=False) -> DNDarray:
     """
     Out of place balance function. More information on the meaning of balance can be found in
diff --git a/heat/core/tests/test_manipulations.py b/heat/core/tests/test_manipulations.py
index 9825d333e9..89aec84cce 100644
--- a/heat/core/tests/test_manipulations.py
+++ b/heat/core/tests/test_manipulations.py
@@ -6,6 +6,118 @@
 
 
 class TestManipulations(TestCase):
+    def test_argsort(self):
+        size = ht.MPI_WORLD.size
+        rank = ht.MPI_WORLD.rank
+        tensor = (
+            torch.arange(size, device=self.device.torch_device).repeat(size).reshape(size, size)
+        )
+
+        data = ht.array(tensor, split=None)
+        result_indices = ht.argsort(data, axis=0, descending=True)
+        exp_indices = torch.argsort(tensor, dim=0, descending=True)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        result_indices = ht.argsort(data, axis=1, descending=True)
+        exp_indices = torch.argsort(tensor, dim=1, descending=True)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        data = ht.array(tensor, split=0)
+
+        exp_indices = torch.tensor([[rank] * size], device=self.device.torch_device)
+        result_indices = ht.argsort(data, descending=True, axis=0)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        exp_indices = (
+            torch.arange(size, device=self.device.torch_device)
+            .reshape(1, size)
+            .argsort(dim=1, descending=True)
+        )
+        result_indices = ht.argsort(data, descending=True, axis=1)
+        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
+
+        indices1 = ht.argsort(data, axis=1, descending=True)
+        indices2 = ht.argsort(data, descending=True)
+        self.assertTrue(ht.equal(indices1, indices2))
+
+        data = ht.array(tensor, split=1)
+
+        indices_axis_zero = torch.arange(
+            size, dtype=torch.int64, device=self.device.torch_device
+        ).reshape(size, 1)
+        result_indices = ht.argsort(data, axis=0, descending=True)
+        # comparison value is only true on CPU
+        if result_indices.larray.is_cuda is False:
+            self.assertTrue(torch.equal(result_indices.larray, indices_axis_zero.int()))
+
+        exp_axis_one = (
+            torch.tensor(size - rank - 1, device=self.device.torch_device)
+            .repeat(size)
+            .reshape(size, 1)
+        )
+        result_indices = ht.argsort(data, descending=True, axis=1)
+        self.assertTrue(torch.equal(result_indices.larray, exp_axis_one.int()))
+
+        tensor = torch.tensor(
+            [
+                [[2, 8, 5], [7, 2, 3]],
+                [[6, 5, 2], [1, 8, 7]],
+                [[9, 3, 0], [1, 2, 4]],
+                [[8, 4, 7], [0, 8, 9]],
+            ],
+            dtype=torch.int32,
+            device=self.device.torch_device,
+        )
+
+        data = ht.array(tensor, split=0)
+        if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=self.device.torch_device
+            )
+        else:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device.torch_device
+            )
+        result_indices = ht.argsort(data, axis=0)
+        first_indices = result_indices[0].larray
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_zero))
+
+        data = ht.array(tensor, split=1)
+        indices_axis_one = torch.tensor(
+            [[0, 1, 1]], dtype=torch.int32, device=self.device.torch_device
+        )
+        result_indices = ht.argsort(data, axis=1)
+        first_indices = result_indices[0].larray[:1]
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_one))
+
+        data = ht.array(tensor, split=2)
+        indices_axis_two = torch.tensor(
+            [[0], [1]], dtype=torch.int32, device=self.device.torch_device
+        )
+        result_indices = ht.argsort(data, axis=2)
+        first_indices = result_indices[0].larray[:, :1]
+        if rank == 0:
+            self.assertTrue(torch.equal(first_indices, indices_axis_two))
+
+        # test exceptions
+        with self.assertRaises(ValueError):
+            ht.argsort(data, axis=3)
+        with self.assertRaises(TypeError):
+            ht.argsort(data, axis="1")
+
+        rank = ht.MPI_WORLD.rank
+        ht.random.seed(1)
+        data = ht.random.randn(100, 1, split=0)
+        indices = ht.argsort(data, axis=0)
+        result = data[indices]
+        counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
+        for i, c in enumerate(counts):
+            for idx in range(c - 1):
+                if rank == i:
+                    self.assertTrue(torch.lt(result.larray[idx], result.larray[idx + 1]).all())
+
     def test_broadcast_arrays(self):
         a = ht.array([[1], [2]])
         b = ht.array([[0, 1]])

From ba90f48e8c696aadb7424896fc2df749ba6af74a Mon Sep 17 00:00:00 2001
From: Tomer Michaeli <tomer.michaeli2@cs.huji.ac.il>
Date: Wed, 4 Oct 2023 12:44:15 +0300
Subject: [PATCH 03/23] fixed small issue in test

---
 heat/core/tests/test_manipulations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/heat/core/tests/test_manipulations.py b/heat/core/tests/test_manipulations.py
index 89aec84cce..b1f9470c5e 100644
--- a/heat/core/tests/test_manipulations.py
+++ b/heat/core/tests/test_manipulations.py
@@ -111,7 +111,7 @@ def test_argsort(self):
         ht.random.seed(1)
         data = ht.random.randn(100, 1, split=0)
         indices = ht.argsort(data, axis=0)
-        result = data[indices]
+        result = data[indices.larray.tolist()]
         counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
         for i, c in enumerate(counts):
             for idx in range(c - 1):

From 2854dd2dfc1053af06ac9de2e8c8963bac4c54d1 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Thu, 16 Jul 2026 10:10:52 +0200
Subject: [PATCH 04/23] - Changed Sort and Argsort to mirror numpy api. - Fixed
 argsort test. - Fixed sort calls for old the sort API

---
 heat/core/manipulations.py       | 55 +++++++++++++++++++++++++++++---
 heat/core/statistics.py          |  4 +--
 heat/utils/data/matrixgallery.py |  2 +-
 tests/core/test_manipulations.py | 36 ++++++++++-----------
 4 files changed, 70 insertions(+), 27 deletions(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index e2890510bc..c3890bc790 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -66,7 +66,15 @@
 ]
 
 
-def argsort(a: DNDarray, axis: int = -1, descending: bool = False) -> DNDarray:
+def argsort(
+    a: DNDarray,
+    axis: int = -1,
+    kind: str | None = None,
+    order: str | list[str] | None = None,
+    *,
+    stable: bool | None = None,
+    descending: bool | None = None
+) -> DNDarray:
     """
     Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
     The sorting is not stable which means that equal elements in the result may have a different ordering than in the
@@ -80,6 +88,12 @@ def argsort(a: DNDarray, axis: int = -1, descending: bool = False) -> DNDarray:
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
+    kind : str, optional
+        Sorting algorithm. Gets ignored. 
+    order : str, list[str], optional
+        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
+    stable : bool, optional
+        Sort stability, currenty not supported.
     descending : bool, optional
         If set to `True`, indices are sorted in descending order.
 
@@ -102,7 +116,7 @@ def argsort(a: DNDarray, axis: int = -1, descending: bool = False) -> DNDarray:
     (array([[0, 1],
             [1, 0]]))
     """
-    _, indices = sort(a=a, axis=axis, descending=descending, out=None)
+    _, indices = sort(a=a, axis=axis, kind=kind, order=order, stable=stable, descending=descending, out=None, return_sort_indices=True)
     return indices
 
 
@@ -2590,7 +2604,17 @@ def shape(a: DNDarray) -> Tuple[int, ...]:
     return a.gshape
 
 
-def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DNDarray] = None):
+def sort(
+    a: DNDarray,
+    axis: int = -1,
+    kind: str | None = None,
+    order: str | list[str] | None = None,
+    *,
+    stable: bool | None = None,
+    descending: bool | None = False, 
+    out: Optional[DNDarray] = None,
+    return_sort_indices: bool | None = None
+):
     """
     Sorts the elements of `a` along the given dimension (by default in ascending order) by their value.
     The sorting is not stable which means that equal elements in the result may have a different ordering than in the
@@ -2605,11 +2629,20 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
+    kind : str, optional
+        Sorting algorithm. Gets ignored. 
+    order : str, list[str], optional
+        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
+    stable : bool, optional
+        Sort stability, currenty not supported.
     descending : bool, optional
         If set to `True`, values are sorted in descending order.
     out : DNDarray, optional
         A location in which to store the results. If provided, it must have a broadcastable shape. If not provided
         or set to `None`, a fresh array is allocated.
+    return_sort_indices: bool, optional
+        Wether to return the indices by which the array was sorted. 
+        If ``out`` is provided, returns the indices if ``True``, otherwise ``None``.
 
     Raises
     ------
@@ -2637,6 +2670,14 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN
         message=r".*__array_wrap__ must accept context and return_scalar arguments.*",
     )
     stride_tricks.sanitize_axis(a.shape, axis)
+    
+    if axis is None:
+        a = flatten(a)
+        axis = 0
+
+    descending = descending or False
+    return_sort_indices = return_sort_indices or False
+    
 
     if not a.is_distributed() or axis != a.split:
         # sorting is not affected by split -> we can just sort along the axis
@@ -2848,12 +2889,16 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN
     )
     if out is not None:
         out.larray = final_result
-        return return_indices
+        if return_sort_indices:
+            return return_indices
+        return None
     else:
         tensor = factories.array(
             final_result, dtype=a.dtype, is_split=a.split, device=a.device, comm=a.comm
         )
-        return tensor, return_indices
+        if return_sort_indices:
+            return tensor, return_indices
+        return tensor
 
 
 def split(x: DNDarray, indices_or_sections: Iterable, axis: int = 0) -> List[DNDarray, ...]:
diff --git a/heat/core/statistics.py b/heat/core/statistics.py
index c9f49b9936..54f86daeaa 100644
--- a/heat/core/statistics.py
+++ b/heat/core/statistics.py
@@ -1543,7 +1543,7 @@ def _create_sketch(
 
         # create a random sample of indices
         indices = manipulations.sort(
-            randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64)
+            randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64), return_sort_indices=True
         )[0]
         sketch = a.swapaxes(0, axis)
         sketch = a[indices, ...].resplit_(None)
@@ -1697,7 +1697,7 @@ def _create_sketch(
         )
 
     # sort data
-    sorted_x, _ = manipulations.sort(x, axis=axis)
+    sorted_x, _ = manipulations.sort(x, axis=axis, return_sort_indices=True)
     del _
     sorted_x = sorted_x.astype(output_dtype)
 
diff --git a/heat/utils/data/matrixgallery.py b/heat/utils/data/matrixgallery.py
index 47b6c0b19c..b56e308b60 100644
--- a/heat/utils/data/matrixgallery.py
+++ b/heat/utils/data/matrixgallery.py
@@ -197,7 +197,7 @@ def random_known_rank(
         raise RuntimeError("rank must not exceed matrix dimensions.")
 
     singular_values = rand(r, dtype=dtype, comm=comm, device=device)
-    singular_values = sort(quantile_function(singular_values), descending=True)[0]
+    singular_values = sort(quantile_function(singular_values), descending=True, return_sort_indices=True)[0]
 
     return random_known_singularvalues(
         m, n, singular_values, split=split, device=device, comm=comm, dtype=dtype
diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index a0308ff2a7..ca97a9cd2d 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -112,12 +112,10 @@ def test_argsort(self):
         ht.random.seed(1)
         data = ht.random.randn(100, 1, split=0)
         indices = ht.argsort(data, axis=0)
-        result = data[indices.larray.tolist()]
-        counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
-        for i, c in enumerate(counts):
-            for idx in range(c - 1):
-                if rank == i:
-                    self.assertTrue(torch.lt(result.larray[idx], result.larray[idx + 1]).all())
+        result = data[indices]
+        
+        arr = ht.resplit(result.flatten(), axis=None)
+        self.assertTrue((arr.larray[:-1] <= arr.larray[1:]).all())
 
     def test_broadcast_arrays(self):
         a = ht.array([[1], [2]])
@@ -2899,12 +2897,12 @@ def test_sort(self):
         )
 
         data = ht.array(tensor, split=None)
-        result, result_indices = ht.sort(data, axis=0, descending=True)
+        result, result_indices = ht.sort(data, axis=0, descending=True, return_sort_indices=True)
         expected, exp_indices = torch.sort(tensor, dim=0, descending=True)
         self.assertTrue(torch.equal(result.larray, expected))
         self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
 
-        result, result_indices = ht.sort(data, axis=1, descending=True)
+        result, result_indices = ht.sort(data, axis=1, descending=True, return_sort_indices=True)
         expected, exp_indices = torch.sort(tensor, dim=1, descending=True)
         self.assertTrue(torch.equal(result.larray, expected))
         self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
@@ -2913,7 +2911,7 @@ def test_sort(self):
 
         exp_axis_zero = torch.arange(size, device=self.device.torch_device).reshape(1, size)
         exp_indices = torch.tensor([[rank] * size], device=self.device.torch_device)
-        result, result_indices = ht.sort(data, descending=True, axis=0)
+        result, result_indices = ht.sort(data, descending=True, axis=0, return_sort_indices=True)
         self.assertTrue(torch.equal(result.larray, exp_axis_zero))
         self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
 
@@ -2922,12 +2920,12 @@ def test_sort(self):
             .reshape(1, size)
             .sort(dim=1, descending=True)
         )
-        result, result_indices = ht.sort(data, descending=True, axis=1)
+        result, result_indices = ht.sort(data, descending=True, axis=1, return_sort_indices=True)
         self.assertTrue(torch.equal(result.larray, exp_axis_one))
         self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
 
-        result1 = ht.sort(data, axis=1, descending=True)
-        result2 = ht.sort(data, descending=True)
+        result1 = ht.sort(data, axis=1, descending=True, return_sort_indices=True)
+        result2 = ht.sort(data, descending=True, return_sort_indices=True)
         self.assertTrue(ht.equal(result1[0], result2[0]))
         self.assertTrue(ht.equal(result1[1], result2[1]))
 
@@ -2939,7 +2937,7 @@ def test_sort(self):
         indices_axis_zero = torch.arange(
             size, dtype=torch.int64, device=self.device.torch_device
         ).reshape(size, 1)
-        result, result_indices = ht.sort(data, axis=0, descending=True)
+        result, result_indices = ht.sort(data, axis=0, descending=True, return_sort_indices=True)
         self.assertTrue(torch.equal(result.larray, exp_axis_zero))
         # comparison value is only true on CPU
         if result_indices.larray.is_cuda is False:
@@ -2950,7 +2948,7 @@ def test_sort(self):
             .repeat(size)
             .reshape(size, 1)
         )
-        result, result_indices = ht.sort(data, descending=True, axis=1)
+        result, result_indices = ht.sort(data, descending=True, axis=1, return_sort_indices=True)
         self.assertTrue(torch.equal(result.larray, exp_axis_one))
         self.assertTrue(torch.equal(result_indices.larray, exp_axis_one.int()))
 
@@ -2972,7 +2970,7 @@ def test_sort(self):
         indices_axis_zero = torch.tensor(
             [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device.torch_device
         )
-        result, result_indices = ht.sort(data, axis=0)
+        result, result_indices = ht.sort(data, axis=0, return_sort_indices=True)
         first = result[0].larray
         first_indices = result_indices[0].larray
         if rank == 0:
@@ -2991,7 +2989,7 @@ def test_sort(self):
         indices_axis_one = torch.tensor(
             [[0, 1, 1]], dtype=torch.int32, device=self.device.torch_device
         )
-        result, result_indices = ht.sort(data, axis=1)
+        result, result_indices = ht.sort(data, axis=1, return_sort_indices=True)
         first = result[0].larray[:1]
         first_indices = result_indices[0].larray[:1]
         if rank == 0:
@@ -3003,7 +3001,7 @@ def test_sort(self):
         indices_axis_two = torch.tensor(
             [[0], [1]], dtype=torch.int32, device=self.device.torch_device
         )
-        result, result_indices = ht.sort(data, axis=2)
+        result, result_indices = ht.sort(data, axis=2, return_sort_indices=True)
         first = result[0].larray[:, :1]
         first_indices = result_indices[0].larray[:, :1]
         if rank == 0:
@@ -3011,7 +3009,7 @@ def test_sort(self):
             self.assertTrue(torch.equal(first_indices, indices_axis_two))
         #
         out = ht.empty_like(data)
-        indices = ht.sort(data, axis=2, out=out)
+        indices = ht.sort(data, axis=2, out=out, return_sort_indices=True)
         self.assertTrue(ht.equal(out, result))
         self.assertTrue(ht.equal(indices, result_indices))
 
@@ -3023,7 +3021,7 @@ def test_sort(self):
         rank = ht.MPI_WORLD.rank
         ht.random.seed(1)
         data = ht.random.randn(100, 1, split=0)
-        result, _ = ht.sort(data, axis=0)
+        result, _ = ht.sort(data, axis=0, return_sort_indices=True)
         counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0)
         for i, c in enumerate(counts):
             for idx in range(c - 1):

From 686185256ab76bacaca623fe2a31e2006bbd8a0d Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Thu, 16 Jul 2026 11:08:12 +0200
Subject: [PATCH 05/23] Fixed Split Advanced indexing error on the argsort test

---
 tests/core/test_manipulations.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index ca97a9cd2d..cd95ba12ff 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -112,9 +112,9 @@ def test_argsort(self):
         ht.random.seed(1)
         data = ht.random.randn(100, 1, split=0)
         indices = ht.argsort(data, axis=0)
-        result = data[indices]
-        
+        result = ht.resplit(data, None)[indices]
         arr = ht.resplit(result.flatten(), axis=None)
+        
         self.assertTrue((arr.larray[:-1] <= arr.larray[1:]).all())
 
     def test_broadcast_arrays(self):

From 21083197c78c53142894f3739b2bde9209d6899c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 16 Jul 2026 09:14:34 +0000
Subject: [PATCH 06/23] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 heat/core/manipulations.py       | 26 +++++++++++++++++---------
 heat/core/statistics.py          |  3 ++-
 heat/utils/data/matrixgallery.py |  4 +++-
 tests/core/test_manipulations.py |  2 +-
 4 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index c3890bc790..19f15342da 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -73,7 +73,7 @@ def argsort(
     order: str | list[str] | None = None,
     *,
     stable: bool | None = None,
-    descending: bool | None = None
+    descending: bool | None = None,
 ) -> DNDarray:
     """
     Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
@@ -89,7 +89,7 @@ def argsort(
         The dimension to sort along.
         Default is the last axis.
     kind : str, optional
-        Sorting algorithm. Gets ignored. 
+        Sorting algorithm. Gets ignored.
     order : str, list[str], optional
         Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
     stable : bool, optional
@@ -116,7 +116,16 @@ def argsort(
     (array([[0, 1],
             [1, 0]]))
     """
-    _, indices = sort(a=a, axis=axis, kind=kind, order=order, stable=stable, descending=descending, out=None, return_sort_indices=True)
+    _, indices = sort(
+        a=a,
+        axis=axis,
+        kind=kind,
+        order=order,
+        stable=stable,
+        descending=descending,
+        out=None,
+        return_sort_indices=True,
+    )
     return indices
 
 
@@ -2611,9 +2620,9 @@ def sort(
     order: str | list[str] | None = None,
     *,
     stable: bool | None = None,
-    descending: bool | None = False, 
+    descending: bool | None = False,
     out: Optional[DNDarray] = None,
-    return_sort_indices: bool | None = None
+    return_sort_indices: bool | None = None,
 ):
     """
     Sorts the elements of `a` along the given dimension (by default in ascending order) by their value.
@@ -2630,7 +2639,7 @@ def sort(
         The dimension to sort along.
         Default is the last axis.
     kind : str, optional
-        Sorting algorithm. Gets ignored. 
+        Sorting algorithm. Gets ignored.
     order : str, list[str], optional
         Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
     stable : bool, optional
@@ -2641,7 +2650,7 @@ def sort(
         A location in which to store the results. If provided, it must have a broadcastable shape. If not provided
         or set to `None`, a fresh array is allocated.
     return_sort_indices: bool, optional
-        Wether to return the indices by which the array was sorted. 
+        Wether to return the indices by which the array was sorted.
         If ``out`` is provided, returns the indices if ``True``, otherwise ``None``.
 
     Raises
@@ -2670,14 +2679,13 @@ def sort(
         message=r".*__array_wrap__ must accept context and return_scalar arguments.*",
     )
     stride_tricks.sanitize_axis(a.shape, axis)
-    
+
     if axis is None:
         a = flatten(a)
         axis = 0
 
     descending = descending or False
     return_sort_indices = return_sort_indices or False
-    
 
     if not a.is_distributed() or axis != a.split:
         # sorting is not affected by split -> we can just sort along the axis
diff --git a/heat/core/statistics.py b/heat/core/statistics.py
index 54f86daeaa..cd6ae25de0 100644
--- a/heat/core/statistics.py
+++ b/heat/core/statistics.py
@@ -1543,7 +1543,8 @@ def _create_sketch(
 
         # create a random sample of indices
         indices = manipulations.sort(
-            randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64), return_sort_indices=True
+            randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64),
+            return_sort_indices=True,
         )[0]
         sketch = a.swapaxes(0, axis)
         sketch = a[indices, ...].resplit_(None)
diff --git a/heat/utils/data/matrixgallery.py b/heat/utils/data/matrixgallery.py
index b56e308b60..779dec9749 100644
--- a/heat/utils/data/matrixgallery.py
+++ b/heat/utils/data/matrixgallery.py
@@ -197,7 +197,9 @@ def random_known_rank(
         raise RuntimeError("rank must not exceed matrix dimensions.")
 
     singular_values = rand(r, dtype=dtype, comm=comm, device=device)
-    singular_values = sort(quantile_function(singular_values), descending=True, return_sort_indices=True)[0]
+    singular_values = sort(
+        quantile_function(singular_values), descending=True, return_sort_indices=True
+    )[0]
 
     return random_known_singularvalues(
         m, n, singular_values, split=split, device=device, comm=comm, dtype=dtype
diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index cd95ba12ff..6bc744804d 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -114,7 +114,7 @@ def test_argsort(self):
         indices = ht.argsort(data, axis=0)
         result = ht.resplit(data, None)[indices]
         arr = ht.resplit(result.flatten(), axis=None)
-        
+
         self.assertTrue((arr.larray[:-1] <= arr.larray[1:]).all())
 
     def test_broadcast_arrays(self):

From c3aa612b153e076aa60e47f7a2961762b06c5d8c Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 20 Jul 2026 10:34:26 +0200
Subject: [PATCH 07/23] - added warning for ignored arguments instead of
 ignoring.

---
 heat/core/manipulations.py | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 19f15342da..0c4859e655 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -2616,19 +2616,17 @@ def shape(a: DNDarray) -> Tuple[int, ...]:
 def sort(
     a: DNDarray,
     axis: int = -1,
-    kind: str | None = None,
-    order: str | list[str] | None = None,
-    *,
-    stable: bool | None = None,
+    *args,
     descending: bool | None = False,
     out: Optional[DNDarray] = None,
-    return_sort_indices: bool | None = None,
+    return_sort_indices: bool = False,
+    **kwargs
 ):
     """
     Sorts the elements of `a` along the given dimension (by default in ascending order) by their value.
     The sorting is not stable which means that equal elements in the result may have a different ordering than in the
     original array.
-    Sorting where `axis==a.split` needs a lot of communication between the processes of MPI.
+    Sorting with `axis==a.split` needs a lot of communication between the processes of MPI.
     Returns a tuple `(values, indices)` with the sorted local results and the indices of the elements in the original data
 
     Parameters
@@ -2638,12 +2636,6 @@ def sort(
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
-    kind : str, optional
-        Sorting algorithm. Gets ignored.
-    order : str, list[str], optional
-        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
-    stable : bool, optional
-        Sort stability, currenty not supported.
     descending : bool, optional
         If set to `True`, values are sorted in descending order.
     out : DNDarray, optional
@@ -2680,12 +2672,17 @@ def sort(
     )
     stride_tricks.sanitize_axis(a.shape, axis)
 
+    for arg in args:
+        warnings.warn(f"[ht.sort] Argument: '{arg}' gets ignored.")
+
+    for k in kwargs.keys():
+        warnings.warn(f"[ht.sort] Keyword Argument: '{k}' gets ignored.")
+
     if axis is None:
         a = flatten(a)
         axis = 0
 
     descending = descending or False
-    return_sort_indices = return_sort_indices or False
 
     if not a.is_distributed() or axis != a.split:
         # sorting is not affected by split -> we can just sort along the axis

From 909517d2da2f2d8341921b1aaa479e812226c837 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 20 Jul 2026 08:34:43 +0000
Subject: [PATCH 08/23] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 heat/core/manipulations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 0c4859e655..c10eca83f2 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -2620,7 +2620,7 @@ def sort(
     descending: bool | None = False,
     out: Optional[DNDarray] = None,
     return_sort_indices: bool = False,
-    **kwargs
+    **kwargs,
 ):
     """
     Sorts the elements of `a` along the given dimension (by default in ascending order) by their value.

From 122d3d09640fca2556798f006913fae7a7a6230d Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Tue, 21 Jul 2026 09:34:19 +0200
Subject: [PATCH 09/23] added ignore warning to argsort

---
 heat/core/manipulations.py | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index c10eca83f2..661a6cf258 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -69,11 +69,9 @@
 def argsort(
     a: DNDarray,
     axis: int = -1,
-    kind: str | None = None,
-    order: str | list[str] | None = None,
-    *,
-    stable: bool | None = None,
+    *args,
     descending: bool | None = None,
+    **kwargs
 ) -> DNDarray:
     """
     Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
@@ -88,12 +86,6 @@ def argsort(
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
-    kind : str, optional
-        Sorting algorithm. Gets ignored.
-    order : str, list[str], optional
-        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
-    stable : bool, optional
-        Sort stability, currenty not supported.
     descending : bool, optional
         If set to `True`, indices are sorted in descending order.
 
@@ -116,12 +108,15 @@ def argsort(
     (array([[0, 1],
             [1, 0]]))
     """
+    for arg in args:
+        warnings.warn(f"[ht.argsort] Argument: '{arg}' gets ignored.")
+
+    for k in kwargs.keys():
+        warnings.warn(f"[ht.argsort] Keyword Argument: '{k}' gets ignored.")
+
     _, indices = sort(
         a=a,
         axis=axis,
-        kind=kind,
-        order=order,
-        stable=stable,
         descending=descending,
         out=None,
         return_sort_indices=True,

From 100c377cb97efd8a57cc7eccf8ec78fb290ff871 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 21 Jul 2026 07:35:31 +0000
Subject: [PATCH 10/23] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 heat/core/manipulations.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 661a6cf258..bc2c2508b3 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -67,11 +67,7 @@
 
 
 def argsort(
-    a: DNDarray,
-    axis: int = -1,
-    *args,
-    descending: bool | None = None,
-    **kwargs
+    a: DNDarray, axis: int = -1, *args, descending: bool | None = None, **kwargs
 ) -> DNDarray:
     """
     Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.

From f23bcf378d27824f455f7e22fd9aa13bbeee756e Mon Sep 17 00:00:00 2001
From: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:39:53 +0200
Subject: [PATCH 11/23] Apply suggestion from @brownbaerchen

---
 heat/core/manipulations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index bc2c2508b3..c6af00e1a3 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -73,7 +73,7 @@ def argsort(
     Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`.
     The sorting is not stable which means that equal elements in the result may have a different ordering than in the
     original array.
-    Sorting where `axis==a.split` needs a lot of communication between the processes of MPI.
+    Sorting with `axis==a.split` needs a lot of communication between the processes.
 
     Parameters
     ----------

From 09063e8b0a3d68468c63cfffa84c09d8a51611b3 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 27 Jul 2026 10:45:25 +0200
Subject: [PATCH 12/23] Added info for ignored parameters in the docstring of
 `sort` and `argsort`

---
 heat/core/manipulations.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index c6af00e1a3..ccddf11bf8 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -82,8 +82,12 @@ def argsort(
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
+    *args: Any
+        Any arguments that are not specified in the function header will be ignored.
     descending : bool, optional
         If set to `True`, indices are sorted in descending order.
+    **kwargs: Any
+        Any keyword arguments that are not specified in the function header will be ignored.
 
     Raises
     ------
@@ -2627,6 +2631,8 @@ def sort(
     axis : int, optional
         The dimension to sort along.
         Default is the last axis.
+    *args: Any
+        Any arguments that are not specified in the function header will be ignored.
     descending : bool, optional
         If set to `True`, values are sorted in descending order.
     out : DNDarray, optional
@@ -2635,6 +2641,8 @@ def sort(
     return_sort_indices: bool, optional
         Wether to return the indices by which the array was sorted.
         If ``out`` is provided, returns the indices if ``True``, otherwise ``None``.
+    **kwargs: Any
+        Any keyword arguments that are not specified in the function header will be ignored.
 
     Raises
     ------

From aafd82356eaad0137ed3c9179434ab8eef9308c9 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Tue, 28 Jul 2026 16:07:04 +0200
Subject: [PATCH 13/23] changed first portion of `argsort` tests to use random
 data

---
 tests/core/test_manipulations.py | 55 +++++---------------------------
 1 file changed, 8 insertions(+), 47 deletions(-)

diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index 76f9087cb2..d32c5263d9 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -10,54 +10,15 @@ class TestManipulations(TestCase):
     def test_argsort(self):
         size = ht.MPI_WORLD.size
         rank = ht.MPI_WORLD.rank
-        tensor = (
-            torch.arange(size, device=self.device.torch_device).repeat(size).reshape(size, size)
-        )
-
-        data = ht.array(tensor, split=None)
-        result_indices = ht.argsort(data, axis=0, descending=True)
-        exp_indices = torch.argsort(tensor, dim=0, descending=True)
-        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
-
-        result_indices = ht.argsort(data, axis=1, descending=True)
-        exp_indices = torch.argsort(tensor, dim=1, descending=True)
-        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
-
-        data = ht.array(tensor, split=0)
-
-        exp_indices = torch.tensor([[rank] * size], device=self.device.torch_device)
-        result_indices = ht.argsort(data, descending=True, axis=0)
-        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
-
-        exp_indices = (
-            torch.arange(size, device=self.device.torch_device)
-            .reshape(1, size)
-            .argsort(dim=1, descending=True)
-        )
-        result_indices = ht.argsort(data, descending=True, axis=1)
-        self.assertTrue(torch.equal(result_indices.larray, exp_indices.int()))
-
-        indices1 = ht.argsort(data, axis=1, descending=True)
-        indices2 = ht.argsort(data, descending=True)
-        self.assertTrue(ht.equal(indices1, indices2))
 
-        data = ht.array(tensor, split=1)
-
-        indices_axis_zero = torch.arange(
-            size, dtype=torch.int64, device=self.device.torch_device
-        ).reshape(size, 1)
-        result_indices = ht.argsort(data, axis=0, descending=True)
-        # comparison value is only true on CPU
-        if result_indices.larray.is_cuda is False:
-            self.assertTrue(torch.equal(result_indices.larray, indices_axis_zero.int()))
-
-        exp_axis_one = (
-            torch.tensor(size - rank - 1, device=self.device.torch_device)
-            .repeat(size)
-            .reshape(size, 1)
-        )
-        result_indices = ht.argsort(data, descending=True, axis=1)
-        self.assertTrue(torch.equal(result_indices.larray, exp_axis_one.int()))
+        data = ht.random.rand(2, 3, 4)
+        for axis in [None, 0, 1, 2]:
+            for descending in [True, False]:
+                for split in [0, 1, 2]:
+                    data.resplit_(split)
+                    result_indices = ht.argsort(data, axis=axis, descending=descending)
+                    exp_indices = np.argsort(data.numpy(), axis=axis, descending=descending)
+                    self.assertTrue(np.allclose(result_indices.numpy(), exp_indices))
 
         tensor = torch.tensor(
             [

From df9ca0a040bba431a8e9ad4184d7c7847d93163f Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Wed, 29 Jul 2026 11:18:27 +0200
Subject: [PATCH 14/23] Added test for `sort, and added more parameter
 combinations for `sort` and `argsort`.

---
 tests/core/test_manipulations.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index d32c5263d9..cac59f9297 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -7,14 +7,29 @@
 
 
 class TestManipulations(TestCase):
+    def test_sort(self):
+        size = ht.MPI_WORLD.size
+        rank = ht.MPI_WORLD.rank
+
+        data = ht.random.rand(2, 3, 4)
+        for axis in [None, 0, 1, 2]:
+            for descending in [None, True, False]:
+                for split in [None, 0, 1, 2]:
+                    data.resplit_(split)
+                    result, idx = ht.sort(data, axis=axis, descending=descending, return_sort_indices=True)
+                    exp = np.sort(data.numpy(), axis=axis, descending=descending)
+                    exp_idx = np.argsort(data.numpy(), axis=axis, descending=descending)
+                    self.assertTrue(np.allclose(result.numpy(), exp))
+                    self.assertTrue(np.allclose(idx.numpy(), exp_idx))
+
     def test_argsort(self):
         size = ht.MPI_WORLD.size
         rank = ht.MPI_WORLD.rank
 
         data = ht.random.rand(2, 3, 4)
         for axis in [None, 0, 1, 2]:
-            for descending in [True, False]:
-                for split in [0, 1, 2]:
+            for descending in [None, True, False]:
+                for split in [None, 0, 1, 2]:
                     data.resplit_(split)
                     result_indices = ht.argsort(data, axis=axis, descending=descending)
                     exp_indices = np.argsort(data.numpy(), axis=axis, descending=descending)

From 1d759ec0014d9c5195dc0dd910ba19e92b3a6e1f Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 08:40:49 +0200
Subject: [PATCH 15/23] - Created own class for sorting testing. - Switched to
 pytest for testing of sorting.

---
 tests/core/test_manipulations.py |  87 -------------------------
 tests/core/test_sorting.py       | 108 +++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 87 deletions(-)
 create mode 100644 tests/core/test_sorting.py

diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py
index cac59f9297..51d123a9d7 100644
--- a/tests/core/test_manipulations.py
+++ b/tests/core/test_manipulations.py
@@ -5,94 +5,7 @@
 import heat as ht
 from heat.testing.basic_test import TestCase
 
-
 class TestManipulations(TestCase):
-    def test_sort(self):
-        size = ht.MPI_WORLD.size
-        rank = ht.MPI_WORLD.rank
-
-        data = ht.random.rand(2, 3, 4)
-        for axis in [None, 0, 1, 2]:
-            for descending in [None, True, False]:
-                for split in [None, 0, 1, 2]:
-                    data.resplit_(split)
-                    result, idx = ht.sort(data, axis=axis, descending=descending, return_sort_indices=True)
-                    exp = np.sort(data.numpy(), axis=axis, descending=descending)
-                    exp_idx = np.argsort(data.numpy(), axis=axis, descending=descending)
-                    self.assertTrue(np.allclose(result.numpy(), exp))
-                    self.assertTrue(np.allclose(idx.numpy(), exp_idx))
-
-    def test_argsort(self):
-        size = ht.MPI_WORLD.size
-        rank = ht.MPI_WORLD.rank
-
-        data = ht.random.rand(2, 3, 4)
-        for axis in [None, 0, 1, 2]:
-            for descending in [None, True, False]:
-                for split in [None, 0, 1, 2]:
-                    data.resplit_(split)
-                    result_indices = ht.argsort(data, axis=axis, descending=descending)
-                    exp_indices = np.argsort(data.numpy(), axis=axis, descending=descending)
-                    self.assertTrue(np.allclose(result_indices.numpy(), exp_indices))
-
-        tensor = torch.tensor(
-            [
-                [[2, 8, 5], [7, 2, 3]],
-                [[6, 5, 2], [1, 8, 7]],
-                [[9, 3, 0], [1, 2, 4]],
-                [[8, 4, 7], [0, 8, 9]],
-            ],
-            dtype=torch.int32,
-            device=self.device.torch_device,
-        )
-
-        data = ht.array(tensor, split=0)
-        if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
-            indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=self.device.torch_device
-            )
-        else:
-            indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device.torch_device
-            )
-        result_indices = ht.argsort(data, axis=0)
-        first_indices = result_indices[0].larray
-        if rank == 0:
-            self.assertTrue(torch.equal(first_indices, indices_axis_zero))
-
-        data = ht.array(tensor, split=1)
-        indices_axis_one = torch.tensor(
-            [[0, 1, 1]], dtype=torch.int32, device=self.device.torch_device
-        )
-        result_indices = ht.argsort(data, axis=1)
-        first_indices = result_indices[0].larray[:1]
-        if rank == 0:
-            self.assertTrue(torch.equal(first_indices, indices_axis_one))
-
-        data = ht.array(tensor, split=2)
-        indices_axis_two = torch.tensor(
-            [[0], [1]], dtype=torch.int32, device=self.device.torch_device
-        )
-        result_indices = ht.argsort(data, axis=2)
-        first_indices = result_indices[0].larray[:, :1]
-        if rank == 0:
-            self.assertTrue(torch.equal(first_indices, indices_axis_two))
-
-        # test exceptions
-        with self.assertRaises(ValueError):
-            ht.argsort(data, axis=3)
-        with self.assertRaises(TypeError):
-            ht.argsort(data, axis="1")
-
-        rank = ht.MPI_WORLD.rank
-        ht.random.seed(1)
-        data = ht.random.randn(100, 1, split=0)
-        indices = ht.argsort(data, axis=0)
-        result = ht.resplit(data, None)[indices]
-        arr = ht.resplit(result.flatten(), axis=None)
-
-        self.assertTrue((arr.larray[:-1] <= arr.larray[1:]).all())
-
     def test_broadcast_arrays(self):
         a = ht.array([[1], [2]])
         b = ht.array([[0, 1]])
diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
new file mode 100644
index 0000000000..387024b30d
--- /dev/null
+++ b/tests/core/test_sorting.py
@@ -0,0 +1,108 @@
+import numpy as np
+import torch
+import pytest
+
+import heat as ht
+
+
+class TestSorting:
+    @pytest.mark.parametrize("device", ["cpu", "gpu"])
+    @pytest.mark.parametrize("split", [None, 0, 1, 2])
+    @pytest.mark.parametrize("descending", [None, True, False])
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2])
+    def test_sort(self, axis, descending, split, device):
+        if not torch.cuda.is_available() and device == "gpu":
+            pytest.skip("No gpu available for testing.")
+
+        if np.lib.NumpyVersion(np.__version__) < '2.5.0' and descending is not None:
+            pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+
+        data = ht.random.rand(2, 3, 4, split=split)
+        result, idx = ht.sort(data, axis=axis, descending=descending, return_sort_indices=True)
+        exp = np.sort(data.numpy(), axis=axis, descending=descending)
+        exp_idx = np.argsort(data.numpy(), axis=axis, descending=descending)
+
+        assert np.allclose(result.numpy(), exp)
+        assert np.allclose(idx.numpy(), exp_idx)
+
+    @pytest.mark.parametrize("device", ["cpu", "gpu"])
+    @pytest.mark.parametrize("split", [None, 0, 1, 2])
+    @pytest.mark.parametrize("descending", [None, True, False])
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2])
+    def test_argsort_random(self, axis, descending, split, device):
+        if not torch.cuda.is_available() and device == "gpu":
+            pytest.skip("No gpu available for testing.")
+
+        if np.lib.NumpyVersion(np.__version__) < '2.5.0' and descending is not None:
+            pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+
+        data = ht.random.rand(2, 3, 4, split=split, device=device)
+        result_indices = ht.argsort(data, axis=axis, descending=descending)
+        exp_indices = np.argsort(data.numpy(), axis=axis, descending=descending)
+        assert np.allclose(result_indices.numpy(), exp_indices)
+
+    @pytest.mark.parametrize("device", ["cpu", "gpu"])
+    def test_argsort_specific(self, device):
+        if not torch.cuda.is_available() and device == "gpu":
+            pytest.skip("No gpu available for testing.")
+
+        size = ht.MPI_WORLD.size
+        rank = ht.MPI_WORLD.rank
+
+        tensor = torch.tensor(
+            [
+                [[2, 8, 5], [7, 2, 3]],
+                [[6, 5, 2], [1, 8, 7]],
+                [[9, 3, 0], [1, 2, 4]],
+                [[8, 4, 7], [0, 8, 9]],
+            ],
+            dtype=torch.int32,
+            device=device,
+        )
+
+        data = ht.array(tensor, split=0)
+        if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=device
+            )
+        else:
+            indices_axis_zero = torch.tensor(
+                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=device
+            )
+        result_indices = ht.argsort(data, axis=0)
+        first_indices = result_indices[0].larray
+        if rank == 0:
+            assert torch.equal(first_indices, indices_axis_zero)
+
+        data = ht.array(tensor, split=1)
+        indices_axis_one = torch.tensor(
+            [[0, 1, 1]], dtype=torch.int32, device=device
+        )
+        result_indices = ht.argsort(data, axis=1)
+        first_indices = result_indices[0].larray[:1]
+        if rank == 0:
+            assert torch.equal(first_indices, indices_axis_one)
+
+        data = ht.array(tensor, split=2)
+        indices_axis_two = torch.tensor(
+            [[0], [1]], dtype=torch.int32, device=device
+        )
+        result_indices = ht.argsort(data, axis=2)
+        first_indices = result_indices[0].larray[:, :1]
+        if rank == 0:
+            assert torch.equal(first_indices, indices_axis_two)
+
+        # test exceptions
+        with pytest.raises(ValueError):
+            ht.argsort(data, axis=3)
+        with pytest.raises(TypeError):
+            ht.argsort(data, axis="1")
+
+        rank = ht.MPI_WORLD.rank
+        ht.random.seed(1)
+        data = ht.random.randn(100, 1, split=0, device=device)
+        indices = ht.argsort(data, axis=0)
+        result = ht.resplit(data, None)[indices]
+        arr = ht.resplit(result.flatten(), axis=None)
+
+        assert (arr.larray[:-1] <= arr.larray[1:]).all()
\ No newline at end of file

From bbfee73d7de60fe6feaf050d3daef6a4ebc13003 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 3 Aug 2026 06:41:09 +0000
Subject: [PATCH 16/23] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/core/test_sorting.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 387024b30d..fd477f1c85 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -105,4 +105,4 @@ def test_argsort_specific(self, device):
         result = ht.resplit(data, None)[indices]
         arr = ht.resplit(result.flatten(), axis=None)
 
-        assert (arr.larray[:-1] <= arr.larray[1:]).all()
\ No newline at end of file
+        assert (arr.larray[:-1] <= arr.larray[1:]).all()

From 937d55147fee67f0b9f427f6d32cf0f1fe028329 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 08:52:20 +0200
Subject: [PATCH 17/23] Use version class on str for comparison

---
 tests/core/test_sorting.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index fd477f1c85..1d70db3454 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -14,7 +14,7 @@ def test_sort(self, axis, descending, split, device):
         if not torch.cuda.is_available() and device == "gpu":
             pytest.skip("No gpu available for testing.")
 
-        if np.lib.NumpyVersion(np.__version__) < '2.5.0' and descending is not None:
+        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0') and descending is not None:
             pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
 
         data = ht.random.rand(2, 3, 4, split=split)
@@ -33,7 +33,7 @@ def test_argsort_random(self, axis, descending, split, device):
         if not torch.cuda.is_available() and device == "gpu":
             pytest.skip("No gpu available for testing.")
 
-        if np.lib.NumpyVersion(np.__version__) < '2.5.0' and descending is not None:
+        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0') and descending is not None:
             pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
 
         data = ht.random.rand(2, 3, 4, split=split, device=device)

From 02e5982a1c09db879922012898fd5e3c2d89baab Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 09:01:00 +0200
Subject: [PATCH 18/23] - Removed descending parameter check. - Removed
 descending parameter from keyword when not supported.

---
 tests/core/test_sorting.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 1d70db3454..5a92843ff1 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -14,13 +14,17 @@ def test_sort(self, axis, descending, split, device):
         if not torch.cuda.is_available() and device == "gpu":
             pytest.skip("No gpu available for testing.")
 
-        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0') and descending is not None:
+        kwargs = {"axis": axis}
+
+        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
             pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+        else:
+            kwargs["descending"] = descending
 
         data = ht.random.rand(2, 3, 4, split=split)
-        result, idx = ht.sort(data, axis=axis, descending=descending, return_sort_indices=True)
-        exp = np.sort(data.numpy(), axis=axis, descending=descending)
-        exp_idx = np.argsort(data.numpy(), axis=axis, descending=descending)
+        result, idx = ht.sort(data, return_sort_indices=True, **kwargs)
+        exp = np.sort(data.numpy(), **kwargs)
+        exp_idx = np.argsort(data.numpy(), **kwargs)
 
         assert np.allclose(result.numpy(), exp)
         assert np.allclose(idx.numpy(), exp_idx)
@@ -33,12 +37,16 @@ def test_argsort_random(self, axis, descending, split, device):
         if not torch.cuda.is_available() and device == "gpu":
             pytest.skip("No gpu available for testing.")
 
-        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0') and descending is not None:
+        kwargs = {"axis": axis}
+
+        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
             pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+        else:
+            kwargs["descending"] = descending
 
         data = ht.random.rand(2, 3, 4, split=split, device=device)
-        result_indices = ht.argsort(data, axis=axis, descending=descending)
-        exp_indices = np.argsort(data.numpy(), axis=axis, descending=descending)
+        result_indices = ht.argsort(data, **kwargs)
+        exp_indices = np.argsort(data.numpy(), **kwargs)
         assert np.allclose(result_indices.numpy(), exp_indices)
 
     @pytest.mark.parametrize("device", ["cpu", "gpu"])

From 12bd457384bcd66000c8cf3642d4ff74c884d1d1 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 09:45:36 +0200
Subject: [PATCH 19/23] - Changed chosen device to the one described by the
 environmental variable.

---
 tests/core/test_sorting.py | 39 +++++++++++++++-----------------------
 1 file changed, 15 insertions(+), 24 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 5a92843ff1..08a4e0a346 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -1,19 +1,18 @@
 import numpy as np
 import torch
 import pytest
-
+import os
 import heat as ht
 
 
 class TestSorting:
-    @pytest.mark.parametrize("device", ["cpu", "gpu"])
+    def __init__(self, *args, **kwargs):
+        self.device = os.environ.get("HEAT_TEST_USE_DEVICE", "cpu")
+
     @pytest.mark.parametrize("split", [None, 0, 1, 2])
     @pytest.mark.parametrize("descending", [None, True, False])
     @pytest.mark.parametrize("axis", [None, 0, 1, 2])
-    def test_sort(self, axis, descending, split, device):
-        if not torch.cuda.is_available() and device == "gpu":
-            pytest.skip("No gpu available for testing.")
-
+    def test_sort(self, axis, descending, split):
         kwargs = {"axis": axis}
 
         if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
@@ -21,7 +20,7 @@ def test_sort(self, axis, descending, split, device):
         else:
             kwargs["descending"] = descending
 
-        data = ht.random.rand(2, 3, 4, split=split)
+        data = ht.random.rand(2, 3, 4, split=split, device=self.device)
         result, idx = ht.sort(data, return_sort_indices=True, **kwargs)
         exp = np.sort(data.numpy(), **kwargs)
         exp_idx = np.argsort(data.numpy(), **kwargs)
@@ -29,14 +28,10 @@ def test_sort(self, axis, descending, split, device):
         assert np.allclose(result.numpy(), exp)
         assert np.allclose(idx.numpy(), exp_idx)
 
-    @pytest.mark.parametrize("device", ["cpu", "gpu"])
     @pytest.mark.parametrize("split", [None, 0, 1, 2])
     @pytest.mark.parametrize("descending", [None, True, False])
     @pytest.mark.parametrize("axis", [None, 0, 1, 2])
-    def test_argsort_random(self, axis, descending, split, device):
-        if not torch.cuda.is_available() and device == "gpu":
-            pytest.skip("No gpu available for testing.")
-
+    def test_argsort_random(self, axis, descending, split):
         kwargs = {"axis": axis}
 
         if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
@@ -44,16 +39,12 @@ def test_argsort_random(self, axis, descending, split, device):
         else:
             kwargs["descending"] = descending
 
-        data = ht.random.rand(2, 3, 4, split=split, device=device)
+        data = ht.random.rand(2, 3, 4, split=split, device=self.device)
         result_indices = ht.argsort(data, **kwargs)
         exp_indices = np.argsort(data.numpy(), **kwargs)
         assert np.allclose(result_indices.numpy(), exp_indices)
 
-    @pytest.mark.parametrize("device", ["cpu", "gpu"])
-    def test_argsort_specific(self, device):
-        if not torch.cuda.is_available() and device == "gpu":
-            pytest.skip("No gpu available for testing.")
-
+    def test_argsort_specific(self):
         size = ht.MPI_WORLD.size
         rank = ht.MPI_WORLD.rank
 
@@ -65,17 +56,17 @@ def test_argsort_specific(self, device):
                 [[8, 4, 7], [0, 8, 9]],
             ],
             dtype=torch.int32,
-            device=device,
+            device=self.device,
         )
 
         data = ht.array(tensor, split=0)
         if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
             indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=device
+                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=self.device
             )
         else:
             indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=device
+                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device
             )
         result_indices = ht.argsort(data, axis=0)
         first_indices = result_indices[0].larray
@@ -84,7 +75,7 @@ def test_argsort_specific(self, device):
 
         data = ht.array(tensor, split=1)
         indices_axis_one = torch.tensor(
-            [[0, 1, 1]], dtype=torch.int32, device=device
+            [[0, 1, 1]], dtype=torch.int32, device=self.device
         )
         result_indices = ht.argsort(data, axis=1)
         first_indices = result_indices[0].larray[:1]
@@ -93,7 +84,7 @@ def test_argsort_specific(self, device):
 
         data = ht.array(tensor, split=2)
         indices_axis_two = torch.tensor(
-            [[0], [1]], dtype=torch.int32, device=device
+            [[0], [1]], dtype=torch.int32, device=self.device
         )
         result_indices = ht.argsort(data, axis=2)
         first_indices = result_indices[0].larray[:, :1]
@@ -108,7 +99,7 @@ def test_argsort_specific(self, device):
 
         rank = ht.MPI_WORLD.rank
         ht.random.seed(1)
-        data = ht.random.randn(100, 1, split=0, device=device)
+        data = ht.random.randn(100, 1, split=0, device=self.device)
         indices = ht.argsort(data, axis=0)
         result = ht.resplit(data, None)[indices]
         arr = ht.resplit(result.flatten(), axis=None)

From 5ea6dd9db24a9ff1c010fbd66db76e6a45e25e51 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 12:44:26 +0200
Subject: [PATCH 20/23] Changed device selection

---
 tests/core/test_sorting.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 08a4e0a346..5198720a50 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -3,11 +3,12 @@
 import pytest
 import os
 import heat as ht
+from heat.testing.basic_test import TestCase
 
 
 class TestSorting:
     def __init__(self, *args, **kwargs):
-        self.device = os.environ.get("HEAT_TEST_USE_DEVICE", "cpu")
+        TestCase.setUpClass()
 
     @pytest.mark.parametrize("split", [None, 0, 1, 2])
     @pytest.mark.parametrize("descending", [None, True, False])
@@ -20,7 +21,7 @@ def test_sort(self, axis, descending, split):
         else:
             kwargs["descending"] = descending
 
-        data = ht.random.rand(2, 3, 4, split=split, device=self.device)
+        data = ht.random.rand(2, 3, 4, split=split)
         result, idx = ht.sort(data, return_sort_indices=True, **kwargs)
         exp = np.sort(data.numpy(), **kwargs)
         exp_idx = np.argsort(data.numpy(), **kwargs)
@@ -39,7 +40,7 @@ def test_argsort_random(self, axis, descending, split):
         else:
             kwargs["descending"] = descending
 
-        data = ht.random.rand(2, 3, 4, split=split, device=self.device)
+        data = ht.random.rand(2, 3, 4, split=split)
         result_indices = ht.argsort(data, **kwargs)
         exp_indices = np.argsort(data.numpy(), **kwargs)
         assert np.allclose(result_indices.numpy(), exp_indices)
@@ -56,17 +57,17 @@ def test_argsort_specific(self):
                 [[8, 4, 7], [0, 8, 9]],
             ],
             dtype=torch.int32,
-            device=self.device,
+            device=ht.get_device().torch_device,
         )
 
         data = ht.array(tensor, split=0)
         if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
             indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=self.device
+                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=ht.get_device().torch_device
             )
         else:
             indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device
+                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=ht.get_device().torch_device
             )
         result_indices = ht.argsort(data, axis=0)
         first_indices = result_indices[0].larray
@@ -75,7 +76,7 @@ def test_argsort_specific(self):
 
         data = ht.array(tensor, split=1)
         indices_axis_one = torch.tensor(
-            [[0, 1, 1]], dtype=torch.int32, device=self.device
+            [[0, 1, 1]], dtype=torch.int32, device=ht.get_device().torch_device
         )
         result_indices = ht.argsort(data, axis=1)
         first_indices = result_indices[0].larray[:1]
@@ -84,7 +85,7 @@ def test_argsort_specific(self):
 
         data = ht.array(tensor, split=2)
         indices_axis_two = torch.tensor(
-            [[0], [1]], dtype=torch.int32, device=self.device
+            [[0], [1]], dtype=torch.int32, device=ht.get_device().torch_device
         )
         result_indices = ht.argsort(data, axis=2)
         first_indices = result_indices[0].larray[:, :1]
@@ -99,7 +100,7 @@ def test_argsort_specific(self):
 
         rank = ht.MPI_WORLD.rank
         ht.random.seed(1)
-        data = ht.random.randn(100, 1, split=0, device=self.device)
+        data = ht.random.randn(100, 1, split=0)
         indices = ht.argsort(data, axis=0)
         result = ht.resplit(data, None)[indices]
         arr = ht.resplit(result.flatten(), axis=None)

From 125a3da7c88eefd7ff597f5de7a7706025aaba36 Mon Sep 17 00:00:00 2001
From: Berkant <51971304+Berkant03@users.noreply.github.com>
Date: Mon, 3 Aug 2026 12:58:10 +0200
Subject: [PATCH 21/23] Update tests/core/test_sorting.py

Co-authored-by: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com>
---
 tests/core/test_sorting.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 5198720a50..16f44ae255 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -16,10 +16,11 @@ def __init__(self, *args, **kwargs):
     def test_sort(self, axis, descending, split):
         kwargs = {"axis": axis}
 
-        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
-            pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
-        else:
-            kwargs["descending"] = descending
+        if descending in [True, False]:
+            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
+                pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+            else:
+                kwargs["descending"] = descending
 
         data = ht.random.rand(2, 3, 4, split=split)
         result, idx = ht.sort(data, return_sort_indices=True, **kwargs)

From c79c8f258f92ff9687c7dacde1aa5b4deea591d8 Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Mon, 3 Aug 2026 13:00:46 +0200
Subject: [PATCH 22/23] Changed descending keyword behaviour in older numpy
 versions.

---
 tests/core/test_sorting.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index 16f44ae255..f903d2cd7b 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -36,10 +36,11 @@ def test_sort(self, axis, descending, split):
     def test_argsort_random(self, axis, descending, split):
         kwargs = {"axis": axis}
 
-        if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
-            pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
-        else:
-            kwargs["descending"] = descending
+        if descending in [True, False]:
+            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
+                pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
+            else:
+                kwargs["descending"] = descending
 
         data = ht.random.rand(2, 3, 4, split=split)
         result_indices = ht.argsort(data, **kwargs)

From 91802fd173fa7235cf11143d64d342b8d15a4efe Mon Sep 17 00:00:00 2001
From: Berkant Palazoglu <berkantpalazoglu03@gmail.com>
Date: Tue, 4 Aug 2026 09:32:09 +0200
Subject: [PATCH 23/23] removed specific test for `argsort`

---
 tests/core/test_sorting.py | 66 ++------------------------------------
 1 file changed, 2 insertions(+), 64 deletions(-)

diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py
index f903d2cd7b..34806e6502 100644
--- a/tests/core/test_sorting.py
+++ b/tests/core/test_sorting.py
@@ -17,7 +17,7 @@ def test_sort(self, axis, descending, split):
         kwargs = {"axis": axis}
 
         if descending in [True, False]:
-            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
+            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion("2.5.0"):
                 pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
             else:
                 kwargs["descending"] = descending
@@ -37,7 +37,7 @@ def test_argsort_random(self, axis, descending, split):
         kwargs = {"axis": axis}
 
         if descending in [True, False]:
-            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion('2.5.0'):
+            if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion("2.5.0"):
                 pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.")
             else:
                 kwargs["descending"] = descending
@@ -46,65 +46,3 @@ def test_argsort_random(self, axis, descending, split):
         result_indices = ht.argsort(data, **kwargs)
         exp_indices = np.argsort(data.numpy(), **kwargs)
         assert np.allclose(result_indices.numpy(), exp_indices)
-
-    def test_argsort_specific(self):
-        size = ht.MPI_WORLD.size
-        rank = ht.MPI_WORLD.rank
-
-        tensor = torch.tensor(
-            [
-                [[2, 8, 5], [7, 2, 3]],
-                [[6, 5, 2], [1, 8, 7]],
-                [[9, 3, 0], [1, 2, 4]],
-                [[8, 4, 7], [0, 8, 9]],
-            ],
-            dtype=torch.int32,
-            device=ht.get_device().torch_device,
-        )
-
-        data = ht.array(tensor, split=0)
-        if torch.cuda.is_available() and data.device == ht.gpu and size < 4:
-            indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 2, 0]], dtype=torch.int32, device=ht.get_device().torch_device
-            )
-        else:
-            indices_axis_zero = torch.tensor(
-                [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=ht.get_device().torch_device
-            )
-        result_indices = ht.argsort(data, axis=0)
-        first_indices = result_indices[0].larray
-        if rank == 0:
-            assert torch.equal(first_indices, indices_axis_zero)
-
-        data = ht.array(tensor, split=1)
-        indices_axis_one = torch.tensor(
-            [[0, 1, 1]], dtype=torch.int32, device=ht.get_device().torch_device
-        )
-        result_indices = ht.argsort(data, axis=1)
-        first_indices = result_indices[0].larray[:1]
-        if rank == 0:
-            assert torch.equal(first_indices, indices_axis_one)
-
-        data = ht.array(tensor, split=2)
-        indices_axis_two = torch.tensor(
-            [[0], [1]], dtype=torch.int32, device=ht.get_device().torch_device
-        )
-        result_indices = ht.argsort(data, axis=2)
-        first_indices = result_indices[0].larray[:, :1]
-        if rank == 0:
-            assert torch.equal(first_indices, indices_axis_two)
-
-        # test exceptions
-        with pytest.raises(ValueError):
-            ht.argsort(data, axis=3)
-        with pytest.raises(TypeError):
-            ht.argsort(data, axis="1")
-
-        rank = ht.MPI_WORLD.rank
-        ht.random.seed(1)
-        data = ht.random.randn(100, 1, split=0)
-        indices = ht.argsort(data, axis=0)
-        result = ht.resplit(data, None)[indices]
-        arr = ht.resplit(result.flatten(), axis=None)
-
-        assert (arr.larray[:-1] <= arr.larray[1:]).all()