Repository: PennyLaneAI/pennylane

Issue:
Add `torch` exception for `sort` to math's `single_dispatch.py`.
### Feature details

Update handling of `sort` in `qml.math` such that the return values are the same across all interfaces. 

The current behaviour is as follows. For `numpy`, `jax`, and `tf`, 
```pycon
>>> a = np.array([1, 3, 4, 2])
>>> qml.math.sort(a)
tensor([1, 2, 3, 4], requires_grad=True)
``` 

But for `torch`, the argument indices are also returned:
```pycon
>>> a = torch.tensor([1, 3, 4, 2])
>>> qml.math.sort(a)
torch.return_types.sort(
values=tensor([1, 2, 3, 4]),
indices=tensor([0, 3, 1, 2]))
```

### Implementation

It should be possible to just add a custom implementation in `math/single_dispatch.py` and then register the new function. For example (taken from the current implementation of `_take_torch`):

```python
def _sort_torch(tensor):
    # Do the sort 
    # return only the sorted array 

ar.register_function("torch", "sort", _sort_torch)
```

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_


Discussion:
Hi @glassnotes I am new to the open source community and wanted to ask if it is alright if I try this as my first issue? 
Hi @RubidgeCarrie , welcome! Absolutely, please go ahead, and let us know if you have questions. :grin: 

Resolving patch:
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -357,6 +357,15 @@ def _scatter_element_add_torch(tensor, index, value):
 ar.register_function("torch", "scatter_element_add", _scatter_element_add_torch)
 
 
+def _sort_torch(tensor):
+    """Update handling of sort to return only values not indices."""
+    sorted_tensor = _i("torch").sort(tensor)
+    return sorted_tensor.values
+
+
+ar.register_function("torch", "sort", _sort_torch)
+
+
 # -------------------------------- JAX --------------------------------- #
 
 
diff --git a/pennylane/transforms/decompositions/two_qubit_unitary.py b/pennylane/transforms/decompositions/two_qubit_unitary.py
--- a/pennylane/transforms/decompositions/two_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/two_qubit_unitary.py
@@ -132,10 +132,7 @@ def _compute_num_cnots(U):
     # To distinguish between 1/2 CNOT cases, we need to look at the eigenvalues
     evs = math.linalg.eigvals(gammaU)
 
-    if math.get_interface(u) == "torch":
-        sorted_evs, _ = math.sort(math.imag(evs))
-    else:
-        sorted_evs = math.sort(math.imag(evs))
+    sorted_evs = math.sort(math.imag(evs))
 
     # Case: 1 CNOT, the trace is 0, and the eigenvalues of gammaU are [-1j, -1j, 1j, 1j]
     # Checking the eigenvalues is needed because of some special 2-CNOT cases that yield
@@ -366,11 +363,7 @@ def _decomposition_2_cnots(U, wires):
     # some reason this case is not handled properly with the full algorithm, so
     # we treat it separately.
 
-    if math.get_interface(u) == "torch":
-        # Torch's sort function returns both the sorted values and the new order
-        sorted_evs, _ = math.sort(math.real(evs))
-    else:
-        sorted_evs = math.sort(math.real(evs))
+    sorted_evs = math.sort(math.real(evs))
 
     if math.allclose(sorted_evs, [-1, -1, 1, 1]):
         interior_decomp = [


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1630,3 +1630,25 @@ def cost_fn(*params):
         grad = qml.grad(cost_fn)(*values)
 
         assert res == {0, 1}
+
+
+test_sort_data = [
+    ([1, 3, 4, 2], [1, 2, 3, 4]),
+    (onp.array([1, 3, 4, 2]), onp.array([1, 2, 3, 4])),
+    (np.array([1, 3, 4, 2]), np.array([1, 2, 3, 4])),
+    (jnp.array([1, 3, 4, 2]), jnp.array([1, 2, 3, 4])),
+    (torch.tensor([1, 3, 4, 2]), torch.tensor([1, 2, 3, 4])),
+    (tf.Variable([1, 3, 4, 2]), tf.Variable([1, 2, 3, 4])),
+    (tf.constant([1, 3, 4, 2]), tf.constant([1, 2, 3, 4])),
+]
+
+
+class TestSortFunction:
+    """Test the sort function works across all interfaces"""
+
+    @pytest.mark.parametrize("input, test_output", test_sort_data)
+    def test_sort(self, input, test_output):
+        """Test the sort method is outputting only sorted values not indices"""
+        result = fn.sort(input)
+
+        assert all(result == test_output)