Repository: PennyLaneAI/pennylane
Pull request #5475 (2024-04-04 20:27:18)

Legacy device API handles Prod observables
**Context:**
The legacy device API has special handling of `Tensor` and `Sum` as observables, but `Prod` is not covered by that.

**Description of the Change:**
- `Prod` observables are expanded if they simplify to a `Sum`
- In `check_validity`, the operands of `Prod` observables are checked.

**Possible Drawbacks:**
`simplify` is called on `Prod` observables many times, which might be inefficient

**Related Shortcut Story:**
[sc-60584]


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/5475?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 `92.30769%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.67%. Comparing base [(`3b0c1b6`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3b0c1b613727f4cb22269c72c74900c160f21f09?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`174e4bf`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5475?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 4 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5475?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/\_device.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5475?src=pr&el=tree&filepath=pennylane%2F_device.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | 90.00% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5475?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    #5475      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         402      402              
  Lines       37540    37300     -240     
==========================================
- Hits        37420    37178     -242     
- Misses        120      122       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5475?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).

astralcai: > Concerning the renaming of the mock devices in the test I dont have an opinion and would refer to someone from core to judge whether changing all of them is a good idea or not.

@Qottmann The renaming was because multiple test cases were using different mock devices that were exact duplicates of one another. I removed the duplicates, and changed all these test cases to use just the one that was left.
astralcai: > Thanks @astralcai
> 
> My questions are mostly answered and I am overall happy with the PR. The only thing left to discuss is the simplify. Why do we need it?

@Qottmann The purpose of this is to expand a nested `Prod` out into either a `Prod` or a `Sum`. For example if the observable was `qml.X(0) @ qml.Y(1) @ qml.Z(2)`, we get a `Prod` with two operands, a `Prod` and a `Z`, or if the observable was `qml.X(0) @ (qml.Y(1) + qml.Z(2))`, we would also want to flatten out the nested structure. I know `simplify` does this, but do you know if there is another less expensive alternative?

Patch:
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 18b0c5338..07a7dea8d 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -42,7 +42,7 @@ from pennylane.measurements import (
 )
 
 from pennylane.operation import Observable, Operation, Tensor, Operator, StatePrepBase
-from pennylane.ops import Hamiltonian, Sum, LinearCombination
+from pennylane.ops import Hamiltonian, Sum, LinearCombination, Prod
 from pennylane.tape import QuantumScript, QuantumTape, expand_tape_state_prep
 from pennylane.wires import WireError, Wires
 from pennylane.queuing import QueuingManager
@@ -744,7 +744,6 @@ class Device(abc.ABC):
             to be applied to the list of evaluated circuit results.
         """
         supports_hamiltonian = self.supports_observable("Hamiltonian")
-
         supports_sum = self.supports_observable("Sum")
         finite_shots = self.shots is not None
         grouping_known = all(
@@ -759,7 +758,12 @@ class Device(abc.ABC):
             isinstance(obs, (Hamiltonian, LinearCombination)) for obs in circuit.observables
         )
         expval_sum_in_obs = any(
-            isinstance(m.obs, Sum) and isinstance(m, ExpectationMP) for m in circuit.measurements
+            (
+                isinstance(m.obs, Sum)
+                or (isinstance(m.obs, Prod) and isinstance(m.obs.simplify(), Sum))
+            )
+            and isinstance(m, ExpectationMP)
+            for m in circuit.measurements
         )
 
         is_shadow = any(isinstance(m, ShadowExpvalMP) for m in circuit.measurements)
@@ -1007,6 +1011,21 @@ class Device(abc.ABC):
                         raise DeviceError(
                             f"Observable {i.name} not supported on device {self.short_name}"
                         )
+
+            elif isinstance(o, qml.ops.Prod):
+
+                supports_prod = self.supports_observable(o.name)
+                if not supports_prod:
+                    raise DeviceError(f"Observable Prod not supported on device {self.short_name}")
+
+                simplified_op = o.simplify()
+                if isinstance(simplified_op, qml.ops.Prod):
+                    for i in o.simplify().operands:
+                        if not self.supports_observable(i.name):
+                            raise DeviceError(
+                                f"Observable {i.name} not supported on device {self.short_name}"
+                            )
+
             else:
                 observable_name = o.name
 
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index c93aec222..9e90ec801 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -20,7 +20,7 @@ from typing import List, Sequence, Callable
 
 import pennylane as qml
 from pennylane.measurements import ExpectationMP, MeasurementProcess
-from pennylane.ops import SProd, Sum
+from pennylane.ops import SProd, Sum, Prod
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.transforms import transform
 
@@ -341,6 +341,8 @@ def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape],
     idxs_coeffs_dict = {}  # {m_hash: [(location_idx, coeff)]}
     for idx, m in enumerate(tape.measurements):
         obs = m.obs
+        if isinstance(obs, Prod) and isinstance(m, ExpectationMP):
+            obs = obs.simplify()
         if isinstance(obs, Sum) and isinstance(m, ExpectationMP):
             for summand in obs.operands:
                 coeff = 1


Test patch:
diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py
index 393d498c5..72bedbbc1 100644
--- a/pennylane/devices/tests/test_measurements.py
+++ b/pennylane/devices/tests/test_measurements.py
@@ -160,8 +160,10 @@ class TestHamiltonianSupport:
     """Separate test to ensure that the device can differentiate Hamiltonian observables."""
 
     @pytest.mark.parametrize("ham_constructor", [qml.ops.Hamiltonian, qml.ops.LinearCombination])
+    @pytest.mark.filterwarnings("ignore::pennylane.PennyLaneDeprecationWarning")
     def test_hamiltonian_diff(self, ham_constructor, device_kwargs, tol):
         """Tests a simple VQE gradient using parameter-shift rules."""
+
         device_kwargs["wires"] = 1
         dev = qml.device(**device_kwargs)
         coeffs = np.array([-0.05, 0.17])
diff --git a/tests/test_device.py b/tests/test_device.py
index 0422e9dfd..22200272b 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -30,36 +30,6 @@ mock_device_paulis_and_hamiltonian = ["Hamiltonian", "PauliX", "PauliY", "PauliZ
 # pylint: disable=abstract-class-instantiated, no-self-use, redefined-outer-name, invalid-name, missing-function-docstring
 
 
-@pytest.fixture(scope="function")
-def mock_device_with_operations(monkeypatch):
-    """A function to create a mock device with non-empty operations"""
-    with monkeypatch.context() as m:
-        m.setattr(Device, "__abstractmethods__", frozenset())
-        m.setattr(Device, "operations", mock_device_paulis)
-        m.setattr(Device, "observables", mock_device_paulis)
-        m.setattr(Device, "short_name", "MockDevice")
-
-        def get_device(wires=1):
-            return Device(wires=wires)
-
-        yield get_device
-
-
-@pytest.fixture(scope="function")
-def mock_device_with_observables(monkeypatch):
-    """A function to create a mock device with non-empty observables"""
-    with monkeypatch.context() as m:
-        m.setattr(Device, "__abstractmethods__", frozenset())
-        m.setattr(Device, "operations", mock_device_paulis)
-        m.setattr(Device, "observables", mock_device_paulis)
-        m.setattr(Device, "short_name", "MockDevice")
-
-        def get_device(wires=1):
-            return Device(wires=wires)
-
-        yield get_device
-
-
 @pytest.fixture(scope="function")
 def mock_device_with_identity(monkeypatch):
     """A function to create a mock device with non-empty observables"""
@@ -203,19 +173,15 @@ def mock_device(monkeypatch):
 
 
 @pytest.fixture(scope="function")
-def mock_device_arbitrary_wires(monkeypatch):
+def mock_device_supporting_prod(monkeypatch):
     with monkeypatch.context() as m:
         m.setattr(Device, "__abstractmethods__", frozenset())
         m.setattr(Device, "_capabilities", mock_device_capabilities)
-        m.setattr(Device, "operations", ["PauliY", "RX", "Rot"])
-        m.setattr(Device, "observables", ["PauliZ"])
+        m.setattr(Device, "operations", ["PauliX", "PauliZ"])
+        m.setattr(Device, "observables", ["PauliX", "PauliZ", "Prod"])
         m.setattr(Device, "short_name", "MockDevice")
-        m.setattr(Device, "expval", lambda self, x, y, z: 0)
-        m.setattr(Device, "var", lambda self, x, y, z: 0)
-        m.setattr(Device, "sample", lambda self, x, y, z: 0)
-        m.setattr(Device, "apply", lambda self, x, y, z: None)
 
-        def get_device(wires):
+        def get_device(wires=1):
             return Device(wires=wires)
 
         yield get_device
@@ -245,11 +211,11 @@ class TestDeviceSupportedLogic:
 
     # pylint: disable=no-self-use, redefined-outer-name
 
-    def test_supports_operation_argument_types(self, mock_device_with_operations):
+    def test_supports_operation_argument_types(self, mock_device_supporting_paulis):
         """Checks that device.supports_operations returns the correct result
         when passed both string and Operation class arguments"""
 
-        dev = mock_device_with_operations()
+        dev = mock_device_supporting_paulis()
 
         assert dev.supports_operation("PauliX")
         assert dev.supports_operation(qml.PauliX)
@@ -257,10 +223,10 @@ class TestDeviceSupportedLogic:
         assert not dev.supports_operation("S")
         assert not dev.supports_operation(qml.CNOT)
 
-    def test_supports_observable_argument_types(self, mock_device_with_observables):
+    def test_supports_observable_argument_types(self, mock_device_supporting_paulis):
         """Checks that device.supports_observable returns the correct result
         when passed both string and Operation class arguments"""
-        dev = mock_device_with_observables()
+        dev = mock_device_supporting_paulis()
 
         assert dev.supports_observable("PauliX")
         assert dev.supports_observable(qml.PauliX)
@@ -309,14 +275,14 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
     """Test the internal functions of the abstract Device class"""
 
     # pylint: disable=unnecessary-dunder-call
-    def test_repr(self, mock_device_with_operations):
+    def test_repr(self, mock_device_supporting_paulis):
         """Tests the __repr__ function"""
-        dev = mock_device_with_operations()
+        dev = mock_device_supporting_paulis()
         assert "<Device device (wires=1, shots=1000) at " in dev.__repr__()
 
-    def test_str(self, mock_device_with_operations):
+    def test_str(self, mock_device_supporting_paulis):
         """Tests the __str__ function"""
-        dev = mock_device_with_operations()
+        dev = mock_device_supporting_paulis()
         string = str(dev)
         assert "Short name: MockDevice" in string
         assert "Package: pennylane" in string
@@ -340,6 +306,43 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
         # Raises an error if queue or observables are invalid
         dev.check_validity(queue, observables)
 
+    def test_check_validity_containing_prod(self, mock_device_supporting_prod):
+        """Tests that the function Device.check_validity works with Prod"""
+
+        dev = mock_device_supporting_prod()
+
+        queue = [
+            qml.PauliX(wires=0),
+            qml.PauliZ(wires=1),
+        ]
+
+        observables = [
+            qml.expval(qml.PauliX(0) @ qml.PauliZ(1)),
+            qml.expval(qml.PauliZ(0) @ (qml.PauliX(1) @ qml.PauliZ(2))),
+        ]
+
+        dev.check_validity(queue, observables)
+
+    def test_prod_containing_unsupported_nested_observables(self, mock_device_supporting_prod):
+        """Tests that the observables nested within Prod are checked for validity"""
+
+        dev = mock_device_supporting_prod()
+
+        queue = [
+            qml.PauliX(wires=0),
+            qml.PauliZ(wires=1),
+        ]
+
+        unsupported_nested_observables = [
+            qml.expval(qml.PauliZ(0) @ (qml.PauliX(1) @ qml.PauliY(2)))
+        ]
+
+        with pytest.raises(
+            DeviceError,
+            match="Observable PauliY not supported",
+        ):
+            dev.check_validity(queue, unsupported_nested_observables)
+
     @pytest.mark.usefixtures("use_legacy_opmath")
     def test_check_validity_on_tensor_support_legacy_opmath(self, mock_device_supporting_paulis):
         """Tests the function Device.check_validity with tensor support capability"""
@@ -429,9 +432,9 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
         with pytest.raises(DeviceError, match="Observable Hadamard not supported on device"):
             dev.check_validity(queue, observables)
 
-    def test_check_validity_on_projector_as_operation(self, mock_device_with_operations):
+    def test_check_validity_on_projector_as_operation(self, mock_device_supporting_paulis):
         """Test that an error is raised if the operation queue contains qml.Projector"""
-        dev = mock_device_with_operations(wires=1)
+        dev = mock_device_supporting_paulis(wires=1)
 
         queue = [qml.PauliX(0), qml.Projector([0], wires=0), qml.PauliZ(0)]
         observables = []
@@ -592,8 +595,8 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
             (Wires([0]), Wires([0]), Wires([0])),
         ],
     )
-    def test_order_wires(self, wires, subset, expected_subset, mock_device_arbitrary_wires):
-        dev = mock_device_arbitrary_wires(wires=wires)
+    def test_order_wires(self, wires, subset, expected_subset, mock_device):
+        dev = mock_device(wires=wires)
         ordered_subset = dev.order_wires(subset_wires=subset)
         assert ordered_subset == expected_subset
 
@@ -606,8 +609,8 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
             (Wires([0]), Wires([2])),
         ],
     )
-    def test_order_wires_raises_value_error(self, wires, subset, mock_device_arbitrary_wires):
-        dev = mock_device_arbitrary_wires(wires=wires)
+    def test_order_wires_raises_value_error(self, wires, subset, mock_device):
+        dev = mock_device(wires=wires)
         with pytest.raises(ValueError, match="Could not find some or all subset wires"):
             _ = dev.order_wires(subset_wires=subset)
 
@@ -658,11 +661,11 @@ class TestInternalFunctions:  # pylint:disable=too-many-public-methods
         assert new_tape.batch_size == tape.batch_size
         assert new_tape.output_dim == tape.output_dim
 
-    def test_default_expand_fn_with_invalid_op(self, mock_device_with_operations, recwarn):
+    def test_default_expand_fn_with_invalid_op(self, mock_device_supporting_paulis, recwarn):
         """Test that default_expand_fn works with an invalid op and some measurement."""
         invalid_tape = qml.tape.QuantumScript([qml.S(0)], [qml.expval(qml.PauliZ(0))])
         expected_tape = qml.tape.QuantumScript([qml.RZ(np.pi / 2, 0)], [qml.expval(qml.PauliZ(0))])
-        dev = mock_device_with_operations(wires=1)
+        dev = mock_device_supporting_paulis(wires=1)
         expanded_tape = dev.expand_fn(invalid_tape, max_expansion=3)
         assert qml.equal(expanded_tape, expected_tape)
         assert len(recwarn) == 0
@@ -799,33 +802,33 @@ class TestOperations:
         with pytest.raises(DeviceError, match="Gate Hadamard not supported on device"):
             dev.execute(queue, observables)
 
-    def test_execute_obs_probs(self, mock_device_with_observables):
+    def test_execute_obs_probs(self, mock_device_supporting_paulis):
         """Tests that the execute function raises an error if probabilities are
         not supported by the device"""
-        dev = mock_device_with_observables()
+        dev = mock_device_supporting_paulis()
         obs = qml.probs(op=qml.PauliZ(0))
         with pytest.raises(NotImplementedError):
             dev.execute([], [obs])
 
-    def test_var(self, mock_device_with_observables):
+    def test_var(self, mock_device_supporting_paulis):
         """Tests that the variance method are not implemented by the device by
         default"""
-        dev = mock_device_with_observables()
+        dev = mock_device_supporting_paulis()
         with pytest.raises(NotImplementedError):
             dev.var(qml.PauliZ, 0, [])
 
-    def test_sample(self, mock_device_with_observables):
+    def test_sample(self, mock_device_supporting_paulis):
         """Tests that the sample method are not implemented by the device by
         default"""
-        dev = mock_device_with_observables()
+        dev = mock_device_supporting_paulis()
         with pytest.raises(NotImplementedError):
             dev.sample(qml.PauliZ, 0, [])
 
     @pytest.mark.parametrize("wires", [None, []])
-    def test_probability(self, mock_device_with_observables, wires):
+    def test_probability(self, mock_device_supporting_paulis, wires):
         """Tests that the probability method are not implemented by the device
         by default"""
-        dev = mock_device_with_observables()
+        dev = mock_device_supporting_paulis()
         with pytest.raises(NotImplementedError):
             dev.probability(wires=wires)
 
@@ -1204,3 +1207,17 @@ class TestGrouping:
 
         assert len(new_qscripts) == 2
         spy.assert_called()
+
+    def test_batch_transform_expands_prod_containing_sums(self, mocker):
+        """Tests that batch_transform expands a Prod with a nested Sum"""
+
+        H = qml.prod(qml.PauliX(0), qml.sum(qml.PauliY(0), qml.PauliZ(0)))
+        qs = qml.tape.QuantumScript(measurements=[qml.expval(H)])
+        spy = mocker.spy(qml.transforms, "sum_expand")
+
+        dev = self.SomeDevice()
+        dev.supports_observable = lambda *args, **kwargs: False
+        new_qscripts, _ = dev.batch_transform(qs)
+
+        assert len(new_qscripts) == 2
+        spy.assert_called()