Repository: PennyLaneAI/pennylane

Issue:
Support native measurement-based snapshots in Default Clifford
### Feature details

The [`default.clifford`](https://docs.pennylane.ai/en/stable/code/api/pennylane.devices.default_clifford.DefaultClifford.html) device allows PennyLane users to quickly simulate [Clifford circuits](https://en.wikipedia.org/wiki/Clifford_gates) in PennyLane using the [Stim](https://github.com/quantumlib/stim/tree/main) library as a backend. The device was [recently added in v0.35](https://pennylane.ai/blog/2024/03/pennylane-release-0.35#a-new-clifford-device-) and is not yet at feature parity with devices like `default.qubit`.

PennyLane provides support for [circuit snapshots](https://docs.pennylane.ai/en/stable/introduction/inspecting_circuits.html#debugging-with-mid-circuit-snapshots), allowing users to inspect circuits at different stages of execution on some simulator devices. Snapshots can can be placed in circuits using [`qml.Snapshot()`](https://docs.pennylane.ai/en/stable/code/api/pennylane.Snapshot.html), which defaults to providing the snapshot of the statevector at that point in the circuit. If a `measurement` argument is provided to `qml.Snapshot()`, then Pennylane [measurements](https://docs.pennylane.ai/en/stable/introduction/measurements.html) like [`qml.expval()`](https://docs.pennylane.ai/en/stable/code/api/pennylane.expval.html) can be inspected.

The `default.clifford` device only supports statevector-based snapshots. The objective of this feature request is to extend support to measurement-based snapshots.

### Implementation

It should be possible to execute the following code:

```python
import pennylane as qml

dev = qml.device("default.clifford")

@qml.qnode(dev)
def circuit():
    qml.Hadamard(0)
    qml.Snapshot(measurement=qml.expval(qml.X(0)))
    qml.S(wires=0)
    return qml.expval(qml.X(0))

qml.snapshots(circuit)()
```

A successful PR must include:

- The requested feature
- Comprehensive tests
- Relevant updates to the documentation
- An entry in the [release notes](https://github.com/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md)
- Passing all CI tests

Check out our [development guide](https://docs.pennylane.ai/en/stable/development/guide.html) for more information, and don't hesitate to reach out to us if you need assistance.

Additional requirements:

- Measurements done during `qml.Snapshot` shouldn't alter the state of the main execution workflow.
- Efficiency is important, the addition of snapshots should not come with much overhead.

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

Support for measurement-based snapshots was provided in `default.mixed` in [this recent PR](https://github.com/PennyLaneAI/pennylane/pull/5552).


Discussion:
Hi, I would like to work on this issue for unitaryhack.

Resolving patch:
diff --git a/pennylane/devices/default_clifford.py b/pennylane/devices/default_clifford.py
--- a/pennylane/devices/default_clifford.py
+++ b/pennylane/devices/default_clifford.py
@@ -23,6 +23,7 @@
 import numpy as np
 
 import pennylane as qml
+from pennylane._device import DeviceError
 from pennylane.measurements import (
     ClassicalShadowMP,
     CountsMP,
@@ -574,23 +575,10 @@ def simulate(
                 # Note: This is a lot faster than doing `stim_ct.append(gate, wires)`
                 stim_circuit.append_from_stim_program_text(f"{gate} {wires}")
             else:
-                if op.name == "GlobalPhase":
+                if isinstance(op, qml.GlobalPhase):
                     global_phase_ops.append(op)
-                if op.name == "Snapshot":
-                    if debugger is not None and debugger.active:
-                        meas = op.hyperparameters["measurement"]
-                        if meas is not None and not isinstance(meas, qml.measurements.StateMP):
-                            raise ValueError(
-                                f"{self.name} does not support arbitrary measurements of a state with snapshots."
-                            )
-                        # Build a temporary simulator for obtaining state
-                        snap_sim = stim.TableauSimulator()
-                        if self.wires is not None:
-                            snap_sim.set_num_qubits(len(self.wires))
-                        snap_sim.do_circuit(stim_circuit)
-                        state = self._measure_state(meas, snap_sim, circuit=circuit)
-                        # Add to the debugger snapshot
-                        debugger.snapshots[op.tag or len(debugger.snapshots)] = state
+                if isinstance(op, qml.Snapshot):
+                    self._apply_snapshot(circuit, stim_circuit, op, global_phase_ops, debugger)
 
         tableau_simulator.do_circuit(stim_circuit)
         global_phase = qml.GlobalPhase(qml.math.sum(op.data[0] for op in global_phase_ops))
@@ -605,6 +593,66 @@ def simulate(
 
         return meas_results[0] if len(meas_results) == 1 else tuple(meas_results)
 
+    @property
+    def _analytical_measurement_map(self):
+        """Maps measurement type to the desired analytic measurement method."""
+        return {
+            DensityMatrixMP: self._measure_density_matrix,
+            StateMP: self._measure_state,  # kwargs -> circuit, global_phase
+            ExpectationMP: self._measure_expectation,  # kwargs -> circuit, stim_circuit
+            VarianceMP: self._measure_variance,
+            VnEntropyMP: self._measure_vn_entropy,  # kwargs -> circuit
+            MutualInfoMP: self._measure_mutual_info,  # kwargs -> circuit
+            PurityMP: self._measure_purity,  # kwargs -> circuit
+            ProbabilityMP: self._measure_probability,  # kwargs -> circuit, stim_circuit
+        }
+
+    @property
+    def _statistical_measurement_map(self):
+        """Maps measurement type to the desired sampling measurement method."""
+        return {
+            ExpectationMP: self._sample_expectation,
+            VarianceMP: self._sample_variance,
+            ClassicalShadowMP: self._sample_classical_shadow,
+            ShadowExpvalMP: self._sample_expval_shadow,
+        }
+
+    def _apply_snapshot(
+        self,
+        circuit: qml.tape.QuantumScript,
+        stim_circuit,
+        operation: qml.Snapshot,
+        global_phase_ops: Sequence[qml.GlobalPhase],
+        debugger=None,
+    ):
+        """Apply a snapshot operation to the stim circuit."""
+        if debugger is not None and debugger.active:
+            meas = operation.hyperparameters["measurement"] or StateMP()
+            measurement_func = self._analytical_measurement_map.get(type(meas), None)
+
+            if measurement_func is None:  # pragma: no cover
+                raise DeviceError(
+                    f"Snapshots of {type(meas)} are not yet supported on default.clifford."
+                )
+
+            # Build a temporary simulator for obtaining state
+            snap_sim = stim.TableauSimulator()
+            if self.wires is not None:
+                snap_sim.set_num_qubits(len(self.wires))
+            snap_sim.do_circuit(stim_circuit)
+            global_phase = qml.GlobalPhase(qml.math.sum(op.data[0] for op in global_phase_ops))
+
+            snap_result = measurement_func(
+                meas,
+                snap_sim,
+                circuit=circuit,
+                stim_circuit=stim_circuit,
+                global_phase=global_phase,
+            )
+
+            # Add to the debugger snapshot
+            debugger.snapshots[operation.tag or len(debugger.snapshots)] = snap_result
+
     @staticmethod
     def _measure_observable_sample(meas_obs, stim_circuit, shots, sample_seed):
         """Compute sample output from a stim circuit for a given Pauli observable"""
@@ -636,17 +684,9 @@ def measure_statistical(self, circuit, stim_circuit, seed=None):
         num_shots = circuit.shots.total_shots
         sample_seed = seed if isinstance(seed, int) else self._rng.integers(2**31 - 1, size=1)[0]
 
-        # maps measurement type to the desired sampling measurement method
-        measurement_map = {
-            ExpectationMP: self._sample_expectation,
-            VarianceMP: self._sample_variance,
-            ClassicalShadowMP: self._sample_classical_shadow,
-            ShadowExpvalMP: self._sample_expval_shadow,
-        }
-
         results = []
         for meas in circuit.measurements:
-            measurement_func = measurement_map.get(type(meas), None)
+            measurement_func = self._statistical_measurement_map.get(type(meas), None)
             if measurement_func is not None:
                 res = measurement_func(meas, stim_circuit, shots=num_shots, seed=sample_seed)
             else:
@@ -677,22 +717,11 @@ def measure_statistical(self, circuit, stim_circuit, seed=None):
 
     def measure_analytical(self, circuit, stim_circuit, tableau_simulator, global_phase):
         """Given a circuit, compute tableau and return the analytical measurement results."""
-        # maps measurement type to the desired analytic measurement method
-        measurement_map = {
-            DensityMatrixMP: self._measure_density_matrix,
-            StateMP: self._measure_state,  # kwargs -> circuit, global_phase
-            ExpectationMP: self._measure_expectation,  # kwargs -> circuit, stim_circuit
-            VarianceMP: self._measure_variance,
-            VnEntropyMP: self._measure_vn_entropy,  # kwargs -> circuit
-            MutualInfoMP: self._measure_mutual_info,  # kwargs -> circuit
-            PurityMP: self._measure_purity,  # kwargs -> circuit
-            ProbabilityMP: self._measure_probability,  # kwargs -> circuit, stim_circuit
-        }
 
         results = []
         for meas in circuit.measurements:
             # signature: measurement_func(meas, tableaus_simulator, **kwargs) -> meas_result
-            measurement_func = measurement_map.get(type(meas), None)
+            measurement_func = self._analytical_measurement_map.get(type(meas), None)
             if measurement_func is None:  # pragma: no cover
                 raise NotImplementedError(
                     f"default.clifford doesn't support the {type(meas)} measurement analytically at the moment."


Test patch:
diff --git a/tests/devices/test_default_clifford.py b/tests/devices/test_default_clifford.py
--- a/tests/devices/test_default_clifford.py
+++ b/tests/devices/test_default_clifford.py
@@ -424,6 +424,54 @@ def test_pl_to_stim(pl_op, stim_op):
     assert wires == " ".join(map(str, stim_op[1]))
 
 
+@pytest.mark.parametrize(
+    ["measurement", "tag"],
+    [
+        (qml.expval(op=qml.Z(1)), None),
+        (qml.expval(op=qml.Y(0) @ qml.X(1)), "expval"),
+        (qml.var(op=qml.X(0)), None),
+        (qml.var(op=qml.X(0) @ qml.Z(1)), "var"),
+        (qml.density_matrix(wires=[1]), None),
+        (qml.density_matrix(wires=[0, 1]), "dm"),
+        (qml.probs(op=qml.Y(0)), None),
+        (qml.probs(op=qml.X(0) @ qml.Y(1)), "probs"),
+        (qml.vn_entropy(wires=[0]), None),
+        (qml.vn_entropy(wires=[1]), "vn_entropy"),
+        (qml.mutual_info(wires0=[1], wires1=[0]), None),
+        (qml.mutual_info(wires0=[0], wires1=[1]), "mi"),
+        (qml.purity(wires=[0]), None),
+        (qml.purity(wires=[1]), "purity"),
+    ],
+)
+def test_snapshot_supported(measurement, tag):
+    """Tests that applying snapshot of measurements is done correctly"""
+
+    def circuit():
+        """Snapshot circuit"""
+        qml.Hadamard(wires=0)
+        qml.Hadamard(wires=1)
+        qml.Snapshot(measurement=qml.expval(qml.Z(0) @ qml.Z(1)))
+        qml.CNOT(wires=[0, 1])
+        qml.Snapshot(measurement=measurement, tag=tag)
+        qml.CZ(wires=[1, 0])
+        qml.Snapshot("meas2", measurement=measurement)
+        return qml.probs(op=qml.Y(1) @ qml.Z(0))
+
+    dev_qubit = qml.device("default.qubit", wires=2)
+    dev_clifford = qml.device("default.clifford", wires=2)
+
+    qnode_qubit = qml.QNode(circuit, device=dev_qubit)
+    qnode_clifford = qml.QNode(circuit, device=dev_clifford)
+
+    snaps_qubit = qml.snapshots(qnode_qubit)()
+    snaps_clifford = qml.snapshots(qnode_clifford)()
+
+    assert len(snaps_qubit) == len(snaps_clifford)
+    for key1, key2 in zip(snaps_qubit, snaps_clifford):
+        assert key1 == key2
+        assert qml.math.allclose(snaps_qubit[key1], snaps_clifford[key2])
+
+
 def test_max_worker_clifford():
     """Test that the execution of multiple tapes is possible with multiprocessing on this device."""
 
@@ -592,17 +640,6 @@ def circuit_exp():
     ):
         circuit_exp()
 
-    @qml.qnode(qml.device("default.clifford"))
-    def circuit_snap():
-        qml.Snapshot(measurement=qml.expval(qml.PauliZ(0)))
-        return qml.state()
-
-    with pytest.raises(
-        ValueError,
-        match="default.clifford does not support arbitrary measurements of a state with snapshots.",
-    ):
-        qml.snapshots(circuit_snap)()
-
     @qml.qnode(qml.device("default.clifford", wires=3, shots=10))
     def circuit_herm():
         qml.Hadamard(wires=[0])