Repository: PennyLaneAI/pennylane
Pull request #5525 (2024-04-16 15:25:59)

_group_measurements in qubit sampling works with Prod, SProd, Sum, Hamiltonian
**Context:**
The `_group_measurements` function in `qubit/sampling.py` groups measurement processes with commuting observables together to reduce the number of executions. Now it does not handle `Prod`, `SProd` or single-term `Hamiltonian` and `Sum`, which leads to inefficiencies with new opmath.

**Description of the Change:**
Removes the branch that places `Prod`, `SProd`, `Hamiltonian` and `Sum` in separate groups. Now `is_pauli_word` is used for everything.

**Benefits:**
Better grouping leads to fewer number of executions.

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5512

**Related Shortcut Stories:**
[sc-61306]


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/5525?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.66%. Comparing base [(`0747c0a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/0747c0ab22a277b1a3356656af33c9b6c2ce2609?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`61edce9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5525?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5525      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         404      406       +2     
  Lines       37852    37580     -272     
==========================================
- Hits        37728    37455     -273     
- Misses        124      125       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5525?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 30fd684eb..6291af1df 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -474,6 +474,9 @@
 * Fixes a bug in `qml.math.kron` that makes torch incompatible with numpy.
   [(#5540)](https://github.com/PennyLaneAI/pennylane/pull/5540)
 
+* Fixes a bug in `_group_measurements` that fails to group measurements with commuting observables when they are operands of `Prod`.
+  [(#5512)](https://github.com/PennyLaneAI/pennylane/issues/5512)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index b6037ceb4..4264b0fa9 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -16,7 +16,7 @@ from typing import List, Union, Tuple
 
 import numpy as np
 import pennylane as qml
-from pennylane.ops import Sum, Hamiltonian, SProd, Prod, LinearCombination
+from pennylane.ops import Sum, Hamiltonian, LinearCombination
 from pennylane.measurements import (
     SampleMeasurement,
     Shots,
@@ -69,11 +69,6 @@ def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, Sh
         elif mp.obs is None:
             mp_no_obs.append(mp)
             mp_no_obs_indices.append(i)
-        elif isinstance(mp.obs, (Hamiltonian, Sum, SProd, Prod)):
-            # Sum, Hamiltonian, SProd, and Prod are treated as valid Pauli words, but
-            # aren't accepted in qml.pauli.group_observables
-            mp_other_obs.append([mp])
-            mp_other_obs_indices.append([i])
         elif qml.pauli.is_pauli_word(mp.obs):
             mp_pauli_obs.append((i, mp))
         else:
@@ -82,13 +77,13 @@ def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, Sh
 
     if mp_pauli_obs:
         i_to_pauli_mp = dict(mp_pauli_obs)
-        ob_groups, group_indices = qml.pauli.group_observables(
+        _, group_indices = qml.pauli.group_observables(
             [mp.obs for mp in i_to_pauli_mp.values()], list(i_to_pauli_mp.keys())
         )
 
         mp_pauli_groups = []
-        for group, indices in zip(ob_groups, group_indices):
-            mp_group = [i_to_pauli_mp[i].__class__(obs=ob) for ob, i in zip(group, indices)]
+        for indices in group_indices:
+            mp_group = [i_to_pauli_mp[i] for i in indices]
             mp_pauli_groups.append(mp_group)
     else:
         mp_pauli_groups, group_indices = [], []


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 f12af5eab..6f8c4c0f8 100644
--- a/tests/devices/default_qubit/test_default_qubit_tracking.py
+++ b/tests/devices/default_qubit/test_default_qubit_tracking.py
@@ -217,8 +217,8 @@ shot_testing_combos = [
             qml.expval(qml.prod(qml.PauliX(0), qml.PauliX(1))),
             qml.expval(qml.prod(qml.PauliX(1), qml.PauliX(2))),
         ],
-        2,
-        20,
+        1,
+        10,
     ),
     # computational basis measurements
     ([qml.probs(wires=(0, 1)), qml.sample(wires=(0, 1))], 1, 10),
@@ -255,31 +255,9 @@ def test_single_expval(mps, expected_exec, expected_shots):
         assert dev.tracker.totals["shots"] == 3 * expected_shots
 
 
-@pytest.mark.usefixtures("use_new_opmath")
-@pytest.mark.xfail(reason="bug in grouping for tracker with new opmath")
-def test_multiple_expval_with_prods():
-    """Can be combined with test below once the bug is fixed - there shouldn't
-    be a difference in behaviour between old and new opmath here"""
-    mps, expected_exec, expected_shots = (
-        [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(0) @ qml.PauliY(1))],
-        1,
-        10,
-    )
-    dev = qml.device("default.qubit")
-    tape = qml.tape.QuantumScript([], mps, shots=10)
-
-    with dev.tracker:
-        dev.execute(tape)
-
-    assert dev.tracker.totals["executions"] == expected_exec
-    assert dev.tracker.totals["simulations"] == 1
-    assert dev.tracker.totals["shots"] == expected_shots
-
-
-@pytest.mark.usefixtures("use_legacy_opmath")
-def test_multiple_expval_with_tensors_legacy_opmath():
+def test_multiple_expval_with_prod():
     mps, expected_exec, expected_shots = (
-        [qml.expval(qml.PauliX(0)), qml.expval(qml.operation.Tensor(qml.PauliX(0), qml.PauliY(1)))],
+        [qml.expval(qml.PauliX(0)), qml.expval(qml.prod(qml.PauliX(0), qml.PauliY(1)))],
         1,
         10,
     )
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 56663c425..9c85b0f49 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -1078,6 +1078,7 @@ class TestHamiltonianSamples:
     """Test that the measure_with_samples function works as expected for
     Hamiltonian and Sum observables"""
 
+    @pytest.mark.usefixtures("use_legacy_and_new_opmath")
     def test_hamiltonian_expval(self):
         """Test that sampling works well for Hamiltonian observables"""
         x, y = np.array(0.67), np.array(0.95)
@@ -1108,6 +1109,7 @@ class TestHamiltonianSamples:
 
     def test_sum_expval(self):
         """Test that sampling works well for Sum observables"""
+
         x, y = np.array(0.67), np.array(0.95)
         ops = [qml.RY(x, wires=0), qml.RZ(y, wires=0)]
         meas = [qml.expval(qml.s_prod(0.8, qml.PauliZ(0)) + qml.s_prod(0.5, qml.PauliX(0)))]
@@ -1134,6 +1136,32 @@ class TestHamiltonianSamples:
         assert np.allclose(res[0], expected, atol=0.01)
         assert np.allclose(res[1], expected, atol=0.01)
 
+    def test_prod_expval(self):
+        """Tests that sampling works for Prod observables"""
+
+        x, y = np.array(0.67), np.array(0.95)
+        ops = [qml.RY(y, wires=0), qml.RX(x, wires=1)]
+        H = qml.prod(qml.PauliX(0), qml.PauliY(1))
+        tape = qml.tape.QuantumScript(
+            ops, measurements=[qml.expval(qml.PauliX(0)), qml.expval(H)], shots=10000
+        )
+        res = simulate(tape, rng=200)
+        expected = [np.sin(y), -np.sin(y) * np.sin(x)]
+        assert np.allclose(res, expected, atol=0.05)
+
+    def test_sprod_expval(self):
+        """Tests that sampling works for SProd observables"""
+
+        y = np.array(0.95)
+        ops = [qml.RY(y, wires=0)]
+        H = qml.s_prod(1.5, qml.PauliX(0))
+        tape = qml.tape.QuantumScript(
+            ops, measurements=[qml.expval(qml.PauliX(0)), qml.expval(H)], shots=10000
+        )
+        res = simulate(tape, rng=200)
+        expected = [np.sin(y), 1.5 * np.sin(y)]
+        assert np.allclose(res, expected, atol=0.05)
+
     def test_multi_wires(self):
         """Test that sampling works for Sums with large numbers of wires"""
         n_wires = 10
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 3b0471a98..3f226f02e 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -618,9 +618,7 @@ class TestSampleMeasurements:
         assert len(result) == len(mps)
 
         # check that samples are reused when possible
-        # 3 groups for expval and var, 1 group for probs and sample, 2 groups each for
-        # Hamiltonian and Sum, and 1 group each for SProd and Prod
-        assert spy.call_count == 10
+        assert spy.call_count == 8
 
     shots_data = [
         [10000, 10000],