Repository: PennyLaneAI/pennylane
Pull request #5833 (2024-06-11 13:51:43)

Add `mcm_method="single-branch-statistics"` for the default Catalyst MCM behaviour
**Context:**
As name says.

**Description of the Change:**
* Add new `mcm_method` `"single-branch-statistics"` used to request the original Catalyst way of performing mid-circuit measurements, wherein only a single branch of the execution tree is explored.
* Raise an error if calling a qnode with `mcm_method="single-branch-statistics"` when not using `qml.qjit`.
* `QNode.__init__` now uses a `MCMConfig` dataclass for MCM config arguments
* Updated `capture` module to work with the updates to `QNode`.
* Added new metaclass `ABCCaptureMeta` in the `capture` module to use as the metaclass for `MeasurementProcess` and `Operator`. This way, we don't need to create metaclasses local to the files to which those classes belong.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


mudit2812: @trbromley Since this is really for catalyst functionality, I'm wondering if I should or should not mention this new `mcm_method` in the `measurements.rst` file.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5833?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`d65f671`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/d65f671326b71bd9e755a157ec913af31888934b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`93a64fc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/93a64fca2cbe4707fd1d9616a2629a9e3e4b9dcc?dropdown=coverage&el=desc&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    #5833      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         415      415              
  Lines       39824    39530     -294     
==========================================
- Hits        39692    39397     -295     
- Misses        132      133       +1     
```



</details>

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

mudit2812: [sc-65684]

Patch:
diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index e837d5603..6c769f7d2 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -538,7 +538,9 @@ PennyLane. For ease of use, we provide the following configuration options to us
 * ``mcm_method``: To set the method used for applying mid-circuit measurements. Use ``mcm_method="deferred"``
   to use the :ref:`deferred measurements principle <deferred_measurements>` or ``mcm_method="one-shot"`` to use
   the :ref:`one-shot transform <one_shot_transform>`. When executing with finite shots, ``mcm_method="one-shot"``
-  will be the default, and ``mcm_method="deferred"`` otherwise.
+  will be the default, and ``mcm_method="deferred"`` otherwise. Additionally, if using :func:`~pennylane.qjit`,
+  ``mcm_method="single-branch-statistics"`` can also be used and will be the default. Using this method, a single
+  branch of the execution tree will be randomly explored.
 
   .. warning::
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c182bb812..410f3beed 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -35,6 +35,7 @@
   These keyword arguments can be used to configure how the device should behave when running circuits with
   mid-circuit measurements.
   [(#5679)](https://github.com/PennyLaneAI/pennylane/pull/5679)
+  [(#5833)](https://github.com/PennyLaneAI/pennylane/pull/5833)
 
   * `postselect_mode="hw-like"` will indicate to devices to discard invalid shots when postselecting
     mid-circuit measurements. Use `postselect_mode="fill-shots"` to unconditionally sample the postselected
@@ -42,7 +43,8 @@
     matches the total number of shots.
   * `mcm_method` will indicate which strategy to use for running circuits with mid-circuit measurements.
     Use `mcm_method="deferred"` to use the deferred measurements principle, or `mcm_method="one-shot"`
-    to execute once for each shot.
+    to execute once for each shot. If using `qml.jit` with the Catalyst compiler, `mcm_method="single-branch-statistics"`
+    is also available. Using this method, a single branch of the execution tree will be randomly explored.
 
 * The `default.tensor` device is introduced to perform tensor network simulation of a quantum circuit.
   [(#5699)](https://github.com/PennyLaneAI/pennylane/pull/5699)
diff --git a/pennylane/capture/__init__.py b/pennylane/capture/__init__.py
index 3cbb39a3d..c740d8570 100644
--- a/pennylane/capture/__init__.py
+++ b/pennylane/capture/__init__.py
@@ -126,7 +126,7 @@ If needed, developers can also override the implementation method of the primiti
         return type.__call__(MyCustomOp, *args, **kwargs)
 """
 from .switches import disable, enable, enabled
-from .capture_meta import CaptureMeta
+from .capture_meta import CaptureMeta, ABCCaptureMeta
 from .primitives import (
     create_operator_primitive,
     create_measurement_obs_primitive,
diff --git a/pennylane/capture/capture_meta.py b/pennylane/capture/capture_meta.py
index fe82b4926..e5f42871b 100644
--- a/pennylane/capture/capture_meta.py
+++ b/pennylane/capture/capture_meta.py
@@ -16,6 +16,7 @@ Defines a metaclass for automatic integration of any ``Operator`` with plxpr pro
 
 See ``explanations.md`` for technical explanations of how this works.
 """
+from abc import ABCMeta
 from inspect import Signature, signature
 
 from .switches import enabled
@@ -84,3 +85,8 @@ class CaptureMeta(type):
             # use bind to construct the class if we want class construction to add it to the jaxpr
             return cls._primitive_bind_call(*args, **kwargs)
         return type.__call__(cls, *args, **kwargs)
+
+
+# pylint: disable=abstract-method
+class ABCCaptureMeta(CaptureMeta, ABCMeta):
+    """A combination of the capture meta and ABCMeta"""
diff --git a/pennylane/capture/capture_qnode.py b/pennylane/capture/capture_qnode.py
index 415253552..6f305f885 100644
--- a/pennylane/capture/capture_qnode.py
+++ b/pennylane/capture/capture_qnode.py
@@ -16,6 +16,7 @@ This submodule defines a capture compatible call to QNodes.
 """
 
 from copy import copy
+from dataclasses import asdict
 from functools import lru_cache, partial
 
 import pennylane as qml
@@ -161,7 +162,7 @@ def qnode_call(qnode: "qml.QNode", *args, **kwargs) -> "qml.typing.Result":
 
     qfunc_jaxpr = jax.make_jaxpr(qfunc)(*args)
     execute_kwargs = copy(qnode.execute_kwargs)
-    mcm_config = execute_kwargs.pop("mcm_config")
+    mcm_config = asdict(execute_kwargs.pop("mcm_config"))
     qnode_kwargs = {"diff_method": qnode.diff_method, **execute_kwargs, **mcm_config}
     qnode_prim = _get_qnode_prim()
 
diff --git a/pennylane/devices/execution_config.py b/pennylane/devices/execution_config.py
index e8dd9816d..5d0dda175 100644
--- a/pennylane/devices/execution_config.py
+++ b/pennylane/devices/execution_config.py
@@ -42,7 +42,7 @@ class MCMConfig:
 
         Note that this hook is automatically called after init via the dataclass integration.
         """
-        if self.mcm_method not in ("deferred", "one-shot", None):
+        if self.mcm_method not in ("deferred", "one-shot", "single-branch-statistics", None):
             raise ValueError(f"Invalid mid-circuit measurements method '{self.mcm_method}'.")
         if self.postselect_mode not in ("hw-like", "fill-shots", None):
             raise ValueError(f"Invalid postselection mode '{self.postselect_mode}'.")
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 3397d95f0..c24c6aee2 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -18,7 +18,7 @@ and measurement samples using AnnotatedQueues.
 """
 import copy
 import functools
-from abc import ABC, ABCMeta, abstractmethod
+from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Optional, Sequence, Tuple, Union
 
@@ -112,12 +112,7 @@ class MeasurementShapeError(ValueError):
     quantum tape."""
 
 
-# pylint: disable=abstract-method
-class ABCCaptureMeta(qml.capture.CaptureMeta, ABCMeta):
-    """A combination of the capture meta and ABCMeta"""
-
-
-class MeasurementProcess(ABC, metaclass=ABCCaptureMeta):
+class MeasurementProcess(ABC, metaclass=qml.capture.ABCCaptureMeta):
     """Represents a measurement process occurring at the end of a
     quantum variational circuit.
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index b5ea0c127..1d68f0a21 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -256,7 +256,7 @@ from numpy.linalg import multi_dot
 from scipy.sparse import coo_matrix, csr_matrix, eye, kron
 
 import pennylane as qml
-from pennylane.capture import CaptureMeta, create_operator_primitive
+from pennylane.capture import ABCCaptureMeta, create_operator_primitive
 from pennylane.math import expand_matrix
 from pennylane.queuing import QueuingManager
 from pennylane.typing import TensorLike
@@ -406,12 +406,7 @@ def _process_data(op):
     return str([id(d) if qml.math.is_abstract(d) else _mod_and_round(d, mod_val) for d in op.data])
 
 
-# pylint: disable=abstract-method
-class CaptureMetaABC(CaptureMeta, abc.ABCMeta):
-    """Mixing together CaptureMeta and ABCMeta so that Operator can use both."""
-
-
-class Operator(abc.ABC, metaclass=CaptureMetaABC):
+class Operator(abc.ABC, metaclass=ABCCaptureMeta):
     r"""Base class representing quantum operators.
 
     Operators are uniquely defined by their name, the wires they act on, their (trainable) parameters,
diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py
index f27d100e7..13fd2ce7c 100644
--- a/pennylane/workflow/execution.py
+++ b/pennylane/workflow/execution.py
@@ -581,7 +581,10 @@ def execute(
         # This is a current limitation of defer_measurements. "hw-like" behaviour is
         # not yet accessible.
         if config.mcm_config.postselect_mode == "hw-like":
-            raise ValueError("Using postselect_mode='hw-like' is not supported with jax-jit.")
+            raise ValueError(
+                "Using postselect_mode='hw-like' is not supported with jax-jit when using "
+                "mcm_method='deferred'."
+            )
         config.mcm_config.postselect_mode = "fill-shots"
 
     if transform_program is None:
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index fa3baab25..42a0c39f0 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -527,10 +527,7 @@ class QNode:
         self.max_expansion = max_expansion
         cache = (max_diff > 1) if cache == "auto" else cache
 
-        if mcm_method not in ("deferred", "one-shot", None):
-            raise ValueError(f"Invalid mid-circuit measurements method '{mcm_method}'.")
-        if postselect_mode not in ("hw-like", "fill-shots", None):
-            raise ValueError(f"Invalid postselection mode '{postselect_mode}'.")
+        mcm_config = qml.devices.MCMConfig(mcm_method=mcm_method, postselect_mode=postselect_mode)
 
         # execution keyword arguments
         self.execute_kwargs = {
@@ -540,7 +537,7 @@ class QNode:
             "max_diff": max_diff,
             "max_expansion": max_expansion,
             "device_vjp": device_vjp,
-            "mcm_config": {"postselect_mode": postselect_mode, "mcm_method": mcm_method},
+            "mcm_config": mcm_config,
         }
 
         if self.expansion_strategy == "device":
@@ -1045,11 +1042,16 @@ class QNode:
         )
         self._tape_cached = using_custom_cache and self.tape.hash in cache
 
+        mcm_config = self.execute_kwargs["mcm_config"]
         finite_shots = _get_device_shots if override_shots is False else override_shots
-        if not finite_shots and self.execute_kwargs["mcm_config"]["mcm_method"] == "one-shot":
-            raise ValueError(
-                "Cannot use the 'one-shot' method for mid-circuit measurements with analytic mode."
-            )
+        if not finite_shots:
+            mcm_config.postselect_mode = None
+            if mcm_config.mcm_method == "one-shot":
+                raise ValueError(
+                    "Cannot use the 'one-shot' method for mid-circuit measurements with analytic mode."
+                )
+        if mcm_config.mcm_method == "single-branch-statistics":
+            raise ValueError("Cannot use mcm_method='single-branch-statistics' without qml.qjit.")
 
         # Add the device program to the QNode program
         if isinstance(self.device, qml.devices.Device):
@@ -1069,7 +1071,7 @@ class QNode:
             full_transform_program.add_transform(
                 qml.devices.preprocess.mid_circuit_measurements,
                 device=self.device,
-                mcm_config=self.execute_kwargs["mcm_config"],
+                mcm_config=mcm_config,
             )
             override_shots = 1
         elif hasattr(self.device, "capabilities"):


Test patch:
diff --git a/tests/capture/test_capture_qnode.py b/tests/capture/test_capture_qnode.py
index ec5387647..ca8f63ffe 100644
--- a/tests/capture/test_capture_qnode.py
+++ b/tests/capture/test_capture_qnode.py
@@ -14,6 +14,7 @@
 """
 Tests for capturing a qnode into jaxpr.
 """
+from dataclasses import asdict
 from functools import partial
 
 # pylint: disable=protected-access
@@ -121,7 +122,7 @@ def test_simple_qnode(x64_mode):
     assert eqn0.params["shots"] == qml.measurements.Shots(None)
     expected_kwargs = {"diff_method": "best"}
     expected_kwargs.update(circuit.execute_kwargs)
-    expected_kwargs.update(expected_kwargs.pop("mcm_config"))
+    expected_kwargs.update(asdict(expected_kwargs.pop("mcm_config")))
     assert eqn0.params["qnode_kwargs"] == expected_kwargs
 
     qfunc_jaxpr = eqn0.params["qfunc_jaxpr"]
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index cf840c477..dbbdfa48e 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -1823,6 +1823,21 @@ class TestMCMConfiguration:
         ):
             _ = f_jit(param)
 
+    def test_single_branch_statistics_error_without_qjit(self):
+        """Test that an error is raised if attempting to use mcm_method="single-branch-statistics
+        without qml.qjit"""
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, mcm_method="single-branch-statistics")
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.measure(0, postselect=1)
+            return qml.sample(wires=0)
+
+        param = np.pi / 4
+        with pytest.raises(ValueError, match="Cannot use mcm_method='single-branch-statistics'"):
+            _ = circuit(param)
+
 
 class TestTapeExpansion:
     """Test that tape expansion within the QNode works correctly"""