add backend.argsort

Author: refraction-ray

CHANGELOG.md

@@ -20,6 +20,8 @@
 
 - Introducing pyproject.toml finally
 
+- Add `argsort` method for backends
+
 ### Fixed
 
 - Fixed `one_hot` in numpy backend.

examples/hamiltonian_building.py

@@ -7,6 +7,7 @@
 import jax
 import numpy as np
 import quimb
+from quimb import tensor as qt
 import scipy
 import tensorflow as tf
 
@@ -90,7 +91,7 @@
 print("hamiltonian building with quimb")
 print("quimb version: ", quimb.__version__)
 
-builder = quimb.tensor.SpinHam1D()
+builder = qt.SpinHam1D()
 # spin operator instead of Pauli matrix
 builder += 4, "Z", "Z"
 builder += -2, "X"

tensorcircuit/backends/abstract_backend.py

@@ -581,6 +581,21 @@ def argmin(self: Any, a: Tensor, axis: int = 0) -> Tensor:
             "Backend '{}' has not implemented `argmin`.".format(self.name)
         )
 
+    def argsort(self: Any, a: Tensor, axis: int = -1) -> Tensor:
+        """
+        return the indices that would sort an array.
+
+        :param a: the tensor to be sorted
+        :type a: Tensor
+        :param axis: the sorted axis, defaults to -1
+        :type axis: int
+        :return: the sorted indices
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `argsort`.".format(self.name)
+        )
+
     def unique_with_counts(self: Any, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         """
         Find the unique elements and their corresponding counts of the given tensor ``a``.

tensorcircuit/backends/jax_backend.py

@@ -387,6 +387,9 @@ def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return jnp.argmin(a, axis=axis)
 
+    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return jnp.argsort(a, axis=axis)
+
     def unique_with_counts(  # type: ignore
         self, a: Tensor, *, size: Optional[int] = None, fill_value: Optional[int] = None
     ) -> Tuple[Tensor, Tensor]:

tensorcircuit/backends/numpy_backend.py

@@ -251,6 +251,9 @@ def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         return np.searchsorted(a, v, side=side)  # type: ignore
 
+    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return np.argsort(a, axis=axis)
+
     def set_random_state(
         self, seed: Optional[int] = None, get_only: bool = False
     ) -> Any:

tensorcircuit/backends/tensorflow_backend.py

@@ -530,6 +530,9 @@ def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return tf.math.argmin(a, axis=axis)
 
+    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return tf.argsort(a, axis=axis)
+
     def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         r = tf.unique_with_counts(a)
         order = tf.argsort(r.y)

tests/test_backends.py

@@ -481,6 +481,26 @@ def test_arg_cmp(backend):
     )
 
 
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_argsort(backend):
+    # Test basic argsort functionality
+    a = tc.array_to_tensor(np.array([3, 1, 2]), dtype="float32")
+    result = tc.backend.argsort(a)
+    expected = np.array([1, 2, 0])  # indices that would sort the array
+    np.testing.assert_allclose(result, expected)
+
+    # Test argsort with 2D array, default axis=-1
+    b = tc.array_to_tensor(np.array([[3, 1, 2], [4, 0, 1]]), dtype="float32")
+    result = tc.backend.argsort(b)
+    expected = np.array([[1, 2, 0], [1, 2, 0]])
+    np.testing.assert_allclose(result, expected)
+
+    # Test argsort with 2D array, axis=0
+    result = tc.backend.argsort(b, axis=0)
+    expected = np.array([[0, 1, 1], [1, 0, 0]])
+    np.testing.assert_allclose(result, expected)
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_tree_map(backend):
     def f(a, b):