Repository: PennyLaneAI/pennylane
Pull request #5583 (2024-04-29 08:20:53)

Fix a sorting order in parameter-shift terms
**Context:**
Currently, `process_shifts` sorts shifts according to the absolute value of the shift.
It uses `kind="stable"`, which guarantees to keep inputs in their input order if they have same absolute value.
While this is stable against repeated calls, the stable order is somewhat arbitrarily depending on how a user/PennyLane
arrived at the input shifts. 

This is visible in `tests/gradients/core/test_general_shift_rules.py`, where the expected output values sometimes have positive before negative shifts with same magnitude, and sometimes the other way around.

**Description of the Change:**
This PR uses `lexsort`, which is like `argsort` but allows us to use the sign of the shift as a secondary sorting key, producing not only stable but also well-defined sorting results.

**Benefits:**
Predictable sorting results beyond the absolute value alone.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-45043]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5583?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.69%. Comparing base [(`a98998e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a98998ec1c25e1615e0984f068805010dc90ec7c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`503ce65`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5583?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5583      +/-   ##
==========================================
- Coverage   99.69%   99.69%   -0.01%     
==========================================
  Files         412      412              
  Lines       38564    38273     -291     
==========================================
- Hits        38447    38155     -292     
- Misses        117      118       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5583?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5ac7447f9..291e91c17 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,6 +6,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* The sorting order of parameter-shift terms is now guaranteed to resolve ties in the absolute value with the sign of the shifts.
+  [(#5582)](https://github.com/PennyLaneAI/pennylane/pull/5582)
+
 <h4>Mid-circuit measurements and dynamic circuits</h4>
 
 * The `dynamic_one_shot` transform can be compiled with `jax.jit`.
@@ -79,7 +82,7 @@
 
 * ``qml.from_qasm_file`` has been removed. The user can open files and load their content using `qml.from_qasm`.
   [(#5659)](https://github.com/PennyLaneAI/pennylane/pull/5659)
-  
+
 * ``qml.load`` has been removed in favour of more specific functions, such as ``qml.from_qiskit``, etc.
   [(#5654)](https://github.com/PennyLaneAI/pennylane/pull/5654)
 
diff --git a/pennylane/gradients/general_shift_rules.py b/pennylane/gradients/general_shift_rules.py
index 4f901b4e7..615000c40 100644
--- a/pennylane/gradients/general_shift_rules.py
+++ b/pennylane/gradients/general_shift_rules.py
@@ -59,6 +59,7 @@ def process_shifts(rule, tol=1e-10, batch_duplicates=True):
 
     - Finally, the terms are sorted according to the absolute value of ``shift``,
       This ensures that a zero-shift term, if it exists, is returned first.
+      For equal absolute values of two shifts, the positive shift is sorted to come first.
     """
     # set all small coefficients, multipliers if present, and shifts to zero.
     rule[np.abs(rule) < tol] = 0
@@ -78,8 +79,9 @@ def process_shifts(rule, tol=1e-10, batch_duplicates=True):
             coeffs = [np.sum(rule[slc, 0]) for slc in matches.T]
             rule = np.hstack([np.stack(coeffs)[:, np.newaxis], unique_mods])
 
-    # sort columns according to abs(shift)
-    return rule[np.argsort(np.abs(rule[:, -1]), kind="stable")]
+    # sort columns according to abs(shift), ties are resolved with the sign,
+    # positive shifts being returned before negative shifts.
+    return rule[np.lexsort((-np.sign(rule[:, -1]), np.abs(rule[:, -1])))]
 
 
 @functools.lru_cache(maxsize=None)


Test patch:
diff --git a/tests/gradients/core/test_general_shift_rules.py b/tests/gradients/core/test_general_shift_rules.py
index ccf246ac2..0efcc5997 100644
--- a/tests/gradients/core/test_general_shift_rules.py
+++ b/tests/gradients/core/test_general_shift_rules.py
@@ -287,7 +287,7 @@ class TestGenerateShiftRule:
         properly simplified when custom shift values are provided"""
         frequencies = (1,)
         generated_terms = generate_shift_rule(frequencies, shifts=(np.pi / 4,), order=2)
-        correct_terms = [[-1, 0], [0.5, -np.pi / 2], [0.5, np.pi / 2]]
+        correct_terms = [[-1, 0], [0.5, np.pi / 2], [0.5, -np.pi / 2]]
         assert np.allclose(generated_terms, correct_terms)
 
     def test_second_order_four_term_shift_rule(self):
@@ -297,8 +297,8 @@ class TestGenerateShiftRule:
         generated_terms = generate_shift_rule(frequencies, order=2)
         correct_terms = [
             [-0.375, 0],
-            [0.25, -np.pi],
             [0.25, np.pi],
+            [0.25, -np.pi],
             [-0.125, -2 * np.pi],
         ]
         assert np.allclose(generated_terms, correct_terms)
@@ -310,12 +310,12 @@ class TestGenerateShiftRule:
         generated_terms = generate_shift_rule(frequencies, order=2)
         correct_terms = [
             [-6, 0],
-            [3.91421356, -np.pi / 4],
             [3.91421356, np.pi / 4],
-            [-1, -np.pi / 2],
+            [3.91421356, -np.pi / 4],
             [-1, np.pi / 2],
-            [0.08578644, -3 * np.pi / 4],
+            [-1, -np.pi / 2],
             [0.08578644, 3 * np.pi / 4],
+            [0.08578644, -3 * np.pi / 4],
         ]
         assert np.allclose(generated_terms, correct_terms)
 
@@ -335,7 +335,7 @@ class TestMultiShiftRule:
         assert np.allclose(res, expected)
 
         res = generate_multi_shift_rule([(1,)], orders=[2], shifts=[(np.pi / 4,)])
-        expected = [[-1, 0], [0.5, -np.pi / 2], [0.5, np.pi / 2]]
+        expected = [[-1, 0], [0.5, np.pi / 2], [0.5, -np.pi / 2]]
         assert np.allclose(res, expected)
 
     def test_two_single_frequency(self):
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_hessian.py b/tests/gradients/parameter_shift/test_parameter_shift_hessian.py
index 4998dfb50..030d670da 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_hessian.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_hessian.py
@@ -170,8 +170,8 @@ class TestCollectRecipes:
     two_term_2nd_order = [(-0.5, 1.0, 0.0), (0.5, 1.0, -np.pi)]
     four_term_2nd_order = [
         (-0.375, 1.0, 0),
-        (0.25, 1.0, -np.pi),
         (0.25, 1.0, np.pi),
+        (0.25, 1.0, -np.pi),
         (-0.125, 1.0, -2 * np.pi),
     ]
 
@@ -1606,8 +1606,8 @@ class TestParamShiftHessianWithKwargs:
         # - 1 for second diagonal.
         assert len(tapes) == 1 + 2 + 4 + 1
         assert np.allclose(tapes[0].get_parameters(), x)
-        assert np.allclose(tapes[1].get_parameters(), x + np.array([-2 * np.pi / 3, 0.0]))
-        assert np.allclose(tapes[2].get_parameters(), x + np.array([2 * np.pi / 3, 0.0]))
+        assert np.allclose(tapes[1].get_parameters(), x + np.array([2 * np.pi / 3, 0.0]))
+        assert np.allclose(tapes[2].get_parameters(), x + np.array([-2 * np.pi / 3, 0.0]))
         assert np.allclose(tapes[-1].get_parameters(), x + np.array([0.0, -np.pi]))
         expected_shifts = np.array([[1, 1], [1, -1], [-1, 1], [-1, -1]]) * (np.pi / 2)
         for _tape, exp_shift in zip(tapes[3:-1], expected_shifts):
@@ -1662,7 +1662,7 @@ class TestParamShiftHessianWithKwargs:
             assert np.allclose(_tape.get_parameters(), x + exp_shift)
 
         # Check that the vanilla diagonal rule is used for the second diagonal entry
-        shift_order = [-1, 1, -2]
+        shift_order = [1, -1, -2]
         for mult, _tape in zip(shift_order, tapes[10:]):
             assert np.allclose(_tape.get_parameters(), x + np.array([0.0, np.pi * mult]))