Add support for mid-circuit measurement program capture

doc/releases/changelog-dev.md

@@ -4,15 +4,15 @@
 
 <h3>New features since last release</h3>
 
+* Mid-circuit measurements can now be captured with `qml.capture` enabled.
+  [(#6015)](https://github.com/PennyLaneAI/pennylane/pull/6015)
+
 * 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)
 
-* Resolved the bug in `qml.ThermalRelaxationError` where there was a typo from `tq` to `tg`.
-  [(#5988)](https://github.com/PennyLaneAI/pennylane/issues/5988)
-
 * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear
   combination of unitaries.
   [(#5756)](https://github.com/PennyLaneAI/pennylane/pull/5756)
@@ -174,6 +174,9 @@
 
 <h4>Community contributions 🥳</h4>
 
+* Resolved the bug in `qml.ThermalRelaxationError` where there was a typo from `tq` to `tg`.
+  [(#5988)](https://github.com/PennyLaneAI/pennylane/issues/5988)
+
 * `DefaultQutritMixed` readout error has been added using parameters `readout_relaxation_probs` and 
   `readout_misclassification_probs` on the `default.qutrit.mixed` device. These parameters add a `~.QutritAmplitudeDamping`  and a `~.TritFlip` channel, respectively,
   after measurement diagonalization. The amplitude damping error represents the potential for

pennylane/capture/primitives.py

@@ -266,13 +266,8 @@ def create_measurement_mcm_primitive(
     primitive = jax.core.Primitive(name + "_mcm")
 
     @primitive.def_impl
-    def _(*mcms, **kwargs):
-        raise NotImplementedError(
-            "mcm measurements do not currently have a concrete implementation"
-        )
-        # need to figure out how to convert a jaxpr int into a measurement value, and pass
-        # that measurment value here.
-        # return type.__call__(measurement_type, obs=mcms, **kwargs)
+    def _(*mcms, single_mcm=True, **kwargs):
+        return type.__call__(measurement_type, obs=mcms[0] if single_mcm else mcms, **kwargs)
 
     abstract_type = _get_abstract_measurement()
 

pennylane/measurements/measurements.py

@@ -172,8 +172,8 @@ def _primitive_bind_call(cls, obs=None, wires=None, eigvals=None, id=None, **kwa
         ):
             return cls._obs_primitive.bind(obs, **kwargs)
         if isinstance(obs, (list, tuple)):
-            return cls._mcm_primitive.bind(*obs, **kwargs)  # iterable of mcms
-        return cls._mcm_primitive.bind(obs, **kwargs)  # single mcm
+            return cls._mcm_primitive.bind(*obs, single_mcm=False, **kwargs)  # iterable of mcms
+        return cls._mcm_primitive.bind(obs, single_mcm=True, **kwargs)  # single mcm
 
     # pylint: disable=unused-argument
     @classmethod

pennylane/measurements/mid_measure.py

@@ -15,6 +15,7 @@
 This module contains the qml.measure measurement.
 """
 import uuid
+from functools import lru_cache
 from typing import Generic, Hashable, Optional, TypeVar, Union
 
 import pennylane as qml
@@ -209,22 +210,56 @@ def func(x):
               samples, leading to unexpected or incorrect results.
 
     """
-    wire = Wires(wires)
-    if len(wire) > 1:
+    if qml.capture.enabled():
+        primitive = _create_mid_measure_primitive()
+        return primitive.bind(wires, reset=reset, postselect=postselect)
+
+    return _measure_impl(wires, reset=reset, postselect=postselect)
+
+
+def _measure_impl(
+    wires: Union[Hashable, Wires], reset: Optional[bool] = False, postselect: Optional[int] = None
+):
+    """Concrete implementation of qml.measure"""
+    wires = Wires(wires)
+    if len(wires) > 1:
         raise qml.QuantumFunctionError(
             "Only a single qubit can be measured in the middle of the circuit"
         )
 
     # Create a UUID and a map between MP and MV to support serialization
     measurement_id = str(uuid.uuid4())[:8]
-    mp = MidMeasureMP(wires=wire, reset=reset, postselect=postselect, id=measurement_id)
-    if qml.capture.enabled():
-        raise NotImplementedError(
-            "Capture cannot currently handle classical output from mid circuit measurements."
-        )
+    mp = MidMeasureMP(wires=wires, reset=reset, postselect=postselect, id=measurement_id)
     return MeasurementValue([mp], processing_fn=lambda v: v)
 
 
+@lru_cache
+def _create_mid_measure_primitive():
+    """Create a primitive corresponding to an mid-circuit measurement type.
+
+    Called when using :func:`~pennylane.measure`.
+
+    Returns:
+        jax.core.Primitive: A new jax primitive corresponding to a mid-circuit
+        measurement.
+
+    """
+    import jax  # pylint: disable=import-outside-toplevel
+
+    mid_measure_p = jax.core.Primitive("measure")
+
+    @mid_measure_p.def_impl
+    def _(wires, reset=False, postselect=None):
+        return _measure_impl(wires, reset=reset, postselect=postselect)
+
+    @mid_measure_p.def_abstract_eval
+    def _(*_, **__):
+        dtype = jax.numpy.int64 if jax.config.jax_enable_x64 else jax.numpy.int32
+        return jax.core.ShapedArray((), dtype)
+
+    return mid_measure_p
+
+
 T = TypeVar("T")
 
 
@@ -266,7 +301,7 @@ def __init__(
     @classmethod
     def _primitive_bind_call(cls, wires=None, reset=False, postselect=None, id=None):
         wires = () if wires is None else wires
-        return cls._wires_primitive.bind(*wires, reset=reset, postselect=postselect)
+        return cls._wires_primitive.bind(*wires, reset=reset, postselect=postselect, id=id)
 
     @classmethod
     def _abstract_eval(

pennylane/workflow/qnode.py

@@ -1169,7 +1169,11 @@ def _impl_call(self, *args, **kwargs) -> qml.typing.Result:
 
         old_interface = self.interface
         if old_interface == "auto":
-            interface = qml.math.get_interface(*args, *list(kwargs.values()))
+            interface = (
+                "jax"
+                if qml.capture.enabled()
+                else qml.math.get_interface(*args, *list(kwargs.values()))
+            )
             self._interface = INTERFACE_MAP[interface]
 
         if self._qfunc_uses_shots_arg: