From 07744502c76d759afccfcac6081b5800e28933e8 Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Thu, 30 May 2024 12:32:51 +0200
Subject: [PATCH 01/10] first commit on sort for more than on name it works for
 samples-block with tests

---
 .../metatensor/operations/_dispatch.py        | 20 +++++++
 .../metatensor/operations/sort.py             | 39 ++++++++++----
 python/metatensor-operations/tests/sort.py    | 54 +++++++++++++++++++
 3 files changed, 104 insertions(+), 9 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/_dispatch.py b/python/metatensor-operations/metatensor/operations/_dispatch.py
index f020a3c17..b59fc9d2d 100644
--- a/python/metatensor-operations/metatensor/operations/_dispatch.py
+++ b/python/metatensor-operations/metatensor/operations/_dispatch.py
@@ -113,6 +113,26 @@ def allclose(
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
+def argsort(values, axis=-1, reverse: bool = False):
+    """
+    Similar to :py:func:`np.argsort`.
+
+    :param labels_values: numpy.array or torch.Tensor
+    :param reverse: if true, order is descending
+
+    :return: indices corresponding to the sorted values in ``labels_values``
+    """
+    if isinstance(values, TorchTensor):
+        return torch.argsort(values, dim=axis, descending=reverse)
+    elif isinstance(values, np.ndarray):
+        tmp = np.argsort(values, axis=axis)
+        if reverse:
+            tmp = tmp[::-1]
+        return tmp
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
 def argsort_labels_values(labels_values, reverse: bool = False):
     """
     Similar to :py:func:`np.argsort`, but sort the rows as one aggregated
diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 9eceeb4db..364cd90ec 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -16,6 +16,7 @@ def _sort_single_gradient_block(
     gradient_block: TensorBlock,
     axes: List[str],
     descending: bool,
+    name: str = "_last_dimension_",
 ) -> TensorBlock:
     """
     Sorts a single gradient tensor block given the tensor block which the gradients are
@@ -95,6 +96,7 @@ def _sort_single_block(
     block: TensorBlock,
     axes: List[str],
     descending: bool,
+    name: str = "_last_dimension_",
 ) -> TensorBlock:
     """
     Sorts a single TensorBlock without the user input checking and sorting of gradients
@@ -102,6 +104,10 @@ def _sort_single_block(
 
     sample_names = block.samples.names
     sample_values = block.samples.values
+    if name != "_last_dimension_" and name not in sample_names:
+        raise ValueError(
+            "`name` must be or '_last_dimension_' " "or one of the sample names"
+        )
 
     component_names: List[List[str]] = []
     components_values = []
@@ -114,20 +120,32 @@ def _sort_single_block(
 
     values = block.values
     if "samples" in axes:
-        sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+        if name == "_last_dimension_":
+            sorted_idx = _dispatch.argsort_labels_values(
+                sample_values, reverse=descending
+            )
+        else:
+            axis = 0
+            for i, v in enumerate(sample_names):
+                if v == name:
+                    axis = i
+                    break
+            sorted_idx = _dispatch.argsort(sample_values[:, axis], reverse=descending)
         sample_values = sample_values[sorted_idx]
         values = values[sorted_idx]
     if "components" in axes:
         for i, _ in enumerate(block.components):
-            sorted_idx = _dispatch.argsort_labels_values(
-                components_values[i], reverse=descending
-            )
+            if name == "_last_dimension_":
+                sorted_idx = _dispatch.argsort_labels_values(
+                    components_values[i], reverse=descending
+                )
             components_values[i] = components_values[i][sorted_idx]
             values = _dispatch.take(values, sorted_idx, axis=i + 1)
     if "properties" in axes:
-        sorted_idx = _dispatch.argsort_labels_values(
-            properties_values, reverse=descending
-        )
+        if name == "_last_dimension_":
+            sorted_idx = _dispatch.argsort_labels_values(
+                properties_values, reverse=descending
+            )
         properties_values = properties_values[sorted_idx]
         values = _dispatch.take(values, sorted_idx, axis=-1)
 
@@ -151,6 +169,7 @@ def sort_block(
     block: TensorBlock,
     axes: Union[str, List[str]] = "all",
     descending: bool = False,
+    name: str = "_last_dimension_",
 ) -> TensorBlock:
     """
     Rearrange the values of a block according to the order given by the sorted metadata
@@ -250,7 +269,7 @@ def sort_block(
                 f"not '{axis}'"
             )
 
-    result_block = _sort_single_block(block, axes_list, descending)
+    result_block = _sort_single_block(block, axes_list, descending, name)
 
     for parameter, gradient in block.gradients():
         if len(gradient.gradients_list()) != 0:
@@ -259,7 +278,7 @@ def sort_block(
         result_block.add_gradient(
             parameter=parameter,
             gradient=_sort_single_gradient_block(
-                block, gradient, axes_list, descending
+                block, gradient, axes_list, descending, name
             ),
         )
 
@@ -271,6 +290,7 @@ def sort(
     tensor: TensorMap,
     axes: Union[str, List[str]] = "all",
     descending: bool = False,
+    name: str = "_last_dimension_",
 ) -> TensorMap:
     """
     Sort the ``tensor`` according to the key values and the blocks for each specified
@@ -365,6 +385,7 @@ def sort(
                 block=tensor.block(tensor.keys[int(i)]),
                 axes=axes_list,
                 descending=descending,
+                name=name,
             )
         )
 
diff --git a/python/metatensor-operations/tests/sort.py b/python/metatensor-operations/tests/sort.py
index 1d26e42f8..58442210b 100644
--- a/python/metatensor-operations/tests/sort.py
+++ b/python/metatensor-operations/tests/sort.py
@@ -260,6 +260,53 @@ def tensor_sorted_descending():
     return TensorMap(keys, [block_2, block_1])
 
 
+@pytest.fixture
+def tensor_two_samples():
+    # samples are descending, components and properties are ascending
+    block_1 = TensorBlock(
+        values=np.array(
+            [
+                [3, 5],
+                [1, 2],
+                [-1, -2],
+                [11, 22],
+                [41, 42],
+            ]
+        ),
+        samples=Labels(
+            ["s", "a"], np.array([[2, 5], [0, 0], [0, 9], [0, 1], [0, 110]])
+        ),
+        components=[],
+        properties=Labels(["p"], np.array([[0], [1]])),
+    )
+    keys = Labels(names=["key_1"], values=np.array([[0]]))
+    # block order is descending
+    return TensorMap(keys, [block_1])
+
+
+@pytest.fixture
+def tensor_two_samples_ascending_a():
+    block_1 = TensorBlock(
+        values=np.array(
+            [
+                [1, 2],
+                [11, 22],
+                [3, 5],
+                [-1, -2],
+                [41, 42],
+            ]
+        ),
+        samples=Labels(
+            ["s", "a"], np.array([[0, 0], [0, 1], [2, 5], [0, 9], [0, 110]])
+        ),
+        components=[],
+        properties=Labels(["p"], np.array([[0], [1]])),
+    )
+    keys = Labels(names=["key_1"], values=np.array([[0]]))
+    # block order is descending
+    return TensorMap(keys, [block_1])
+
+
 def test_sort(tensor, tensor_sorted_ascending):
     metatensor.allclose_block_raise(
         metatensor.sort_block(tensor.block(0)), tensor_sorted_ascending.block(1)
@@ -294,3 +341,10 @@ def test_raise_error(tensor, tensor_sorted_ascending):
     )
     with pytest.raises(ValueError, match=error_message):
         metatensor.operations.sort(tensor, axes=[5])
+
+
+def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
+    metatensor.allclose_block_raise(
+        metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a"),
+        tensor_two_samples_ascending_a.block(0),
+    )

From fd38ead3bbf6e8b4e4b35d3814028e608355af2f Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Thu, 30 May 2024 16:53:34 +0200
Subject: [PATCH 02/10] added sort of components properties and examples

---
 .../metatensor/operations/sort.py             | 88 ++++++++++++++++---
 python/metatensor-operations/tests/sort.py    | 50 +++++++++++
 2 files changed, 125 insertions(+), 13 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 364cd90ec..dbf69d9e4 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -16,7 +16,7 @@ def _sort_single_gradient_block(
     gradient_block: TensorBlock,
     axes: List[str],
     descending: bool,
-    name: str = "_last_dimension_",
+    name: str = "-1",
 ) -> TensorBlock:
     """
     Sorts a single gradient tensor block given the tensor block which the gradients are
@@ -96,7 +96,7 @@ def _sort_single_block(
     block: TensorBlock,
     axes: List[str],
     descending: bool,
-    name: str = "_last_dimension_",
+    name: str = "-1",
 ) -> TensorBlock:
     """
     Sorts a single TensorBlock without the user input checking and sorting of gradients
@@ -104,10 +104,8 @@ def _sort_single_block(
 
     sample_names = block.samples.names
     sample_values = block.samples.values
-    if name != "_last_dimension_" and name not in sample_names:
-        raise ValueError(
-            "`name` must be or '_last_dimension_' " "or one of the sample names"
-        )
+    if name != "-1" and name not in sample_names:
+        raise ValueError("`name` must be or '-1' " "or one of the sample names")
 
     component_names: List[List[str]] = []
     components_values = []
@@ -120,32 +118,50 @@ def _sort_single_block(
 
     values = block.values
     if "samples" in axes:
-        if name == "_last_dimension_":
+        if name == "-1":
             sorted_idx = _dispatch.argsort_labels_values(
                 sample_values, reverse=descending
             )
         else:
             axis = 0
-            for i, v in enumerate(sample_names):
+            for iv, v in enumerate(sample_names):
                 if v == name:
-                    axis = i
+                    axis = iv
                     break
             sorted_idx = _dispatch.argsort(sample_values[:, axis], reverse=descending)
         sample_values = sample_values[sorted_idx]
         values = values[sorted_idx]
     if "components" in axes:
         for i, _ in enumerate(block.components):
-            if name == "_last_dimension_":
+            if name == "-1":
                 sorted_idx = _dispatch.argsort_labels_values(
                     components_values[i], reverse=descending
                 )
+            else:
+                axis = 0
+                for ic, c in enumerate(component_names[i]):
+                    if c == name:
+                        axis = ic
+                        break
+                sorted_idx = _dispatch.argsort(
+                    components_values[i][:, axis], reverse=descending
+                )
             components_values[i] = components_values[i][sorted_idx]
             values = _dispatch.take(values, sorted_idx, axis=i + 1)
     if "properties" in axes:
-        if name == "_last_dimension_":
+        if name == "-1":
             sorted_idx = _dispatch.argsort_labels_values(
                 properties_values, reverse=descending
             )
+        else:
+            axis = 0
+            for ip, p in enumerate(property_names):
+                if p == name:
+                    axis = ip
+                    break
+            sorted_idx = _dispatch.argsort(
+                properties_values[:, axis], reverse=descending
+            )
         properties_values = properties_values[sorted_idx]
         values = _dispatch.take(values, sorted_idx, axis=-1)
 
@@ -169,7 +185,7 @@ def sort_block(
     block: TensorBlock,
     axes: Union[str, List[str]] = "all",
     descending: bool = False,
-    name: str = "_last_dimension_",
+    name: str = "-1",
 ) -> TensorBlock:
     """
     Rearrange the values of a block according to the order given by the sorted metadata
@@ -184,6 +200,9 @@ def sort_block(
 
     :param descending: if false, the order is ascending
 
+    :param name: name of `axes` to be used for the sorting. Default == "-1"
+         means sort along the last axis
+
     :return: sorted tensor block
 
     >>> import numpy as np
@@ -220,6 +239,23 @@ def sort_block(
     >>> sorted_block = metatensor.sort_block(block)
     >>> np.all(sorted_block.values == block_sorted_stepwise.values)
     True
+    >>> # You can also choose along which axis of "samples“ you sort
+    >>> block2 = TensorBlock(
+    ...     values=np.arange(12).reshape(4, 3),
+    ...     samples=Labels(["system", "atom"], np.array([[0, 2], [1, 0], [2,5],[2,1]])),
+    ...     components=[],
+    ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
+    ... )
+    >>> block_sorted_2_sample = metatensor.sort_block(
+    ...     block2, axes=["samples"],name="atom"
+    ... )
+    >>> # samples (first dimension of the array) are sorted
+    >>> block_sorted_2_sample.values
+    array([[4, 5, 6],
+           [10,11,12],
+           [1, 2, 3],
+           [7, 8, 9],
+           ])
     >>> # This function can also sort gradients:
     >>> sorted_block.add_gradient(
     ...     parameter="g",
@@ -250,9 +286,21 @@ def sort_block(
     if isinstance(axes, str):
         if axes == "all":
             axes_list = ["samples", "components", "properties"]
+            if name != "-1":
+                raise ValueError(
+                    "'name' is allowed only if 'axes' is one of"
+                    "'samples', 'components','properties' but"
+                    "'axes'=='all'"
+                )
         else:
             axes_list = [axes]
     elif isinstance(axes, list):
+        if name != "-1":
+            raise ValueError(
+                "'name' is allowed only if 'axes' is one of"
+                "'samples', 'components','properties' but"
+                "'axes' is a List"
+            )
         axes_list = axes
     else:
         if torch_jit_is_scripting():
@@ -290,7 +338,7 @@ def sort(
     tensor: TensorMap,
     axes: Union[str, List[str]] = "all",
     descending: bool = False,
-    name: str = "_last_dimension_",
+    name: str = "-1",
 ) -> TensorMap:
     """
     Sort the ``tensor`` according to the key values and the blocks for each specified
@@ -305,6 +353,8 @@ def sort(
         Possible values are ``'keys'``, ``'samples'``, ``'components'``,
         ``'properties'`` and ``'all'`` to sort everything.
     :param descending: if false, the order is ascending
+    :param name: name of `axes` to be used for the sorting. Default == "-1"
+         means sort along the last axis
     :return: sorted tensor map
 
     >>> import numpy as np
@@ -336,6 +386,12 @@ def sort(
         if axes == "all":
             axes_list = ["samples", "components", "properties"]
             sort_keys = True
+            if name != "-1":
+                raise ValueError(
+                    "'name' is allowed only if 'axes' is one of"
+                    "'samples', 'components','properties' but"
+                    "'axes'=='all'"
+                )
         elif axes == "keys":
             axes_list = torch_jit_annotate(List[str], [])
             sort_keys = True
@@ -344,6 +400,12 @@ def sort(
             sort_keys = False
 
     elif isinstance(axes, list):
+        if name != "-1":
+            raise ValueError(
+                "'name' is allowed only if 'axes' is one of"
+                "'samples', 'components','properties' but"
+                "'axes' is a List"
+            )
         axes_list = axes
 
         if "keys" in axes_list:
diff --git a/python/metatensor-operations/tests/sort.py b/python/metatensor-operations/tests/sort.py
index 58442210b..5a9b71165 100644
--- a/python/metatensor-operations/tests/sort.py
+++ b/python/metatensor-operations/tests/sort.py
@@ -307,6 +307,29 @@ def tensor_two_samples_ascending_a():
     return TensorMap(keys, [block_1])
 
 
+@pytest.fixture
+def tensor_two_samples_descending_a():
+    block_1 = TensorBlock(
+        values=np.array(
+            [
+                [41, 42],
+                [-1, -2],
+                [3, 5],
+                [11, 22],
+                [1, 2],
+            ]
+        ),
+        samples=Labels(
+            ["s", "a"], np.array([[0, 110], [0, 9], [2, 5], [0, 1], [0, 0]])
+        ),
+        components=[],
+        properties=Labels(["p"], np.array([[0], [1]])),
+    )
+    keys = Labels(names=["key_1"], values=np.array([[0]]))
+    # block order is descending
+    return TensorMap(keys, [block_1])
+
+
 def test_sort(tensor, tensor_sorted_ascending):
     metatensor.allclose_block_raise(
         metatensor.sort_block(tensor.block(0)), tensor_sorted_ascending.block(1)
@@ -342,9 +365,36 @@ def test_raise_error(tensor, tensor_sorted_ascending):
     with pytest.raises(ValueError, match=error_message):
         metatensor.operations.sort(tensor, axes=[5])
 
+    error_message = (
+        "'name' is allowed only if 'axes' is one of"
+        "'samples', 'components','properties' but"
+        "'axes' is a List"
+    )
+    with pytest.raises(ValueError, match=error_message):
+        metatensor.operations.sort(tensor, axes=["samples", "components"], name="s")
+
+    error_message = (
+        "'name' is allowed only if 'axes' is one of"
+        "'samples', 'components','properties' but"
+        "'axes'=='all'"
+    )
+    with pytest.raises(ValueError, match=error_message):
+        metatensor.operations.sort(tensor, axes="all", name="s")
+
 
 def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
     metatensor.allclose_block_raise(
         metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a"),
         tensor_two_samples_ascending_a.block(0),
     )
+
+
+def test_sort_two_sample_descending(
+    tensor_two_samples, tensor_two_samples_descending_a
+):
+    metatensor.allclose_block_raise(
+        metatensor.sort_block(
+            tensor_two_samples.block(0), axes="samples", name="a", descending=True
+        ),
+        tensor_two_samples_descending_a.block(0),
+    )

From e8b1cf427c901cda02205cfbc0ecc9288d27ceeb Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Thu, 30 May 2024 19:08:40 +0200
Subject: [PATCH 03/10] first attempt at gradient

---
 .../metatensor/operations/sort.py               | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index dbf69d9e4..c7fb0f95a 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -44,9 +44,20 @@ def _sort_single_gradient_block(
         # the parent block
         block_sample_values = block.samples.values
         # sample index -> sample labels
-        sorted_idx = _dispatch.argsort_labels_values(
-            block_sample_values, reverse=descending
-        )
+        if name == "-1":
+            sorted_idx = _dispatch.argsort_labels_values(
+                block_sample_values, reverse=descending
+            )
+        else:
+            axis = 0
+            for iv, v in enumerate(sample_names):
+                if v == name:
+                    axis = iv
+                    break
+            sorted_idx = _dispatch.argsort(
+                block_sample_values[:, axis], reverse=descending
+            )
+
         # obtain inverse mapping sample labels -> sample index
         sorted_idx_inverse = _dispatch.empty_like(sorted_idx, shape=(len(sorted_idx),))
         sorted_idx_inverse[sorted_idx] = _dispatch.int_array_like(

From 2a6d41130a236fc5ead35e20aefab3390630559c Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Fri, 31 May 2024 12:33:42 +0200
Subject: [PATCH 04/10] fix gradient (sample) + tests

---
 .../metatensor/operations/sort.py             | 11 ++-
 python/metatensor-operations/tests/sort.py    | 73 ++++++++++++++++++-
 2 files changed, 77 insertions(+), 7 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index c7fb0f95a..98930cced 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -50,7 +50,7 @@ def _sort_single_gradient_block(
             )
         else:
             axis = 0
-            for iv, v in enumerate(sample_names):
+            for iv, v in enumerate(block.samples.names):
                 if v == name:
                     axis = iv
                     break
@@ -70,10 +70,13 @@ def _sort_single_gradient_block(
             _dispatch.to_index_array(sample_values[:, 0])
         ]
 
+        # I AM NOT SURE WHAT THESE LINES DO THE OLD TEST PASS EVEN
+        # IF THEY ARE COMMENTED
+
         # sort the samples in gradient regularly moving the rows considering all columns
-        sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
-        sample_values = sample_values[sorted_idx]
-        values = values[sorted_idx]
+        # sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+        # sample_values = sample_values[sorted_idx]
+        # values = values[sorted_idx]
     if "components" in axes:
         for i, _ in enumerate(gradient_block.components):
             sorted_idx = _dispatch.argsort_labels_values(
diff --git a/python/metatensor-operations/tests/sort.py b/python/metatensor-operations/tests/sort.py
index 5a9b71165..9c8138d32 100644
--- a/python/metatensor-operations/tests/sort.py
+++ b/python/metatensor-operations/tests/sort.py
@@ -279,6 +279,40 @@ def tensor_two_samples():
         components=[],
         properties=Labels(["p"], np.array([[0], [1]])),
     )
+    block_1.add_gradient(
+        parameter="g",
+        gradient=TensorBlock(
+            values=np.array(
+                [
+                    [[1, 2], [3, 4]],
+                    [[5, 6], [7, 8]],
+                    [[9, 10], [11, 12]],
+                    [[13, 14], [15, 16]],
+                    [[17, 18], [19, 20]],
+                    [[21, 22], [23, 24]],
+                    [[25, 26], [27, 28]],
+                ]
+            ),
+            samples=Labels(
+                ["sample", "g"],
+                np.array(
+                    [
+                        [2, 1],
+                        [1, 1],
+                        [0, 1],
+                        [0, 2],
+                        [3, 1],
+                        [4, 1],
+                        [4, 0],
+                    ]
+                ),
+            ),
+            components=[
+                Labels(["c"], np.array([[1], [0]])),
+            ],
+            properties=block_1.properties,
+        ),
+    )
     keys = Labels(names=["key_1"], values=np.array([[0]]))
     # block order is descending
     return TensorMap(keys, [block_1])
@@ -302,6 +336,40 @@ def tensor_two_samples_ascending_a():
         components=[],
         properties=Labels(["p"], np.array([[0], [1]])),
     )
+    block_1.add_gradient(
+        parameter="g",
+        gradient=TensorBlock(
+            values=np.array(
+                [
+                    [[1, 2], [3, 4]],
+                    [[5, 6], [7, 8]],
+                    [[9, 10], [11, 12]],
+                    [[13, 14], [15, 16]],
+                    [[17, 18], [19, 20]],
+                    [[21, 22], [23, 24]],
+                    [[25, 26], [27, 28]],
+                ]
+            ),
+            samples=Labels(
+                ["sample", "g"],
+                np.array(
+                    [
+                        [3, 1],
+                        [0, 1],
+                        [2, 1],
+                        [2, 2],
+                        [1, 1],
+                        [4, 1],
+                        [4, 0],
+                    ]
+                ),
+            ),
+            components=[
+                Labels(["c"], np.array([[1], [0]])),
+            ],
+            properties=block_1.properties,
+        ),
+    )
     keys = Labels(names=["key_1"], values=np.array([[0]]))
     # block order is descending
     return TensorMap(keys, [block_1])
@@ -392,9 +460,8 @@ def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
 def test_sort_two_sample_descending(
     tensor_two_samples, tensor_two_samples_descending_a
 ):
+    t = metatensor.remove_gradients(tensor_two_samples)
     metatensor.allclose_block_raise(
-        metatensor.sort_block(
-            tensor_two_samples.block(0), axes="samples", name="a", descending=True
-        ),
+        metatensor.sort_block(t.block(0), axes="samples", name="a", descending=True),
         tensor_two_samples_descending_a.block(0),
     )

From d3a43cd40b71b566b82371b0f926d27db309478a Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 09:37:04 +0200
Subject: [PATCH 05/10] added metatensor-torch tests

---
 python/metatensor-operations/metatensor/operations/_dispatch.py | 2 +-
 python/metatensor-torch/tests/operations/sort.py                | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/metatensor-operations/metatensor/operations/_dispatch.py b/python/metatensor-operations/metatensor/operations/_dispatch.py
index b59fc9d2d..03d9ac087 100644
--- a/python/metatensor-operations/metatensor/operations/_dispatch.py
+++ b/python/metatensor-operations/metatensor/operations/_dispatch.py
@@ -113,7 +113,7 @@ def allclose(
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-def argsort(values, axis=-1, reverse: bool = False):
+def argsort(values, axis: int = -1, reverse: bool = False):
     """
     Similar to :py:func:`np.argsort`.
 
diff --git a/python/metatensor-torch/tests/operations/sort.py b/python/metatensor-torch/tests/operations/sort.py
index 35d423372..adf546463 100644
--- a/python/metatensor-torch/tests/operations/sort.py
+++ b/python/metatensor-torch/tests/operations/sort.py
@@ -3,6 +3,7 @@
 import pytest
 import torch
 from packaging import version
+import pytest
 
 import metatensor.torch
 from metatensor.torch import Labels, TensorBlock, TensorMap

From 33976fd98ea06220083527c718ef0f5413bc35ed Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 10:56:37 +0200
Subject: [PATCH 06/10] fix docs

---
 python/metatensor-operations/metatensor/operations/sort.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 98930cced..13aaba68b 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -119,7 +119,7 @@ def _sort_single_block(
     sample_names = block.samples.names
     sample_values = block.samples.values
     if name != "-1" and name not in sample_names:
-        raise ValueError("`name` must be or '-1' " "or one of the sample names")
+        raise ValueError("`name` must be or '-1' or one of the sample names")
 
     component_names: List[List[str]] = []
     components_values = []
@@ -309,7 +309,7 @@ def sort_block(
         else:
             axes_list = [axes]
     elif isinstance(axes, list):
-        if name != "-1":
+        if len(axes) > 1 and name != "-1":
             raise ValueError(
                 "'name' is allowed only if 'axes' is one of"
                 "'samples', 'components','properties' but"
@@ -414,7 +414,7 @@ def sort(
             sort_keys = False
 
     elif isinstance(axes, list):
-        if name != "-1":
+        if len(axes) > 1 and name != "-1":
             raise ValueError(
                 "'name' is allowed only if 'axes' is one of"
                 "'samples', 'components','properties' but"

From 1e2d9e6177d3c09d3eaf9a06a93e5260638cbb57 Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 11:20:28 +0200
Subject: [PATCH 07/10] fix docs-

---
 .../metatensor-operations/metatensor/operations/sort.py   | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 13aaba68b..3ad3160f6 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -265,10 +265,10 @@ def sort_block(
     ... )
     >>> # samples (first dimension of the array) are sorted
     >>> block_sorted_2_sample.values
-    array([[4, 5, 6],
-           [10,11,12],
-           [1, 2, 3],
-           [7, 8, 9],
+    array([[3, 4, 5],
+           [9,10,11],
+           [0, 1, 2],
+           [6, 7, 8],
            ])
     >>> # This function can also sort gradients:
     >>> sorted_block.add_gradient(

From 63538eb4776e07ff1401c9801f86e0f61fb2ac14 Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 11:24:23 +0200
Subject: [PATCH 08/10] linter

---
 .../metatensor-operations/metatensor/operations/sort.py  | 9 ++++++---
 python/metatensor-torch/tests/operations/sort.py         | 1 -
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 3ad3160f6..0300c5bc9 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -74,7 +74,8 @@ def _sort_single_gradient_block(
         # IF THEY ARE COMMENTED
 
         # sort the samples in gradient regularly moving the rows considering all columns
-        # sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+        # sorted_idx = _dispatch.argsort_labels_values(sample_values,
+        # reverse=descending)
         # sample_values = sample_values[sorted_idx]
         # values = values[sorted_idx]
     if "components" in axes:
@@ -256,12 +257,14 @@ def sort_block(
     >>> # You can also choose along which axis of "samples“ you sort
     >>> block2 = TensorBlock(
     ...     values=np.arange(12).reshape(4, 3),
-    ...     samples=Labels(["system", "atom"], np.array([[0, 2], [1, 0], [2,5],[2,1]])),
+    ...     samples=Labels(
+    ...         ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
+    ...     ),
     ...     components=[],
     ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
     ... )
     >>> block_sorted_2_sample = metatensor.sort_block(
-    ...     block2, axes=["samples"],name="atom"
+    ...     block2, axes=["samples"], name="atom"
     ... )
     >>> # samples (first dimension of the array) are sorted
     >>> block_sorted_2_sample.values
diff --git a/python/metatensor-torch/tests/operations/sort.py b/python/metatensor-torch/tests/operations/sort.py
index adf546463..35d423372 100644
--- a/python/metatensor-torch/tests/operations/sort.py
+++ b/python/metatensor-torch/tests/operations/sort.py
@@ -3,7 +3,6 @@
 import pytest
 import torch
 from packaging import version
-import pytest
 
 import metatensor.torch
 from metatensor.torch import Labels, TensorBlock, TensorMap

From 00a206136dfc5e85b63217d9eeb213f4b1105a59 Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 12:03:04 +0200
Subject: [PATCH 09/10] doc

---
 .../metatensor-operations/metatensor/operations/sort.py  | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 0300c5bc9..44cfd51a3 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -268,11 +268,10 @@ def sort_block(
     ... )
     >>> # samples (first dimension of the array) are sorted
     >>> block_sorted_2_sample.values
-    array([[3, 4, 5],
-           [9,10,11],
-           [0, 1, 2],
-           [6, 7, 8],
-           ])
+    array([[ 3,  4,  5],
+           [ 9, 10, 11],
+           [ 0,  1,  2],
+           [ 6,  7,  8]])
     >>> # This function can also sort gradients:
     >>> sorted_block.add_gradient(
     ...     parameter="g",

From 86de8f47e95fff8357cd40e200136e20400518ca Mon Sep 17 00:00:00 2001
From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com>
Date: Sat, 1 Jun 2024 12:27:47 +0200
Subject: [PATCH 10/10] gradient

---
 .../metatensor/operations/sort.py             | 47 +++++++++----------
 python/metatensor-operations/tests/sort.py    | 20 +++++---
 2 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/python/metatensor-operations/metatensor/operations/sort.py b/python/metatensor-operations/metatensor/operations/sort.py
index 44cfd51a3..010fb6736 100644
--- a/python/metatensor-operations/metatensor/operations/sort.py
+++ b/python/metatensor-operations/metatensor/operations/sort.py
@@ -70,14 +70,10 @@ def _sort_single_gradient_block(
             _dispatch.to_index_array(sample_values[:, 0])
         ]
 
-        # I AM NOT SURE WHAT THESE LINES DO THE OLD TEST PASS EVEN
-        # IF THEY ARE COMMENTED
-
         # sort the samples in gradient regularly moving the rows considering all columns
-        # sorted_idx = _dispatch.argsort_labels_values(sample_values,
-        # reverse=descending)
-        # sample_values = sample_values[sorted_idx]
-        # values = values[sorted_idx]
+        sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+        sample_values = sample_values[sorted_idx]
+        values = values[sorted_idx]
     if "components" in axes:
         for i, _ in enumerate(gradient_block.components):
             sorted_idx = _dispatch.argsort_labels_values(
@@ -254,24 +250,6 @@ def sort_block(
     >>> sorted_block = metatensor.sort_block(block)
     >>> np.all(sorted_block.values == block_sorted_stepwise.values)
     True
-    >>> # You can also choose along which axis of "samples“ you sort
-    >>> block2 = TensorBlock(
-    ...     values=np.arange(12).reshape(4, 3),
-    ...     samples=Labels(
-    ...         ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
-    ...     ),
-    ...     components=[],
-    ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
-    ... )
-    >>> block_sorted_2_sample = metatensor.sort_block(
-    ...     block2, axes=["samples"], name="atom"
-    ... )
-    >>> # samples (first dimension of the array) are sorted
-    >>> block_sorted_2_sample.values
-    array([[ 3,  4,  5],
-           [ 9, 10, 11],
-           [ 0,  1,  2],
-           [ 6,  7,  8]])
     >>> # This function can also sort gradients:
     >>> sorted_block.add_gradient(
     ...     parameter="g",
@@ -297,7 +275,24 @@ def sort_block(
     <BLANKLINE>
            [[ 6,  7,  8],
             [ 9, 10, 11]]])
-
+    >>> # You can also choose along which axis of "samples“ you sort
+    >>> block2 = TensorBlock(
+    ...     values=np.arange(12).reshape(4, 3),
+    ...     samples=Labels(
+    ...         ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
+    ...     ),
+    ...     components=[],
+    ...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
+    ... )
+    >>> block_sorted_2_sample = metatensor.sort_block(
+    ...     block2, axes=["samples"], name="atom"
+    ... )
+    >>> # samples (first dimension of the array) are sorted
+    >>> block_sorted_2_sample.values
+    array([[ 3,  4,  5],
+           [ 9, 10, 11],
+           [ 0,  1,  2],
+           [ 6,  7,  8]])
     """
     if isinstance(axes, str):
         if axes == "all":
diff --git a/python/metatensor-operations/tests/sort.py b/python/metatensor-operations/tests/sort.py
index 9c8138d32..a15c57b7d 100644
--- a/python/metatensor-operations/tests/sort.py
+++ b/python/metatensor-operations/tests/sort.py
@@ -341,26 +341,26 @@ def tensor_two_samples_ascending_a():
         gradient=TensorBlock(
             values=np.array(
                 [
-                    [[1, 2], [3, 4]],
                     [[5, 6], [7, 8]],
+                    [[17, 18], [19, 20]],
                     [[9, 10], [11, 12]],
                     [[13, 14], [15, 16]],
-                    [[17, 18], [19, 20]],
-                    [[21, 22], [23, 24]],
+                    [[1, 2], [3, 4]],
                     [[25, 26], [27, 28]],
+                    [[21, 22], [23, 24]],
                 ]
             ),
             samples=Labels(
                 ["sample", "g"],
                 np.array(
                     [
-                        [3, 1],
                         [0, 1],
+                        [1, 1],
                         [2, 1],
                         [2, 2],
-                        [1, 1],
-                        [4, 1],
+                        [3, 1],
                         [4, 0],
+                        [4, 1],
                     ]
                 ),
             ),
@@ -451,6 +451,14 @@ def test_raise_error(tensor, tensor_sorted_ascending):
 
 
 def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
+    print(
+        "ii",
+        metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a")
+        .gradient("g")
+        .samples,
+    )
+    print("jj", tensor_two_samples_ascending_a.block(0).gradient("g").samples)
+
     metatensor.allclose_block_raise(
         metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a"),
         tensor_two_samples_ascending_a.block(0),