More sophisticated measurement validation for default-qubit

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,

pennylane/devices/default_qubit.py

@@ -60,30 +60,6 @@
 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 @@ def preprocess(
             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)

pennylane/devices/preprocess.py

@@ -479,7 +479,9 @@ def sample_measurements(m):
         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}")

tests/devices/default_qubit/test_default_qubit_preprocessing.py

@@ -20,6 +20,7 @@
 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 @@ def has_matrix(self):
         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 @@ def has_matrix(self):
         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."""

tests/devices/qubit/test_measure.py

@@ -179,12 +179,9 @@ def test_op_math_observable_jit_compatible(self):
 
         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))

tests/ops/op_math/test_prod.py

@@ -1522,7 +1522,7 @@ def my_circ():
             qml.PauliX(0)
             return qml.expval(prod_op)
 
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(qml.DeviceError):
             my_circ()
 
     def test_operation_integration(self):

tests/transforms/core/test_transform_dispatcher.py

@@ -639,8 +639,8 @@ def test_device_transform(self, valid_transform):
         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
 

tests/workflow/test_construct_batch.py

@@ -105,7 +105,7 @@ def circuit():
         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 @@ def circuit(x):
             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),