Repository: PennyLaneAI/pennylane
Pull request #7219 (2025-04-07 16:31:06)

Capture parametric MCMs in plxpr
Current implementation summary:
- We leave `ftqc.measure_z` completely equivalent to `qml.measure`, so it captures the standard `measure` primitive for PennyLane

```
qml.capture.enable()

@qml.qnode(dev)
def circ1():
    qml.measure(0)
    return qml.expval(Z(0))
    
@qml.qnode(dev)
def circ2():
    measure_z(0)
    return qml.expval(Z(0))
```

```
>>> jaxpr = jax.make_jaxpr(circ1)()
>>> jaxpr.eqns[0].params["qfunc_jaxpr"].eqns
[_:i64[] = measure[postselect=None reset=False] 0,
 a:AbstractOperator() = PauliZ[n_wires=1] 0,
 a:AbstractMeasurement(n_wires=None) = expval_obs b]
```

```
>>> jaxpr = jax.make_jaxpr(circ2)()
>>> jaxpr.eqns[0].params["qfunc_jaxpr"].eqns
[_:i64[] = measure[postselect=None reset=False] 0,
 a:AbstractOperator() = PauliZ[n_wires=1] 0,
 a:AbstractMeasurement(n_wires=None) = expval_obs b]
```

- We capture `measure_x`, `measure_y` and `measure_arbitrary_basis` all with the same parametrized measurement primitive, `p_measure`. This is under the assumption that the distinction is primarily a user-facing one for reading code and drawing circuits, as well as more specific diagonalization, and none of those things are relevant at the point where we are capturing for catalyst, so they may as well all just be one thing that is handled in a uniform manner. So when captured, these circuits are identical:

```
@qml.qnode(dev)
def circ1():
    measure_x(0)
    return qml.expval(Z(0))

@qml.qnode(dev)
def circ2():
    measure_arbitrary_basis(0, angle=0, plane="XY")
    return qml.expval(Z(0))
```

```
>>> jaxpr = jax.make_jaxpr(circ1)()
>>> jaxpr.eqns[0].params["qfunc_jaxpr"].eqns
[_:i64[] = p_measure[angle=0 plane=XY postselect=None reset=False] 0,
 a:AbstractOperator() = PauliZ[n_wires=1] 0,
 a:AbstractMeasurement(n_wires=None) = expval_obs b]
```

```
>>> jaxpr = jax.make_jaxpr(circ2)()
>>> jaxpr.eqns[0].params["qfunc_jaxpr"].eqns
[_:i64[] = p_measure[angle=0 plane=XY postselect=None reset=False] 0,
 a:AbstractOperator() = PauliZ[n_wires=1] 0,
 a:AbstractMeasurement(n_wires=None) = expval_obs b]
```

- We capture `cond_measure` by putting it inside a `cond` primitive using `qml.cond`. This wouldn't work without program capture on, but my understanding is that this will be more flexible with plxpr. We can make a new primitive if that's not the case. The essential thing here is that we need to be able to use the value returned by whichever function is called going forward, so `cond` needs to have the option to return something, i.e. in the following we need `a:i64[]` to be the value returned by whichever of the measurement functions we end up executing inside cond:

```
@qml.qnode(dev)
def circ():
    m1 = measure_x(0)
    m2 = cond_measure(m1, measure_x, measure_y)(2)
    cond_measure(m2, measure_x, measure_y)(3)
    return qml.expval(Z(0))
```

```
>>> jaxpr = jax.make_jaxpr(circ)()
>>> jaxpr.eqns[0].params["qfunc_jaxpr"].eqns[1]
a:i64[] = cond[
  args_slice=slice(2, None, None)
  consts_slices=[slice(2, 2, None), slice(2, 2, None)]
  jaxpr_branches=[{ lambda ; a:i64[]. let
    b:i64[] = p_measure[angle=0 plane=XY postselect=None reset=False] a
  in (b,) }, { lambda ; a:i64[]. let
    b:i64[] = p_measure[
      angle=1.5707963267948966
      plane=XY
      postselect=None
      reset=False
    ] a
  in (b,) }]
] b True 2
```

- We make a new function, `make_graph_state`, that either creates a GraphStatePrep operation (if `capture` is off), or decomposes it and captures the operations generated in the decomposition (if `capture` is on). This circumvents JAX's inability to handle a networkx graph. Maybe we want to maintain some more efficient way of storing this in a single primitive in the future, but for now this seems simple and acceptable.

```
import networkx as nx
from pennylane.ftqc import make_graph_state

@qml.qnode(dev)
def circ():
    make_graph_state(nx.grid_graph((4,)), wires=[1, 2, 3, 4])
    return qml.expval(X(4))
``` 

```
>>> jax.make_jaxpr(circ.func)()
{ lambda ; . let
    _:AbstractOperator() = Hadamard[n_wires=1] 1
    _:AbstractOperator() = Hadamard[n_wires=1] 2
    _:AbstractOperator() = Hadamard[n_wires=1] 3
    _:AbstractOperator() = Hadamard[n_wires=1] 4
    _:AbstractOperator() = CZ[n_wires=2] 1 2
    _:AbstractOperator() = CZ[n_wires=2] 2 3
    _:AbstractOperator() = CZ[n_wires=2] 3 4
    a:AbstractOperator() = PauliX[n_wires=1] 4
    b:AbstractMeasurement(n_wires=None) = expval_obs a
  in (b,) }
```

```
>>> qml.capture.disable()
>>> qml.workflow.construct_tape(circ)().operations
[GraphStatePrep(Hadamard, CZ)]
```
---------------------------------

**Context:**

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


[sc-88520]

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/7219?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 `92.72727%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 99.64%. Comparing base [(`a2ee4b3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a2ee4b35665588be03b6dc9013c7867c1997175d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`2adb61c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/2adb61c7a59fb9e16271231bcd4a9a59c5dea43d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7219?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/ftqc/primitives.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7219?src=pr&el=tree&filepath=pennylane%2Fftqc%2Fprimitives.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2Z0cWMvcHJpbWl0aXZlcy5weQ==) | 0.00% | [3 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7219?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/ftqc/parametric\_midmeasure.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7219?src=pr&el=tree&filepath=pennylane%2Fftqc%2Fparametric_midmeasure.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2Z0cWMvcGFyYW1ldHJpY19taWRtZWFzdXJlLnB5) | 97.56% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7219?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    #7219      +/-   ##
==========================================
- Coverage   99.65%   99.64%   -0.01%     
==========================================
  Files         526      527       +1     
  Lines       50317    50353      +36     
==========================================
+ Hits        50142    50174      +32     
- Misses        175      179       +4     
```

</details>

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

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
mlxd: > Maybe we want to maintain some more efficient way of storing this in a single primitive in the future, but for now this seems simple and acceptable.

All good with the decomp for now. Since we only care about index and connectivity for the graph (at this stage) we could just replace this entire thing with the sparse matrices, and feed JAX 3 vectors for the sparse format. But, let's revisit that point

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index be132e449..0d3749676 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -127,6 +127,10 @@
 * Alias for Identity (`I`) is now accessible from `qml.ops`.
   [(#7200)](https://github.com/PennyLaneAI/pennylane/pull/7200)
 
+* The `ftqc` module `measure_arbitrary_basis`, `measure_x` and `measure_y` functions
+  can now be captured when program capture is enabled.
+  [(#7219)](https://github.com/PennyLaneAI/pennylane/pull/7219)
+
 * `Operator.num_wires` now defaults to `None` to indicate that the operator can be on
   any number of wires.
   [(#7312)](https://github.com/PennyLaneAI/pennylane/pull/7312)
diff --git a/pennylane/ftqc/__init__.py b/pennylane/ftqc/__init__.py
index 6817153ea..868dc1e88 100644
--- a/pennylane/ftqc/__init__.py
+++ b/pennylane/ftqc/__init__.py
@@ -39,7 +39,7 @@ from .parametric_midmeasure import (
 )
 from .conditional_measure import cond_measure
 from .lattice import Lattice, generate_lattice
-from .graph_state_preparation import GraphStatePrep
+from .graph_state_preparation import GraphStatePrep, make_graph_state
 from .qubit_graph import QubitGraph
 from .utils import QubitMgr
 
@@ -60,4 +60,5 @@ __all__ = [
     "diagonalize_mcms",
     "convert_to_mbqc_gateset",
     "generate_lattice",
+    "make_graph_state",
 ]
diff --git a/pennylane/ftqc/conditional_measure.py b/pennylane/ftqc/conditional_measure.py
index 80ab417a0..ad83aa318 100644
--- a/pennylane/ftqc/conditional_measure.py
+++ b/pennylane/ftqc/conditional_measure.py
@@ -17,9 +17,10 @@ Contains the condition transform.
 from functools import wraps
 from typing import Callable, Union
 
-import pennylane as qml
+from pennylane import capture
 from pennylane.measurements import MeasurementValue, MidMeasureMP
-from pennylane.ops.op_math.condition import CondCallable, Conditional
+from pennylane.ops.op_math.condition import CondCallable, Conditional, cond
+from pennylane.queuing import QueuingManager
 
 
 def cond_measure(
@@ -31,11 +32,6 @@ def cond_measure(
     supplied expression. This conditional expression may involve the results of other mid-circuit
     qubit measurements.
 
-    .. note::
-
-        This function is currently not compatible with :func:`~.qjit`, or with
-        :func:`.pennylane.capture.enabled`.
-
     Args:
         condition (Union[.MeasurementValue, bool]): a conditional expression that may involve a mid-circuit
            measurement value (see :func:`.pennylane.measure`).
@@ -49,7 +45,6 @@ def cond_measure(
         wire, and they must have the same settings for `reset` and `postselection`. The two
         branches can differ only in regard to the measurement basis of the applied measurement.
 
-
     Returns:
         function: A new function that applies the conditional measurements. The returned
         function takes the same input arguments as ``true_fn`` and ``false_fn``.
@@ -94,10 +89,8 @@ def cond_measure(
         While such statements may not result in errors, they may result in
         incorrect behaviour.
     """
-    if qml.capture.enabled():
-        raise NotImplementedError(
-            "The `cond_measure` function is not compatible with program capture"
-        )
+    if capture.enabled():
+        cond(condition, true_fn, false_fn)
 
     if not isinstance(condition, MeasurementValue):
         # The condition is not a mid-circuit measurement - we can simplify immediately
@@ -112,7 +105,7 @@ def cond_measure(
         @wraps(true_fn)
         def wrapper(*args, **kwargs):
 
-            with qml.QueuingManager.stop_recording():
+            with QueuingManager.stop_recording():
                 true_meas_return = true_fn(*args, **kwargs)
                 false_meas_return = false_fn(*args, **kwargs)
 
diff --git a/pennylane/ftqc/graph_state_preparation.py b/pennylane/ftqc/graph_state_preparation.py
index 19a9eb29d..bb5e8b8ba 100644
--- a/pennylane/ftqc/graph_state_preparation.py
+++ b/pennylane/ftqc/graph_state_preparation.py
@@ -25,6 +25,18 @@ from pennylane.wires import Wires
 from .qubit_graph import QubitGraph
 
 
+def make_graph_state(graph, wires, one_qubit_ops=qml.H, two_qubit_ops=qml.CZ):
+    """A program-capture compatible way to create a GraphStatePrep template.
+    We can't capture the graph object in plxpr, so instead, if capture is enabled,
+    we capture the operations generated in computing the decomposition."""
+    if qml.capture.enabled():
+        GraphStatePrep.compute_decomposition(wires, graph, one_qubit_ops, two_qubit_ops)
+    else:
+        GraphStatePrep(
+            graph=graph, wires=wires, one_qubit_ops=one_qubit_ops, two_qubit_ops=two_qubit_ops
+        )
+
+
 class GraphStatePrep(Operation):
     r"""
     Encode a graph state with a single graph operation applied on each qubit, and an entangling
diff --git a/pennylane/ftqc/parametric_midmeasure.py b/pennylane/ftqc/parametric_midmeasure.py
index 0c719d8eb..c3807e599 100644
--- a/pennylane/ftqc/parametric_midmeasure.py
+++ b/pennylane/ftqc/parametric_midmeasure.py
@@ -17,16 +17,60 @@ mid-circuit measurements with a parameterized measurement axis."""
 
 import uuid
 from collections.abc import Hashable
+from functools import lru_cache
 from typing import Iterable, Optional, Union
 
 import numpy as np
 
-import pennylane as qml
+from pennylane import capture
 from pennylane.drawer.tape_mpl import _add_operation_to_drawer
+from pennylane.exceptions import QuantumFunctionError
 from pennylane.measurements.mid_measure import MeasurementValue, MidMeasureMP, measure
+from pennylane.ops.op_math import Conditional, adjoint
+from pennylane.ops.qubit import RX, RY, H, PhaseShift, S
+from pennylane.queuing import QueuingManager
+from pennylane.transforms import transform
 from pennylane.wires import Wires
 
 
+@lru_cache(maxsize=1)
+def _create_parametrized_mid_measure_primitive():
+    """Create a primitive corresponding to a parametrized mid-circuit measurement type.
+
+    Called when using a parametrized mid-circuit measurement, such as
+    :func:`~pennylane.measure_arbitrary_basis`.
+
+    Returns:
+        jax.core.Primitive: A new jax primitive corresponding to a mid-circuit
+        measurement.
+
+    """
+    # pylint: disable=import-outside-toplevel
+    import jax
+
+    from pennylane.capture.custom_primitives import NonInterpPrimitive
+
+    measure_in_basis_p = NonInterpPrimitive("measure_in_basis")
+
+    @measure_in_basis_p.def_impl
+    def _(wires, angle=0.0, plane="ZX", reset=False, postselect=None):
+        return _measure_impl(
+            wires,
+            measurement_class=ParametricMidMeasureMP,
+            angle=angle,
+            plane=plane,
+            reset=reset,
+            postselect=postselect,
+        )
+
+    @measure_in_basis_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 measure_in_basis_p
+
+
 def measure_arbitrary_basis(
     wires: Union[Hashable, Wires],
     angle: float,
@@ -126,8 +170,14 @@ def measure_arbitrary_basis(
         result will return a binary sequence of samples.
         See :ref:`here <mid_circuit_measurements_statistics>` for more details.
     """
+    if len(Wires(wires)) > 1:
+        raise QuantumFunctionError(
+            "Only a single qubit can be measured in the middle of the circuit"
+        )
 
-    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+    if capture.enabled():
+        primitive = _create_parametrized_mid_measure_primitive()
+        return primitive.bind(angle, wires, plane=plane, reset=reset, postselect=postselect)
 
     return _measure_impl(
         wires, ParametricMidMeasureMP, angle=angle, plane=plane, reset=reset, postselect=postselect
@@ -171,8 +221,14 @@ def measure_x(
         QuantumFunctionError: if multiple wires were specified
 
     """
+    if len(Wires(wires)) > 1:
+        raise QuantumFunctionError(
+            "Only a single qubit can be measured in the middle of the circuit"
+        )
 
-    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+    if capture.enabled():
+        primitive = _create_parametrized_mid_measure_primitive()
+        return primitive.bind(0.0, wires, plane="XY", reset=reset, postselect=postselect)
 
     return _measure_impl(wires, XMidMeasureMP, reset=reset, postselect=postselect)
 
@@ -214,8 +270,14 @@ def measure_y(
         QuantumFunctionError: if multiple wires were specified
 
     """
+    if len(Wires(wires)) > 1:
+        raise QuantumFunctionError(
+            "Only a single qubit can be measured in the middle of the circuit"
+        )
 
-    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+    if capture.enabled():
+        primitive = _create_parametrized_mid_measure_primitive()
+        return primitive.bind(np.pi / 2, wires, plane="XY", reset=reset, postselect=postselect)
 
     return _measure_impl(wires, YMidMeasureMP, reset=reset, postselect=postselect)
 
@@ -249,6 +311,7 @@ def measure_z(
         QuantumFunctionError: if multiple wires were specified
 
     """
+    # capture is already handled inside qml.measure
     return measure(wires, reset=reset, postselect=postselect)
 
 
@@ -259,10 +322,6 @@ def _measure_impl(
 ):
     """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())
@@ -341,6 +400,16 @@ class ParametricMidMeasureMP(MidMeasureMP):
 
         return hash(fingerprint)
 
+    # pylint: disable=too-many-positional-arguments
+    @classmethod
+    def _primitive_bind_call(
+        cls, angle=0.0, wires=None, plane="ZX", reset=False, postselect=None, id=None
+    ):
+        wires = () if wires is None else wires
+        return cls._wires_primitive.bind(
+            *wires, angle=angle, plane=plane, reset=reset, postselect=postselect, id=id
+        )
+
     def __repr__(self):
         """Representation of this class."""
         return f"{self._shortname}_{self.plane.lower()}(wires={self.wires.tolist()}, angle={self.angle})"
@@ -353,11 +422,11 @@ class ParametricMidMeasureMP(MidMeasureMP):
     def diagonalizing_gates(self):
         """Decompose to a diagonalizing gate and a standard MCM in the computational basis"""
         if self.plane == "XY":
-            return [qml.PhaseShift(-self.angle, self.wires), qml.H(self.wires)]
+            return [PhaseShift(-self.angle, self.wires), H(self.wires)]
         if self.plane == "ZX":
-            return [qml.RY(-self.angle, self.wires)]
+            return [RY(-self.angle, self.wires)]
         if self.plane == "YZ":
-            return [qml.RX(-self.angle, self.wires)]
+            return [RX(-self.angle, self.wires)]
 
         raise NotImplementedError(
             f"{self.plane} plane not implemented. Available plans are 'XY' 'ZX' and 'YZ'."
@@ -453,7 +522,7 @@ class XMidMeasureMP(ParametricMidMeasureMP):
 
     def diagonalizing_gates(self):
         """Decompose to a diagonalizing gate and a standard MCM in the computational basis"""
-        return [qml.H(self.wires)]
+        return [H(self.wires)]
 
 
 class YMidMeasureMP(ParametricMidMeasureMP):
@@ -520,7 +589,7 @@ class YMidMeasureMP(ParametricMidMeasureMP):
     def diagonalizing_gates(self):
         """Decompose to a diagonalizing gate and a standard MCM in the computational basis"""
         # alternatively we could apply (Z, S) instead of adjoint(S)
-        return [qml.adjoint(qml.S(self.wires)), qml.H(self.wires)]
+        return [adjoint(S(self.wires)), H(self.wires)]
 
 
 @_add_operation_to_drawer.register
@@ -551,7 +620,7 @@ def null_postprocessing(results):
     return results[0]
 
 
-@qml.transform
+@transform
 def diagonalize_mcms(tape):
     """Diagonalize any mid-circuit measurements in a parameterized basis into the computational basis.
 
@@ -657,14 +726,14 @@ def diagonalize_mcms(tape):
             new_operations.extend(diag_gates)
 
             # add computational basis MCM to tape
-            with qml.QueuingManager.stop_recording():
+            with QueuingManager.stop_recording():
                 new_mp = MidMeasureMP(op.wires, reset=op.reset, postselect=op.postselect, id=op.id)
             new_operations.append(new_mp)
 
             # track mapping from original to computational basis MCMs
             mps_mapping[op] = new_mp
 
-        elif isinstance(op, qml.ops.Conditional):
+        elif isinstance(op, Conditional):
 
             # from MCM mapping, map any MCMs in the condition if needed
             mps = [mps_mapping.get(op, op) for op in op.meas_val.measurements]
@@ -681,14 +750,14 @@ def diagonalize_mcms(tape):
                 expr_true = MeasurementValue(mps, processing_fn=true_cond.meas_val.processing_fn)
                 expr_false = MeasurementValue(mps, processing_fn=false_cond.meas_val.processing_fn)
 
-                with qml.QueuingManager.stop_recording():
+                with QueuingManager.stop_recording():
                     diag_gates_true = [
-                        qml.ops.Conditional(expr=expr_true, then_op=gate)
+                        Conditional(expr=expr_true, then_op=gate)
                         for gate in true_cond.diagonalizing_gates()
                     ]
 
                     diag_gates_false = [
-                        qml.ops.Conditional(expr=expr_false, then_op=gate)
+                        Conditional(expr=expr_false, then_op=gate)
                         for gate in false_cond.diagonalizing_gates()
                     ]
 
@@ -707,8 +776,8 @@ def diagonalize_mcms(tape):
                 processing_fn = op.meas_val.processing_fn
                 expr = MeasurementValue(mps, processing_fn=processing_fn)
 
-                with qml.QueuingManager.stop_recording():
-                    new_cond = qml.ops.Conditional(expr=expr, then_op=op.base)
+                with QueuingManager.stop_recording():
+                    new_cond = Conditional(expr=expr, then_op=op.base)
                 new_operations.append(new_cond)
 
         else:
diff --git a/pennylane/ftqc/primitives.py b/pennylane/ftqc/primitives.py
new file mode 100644
index 000000000..fd65d94dc
--- /dev/null
+++ b/pennylane/ftqc/primitives.py
@@ -0,0 +1,25 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This submodule offers all the non-operator/ measurement custom primitives
+created in the ftqc module.
+"""
+
+from .parametric_midmeasure import _create_parametrized_mid_measure_primitive
+
+measure_in_basis_prim = _create_parametrized_mid_measure_primitive()
+
+__all__ = [
+    "measure_in_basis_prim",
+]


Test patch:
diff --git a/tests/ftqc/test_cond_measure.py b/tests/ftqc/test_cond_measure.py
index 1e74ede87..78a3bc169 100644
--- a/tests/ftqc/test_cond_measure.py
+++ b/tests/ftqc/test_cond_measure.py
@@ -178,16 +178,26 @@ class TestValidation:
             cond_measure(m, partial(measure_y, wires=0), partial(measure_x, wires=1))()
 
     @pytest.mark.jax
-    def test_program_capture_raises_error(self):
-        """Test that a clear error is raised when executing with program capture"""
-        try:
-            qml.capture.enable()
-
-            with pytest.raises(NotImplementedError, match="not compatible with program capture"):
-                m = qml.measure(0)
-                cond_measure(m, measure_x, measure_y)(0)
-        finally:
-            qml.capture.disable()
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    def test_program_capture(self):
+        """Test that program capture works as expected with cond_measure"""
+        import jax
+
+        def func():
+            m = qml.measure(0)
+            cond_measure(m, measure_x, measure_y)(0)
+
+        plxpr = jax.make_jaxpr(func)()
+
+        cond_eq = plxpr.eqns[1]
+        assert "cond" in str(cond_eq)
+        cond_branches = cond_eq.params["jaxpr_branches"]
+        assert len(cond_branches) == 2
+        for branch, angle in zip(cond_branches, [0, 1.57]):
+            branch_str = str(branch)
+            assert "measure_in_basis" in branch_str
+            assert "plane=XY" in branch_str
+            assert str(angle) in branch_str
 
 
 class TestWorkflows:
diff --git a/tests/ftqc/test_ftqc_graph_state_prep.py b/tests/ftqc/test_ftqc_graph_state_prep.py
index f936cf26b..b25421782 100644
--- a/tests/ftqc/test_ftqc_graph_state_prep.py
+++ b/tests/ftqc/test_ftqc_graph_state_prep.py
@@ -19,7 +19,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.ftqc import GraphStatePrep, QubitGraph, generate_lattice
+from pennylane.ftqc import GraphStatePrep, QubitGraph, generate_lattice, make_graph_state
 from pennylane.ops.functions import assert_valid
 
 
@@ -289,6 +289,93 @@ class TestGraphStatePrep:
             )
 
 
+class TestMakeGraphState:
+    """Test the helper-function make_graph_state that allows graph state preparation
+    to be capture compatible"""
+
+    def test_make_graph_state_no_capture(self, mocker):
+        """Test that the default behaviour for make_graph_state works as expected"""
+
+        lattice = generate_lattice([2, 2], "square")
+        q_graph = QubitGraph(lattice.graph)
+
+        spy = mocker.spy(qml.ftqc.graph_state_preparation, "GraphStatePrep")
+
+        with qml.queuing.AnnotatedQueue() as q:
+            make_graph_state(q_graph, [0, 1, 2, 3])
+
+        assert len(q) == 1
+        assert isinstance(q.queue[0], GraphStatePrep)
+        spy.assert_called_with(
+            graph=q_graph, one_qubit_ops=qml.H, two_qubit_ops=qml.CZ, wires=[0, 1, 2, 3]
+        )
+
+    def test_passing_operations_no_capture(self, mocker):
+        """Test that gates for the graph state can be specified in make_graph_state,
+        and are passed to GraphStatePrep"""
+
+        lattice = generate_lattice([2, 2], "square")
+        q_graph = QubitGraph(lattice.graph)
+
+        spy = mocker.spy(qml.ftqc.graph_state_preparation, "GraphStatePrep")
+
+        with qml.queuing.AnnotatedQueue() as q:
+            make_graph_state(q_graph, [0, 1, 2, 3], one_qubit_ops=qml.X, two_qubit_ops=qml.CNOT)
+
+        assert len(q) == 1
+        assert isinstance(q.queue[0], GraphStatePrep)
+        spy.assert_called_with(
+            graph=q_graph, one_qubit_ops=qml.X, two_qubit_ops=qml.CNOT, wires=[0, 1, 2, 3]
+        )
+
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    def test_make_graph_state_with_capture(self, mocker):
+        """Test that make_graph_state adds the decomposed graph state to the plxpr"""
+        import jax
+
+        spy = mocker.spy(qml.ftqc.graph_state_preparation.GraphStatePrep, "compute_decomposition")
+
+        lattice = generate_lattice([2, 2], "square")
+        q_graph = QubitGraph(lattice.graph)
+        wires = [0, 1, 2, 3]
+
+        def func():
+            make_graph_state(q_graph, wires)
+
+        plxpr = jax.make_jaxpr(func)()
+
+        # the operator queue looks correct
+        assert len(plxpr.eqns) == 8
+        for eq in plxpr.eqns[:4]:
+            assert "Hadamard" in str(eq)
+        for eq in plxpr.eqns[4:]:
+            assert "CZ" in str(eq)
+
+        # the queue was generated by the expected method
+        spy.assert_called_with([0, 1, 2, 3], q_graph, qml.H, qml.CZ)
+
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    def test_passing_operations_with_capture(self, mocker):
+        """Test that gates for the graph state can be specified in make_graph_state,
+        and are passed to compute_decomposition when capture is enabled"""
+        import jax
+
+        lattice = generate_lattice([2, 2], "square")
+        q_graph = QubitGraph(lattice.graph)
+
+        spy = mocker.spy(qml.ftqc.graph_state_preparation.GraphStatePrep, "compute_decomposition")
+
+        def func():
+            make_graph_state(q_graph, [0, 1, 2, 3], one_qubit_ops=qml.X, two_qubit_ops=qml.CNOT)
+
+        plxpr = jax.make_jaxpr(func)()
+
+        assert len(plxpr.eqns) == 8
+        spy.assert_called_with([0, 1, 2, 3], q_graph, qml.X, qml.CNOT)
+
+
 class TestGraphStateInvariantUnderInternalGraphOrdering:
     """Additional tests of graph-state preparation to ensure that the resulting state is invariant
     under the internal ordering of the nodes and edges in the graph.
diff --git a/tests/ftqc/test_parametric_mid_measure.py b/tests/ftqc/test_parametric_mid_measure.py
index 070dcff83..d1209e6d4 100644
--- a/tests/ftqc/test_parametric_mid_measure.py
+++ b/tests/ftqc/test_parametric_mid_measure.py
@@ -463,6 +463,20 @@ class TestMeasureFunctions:
         ):
             func([0, 1])
 
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    @pytest.mark.parametrize(
+        "func", [partial(measure_arbitrary_basis, angle=-0.8, plane="XY"), measure_x, measure_y]
+    )
+    def test_error_is_raised_if_too_many_wires_capture(self, func):
+        """Test that a QuanutmFunctionError is raised if too many wires are passed when using capture"""
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Only a single qubit can be measured in the middle of the circuit",
+        ):
+            func([0, 1])
+
     @pytest.mark.parametrize("reset", [True, False])
     @pytest.mark.parametrize("postselect", [0, 1, None])
     def test_measure_z_dispatches_to_measure(self, reset, postselect):
@@ -480,6 +494,100 @@ class TestMeasureFunctions:
         assert mp.postselect == postselect
         assert isinstance(mp, MidMeasureMP)
 
+    # pylint: disable=too-many-positional-arguments, too-many-arguments
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    @pytest.mark.parametrize(
+        "meas_func, angle, plane", [(measure_x, 0.0, "XY"), (measure_y, 1.5707, "XY")]
+    )
+    @pytest.mark.parametrize(
+        "wire, reset, postselect", ((2, True, None), (3, False, 0), (0, True, 1))
+    )
+    def test_x_and_y_with_program_capture(self, meas_func, angle, plane, wire, reset, postselect):
+        """Test that the measure_ functions are captured as expected"""
+        import jax
+
+        def circ():
+            m = meas_func(wire, reset=reset, postselect=postselect)
+            qml.cond(m, qml.X, qml.Y)(0)
+            return qml.expval(qml.Z(2))
+
+        plxpr = jax.make_jaxpr(circ)()
+        captured_measurement = str(plxpr.eqns[0])
+
+        # measurement is captured as epxected
+        assert "measure_in_basis" in captured_measurement
+        assert f"plane={plane}" in captured_measurement
+        assert f"postselect={postselect}" in captured_measurement
+        assert f"reset={reset}" in captured_measurement
+
+        # last section is parameters
+        dynamic_params = captured_measurement.rsplit("] ", maxsplit=1)[-1]
+        assert str(angle) in dynamic_params
+        assert str(wire) in dynamic_params
+
+        # measurement value is assigned and passed forward
+        conditional = str(plxpr.eqns[1])
+        assert "cond" in conditional
+        assert captured_measurement[:7] == "a:i64[]"
+        assert "lambda ; a:i64[]" in conditional
+
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    @pytest.mark.parametrize("angle, plane", [(1.23, "XY"), (1.5707, "YZ"), (-0.34, "ZX")])
+    @pytest.mark.parametrize(
+        "wire, reset, postselect", ((2, True, None), (3, False, 0), (0, True, 1))
+    )
+    def test_arbitrary_basis_with_program_capture(self, angle, plane, wire, reset, postselect):
+        """Test that the measure_ functions are captured as expected"""
+        import jax
+        import networkx as nx
+
+        def circ():
+            m = measure_arbitrary_basis(
+                wire, angle=angle, plane=plane, reset=reset, postselect=postselect
+            )
+            qml.cond(m, qml.X, qml.Y)(0)
+            qml.ftqc.make_graph_state(nx.grid_graph((4,)), [0, 1, 2, 3])
+            return qml.expval(qml.Z(2))
+
+        plxpr = jax.make_jaxpr(circ)()
+        captured_measurement = str(plxpr.eqns[0])
+
+        # measurement is captured as epxected
+        assert "measure_in_basis" in captured_measurement
+        assert f"plane={plane}" in captured_measurement
+        assert f"postselect={postselect}" in captured_measurement
+        assert f"reset={reset}" in captured_measurement
+
+        # last section is parameters
+        dynamic_params = captured_measurement.rsplit("] ", maxsplit=1)[-1]
+        assert str(angle) in dynamic_params
+        assert str(wire) in dynamic_params
+
+        # measurement value is assigned and passed forward
+        conditional = str(plxpr.eqns[1])
+        assert "cond" in conditional
+        assert captured_measurement[:7] == "a:i64[]"
+        assert "lambda ; a:i64[]" in conditional
+
+    @pytest.mark.jax
+    @pytest.mark.usefixtures("enable_disable_plxpr")
+    @pytest.mark.parametrize(
+        "func, kwargs",
+        [
+            (measure_x, {"wires": 2}),
+            (measure_y, {"wires": 2}),
+            (measure_arbitrary_basis, {"wires": 2, "angle": 1.2, "plane": "XY"}),
+        ],
+    )
+    def test_calling_functions_with_capture_enabled(self, func, kwargs):
+        """Test that the functions can still be called and return a measurement value
+        with capture enabled."""
+
+        m = func(**kwargs)
+        assert isinstance(m, MeasurementValue)
+
 
 class TestDrawParametricMidMeasure:
     @pytest.mark.matplotlib
diff --git a/tests/ftqc/test_primitives.py b/tests/ftqc/test_primitives.py
new file mode 100644
index 000000000..7668d6622
--- /dev/null
+++ b/tests/ftqc/test_primitives.py
@@ -0,0 +1,28 @@
+# Copyright 2018-2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for the custom capture primitives for the ftqc module"""
+import pytest
+
+pytest.importorskip("jax")
+
+# pylint: disable=wrong-import-position
+from pennylane.capture.custom_primitives import NonInterpPrimitive
+from pennylane.ftqc.primitives import measure_in_basis_prim
+
+
+def test_importing_primitive():
+    """Test that the measure_in_basis_prim is accessible from pennylane.ftqc.primitives.
+    This is mostly for CodeCov."""
+
+    assert isinstance(measure_in_basis_prim, NonInterpPrimitive)