Fix a sorting order in parameter-shift terms

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)
 

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)

tests/gradients/core/test_general_shift_rules.py

@@ -287,7 +287,7 @@ def test_second_order_two_term_shift_rule_custom_shifts(self):
  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 @@ def test_second_order_four_term_shift_rule(self):
  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 @@ def test_second_order_non_equidistant_shift_rule(self):
  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 @@ def test_single_parameter(self):
  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):

tests/gradients/parameter_shift/test_parameter_shift_hessian.py

@@ -170,8 +170,8 @@ class DummyOp(qml.RX):
  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 @@ def circuit(x):
  # - 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 @@ def circuit(x):
  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]))