Repository: PennyLaneAI/pennylane
Pull request #5559 (2024-04-23 15:07:27)

Follow-up bug fixes for `_group_measurements`
**Context:**
PR https://github.com/PennyLaneAI/pennylane/pull/5525 led to two bugs:
- `group_observables` uses `are_identity_pauli_words`, which does not work for legacy `Hamiltonian`.
- `get_num_shots_and_executions` incorrectly counts number of executions for `Sum`

**Description of the Change:**
- Converts single-term `Hamiltonian` to `SProd`
- Use the grouping information of `Sum` to correctly count number of executions.

**Related Shortcut Issues:**
[sc-61807]
[sc-61812]
Fixes https://github.com/PennyLaneAI/pennylane/issues/5561

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `94.11765%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.68%. Comparing base [(`a8bba50`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a8bba509e16a8582d57a0baa978ddb89a6d4d75a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`801d6e5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 3 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/qubit/sampling.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?src=pr&el=tree&filepath=pennylane%2Fdevices%2Fqubit%2Fsampling.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2FtcGxpbmcucHk=) | 90.00% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #5559      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         410      410              
  Lines       38230    37953     -277     
==========================================
- Hits        38113    37834     -279     
- Misses        117      119       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5559?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/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index e18b5dcec..14b14c974 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -99,13 +99,27 @@ def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, Sh
     return all_mp_groups, all_indices
 
 
-def _get_num_shots_for_expval_H(obs):
+def _get_num_executions_for_expval_H(obs):
     indices = obs.grouping_indices
     if indices:
         return len(indices)
     return sum(int(not isinstance(o, qml.Identity)) for o in obs.terms()[1])
 
 
+def _get_num_executions_for_sum(obs):
+
+    if obs.grouping_indices:
+        return len(obs.grouping_indices)
+
+    if not obs.pauli_rep:
+        return sum(int(not isinstance(o, qml.Identity)) for o in obs.terms()[1])
+
+    _, ops = obs.terms()
+    with qml.QueuingManager.stop_recording():
+        op_groups = qml.pauli.group_observables(ops)
+    return len(op_groups)
+
+
 # pylint: disable=no-member
 def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]:
     """Get the total number of qpu executions and shots.
@@ -125,14 +139,15 @@ def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]:
         if isinstance(group[0], ExpectationMP) and isinstance(
             group[0].obs, (qml.ops.Hamiltonian, qml.ops.LinearCombination)
         ):
-            H_executions = _get_num_shots_for_expval_H(group[0].obs)
+            H_executions = _get_num_executions_for_expval_H(group[0].obs)
             num_executions += H_executions
             if tape.shots:
                 num_shots += tape.shots.total_shots * H_executions
         elif isinstance(group[0], ExpectationMP) and isinstance(group[0].obs, qml.ops.Sum):
-            num_executions += len(group[0].obs)
+            sum_executions = _get_num_executions_for_sum(group[0].obs)
+            num_executions += sum_executions
             if tape.shots:
-                num_shots += tape.shots.total_shots * len(group[0].obs)
+                num_shots += tape.shots.total_shots * sum_executions
         elif isinstance(group[0], (ClassicalShadowMP, ShadowExpvalMP)):
             num_executions += tape.shots.total_shots
             if tape.shots:
diff --git a/pennylane/pauli/grouping/group_observables.py b/pennylane/pauli/grouping/group_observables.py
index 98d9e4025..1031de1d0 100644
--- a/pennylane/pauli/grouping/group_observables.py
+++ b/pennylane/pauli/grouping/group_observables.py
@@ -246,6 +246,14 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
     if coefficients is None:
         return partitioned_paulis
 
+    partitioned_coeffs = _partition_coeffs(partitioned_paulis, observables, coefficients)
+
+    return partitioned_paulis, partitioned_coeffs
+
+
+def _partition_coeffs(partitioned_paulis, observables, coefficients):
+    """Partition the coefficients according to the Pauli word groupings."""
+
     partitioned_coeffs = [
         qml.math.cast_like([0] * len(g), coefficients) for g in partitioned_paulis
     ]
@@ -259,6 +267,12 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
         for pauli_word in partition:
             # find index of this pauli word in remaining original observables,
             for ind, observable in enumerate(observables):
+                if isinstance(observable, qml.ops.Hamiltonian):
+                    # Converts single-term Hamiltonian to SProd because
+                    # are_identical_pauli_words cannot handle Hamiltonian
+                    coeffs, ops = observable.terms()
+                    # Assuming the Hamiltonian has only one term
+                    observable = qml.s_prod(coeffs[0], ops[0])
                 if are_identical_pauli_words(pauli_word, observable):
                     indices.append(coeff_indices[ind])
                     observables.pop(ind)
@@ -273,4 +287,4 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
     if isinstance(coefficients, list):
         partitioned_coeffs = [list(p) for p in partitioned_coeffs]
 
-    return partitioned_paulis, partitioned_coeffs
+    return partitioned_coeffs


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit_tracking.py b/tests/devices/default_qubit/test_default_qubit_tracking.py
index 6f8c4c0f8..2d2c5e81d 100644
--- a/tests/devices/default_qubit/test_default_qubit_tracking.py
+++ b/tests/devices/default_qubit/test_default_qubit_tracking.py
@@ -211,7 +211,14 @@ shot_testing_combos = [
         20,
     ),
     # op arithmetic test cases
-    ([qml.expval(qml.sum(qml.PauliX(0), qml.PauliX(1)))], 2, 20),
+    ([qml.expval(qml.sum(qml.PauliX(0), qml.PauliY(0)))], 2, 20),
+    ([qml.expval(qml.sum(qml.PauliX(0), qml.PauliX(0) @ qml.PauliX(1)))], 1, 10),
+    ([qml.expval(qml.sum(qml.PauliX(0), qml.Hadamard(0)))], 2, 20),
+    (
+        [qml.expval(qml.sum(qml.PauliX(0), qml.PauliY(1) @ qml.PauliX(1), grouping_type="qwc"))],
+        1,
+        10,
+    ),
     (
         [
             qml.expval(qml.prod(qml.PauliX(0), qml.PauliX(1))),
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
index aeeb9c8ae..6c9aaca55 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
@@ -726,24 +726,24 @@ class TestHigherOrderDerivatives:
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
-@pytest.mark.parametrize("use_new_op_math", (True, False))
+@pytest.mark.usefixtures("use_legacy_and_new_opmath")
 class TestHamiltonianWorkflows:
     """Test that tapes ending with expectations
     of Hamiltonians provide correct results and gradients"""
 
     @pytest.fixture
-    def cost_fn(self, execute_kwargs, shots, device, use_new_op_math):
+    def cost_fn(self, execute_kwargs, shots, device):
         """Cost function for gradient tests"""
 
         def _cost_fn(weights, coeffs1, coeffs2):
             obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
             H1 = qml.Hamiltonian(coeffs1, obs1)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H1 = qml.pauli.pauli_sentence(H1).operation()
 
             obs2 = [qml.PauliZ(0)]
             H2 = qml.Hamiltonian(coeffs2, obs2)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H2 = qml.pauli.pauli_sentence(H2).operation()
 
             with qml.queuing.AnnotatedQueue() as q:
@@ -786,12 +786,10 @@ class TestHamiltonianWorkflows:
             ]
         )
 
-    def test_multiple_hamiltonians_not_trainable(
-        self, execute_kwargs, cost_fn, shots, use_new_op_math
-    ):
+    def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with no trainable parameters."""
 
-        if execute_kwargs["gradient_fn"] == "adjoint" and not use_new_op_math:
+        if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath():
             pytest.skip("adjoint differentiation does not suppport hamiltonians.")
 
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False)
@@ -814,11 +812,11 @@ class TestHamiltonianWorkflows:
         else:
             assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
-    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots, use_new_op_math):
+    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with trainable parameters."""
         if execute_kwargs["gradient_fn"] == "adjoint":
             pytest.skip("trainable hamiltonians not supported with adjoint")
-        if use_new_op_math:
+        if qml.operation.active_new_opmath():
             pytest.skip("parameter shift derivatives do not yet support sums.")
 
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index 770d8112d..87d4773e6 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -701,24 +701,24 @@ class TestHigherOrderDerivatives:
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
-@pytest.mark.parametrize("use_new_op_math", (True, False))
+@pytest.mark.usefixtures("use_legacy_and_new_opmath")
 class TestHamiltonianWorkflows:
     """Test that tapes ending with expectations
     of Hamiltonians provide correct results and gradients"""
 
     @pytest.fixture
-    def cost_fn(self, execute_kwargs, shots, device, use_new_op_math):
+    def cost_fn(self, execute_kwargs, shots, device):
         """Cost function for gradient tests"""
 
         def _cost_fn(weights, coeffs1, coeffs2):
             obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
             H1 = qml.Hamiltonian(coeffs1, obs1)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H1 = qml.pauli.pauli_sentence(H1).operation()
 
             obs2 = [qml.PauliZ(0)]
             H2 = qml.Hamiltonian(coeffs2, obs2)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H2 = qml.pauli.pauli_sentence(H2).operation()
 
             with qml.queuing.AnnotatedQueue() as q:
@@ -764,12 +764,10 @@ class TestHamiltonianWorkflows:
             ]
         )
 
-    def test_multiple_hamiltonians_not_trainable(
-        self, execute_kwargs, cost_fn, shots, use_new_op_math
-    ):
+    def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with no trainable parameters."""
 
-        if execute_kwargs["gradient_fn"] == "adjoint" and not use_new_op_math:
+        if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath():
             pytest.skip("adjoint differentiation does not suppport hamiltonians.")
 
         coeffs1 = jnp.array([0.1, 0.2, 0.3])
@@ -792,11 +790,11 @@ class TestHamiltonianWorkflows:
         else:
             assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
-    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots, use_new_op_math):
+    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with trainable parameters."""
         if execute_kwargs["gradient_fn"] == "adjoint":
             pytest.skip("trainable hamiltonians not supported with adjoint")
-        if use_new_op_math:
+        if qml.operation.active_new_opmath():
             pytest.skip("parameter shift derivatives do not yet support sums.")
 
         coeffs1 = jnp.array([0.1, 0.2, 0.3])
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
index e5cb1f775..3dd6e293b 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
@@ -707,24 +707,24 @@ class TestHigherOrderDerivatives:
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
-@pytest.mark.parametrize("use_new_op_math", (True, False))
+@pytest.mark.usefixtures("use_legacy_and_new_opmath")
 class TestHamiltonianWorkflows:
     """Test that tapes ending with expectations
     of Hamiltonians provide correct results and gradients"""
 
     @pytest.fixture
-    def cost_fn(self, execute_kwargs, shots, device, use_new_op_math):
+    def cost_fn(self, execute_kwargs, shots, device):
         """Cost function for gradient tests"""
 
         def _cost_fn(weights, coeffs1, coeffs2):
             obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
             H1 = qml.Hamiltonian(coeffs1, obs1)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H1 = qml.pauli.pauli_sentence(H1).operation()
 
             obs2 = [qml.PauliZ(0)]
             H2 = qml.Hamiltonian(coeffs2, obs2)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H2 = qml.pauli.pauli_sentence(H2).operation()
 
             with qml.queuing.AnnotatedQueue() as q:
@@ -767,12 +767,10 @@ class TestHamiltonianWorkflows:
             ]
         )
 
-    def test_multiple_hamiltonians_not_trainable(
-        self, execute_kwargs, cost_fn, shots, use_new_op_math
-    ):
+    def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with no trainable parameters."""
 
-        if execute_kwargs["gradient_fn"] == "adjoint" and not use_new_op_math:
+        if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath():
             pytest.skip("adjoint differentiation does not suppport hamiltonians.")
 
         device_vjp = execute_kwargs.get("device_vjp", False)
@@ -791,11 +789,11 @@ class TestHamiltonianWorkflows:
         expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2]
         assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0)
 
-    def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots, use_new_op_math):
+    def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots):
         """Test hamiltonian with trainable parameters."""
         if execute_kwargs["gradient_fn"] == "adjoint":
             pytest.skip("trainable hamiltonians not supported with adjoint")
-        if use_new_op_math:
+        if qml.operation.active_new_opmath():
             pytest.skip("parameter shift derivatives do not yet support sums.")
 
         coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64)
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
index 0618ab8bc..3fdd04179 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
@@ -727,24 +727,24 @@ class TestHigherOrderDerivatives:
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
-@pytest.mark.parametrize("use_new_op_math", (True, False))
+@pytest.mark.usefixtures("use_legacy_and_new_opmath")
 class TestHamiltonianWorkflows:
     """Test that tapes ending with expectations
     of Hamiltonians provide correct results and gradients"""
 
     @pytest.fixture
-    def cost_fn(self, execute_kwargs, shots, device, use_new_op_math):
+    def cost_fn(self, execute_kwargs, shots, device):
         """Cost function for gradient tests"""
 
         def _cost_fn(weights, coeffs1, coeffs2):
             obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
             H1 = qml.Hamiltonian(coeffs1, obs1)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H1 = qml.pauli.pauli_sentence(H1).operation()
 
             obs2 = [qml.PauliZ(0)]
             H2 = qml.Hamiltonian(coeffs2, obs2)
-            if use_new_op_math:
+            if qml.operation.active_new_opmath():
                 H2 = qml.pauli.pauli_sentence(H2).operation()
 
             with qml.queuing.AnnotatedQueue() as q:
@@ -795,12 +795,10 @@ class TestHamiltonianWorkflows:
             ]
         )
 
-    def test_multiple_hamiltonians_not_trainable(
-        self, execute_kwargs, cost_fn, shots, use_new_op_math
-    ):
+    def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with no trainable parameters."""
 
-        if execute_kwargs["gradient_fn"] == "adjoint" and not use_new_op_math:
+        if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath():
             pytest.skip("adjoint differentiation does not suppport hamiltonians.")
 
         coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False)
@@ -823,11 +821,11 @@ class TestHamiltonianWorkflows:
         else:
             assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
-    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots, use_new_op_math):
+    def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots):
         """Test hamiltonian with trainable parameters."""
         if execute_kwargs["gradient_fn"] == "adjoint":
             pytest.skip("trainable hamiltonians not supported with adjoint")
-        if use_new_op_math:
+        if qml.operation.active_new_opmath():
             pytest.skip("parameter shift derivatives do not yet support sums.")
 
         coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True)