Repository: PennyLaneAI/pennylane
Pull request #6953 (2025-02-11 22:33:51)

Add measure functions for MCMs in X, Y and arbitrary measurement basis
**Note:** I also added docs rendering in this PR (I meant to in the previous one), so we should include the docs rendering for the `ftqc` module in the review process

-----------------

**Context:**
In the existing architecture of PennyLane, if a MCM returns a value that we would like to manipulate/use following the measurement, the MCM needs to be stored in a `MeasurementValue`. For the computational basis, `qml.measure` is the function that creates an MCM, links it to a `MeasurementValue`, and returns the `MeasurementValue` for use.

We've added a base class for creating MCMs in bases other than the computational basis to the ftqc module, but it doesn't have this functionality yet. 

**Description of the Change:**
We add `ftqc.measure_arbitrary_basis`, `ftqc.measure_x`, and `ftqc.measure_y`, analogous to `qml.measure`.
Additionally, we add classes `XMidMeasureMP` and `YMidMeasureMP` (used in the x and y measurement functions), as subclasses of `ParametrizedMidMeasureMP`.

**Benefits:**
- the `measure` functions allow us to easily return a measurement value and use or manipulate is, as in:

```
    m = measure_y(0)
    qml.cond(m, qml.RX)(2.345, 1)
```

- designated classes for X and Y basis measurements will simplify logic in creating MBQC circuits, which rely mostly on X and Y basis MPs
- designated classes for X and Y basis measurements allows them to be labelled clearly as X or Y (instead of XY-plane) measurements will improve circuit visualization


[sc-83347]


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6953?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.59%. Comparing base [(`ed30bf4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ed30bf4721d40daecd876354e7c8552f7a532429?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`97102d3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/97102d3ca4b85381b973ca4e7e6e43b5d9d30039?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    #6953   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         481      481           
  Lines       45783    45841   +58     
=======================================
+ Hits        45598    45656   +58     
  Misses        185      185           
```

</details>

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


Patch:
diff --git a/doc/code/qml_ftqc.rst b/doc/code/qml_ftqc.rst
index 3fdd187b1..fc1858855 100644
--- a/doc/code/qml_ftqc.rst
+++ b/doc/code/qml_ftqc.rst
@@ -1,12 +1,12 @@
-qml.ftqc
-========
-
-.. currentmodule:: pennylane.ftqc
-
-.. warning::
-
-    This module is currently experimental and will not maintain API stability between releases.
-
-.. automodapi:: pennylane.ftqc
-    :no-heading:
-    :include-all-objects:
+qml.ftqc
+========
+
+.. currentmodule:: pennylane.ftqc
+
+.. warning::
+
+    This module is currently experimental and will not maintain API stability between releases.
+
+.. automodapi:: pennylane.ftqc
+    :no-heading:
+    :include-all-objects:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 388202a98..2266d2de0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -197,12 +197,20 @@
   [(#6954)](https://github.com/PennyLaneAI/pennylane/pull/6954)
 
 * A `ParametrizedMidMeasure` class is added to represent a mid-circuit measurement in an arbitrary
-  measurement basis in the XY, YZ or ZX plane. 
+  measurement basis in the XY, YZ or ZX plane. Subclasses `XMidMeasureMP` and `YMidMeasureMP` represent
+  X-basis and Y-basis measurements. These classes are part of the experimental `ftqc` module.
   [(#6938)](https://github.com/PennyLaneAI/pennylane/pull/6938)
+  [(#6953)](https://github.com/PennyLaneAI/pennylane/pull/6953)
 
 * A `diagonalize_mcms` transform is added that diagonalizes any `ParametrizedMidMeasure`, for devices 
   that only natively support mid-circuit measurements in the computational basis.
   [(#6938)](https://github.com/PennyLaneAI/pennylane/pull/6938)
+
+* Measurement functions `measure_x`, `measure_y` and `measure_arbitrary_basis` are added in the experimental `ftqc` module. These functions 
+  apply a mid-circuit measurement and return a `MeasurementValue`. They are analogous to `qml.measure` for 
+  the computational basis, but instead measure in the X-basis, Y-basis, or an arbitrary basis, respectively.
+  Function `qml.ftqc.measure_z` is also added as an alias for `qml.measure`.
+  [(#6953)](https://github.com/PennyLaneAI/pennylane/pull/6953)
   
 * `null.qubit` can now execute jaxpr.
   [(#6924)](https://github.com/PennyLaneAI/pennylane/pull/6924)
diff --git a/pennylane/ftqc/__init__.py b/pennylane/ftqc/__init__.py
index 951555808..abd8a2124 100644
--- a/pennylane/ftqc/__init__.py
+++ b/pennylane/ftqc/__init__.py
@@ -25,10 +25,24 @@ Modules
     :toctree: api
 
 """
-from pennylane import ExperimentalWarning
-from .parametric_midmeasure import ParametricMidMeasureMP, diagonalize_mcms
+from .parametric_midmeasure import (
+    ParametricMidMeasureMP,
+    XMidMeasureMP,
+    YMidMeasureMP,
+    measure_arbitrary_basis,
+    measure_x,
+    measure_y,
+    measure_z,
+    diagonalize_mcms,
+)
 
 __all__ = [
     "ParametricMidMeasureMP",
+    "XMidMeasureMP",
+    "YMidMeasureMP",
+    "measure_arbitrary_basis",
+    "measure_x",
+    "measure_y",
+    "measure_z",
     "diagonalize_mcms",
 ]
diff --git a/pennylane/ftqc/parametric_midmeasure.py b/pennylane/ftqc/parametric_midmeasure.py
index e4eef097e..8ead582bb 100644
--- a/pennylane/ftqc/parametric_midmeasure.py
+++ b/pennylane/ftqc/parametric_midmeasure.py
@@ -15,14 +15,261 @@
 """This module contains the classes and functions for creating and diagonalizing
 mid-circuit measurements with a parameterized measurement axis."""
 
-from typing import Optional
+import uuid
+from collections.abc import Hashable
+from typing import Iterable, Optional, Union
+
+import numpy as np
 
 import pennylane as qml
 from pennylane.drawer.tape_mpl import _add_operation_to_drawer
-from pennylane.measurements.mid_measure import MeasurementValue, MidMeasureMP
+from pennylane.measurements.mid_measure import MeasurementValue, MidMeasureMP, measure
 from pennylane.wires import Wires
 
 
+def measure_arbitrary_basis(
+    wires: Union[Hashable, Wires],
+    angle: float,
+    plane: str,
+    reset: bool = False,
+    postselect: Optional[int] = None,
+):
+    r"""Perform a mid-circuit measurement in the basis defined by the plane and angle on the
+    supplied qubit.
+
+    The measurements are performed using the 0, 1 convention rather than the ±1 convention.
+
+    If a device doesn't support mid-circuit measurements natively, then the desired ``mcm_method`` for
+    executing mid-circuit measurements should be passed to the QNode.
+
+    .. warning::
+        Measurements should be diagonalized before execution for any device that only natively supports
+        mid-circuit measurements in the computational basis. To diagonalize, the :func:`diagonalize_mcms <pennylane.ftqc.diagonalize_mcms>`
+        transform can be applied.
+
+        Skipping diagonalization for a circuit containing parametric mid-circuit measurements may result
+        in a completed execution with incorrect results.
+
+    Args:
+        wires (Wires): The wire to measure.
+        angle (float): The angle of rotation defining the axis, specified in radians.
+        plane (str): The plane the measurement basis lies in. Options are "XY", "YZ" and "ZX"
+        reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
+            state after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
+
+    Returns:
+        MeasurementValue: The mid-circuit measurement result linked to the created ``MidMeasureMP``.
+
+    Raises:
+        QuantumFunctionError: if multiple wires were specified
+
+    .. note::
+        Reset behaviour will depend on the execution method for mid-circuit measurements,
+        and may not work for all configurations.
+
+    **Example:**
+
+    .. code-block:: python3
+
+        import pennylane as qml
+        from pennylane.ftqc import diagonalize_mcms, measure_arbitrary_basis
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method="tree-traversal")
+        def func(x, y):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            m_0 = measure_arbitrary_basis(1, angle=np.pi/3, plane="XY")
+
+            qml.cond(m_0, qml.RY)(y, wires=0)
+            return qml.probs(wires=[0])
+
+    Executing this QNode:
+
+    >>> pars = np.array([0.643, 0.246])
+    >>> func(*pars)
+    array([0.91237915, 0.08762085])
+
+    .. details::
+        :title: Plane and angle
+
+        The plane and angle are related to the axis of measurement by the following formulas:
+
+        .. math:: M_{XY}(\phi) =\frac{1}{\sqrt{2}} (|0\rangle + e^{i\phi} |1\rangle),
+
+        .. math:: M_{YZ}(\theta) =\cos{\frac{\theta}{2}}|0\rangle + i \sin{\frac{\theta}{2}} |1\rangle,\text{ and}
+
+        .. math:: M_{ZX}(\theta) = \cos{\frac{\theta}{2}}|0\rangle + \sin{\frac{\theta}{2}} |1\rangle
+
+        where, in terms of `spherical coordinates <https://en.wikipedia.org/wiki/Spherical_coordinate_system>`_ in
+        the physics convention, the angles :math:`\phi` and :math:`\theta` are the azimuthal and polar angles,
+        respectively.
+
+    .. details::
+        :title: Using mid-circuit measurements
+
+        Measurement outcomes can be used to conditionally apply operations, and measurement
+        statistics can be gathered and returned by a quantum function. Measurement outcomes can
+        also be manipulated using arithmetic operators like ``+``, ``-``, ``*``, ``/``, etc. with
+        other mid-circuit measurements or scalars.
+
+        See the :func:`qml.measure <pennylane.measurements.measure>` function
+        for details on the available arithmetic operators for mid-circuit measurement results.
+
+        Mid-circuit measurement results can be processed with the usual measurement functions such as
+        :func:`~.expval`. For QNodes with finite shots, :func:`~.sample` applied to a mid-circuit measurement
+        result will return a binary sequence of samples.
+        See :ref:`here <mid_circuit_measurements_statistics>` for more details.
+    """
+
+    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+
+    return _measure_impl(
+        wires, ParametricMidMeasureMP, angle=angle, plane=plane, reset=reset, postselect=postselect
+    )
+
+
+def measure_x(
+    wires: Union[Hashable, Wires],
+    reset: bool = False,
+    postselect: Optional[int] = None,
+):
+    r"""Perform a mid-circuit measurement in the X basis. The measurements are performed using the 0, 1
+    convention rather than the ±1 convention.
+
+    For more details on the results of mid-circuit measurements and how to use them,
+    see :func:`qml.measure <pennylane.measure>`.
+
+    For more details on mid-circuit measurements in an arbitrary basis (besides the computational basis),
+    see :func:`measure_arbitrary_basis <pennylane.ftqc.measure_arbitrary_basis>`.
+
+    .. warning::
+        Measurements should be diagonalized before execution for any device that only natively supports
+        mid-circuit measurements in the computational basis. To diagonalize, the :func:`diagonalize_mcms <pennylane.ftqc.diagonalize_mcms>`
+        transform can be applied.
+
+        Skipping diagonalization for a circuit containing parametric mid-circuit measurements may result
+        in a completed execution with incorrect results.
+
+    Args:
+        wires (Wires): The wire to measure.
+        reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
+            state after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
+
+    Returns:
+        MeasurementValue: The mid-circuit measurement result linked to the created ``MidMeasureMP``.
+
+    Raises:
+        QuantumFunctionError: if multiple wires were specified
+
+    """
+
+    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+
+    return _measure_impl(wires, XMidMeasureMP, reset=reset, postselect=postselect)
+
+
+def measure_y(
+    wires: Union[Hashable, Wires],
+    reset: bool = False,
+    postselect: Optional[int] = None,
+):
+    r"""Perform a mid-circuit measurement in the Y basis. The measurements are performed using the 0, 1
+    convention rather than the ±1 convention.
+
+    For more details on the results of mid-circuit measurements and how to use them,
+    see :func:`qml.measure <pennylane.measure>`.
+
+    For more details on mid-circuit measurements in an arbitrary basis (besides the computational basis),
+    see :func:`measure_arbitrary_basis <pennylane.ftqc.measure_arbitrary_basis>`.
+
+    .. warning::
+        Measurements should be diagonalized before execution for any device that only natively supports
+        mid-circuit measurements in the computational basis. To diagonalize, the :func:`diagonalize_mcms <pennylane.ftqc.diagonalize_mcms>`
+        transform can be applied.
+
+        Skipping diagonalization for a circuit containing parametric mid-circuit measurements may result
+        in a completed execution with incorrect results.
+
+    Args:
+        wires (Wires): The wire to measure.
+        reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
+            state after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
+
+    Returns:
+        MeasurementValue: The mid-circuit measurement result linked to the created ``MidMeasureMP``.
+
+    Raises:
+        QuantumFunctionError: if multiple wires were specified
+
+    """
+
+    # ToDo: if capture is enabled, create and bind primitive here and return primitive instead (subsequent PR)
+
+    return _measure_impl(wires, YMidMeasureMP, reset=reset, postselect=postselect)
+
+
+def measure_z(
+    wires: Union[Hashable, Wires],
+    reset: bool = False,
+    postselect: Optional[int] = None,
+):
+    r"""Perform a mid-circuit measurement in the Z basis. The measurements are performed using the 0, 1
+    convention rather than the ±1 convention.
+
+    .. note::
+        This function dispatches to :func:`qml.measure <pennylane.measure>`
+
+    For more details on the results of mid-circuit measurements and how to use them,
+    see :func:`qml.measure <pennylane.measure>`.
+
+    Args:
+        wires (Wires): The wire to measure.
+        reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
+            state after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
+
+    Returns:
+        MeasurementValue: The mid-circuit measurement result linked to the created ``MidMeasureMP``.
+
+    Raises:
+        QuantumFunctionError: if multiple wires were specified
+
+    """
+    return measure(wires, reset=reset, postselect=postselect)
+
+
+def _measure_impl(
+    wires: Union[Hashable, Wires],
+    measurement_class=MidMeasureMP,
+    **kwargs,
+):
+    """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())
+    mp = measurement_class(wires=wires, id=measurement_id, **kwargs)
+    return MeasurementValue([mp], processing_fn=lambda v: v)
+
+
 class ParametricMidMeasureMP(MidMeasureMP):
     """Parametric mid-circuit measurement. The basis for the measurement is parametrized by
     a plane ("XY", "YZ" or "ZX"), and an angle within the plane.
@@ -116,16 +363,18 @@ class ParametricMidMeasureMP(MidMeasureMP):
             f"{self.plane} plane not implemented. Available plans are 'XY' 'ZX' and 'YZ'."
         )
 
-    def label(self, decimals=None, base_label=None, cache=None):  # pylint: disable=unused-argument
+    def label(
+        self, decimals: int = None, base_label: Iterable[str] = None, cache: dict = None
+    ):  # pylint: disable=unused-argument
         r"""How the mid-circuit measurement is represented in diagrams and drawings.
 
         Args:
-            decimals=None (Int): If ``None``, no parameters are included. Else,
-                how to round the parameters.
-            base_label=None (Iterable[str]): overwrite the non-parameter component of the label.
+            decimals: If ``None``, no parameters are included. Else,
+                how to round the parameters. Defaults to None.
+            base_label: overwrite the non-parameter component of the label.
                 Required to match general call signature. Not used.
-            cache=None (dict): dictionary that carries information between label calls
-                in the same drawing. Required to match general call signature. Not used.
+            cache: dictionary that carries information between label calls in the
+                same drawing. Required to match general call signature. Not used.
 
         Returns:
             str: label to use in drawings
@@ -146,9 +395,142 @@ class ParametricMidMeasureMP(MidMeasureMP):
         return _label
 
 
+class XMidMeasureMP(ParametricMidMeasureMP):
+    """A subclass of ParametricMidMeasureMP that uses the X measurement basis
+    (angle=0, plane="XY"). For labels and visualizations, this will be represented
+    as a X measurement. It is otherwise identical to the parent class."""
+
+    _shortname = "measure_x"
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        wires: Optional[Wires],
+        reset: Optional[bool] = False,
+        postselect: Optional[int] = None,
+        id: Optional[str] = None,
+    ):
+        super().__init__(
+            wires=Wires(wires), angle=0, plane="XY", reset=reset, postselect=postselect, id=id
+        )
+
+    def _flatten(self):
+        metadata = (
+            ("wires", self.raw_wires),
+            ("reset", self.reset),
+            ("id", self.id),
+        )
+        return (None, None), metadata
+
+    def __repr__(self):
+        """Representation of this class."""
+        return f"{self._shortname}(wires={self.wires.tolist()})"
+
+    def label(
+        self, decimals: int = None, base_label: Iterable[str] = None, cache: dict = None
+    ):  # pylint: disable=unused-argument
+        r"""How the mid-circuit measurement is represented in diagrams and drawings.
+
+        Args:
+            decimals: If ``None``, no parameters are included. Else, how to round
+                the parameters. Required to match general call signature. Not used.
+            base_label: overwrite the non-parameter component of the label.
+                Required to match general call signature. Not used.
+            cache: dictionary that carries information between label calls in the
+                same drawing. Required to match general call signature. Not used.
+
+        Returns:
+            str: label to use in drawings
+        """
+        _label = "┤↗ˣ"
+
+        if self.postselect is not None:
+            _label += "₁" if self.postselect == 1 else "₀"
+
+        _label += "├" if not self.reset else "│  │0⟩"
+
+        return _label
+
+    def diagonalizing_gates(self):
+        """Decompose to a diagonalizing gate and a standard MCM in the computational basis"""
+        return [qml.H(self.wires)]
+
+
+class YMidMeasureMP(ParametricMidMeasureMP):
+    """A subclass of ParametricMidMeasureMP that uses the Y measurement basis
+    (angle=pi/2, plane="XY"). For labels and visualizations, this will be represented
+    as a Y measurement. It is otherwise identical to the parent class."""
+
+    _shortname = "measure_y"
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        wires: Optional[Wires],
+        reset: Optional[bool] = False,
+        postselect: Optional[int] = None,
+        id: Optional[str] = None,
+    ):
+        super().__init__(
+            wires=Wires(wires),
+            angle=np.pi / 2,
+            plane="XY",
+            reset=reset,
+            postselect=postselect,
+            id=id,
+        )
+
+    def _flatten(self):
+        metadata = (
+            ("wires", self.raw_wires),
+            ("reset", self.reset),
+            ("id", self.id),
+        )
+        return (None, None), metadata
+
+    def __repr__(self):
+        """Representation of this class."""
+        return f"{self._shortname}(wires={self.wires.tolist()})"
+
+    def label(
+        self, decimals: int = None, base_label: str = None, cache: dict = None
+    ):  # pylint: disable=unused-argument
+        r"""How the mid-circuit measurement is represented in diagrams and drawings.
+
+        Args:
+            decimals: If ``None``, no parameters are included. Else, how to round
+                the parameters. Required to match general call signature. Not used.
+            base_label: overwrite the non-parameter component of the label.
+                Required to match general call signature. Not used.
+            cache: dictionary that carries information between label calls in the
+                same drawing. Required to match general call signature. Not used.
+
+        Returns:
+            str: label to use in drawings
+        """
+        _label = "┤↗ʸ"
+
+        if self.postselect is not None:
+            _label += "₁" if self.postselect == 1 else "₀"
+
+        _label += "├" if not self.reset else "│  │0⟩"
+
+        return _label
+
+    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)]
+
+
 @_add_operation_to_drawer.register
 def _(op: ParametricMidMeasureMP, drawer, layer, _):
-    text = op.plane
+    if isinstance(op, XMidMeasureMP):
+        text = "X"
+    elif isinstance(op, YMidMeasureMP):
+        text = "Y"
+    else:
+        text = op.plane
     drawer.measure(layer, op.wires[0], text=text)  # assume one wire
 
     if op.reset:
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 7e194bd78..644820d19 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -227,7 +227,7 @@ def _measure_impl(
         )
 
     # Create a UUID and a map between MP and MV to support serialization
-    measurement_id = str(uuid.uuid4())[:8]
+    measurement_id = str(uuid.uuid4())
     mp = MidMeasureMP(wires=wires, reset=reset, postselect=postselect, id=measurement_id)
     return MeasurementValue([mp], processing_fn=lambda v: v)
 


Test patch:
diff --git a/tests/ftqc/test_parametric_mid_measure.py b/tests/ftqc/test_parametric_mid_measure.py
index bb9cf2710..b528a1d37 100644
--- a/tests/ftqc/test_parametric_mid_measure.py
+++ b/tests/ftqc/test_parametric_mid_measure.py
@@ -20,7 +20,15 @@ import pytest
 
 import pennylane as qml
 from pennylane.devices.qubit import measure as apply_qubit_measurement
-from pennylane.ftqc import ParametricMidMeasureMP, diagonalize_mcms
+from pennylane.ftqc import (
+    ParametricMidMeasureMP,
+    XMidMeasureMP,
+    YMidMeasureMP,
+    diagonalize_mcms,
+    measure_arbitrary_basis,
+    measure_x,
+    measure_y,
+)
 from pennylane.measurements import MeasurementValue, MidMeasureMP
 from pennylane.wires import Wires
 
@@ -28,7 +36,7 @@ from pennylane.wires import Wires
 
 
 class TestParametricMidMeasure:
-    """Tests for the measure function"""
+    """Tests for the parametric mid-circuit measurement class in an arbitrary basis"""
 
     def test_hash(self):
         """Test that the hash for `ParametricMidMeasureMP` is defined correctly."""
@@ -201,64 +209,364 @@ class TestParametricMidMeasure:
             mp.diagonalizing_gates()
 
 
-class TestDrawParametricMidMeasure:
+class TestMidMeasureXAndY:
+    """Tests for the mid-circuit measurement class in the X and Y basis"""
+
+    @pytest.mark.parametrize("mp_class, angle", [(XMidMeasureMP, 0), (YMidMeasureMP, np.pi / 2)])
+    def test_attributes(self, mp_class, angle):
+        """Test that the XMidMeasure and YMidMeasure have the expected attributes for a ParametricMidMeasureMP"""
+
+        mp = mp_class(Wires(0))
+
+        assert isinstance(mp, ParametricMidMeasureMP)
+        assert mp.angle == angle
+        assert mp.plane == "XY"
+
+    @pytest.mark.parametrize("mp_class, angle", [(XMidMeasureMP, 0), (YMidMeasureMP, np.pi / 2)])
+    def test_hash(self, mp_class, angle):
+        """Test that the hash for XMidMeasureMP and YMidMeasureMP are defined correctly."""
+        m1 = mp_class(Wires(0), id="m1")
+        m2 = mp_class(Wires(2), id="m1")
+        m3 = mp_class(Wires(0), id="m2")
+        m4 = ParametricMidMeasureMP(Wires(0), angle=angle, id="m1", plane="XY")
+        m5 = mp_class(Wires(0), id="m1")
+
+        assert m1.hash != m2.hash
+        assert m1.hash != m3.hash
+        assert m1.hash != m4.hash
+        assert m1.hash == m5.hash
+
+    @pytest.mark.parametrize("mp_class", [XMidMeasureMP, YMidMeasureMP])
+    def test_flatten_unflatten(self, mp_class):
+        """Test that we can flatten and unflatten the ParametricMidMeasureMP"""
+
+        op = mp_class(Wires(0), id="m1")
+        data, metadata = op._flatten()  # pylint: disable = protected-access
+
+        assert hash(metadata)  # metadata must be hashable
+
+        unflattened_op = mp_class._unflatten(data, metadata)  # pylint: disable = protected-access
+        assert op.hash == unflattened_op.hash
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("mp_class", [XMidMeasureMP, YMidMeasureMP])
+    def test_flatten_unflatten_jax(self, mp_class):
+        """Test that jax.tree_util can flatten and unflatten the ParametricMidMeasureMP"""
+
+        import jax
+
+        op = mp_class(Wires(0), id="m1")
+
+        leaves, struct = jax.tree_util.tree_flatten(op)
+        unflattened_op = jax.tree_util.tree_unflatten(struct, leaves)
 
+        assert op.hash == unflattened_op.hash
+
+    @pytest.mark.parametrize(
+        "wire, expected",
+        [(0, "measure_x(wires=[0])"), (1, "measure_x(wires=[1])"), ("a", "measure_x(wires=['a'])")],
+    )
+    def test_repr_x(self, wire, expected):
+        """Test the repr for XMidMeasureMP is correct"""
+        mp = XMidMeasureMP(wires=Wires([wire]))
+        assert repr(mp) == expected
+
+    @pytest.mark.parametrize(
+        "wire, expected",
+        [(0, "measure_y(wires=[0])"), (1, "measure_y(wires=[1])"), ("a", "measure_y(wires=['a'])")],
+    )
+    def test_repr_y(self, wire, expected):
+        """Test the repr for YMidMeasureMP is correct"""
+        mp = YMidMeasureMP(wires=Wires([wire]))
+        assert repr(mp) == expected
+
+    @pytest.mark.parametrize(
+        "postselect, reset, expected",
+        [
+            (None, False, "┤↗ˣ├"),
+            (None, True, "┤↗ˣ│  │0⟩"),
+            (0, False, "┤↗ˣ₀├"),
+            (0, True, "┤↗ˣ₀│  │0⟩"),
+            (1, False, "┤↗ˣ₁├"),
+            (1, True, "┤↗ˣ₁│  │0⟩"),
+        ],
+    )
+    def test_label_x(self, postselect, reset, expected):
+        """Test that the label for a XMidMeasureMP is correct"""
+        mp = XMidMeasureMP(Wires([0]), postselect=postselect, reset=reset)
+        label = mp.label()
+        assert label == expected
+
+    @pytest.mark.parametrize(
+        "postselect, reset, expected",
+        [
+            (None, False, "┤↗ʸ├"),
+            (None, True, "┤↗ʸ│  │0⟩"),
+            (0, False, "┤↗ʸ₀├"),
+            (0, True, "┤↗ʸ₀│  │0⟩"),
+            (1, False, "┤↗ʸ₁├"),
+            (1, True, "┤↗ʸ₁│  │0⟩"),
+        ],
+    )
+    def test_label_y(self, postselect, reset, expected):
+        """Test that the label for a YMidMeasureMP is correct"""
+        mp = YMidMeasureMP(Wires([0]), postselect=postselect, reset=reset)
+        label = mp.label()
+        assert label == expected
+
+    def test_diagonalizing_gates_x(self):
+        """Test that diagonalizing a XMidMeasureMP and measuring in the computational
+        basis corresponds to the expected observable"""
+
+        dev = qml.device("default.qubit")
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method="tree-traversal")
+        def circ(state):
+            qml.StatePrep(state, wires=0)
+            mp = XMidMeasureMP([0])
+            assert mp.has_diagonalizing_gates
+            return qml.expval(qml.Z(0))
+
+        rng = np.random.default_rng(seed=111)
+        input_state = rng.random(2) + 1j * rng.random(2)
+        input_state = input_state / np.linalg.norm(input_state)
+
+        res = circ(input_state)
+
+        expected_res = apply_qubit_measurement(qml.expval(qml.X(0)), input_state)
+
+        assert np.allclose(res, expected_res)
+        assert XMidMeasureMP([0]).diagonalizing_gates() == [qml.H(0)]
+
+    def test_diagonalizing_gates_y(self):
+        """Test that diagonalizing a YMidMeasureMP and measuring in the computational
+        basis corresponds to the expected observable"""
+
+        dev = qml.device("default.qubit")
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method="tree-traversal")
+        def circ(state):
+            qml.StatePrep(state, wires=0)
+            mp = YMidMeasureMP([0])
+            assert mp.has_diagonalizing_gates
+            return qml.expval(qml.Z(0))
+
+        rng = np.random.default_rng(seed=111)
+        input_state = rng.random(2) + 1j * rng.random(2)
+        input_state = input_state / np.linalg.norm(input_state)
+
+        res = circ(input_state)
+
+        expected_res = apply_qubit_measurement(qml.expval(qml.Y(0)), input_state)
+
+        assert np.allclose(res, expected_res)
+        assert YMidMeasureMP([0]).diagonalizing_gates() == [qml.adjoint(qml.S(0)), qml.H(0)]
+
+
+class TestMeasureFunctions:
+    """Test that the measure functions (measure_arbitrary_basis, measure_x, and
+    measure_y) behave as expected"""
+
+    @pytest.mark.parametrize("wire", [0, 2])
+    @pytest.mark.parametrize("angle", [1.2, -0.4])
+    @pytest.mark.parametrize("plane", ["XY", "YZ", "XZ"])
+    @pytest.mark.parametrize("reset", [True, False])
+    @pytest.mark.parametrize("postselect", [None, 0, 1])
+    def test_measure_arbitrary_basis(  # pylint: disable=too-many-arguments, too-many-positional-arguments
+        self, wire, angle, plane, reset, postselect
+    ):
+        """Test that measure_arbitrary_basis queues the expected ParametricMidMeasureMP
+        and returns a linked MeasurementValue"""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            m = measure_arbitrary_basis(
+                wire, angle=angle, plane=plane, reset=reset, postselect=postselect
+            )
+
+        # a single op is queued
+        tape = qml.tape.QuantumScript.from_queue(q)
+        assert len(tape.operations) == 1
+        op = tape.operations[0]
+
+        # that op is stored as the measurement on the returned MeasurementValue
+        assert isinstance(m, MeasurementValue)
+        assert m.measurements == [op]
+
+        # the op is the expected ParametricMidMeasureMP
+        assert isinstance(op, ParametricMidMeasureMP)
+        assert op.angle == angle
+        assert op.wires == Wires([wire])
+        assert op.plane == plane
+        assert op.reset == reset
+        assert op.postselect == postselect
+
+    @pytest.mark.parametrize("wire", [0, 2])
+    @pytest.mark.parametrize("reset", [True, False])
+    @pytest.mark.parametrize("postselect", [None, 0, 1])
+    def test_measure_x(self, wire, reset, postselect):
+        """Test that measure_arbitrary_basis queues the expected XMidMeasureMP
+        and returns a linked MeasurementValue"""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            m = measure_x(wire, reset=reset, postselect=postselect)
+
+        # a single op is queued
+        tape = qml.tape.QuantumScript.from_queue(q)
+        assert len(tape.operations) == 1
+        op = tape.operations[0]
+
+        # that op is stored as the measurement on the returned MeasurementValue
+        assert isinstance(m, MeasurementValue)
+        assert m.measurements == [op]
+
+        # the op is the expected ParametricMidMeasureMP
+        assert isinstance(op, XMidMeasureMP)
+        assert op.wires == Wires([wire])
+        assert op.reset == reset
+        assert op.postselect == postselect
+
+    @pytest.mark.parametrize("wire", [0, 2])
+    @pytest.mark.parametrize("reset", [True, False])
+    @pytest.mark.parametrize("postselect", [None, 0, 1])
+    def test_measure_y(self, wire, reset, postselect):
+        """Test that measure_arbitrary_basis queues the expected YMidMeasureMP
+        and returns a linked MeasurementValue"""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            m = measure_y(wire, reset=reset, postselect=postselect)
+
+        # a single op is queued
+        tape = qml.tape.QuantumScript.from_queue(q)
+        assert len(tape.operations) == 1
+        op = tape.operations[0]
+
+        # that op is stored as the measurement on the returned MeasurementValue
+        assert isinstance(m, MeasurementValue)
+        assert m.measurements == [op]
+
+        # the op is the expected ParametricMidMeasureMP
+        assert isinstance(op, YMidMeasureMP)
+        assert op.wires == Wires([wire])
+        assert op.reset == reset
+        assert op.postselect == postselect
+
+    @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(self, func):
+        """Test that a QuanutmFunctionError is raised if too many wires are passed"""
+
+        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):
+        """Test that the measurement in the computational basis calls the standard
+        PennyLane MCM class, and passes through the arguments"""
+
+        m = qml.ftqc.measure_z(0, reset=reset, postselect=postselect)
+
+        assert isinstance(m, MeasurementValue)
+
+        assert len(m.measurements) == 1
+
+        mp = m.measurements[0]
+        assert mp.reset == reset
+        assert mp.postselect == postselect
+        assert isinstance(mp, MidMeasureMP)
+
+
+class TestDrawParametricMidMeasure:
     @pytest.mark.matplotlib
-    def test_draw_mpl_label(self):
+    @pytest.mark.parametrize(
+        "mp_class, expected_label",
+        [(ParametricMidMeasureMP, "XY"), (XMidMeasureMP, "X"), (YMidMeasureMP, "Y")],
+    )
+    def test_draw_mpl_label(self, mp_class, expected_label):
         """Test that the plane label is added to the MCM in a mpl drawing"""
 
         from matplotlib import pyplot as plt
 
         dev = qml.device("default.qubit", wires=2)
 
+        if mp_class == ParametricMidMeasureMP:
+            args = {"wires": Wires([0]), "angle": np.pi / 4, "plane": "XY"}
+        else:
+            args = {"wires": Wires([0])}
+
         @qml.qnode(dev)
         def circ():
-            ParametricMidMeasureMP(Wires([0]), angle=np.pi / 4, plane="XY")
+            mp_class(**args)
             return qml.expval(qml.Z(0))
 
         _, ax = qml.draw_mpl(circ)()
-
         assert len(ax.texts) == 2  # one wire label, 1 box label on the MCM
-
         assert ax.texts[0].get_text() == "0"
-        assert ax.texts[1].get_text() == "XY"
+        assert ax.texts[1].get_text() == expected_label
 
         plt.close()
 
+    @pytest.mark.parametrize(
+        "mp_class, expected_label",
+        [(ParametricMidMeasureMP, "XY"), (XMidMeasureMP, "X"), (YMidMeasureMP, "Y")],
+    )
     @pytest.mark.matplotlib
-    def test_draw_mpl_reset(self):
+    def test_draw_mpl_reset(self, mp_class, expected_label):
         """Test that the reset is added after the MCM as expected in a mpl drawing"""
 
         from matplotlib import pyplot as plt
 
         dev = qml.device("default.qubit", wires=2)
 
+        if mp_class == ParametricMidMeasureMP:
+            args = {"wires": Wires([0]), "angle": np.pi / 4, "plane": "XY"}
+        else:
+            args = {"wires": Wires([0])}
+
         @qml.qnode(dev)
         def circ():
-            ParametricMidMeasureMP(Wires([0]), angle=np.pi / 4, plane="XY", reset=True)
+            mp_class(**args, reset=True)
             return qml.expval(qml.Z(0))
 
         _, ax = qml.draw_mpl(circ)()
-
         assert len(ax.texts) == 3  # one wire label, 1 box label on the MCM, one reset box
 
         assert ax.texts[0].get_text() == "0"
-        assert ax.texts[1].get_text() == "XY"
+        assert ax.texts[1].get_text() == expected_label
         assert ax.texts[2].get_text() == "|0⟩"
 
         plt.close()
 
-    def test_text_drawer(self):
+    @pytest.mark.parametrize(
+        "mp_class, expected_string",
+        [
+            (ParametricMidMeasureMP, "0: ──┤↗ˣʸ(0.79)├─┤  <Z>"),
+            (XMidMeasureMP, "0: ──┤↗ˣ├─┤  <Z>"),
+            (YMidMeasureMP, "0: ──┤↗ʸ├─┤  <Z>"),
+        ],
+    )
+    def test_text_drawer(self, mp_class, expected_string):
         """Test that the text drawer works as expected"""
 
         dev = qml.device("default.qubit", wires=2)
 
+        if mp_class == ParametricMidMeasureMP:
+            args = {"wires": Wires([0]), "angle": np.pi / 4, "plane": "XY"}
+        else:
+            args = {"wires": Wires([0])}
+
         @qml.qnode(dev, mcm_method="tree-traversal")
         def circ():
-            ParametricMidMeasureMP(Wires([0]), angle=np.pi / 4, plane="XY")
+            mp_class(**args)
             return qml.expval(qml.Z(0))
 
-        assert qml.draw(circ)() == "0: ──┤↗ˣʸ(0.79)├─┤  <Z>"
+        assert qml.draw(circ)() == expected_string
 
 
 class TestDiagonalizeMCMs:
@@ -294,6 +602,28 @@ class TestDiagonalizeMCMs:
         assert not isinstance(new_tape.operations[3], ParametricMidMeasureMP)
         assert new_tape.operations[3].wires == tape.operations[1].wires
 
+    @pytest.mark.parametrize("postselect", [None, 0, 1])
+    def test_diagonalize_mcm_transform_preserves_postselect(self, postselect):
+        """Test that the diagonalize_mcm transform preserves postselet on a diagonalized MCM"""
+
+        op = ParametricMidMeasureMP(Wires([0]), angle=np.pi, plane="XY", postselect=postselect)
+        tape = qml.tape.QuantumScript([op])
+        (new_tape,), _ = diagonalize_mcms(tape)
+
+        assert isinstance(new_tape.operations[-1], MidMeasureMP)
+        assert new_tape.operations[-1].postselect == postselect
+
+    @pytest.mark.parametrize("reset", [True, False])
+    def test_diagonalize_mcm_transform_preserves_reset(self, reset):
+        """Test that the diagonalize_mcm transform preserves reset on a diagonalized MCM"""
+
+        op = ParametricMidMeasureMP(Wires([0]), angle=np.pi, plane="XY", reset=reset)
+        tape = qml.tape.QuantumScript([op])
+        (new_tape,), _ = diagonalize_mcms(tape)
+
+        assert isinstance(new_tape.operations[-1], MidMeasureMP)
+        assert new_tape.operations[-1].reset == reset
+
     def test_diagonalize_mcm_in_cond(self):
         """Test that the diagonalize_mcm transform works as expected on a tape
         containing a conditional with a ParametricMidMeasureMP"""
@@ -445,8 +775,16 @@ class TestDiagonalizeMCMs:
 
 
 class TestWorkflows:
+    @pytest.mark.parametrize(
+        "rot_gate, measurement_fn",
+        [
+            (partial(qml.RX, 2.345), measure_y),
+            (partial(qml.RY, -2.345), measure_x),
+            (partial(qml.RX, 2.345), partial(measure_arbitrary_basis, angle=np.pi / 2, plane="XY")),
+        ],
+    )
     @pytest.mark.parametrize("mcm_method, shots", [("tree-traversal", None), ("one-shot", 10000)])
-    def test_execution(self, mcm_method, shots):
+    def test_simple_execution(self, rot_gate, measurement_fn, mcm_method, shots):
         """Test that we can execute a QNode with a ParametricMidMeasureMP and produce
         an accurate result"""
 
@@ -455,11 +793,13 @@ class TestWorkflows:
         @diagonalize_mcms
         @qml.qnode(dev, mcm_method=mcm_method)
         def circ():
-            qml.RX(2.345, 0)
-            ParametricMidMeasureMP(0, angle=np.pi / 2, plane="XY")
+            rot_gate(0)
+            measurement_fn(0)
             return qml.expval(qml.Z(0))
 
         if shots:
+            # the result is on the order of 1 (-0.7), and an uncertainty ~1.5-2 orders of magnitude
+            # smaller than the result is sufficiently accurate for a shots-based measurement
             assert np.isclose(circ(), -np.sin(2.345), atol=0.03)
         else:
             assert np.isclose(circ(), -np.sin(2.345))
@@ -483,3 +823,37 @@ class TestWorkflows:
             return qml.expval(qml.Z(0))
 
         assert np.isclose(circ(), -np.sin(2.345))
+
+    @pytest.mark.parametrize(
+        "rot_gate, measurement_fn",
+        [
+            (partial(qml.RX, np.pi / 2), measure_y),
+            (partial(qml.RY, -np.pi / 2), measure_x),
+            (
+                partial(qml.RX, np.pi / 2),
+                partial(measure_arbitrary_basis, angle=np.pi / 2, plane="XY"),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("mcm_method, shots", [("tree-traversal", None), ("one-shot", 10000)])
+    def test_condition_of_cond(self, rot_gate, measurement_fn, mcm_method, shots):
+        """Test that we can execute a QNode with a ParametricMidMeasureMP as the condition of a conditional,
+        and produce an accurate result"""
+
+        dev = qml.device("default.qubit", shots=shots)
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method=mcm_method)
+        def circ():
+            rot_gate(0)
+            m = measurement_fn(0)  # always 1
+
+            qml.cond(m, qml.RX)(2.345, 1)
+            return qml.expval(qml.Z(1)), qml.expval(qml.Z(0))
+
+        if shots:
+            # both results are on the order of 1 (-0.7, -1), and an uncertainty ~1.5-2 orders
+            # of magnitude smaller than the result is sufficiently accurate for a shots-based measurement
+            assert np.allclose(circ(), [np.cos(2.345), -1], atol=0.03)
+        else:
+            assert np.allclose(circ(), [np.cos(2.345), -1])