_group_measurements in qubit sampling works with Prod, SProd, Sum, Hamiltonian

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):

pennylane/devices/qubit/sampling.py

@@ -16,7 +16,7 @@
 
 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 = [], []

tests/devices/default_qubit/test_default_qubit_tracking.py

@@ -217,8 +217,8 @@ def circuit_3(y):
             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,
     )

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 @@ def test_hamiltonian_expval_shot_vector(self):
 
     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 @@ def test_sum_expval_shot_vector(self):
         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

tests/devices/qubit/test_simulate.py

@@ -618,9 +618,7 @@ def test_shots_reuse(self, mocker):
         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],