Repository: PennyLaneAI/pennylane
Pull request #5890 (2024-06-20 19:51:10)

More sophisticated measurement validation for default-qubit
**Description of the Change:**
Add measurement-specific observable validations to measurement validation

[sc-64700]


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/5890?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 `96.77419%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 99.63%. Comparing base [(`8fec135`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/8fec135c4239fb39afc3391693ffe3b344f93d53?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`51b0996`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/51b09969bdfd9bd31994f5bd7520bd10062f75d8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5890?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/default\_qubit.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5890?src=pr&el=tree&filepath=pennylane%2Fdevices%2Fdefault_qubit.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | 96.77% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5890?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    #5890      +/-   ##
==========================================
- Coverage   99.67%   99.63%   -0.04%     
==========================================
  Files         422      425       +3     
  Lines       40651    40484     -167     
==========================================
- Hits        40518    40338     -180     
- Misses        133      146      +13     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5890?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 db929d3a5..a500b27c0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -3,6 +3,7 @@
 # Release 0.38.0-dev (development release)
 
 <h3>New features since last release</h3>
+
 * A new method `process_density_matrix` has been added to the `ProbabilityMP` and `DensityMatrixMP` classes, allowing for more efficient handling of quantum density matrices, particularly with batch processing support. This method simplifies the calculation of probabilities from quantum states represented as density matrices.
   [(#5830)](https://github.com/PennyLaneAI/pennylane/pull/5830)
 
@@ -23,6 +24,9 @@
 * `QuantumScript.hash` is now cached, leading to performance improvements.
   [(#5919)](https://github.com/PennyLaneAI/pennylane/pull/5919)
 
+* Observable validation for `default.qubit` is now based on execution mode (analytic vs. finite shots) and measurement type (sample measurement vs. state measurement).
+  [(#5890)](https://github.com/PennyLaneAI/pennylane/pull/5890)
+
 <h3>Breaking changes 💔</h3>
 
 * `QuantumScript.interface` has been removed.
@@ -43,6 +47,8 @@
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
+
+Astral Cai,
 Yushao Chen,
 Christina Lee,
 William Maxwell,
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 865bedf6b..aebbe4bbe 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -60,30 +60,6 @@ QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
 PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
 
 
-observables = {
-    "PauliX",
-    "PauliY",
-    "PauliZ",
-    "Hadamard",
-    "Hermitian",
-    "Identity",
-    "Projector",
-    "SparseHamiltonian",
-    "Hamiltonian",
-    "LinearCombination",
-    "Sum",
-    "SProd",
-    "Prod",
-    "Exp",
-    "Evolution",
-}
-
-
-def observable_stopping_condition(obs: qml.operation.Operator) -> bool:
-    """Specifies whether or not an observable is accepted by DefaultQubit."""
-    return obs.name in observables
-
-
 def stopping_condition(op: qml.operation.Operator) -> bool:
     """Specify whether or not an Operator object is supported by the device."""
     if op.name == "QFT" and len(op.wires) >= 6:
@@ -103,16 +79,74 @@ def stopping_condition_shots(op: qml.operation.Operator) -> bool:
     return isinstance(op, (Conditional, MidMeasureMP)) or stopping_condition(op)
 
 
+def observable_accepts_sampling(obs: qml.operation.Operator) -> bool:
+    """Verifies whether an observable supports sample measurement"""
+
+    if isinstance(obs, qml.ops.CompositeOp):
+        return all(observable_accepts_sampling(o) for o in obs.operands)
+
+    if isinstance(obs, qml.ops.SymbolicOp):
+        return observable_accepts_sampling(obs.base)
+
+    if isinstance(obs, qml.ops.Hamiltonian):
+        return all(observable_accepts_sampling(o) for o in obs.ops)
+
+    if isinstance(obs, qml.operation.Tensor):
+        return all(observable_accepts_sampling(o) for o in obs.obs)
+
+    return obs.has_diagonalizing_gates
+
+
+def observable_accepts_analytic(obs: qml.operation.Operator, is_expval=False) -> bool:
+    """Verifies whether an observable supports analytic measurement"""
+
+    if isinstance(obs, qml.ops.CompositeOp):
+        return all(observable_accepts_analytic(o, is_expval) for o in obs.operands)
+
+    if isinstance(obs, qml.ops.SymbolicOp):
+        return observable_accepts_analytic(obs.base, is_expval)
+
+    if isinstance(obs, qml.ops.Hamiltonian):
+        return all(observable_accepts_analytic(o, is_expval) for o in obs.ops)
+
+    if isinstance(obs, qml.operation.Tensor):
+        return all(observable_accepts_analytic(o, is_expval) for o in obs.obs)
+
+    if is_expval and isinstance(obs, (qml.ops.SparseHamiltonian, qml.ops.Hermitian)):
+        return True
+
+    return obs.has_diagonalizing_gates
+
+
 def accepted_sample_measurement(m: qml.measurements.MeasurementProcess) -> bool:
-    """Specifies whether or not a measurement is accepted when sampling."""
-    return isinstance(
+    """Specifies whether a measurement is accepted when sampling."""
+
+    if not isinstance(
         m,
         (
             qml.measurements.SampleMeasurement,
             qml.measurements.ClassicalShadowMP,
             qml.measurements.ShadowExpvalMP,
         ),
-    )
+    ):
+        return False
+
+    if m.obs is not None:
+        return observable_accepts_sampling(m.obs)
+
+    return True
+
+
+def accepted_analytic_measurement(m: qml.measurements.MeasurementProcess) -> bool:
+    """Specifies whether a measurement is accepted when analytic."""
+
+    if not isinstance(m, qml.measurements.StateMeasurement):
+        return False
+
+    if m.obs is not None:
+        return observable_accepts_analytic(m.obs, isinstance(m, qml.measurements.ExpectationMP))
+
+    return True
 
 
 def null_postprocessing(results):
@@ -514,10 +548,10 @@ class DefaultQubit(Device):
             name=self.name,
         )
         transform_program.add_transform(
-            validate_measurements, sample_measurements=accepted_sample_measurement, name=self.name
-        )
-        transform_program.add_transform(
-            validate_observables, stopping_condition=observable_stopping_condition, name=self.name
+            validate_measurements,
+            analytic_measurements=accepted_analytic_measurement,
+            sample_measurements=accepted_sample_measurement,
+            name=self.name,
         )
         if config.mcm_config.mcm_method == "tree-traversal":
             transform_program.add_transform(qml.transforms.broadcast_expand)
diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py
index 3704d8c68..2f38f1f9a 100644
--- a/pennylane/devices/preprocess.py
+++ b/pennylane/devices/preprocess.py
@@ -479,7 +479,9 @@ def validate_measurements(
         and not isinstance(meas := op.hyperparameters["measurement"], qml.measurements.StateMP)
     ]
 
-    if tape.shots:
+    shots = qml.measurements.Shots(tape.shots)
+
+    if shots.total_shots is not None:
         for m in chain(snapshot_measurements, tape.measurements):
             if not sample_measurements(m):
                 raise DeviceError(f"Measurement {m} not accepted with finite shots on {name}")


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
index e8610d922..cb1e87070 100644
--- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py
+++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
@@ -20,6 +20,7 @@ import pennylane as qml
 from pennylane import numpy as pnp
 from pennylane.devices import DefaultQubit, ExecutionConfig
 from pennylane.devices.default_qubit import stopping_condition
+from pennylane.operation import classproperty
 
 
 class NoMatOp(qml.operation.Operation):
@@ -45,6 +46,16 @@ class NoMatNoDecompOp(qml.operation.Operation):
         return False
 
 
+# pylint: disable=too-few-public-methods
+class HasDiagonalizingGatesOp(qml.operation.Operator):
+    """Dummy observable that has diagonalizing gates."""
+
+    # pylint: disable=arguments-renamed,invalid-overridden-method,no-self-argument
+    @classproperty
+    def has_diagonalizing_gates(cls):
+        return True
+
+
 def test_snapshot_multiprocessing_execute():
     """DefaultQubit cannot execute tapes with Snapshot if `max_workers` is not `None`"""
     dev = qml.device("default.qubit", max_workers=2)
@@ -288,6 +299,58 @@ class TestPreprocessing:
         batch, _ = program((tape4,))
         assert batch[0].circuit == tape4.circuit
 
+    @pytest.mark.parametrize(
+        "shots, measurements, supported",
+        [
+            # Supported measurements in analytic mode
+            (None, [qml.state()], True),
+            (None, [qml.expval(qml.X(0))], True),
+            (None, [qml.expval(qml.RX(0.123, 0))], False),
+            (None, [qml.expval(qml.SparseHamiltonian(qml.X.compute_sparse_matrix(), 0))], True),
+            (None, [qml.expval(qml.Hermitian(np.diag([1, 2]), wires=0))], True),
+            (None, [qml.var(qml.SparseHamiltonian(qml.X.compute_sparse_matrix(), 0))], False),
+            (None, [qml.expval(qml.X(0) @ qml.Hermitian(np.diag([1, 2]), wires=1))], True),
+            (
+                None,
+                [
+                    qml.expval(
+                        qml.Hamiltonian(
+                            [0.1, 0.2],
+                            [qml.Z(0), qml.SparseHamiltonian(qml.X.compute_sparse_matrix(), 1)],
+                        )
+                    )
+                ],
+                True,
+            ),
+            (None, [qml.expval(qml.Hamiltonian([0.1, 0.2], [qml.RZ(0.234, 0), qml.X(0)]))], False),
+            (
+                None,
+                [qml.expval(qml.Hamiltonian([1, 1], [qml.Z(0), HasDiagonalizingGatesOp(1)]))],
+                True,
+            ),
+            # Supported measurements in finite shots mode
+            (100, [qml.state()], False),
+            (100, [qml.expval(qml.X(0))], True),
+            (100, [qml.expval(qml.RX(0.123, 0))], False),
+            (100, [qml.expval(qml.X(0) @ qml.RX(0.123, 1))], False),
+            (100, [qml.expval(qml.X(0) @ qml.Y(1))], True),
+            (100, [qml.expval(qml.Hamiltonian([0.1, 0.2], [qml.Z(0), qml.X(1)]))], True),
+            (100, [qml.expval(qml.Hamiltonian([0.1, 0.2], [qml.RZ(0.123, 0), qml.X(1)]))], False),
+        ],
+    )
+    def test_validate_measurements(self, shots, measurements, supported):
+        """Tests that preprocess correctly validates measurements."""
+
+        device = qml.device("default.qubit")
+        tape = qml.tape.QuantumScript(measurements=measurements, shots=shots)
+        program, _ = device.preprocess()
+
+        if not supported:
+            with pytest.raises(qml.DeviceError):
+                program([tape])
+        else:
+            program([tape])
+
 
 class TestPreprocessingIntegration:
     """Test preprocess produces output that can be executed by the device."""
diff --git a/tests/devices/qubit/test_measure.py b/tests/devices/qubit/test_measure.py
index 332c5971a..d0c618311 100644
--- a/tests/devices/qubit/test_measure.py
+++ b/tests/devices/qubit/test_measure.py
@@ -179,12 +179,9 @@ class TestMeasurements:
 
         dev = qml.device("default.qubit", wires=4)
 
-        O1 = qml.X(0)
-        O2 = qml.X(0)
-
         @qml.qnode(dev, interface="jax")
         def qnode(t1, t2):
-            return qml.expval(qml.prod(O1, qml.RX(t1, 0), O2, qml.RX(t2, 0)))
+            return qml.expval((t1 * qml.X(0)) @ (t2 * qml.Y(1)))
 
         t1, t2 = 0.5, 1.0
         assert qml.math.allclose(qnode(t1, t2), jax.jit(qnode)(t1, t2))
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index 80286e7d5..13302dc3e 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -1522,7 +1522,7 @@ class TestIntegration:
             qml.PauliX(0)
             return qml.expval(prod_op)
 
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(qml.DeviceError):
             my_circ()
 
     def test_operation_integration(self):
diff --git a/tests/transforms/core/test_transform_dispatcher.py b/tests/transforms/core/test_transform_dispatcher.py
index be9a1611c..94f7d74e7 100644
--- a/tests/transforms/core/test_transform_dispatcher.py
+++ b/tests/transforms/core/test_transform_dispatcher.py
@@ -639,8 +639,8 @@ class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
         assert isinstance(program, qml.transforms.core.TransformProgram)
         assert isinstance(new_program, qml.transforms.core.TransformProgram)
 
-        assert len(program) == 5
-        assert len(new_program) == 6
+        assert len(program) == 4
+        assert len(new_program) == 5
 
         assert new_program[-1].transform is valid_transform
 
diff --git a/tests/workflow/test_construct_batch.py b/tests/workflow/test_construct_batch.py
index 3bf0a617c..8e8400b2c 100644
--- a/tests/workflow/test_construct_batch.py
+++ b/tests/workflow/test_construct_batch.py
@@ -105,7 +105,7 @@ class TestTransformProgramGetter:
         p_none = get_transform_program(circuit, None)
         assert p_dev == p_default
         assert p_none == p_dev
-        assert len(p_dev) == 9
+        assert len(p_dev) == 8
         config = qml.devices.ExecutionConfig(interface=getattr(circuit, "interface", None))
         assert p_dev == p_grad + dev.preprocess(config)[0]
 
@@ -140,7 +140,7 @@ class TestTransformProgramGetter:
             return qml.expval(qml.PauliZ(0))
 
         full_prog = get_transform_program(circuit)
-        assert len(full_prog) == 13
+        assert len(full_prog) == 12
 
         config = qml.devices.ExecutionConfig(
             interface=getattr(circuit, "interface", None),