Repository: PennyLaneAI/pennylane
Pull request #7815 (2025-07-08 18:47:48)

[Python compiler] Add MBQC dialect
**Context:** Catalyst currently has an `mbqc` MLIR dialect defined in

* [mlir/include/MBQC/IR/MBQCDialect.td](https://github.com/PennyLaneAI/catalyst/blob/main/mlir/include/MBQC/IR/MBQCDialect.td)
* [mlir/include/MBQC/IR/MBQCOps.td](https://github.com/PennyLaneAI/catalyst/blob/main/mlir/include/MBQC/IR/MBQCOps.td)

We need an xDSL equivalent of this dialect defined for the Python compiler.

**Description of the Change:** Adds the `mbqc` xDSL dialect to the Python compiler, exactly as it is currently defined in Catalyst.

[sc-94750]


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.

Patch:
diff --git a/doc/releases/changelog-0.42.0.md b/doc/releases/changelog-0.42.0.md
index 5ca556e31..0b91aaa95 100644
--- a/doc/releases/changelog-0.42.0.md
+++ b/doc/releases/changelog-0.42.0.md
@@ -1195,6 +1195,10 @@ may move operations across a `Snapshot`.
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes a bug where normalization in `qml.StatePrep` with `normalize=True` was skipped if
+  `validate_norm` is set to `False`.
+  [(#7835)](https://github.com/PennyLaneAI/pennylane/pull/7835) 
+
 * Fixes broken support of `qml.matrix` for a `QNode` when using mixed Torch GPU & CPU data for parametric tensors.
   [(#7775)](https://github.com/PennyLaneAI/pennylane/pull/7775) 
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 25e46a2f4..a54479a56 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,6 +6,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* The `mbqc` xDSL dialect has been added to the Python compiler, which is used to represent
+  measurement-based quantum-computing instructions in the xDSL framework.
+  [(#7815)](https://github.com/PennyLaneAI/pennylane/pull/7815)
+
 <h3>Labs: a place for unified and rapid prototyping of research software 🧪</h3>
 
 * Added state of the art resources for the `ResourceSelectPauliRot` template and the
@@ -30,5 +34,6 @@
 
 This release contains contributions from (in alphabetical order):
 
+Joey Carter,
 Andrija Paurevic,
 Jay Soni,
diff --git a/pennylane/compiler/python_compiler/mbqc_dialect.py b/pennylane/compiler/python_compiler/mbqc_dialect.py
new file mode 100644
index 000000000..418b2dd9e
--- /dev/null
+++ b/pennylane/compiler/python_compiler/mbqc_dialect.py
@@ -0,0 +1,126 @@
+# 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 module contains the definition of the MBQC dialect for the Python compiler.
+
+The MBQC dialect is a set of operations and types used to represent measurement-based
+quantum-computing instructions in the xDSL framework.
+
+It was initially generated by xDSL (using the ``xdsl-tblgen`` tool) starting from the
+catalyst/mlir/include/MBQC/IR/MBQCDialect.td file in the catalyst repository.
+"""
+
+from typing import Optional, TypeAlias
+
+from xdsl.dialects.builtin import I32, Float64Type, IntegerAttr, IntegerType
+from xdsl.ir import Dialect, EnumAttribute, Operation, SSAValue
+from xdsl.irdl import (
+    IRDLOperation,
+    irdl_attr_definition,
+    irdl_op_definition,
+    operand_def,
+    opt_prop_def,
+    prop_def,
+    result_def,
+)
+from xdsl.utils.exceptions import VerifyException
+from xdsl.utils.str_enum import StrEnum  # StrEnum is standard in Python>=3.11
+
+from pennylane.compiler.python_compiler.quantum_dialect import QubitType
+
+QubitSSAValue: TypeAlias = SSAValue[QubitType]
+
+
+class MeasurementPlaneEnum(StrEnum):
+    """Enum containing supported measurement-plane attributes"""
+
+    # pylint: disable=too-few-public-methods
+
+    XY = "XY"
+    YZ = "YZ"
+    ZX = "ZX"
+
+
+@irdl_attr_definition
+class MeasurementPlaneAttr(EnumAttribute[MeasurementPlaneEnum]):
+    """Planes in the Bloch sphere representation with support for arbitrary-basis measurements"""
+
+    # pylint: disable=too-few-public-methods
+
+    name = "mbqc.measurement_plane"
+
+
+@irdl_op_definition
+class MeasureInBasisOp(IRDLOperation):
+    """A parametric single-qubit projective measurement in an arbitrary basis."""
+
+    # pylint: disable=too-few-public-methods
+
+    name = "mbqc.measure_in_basis"
+
+    in_qubit = operand_def(QubitType)
+
+    plane = prop_def(MeasurementPlaneAttr)
+
+    angle = operand_def(Float64Type())
+
+    postselect = opt_prop_def(IntegerAttr[I32])
+
+    mres = result_def(IntegerType(1))
+
+    out_qubit = result_def(QubitType)
+
+    assembly_format = """
+            `[` $plane `,` $angle `]` $in_qubit (`postselect` $postselect^)? attr-dict `:` type(results)
+        """
+
+    def __init__(
+        self,
+        in_qubit: QubitSSAValue | Operation,
+        plane: MeasurementPlaneAttr,
+        angle: SSAValue[Float64Type],
+        postselect: Optional[int | IntegerAttr] = None,
+    ):
+        properties = {"plane": plane}
+
+        if isinstance(postselect, int):
+            postselect = IntegerAttr.from_int_and_width(postselect, 32)
+
+        if postselect is not None:
+            properties["postselect"] = postselect
+
+        super().__init__(
+            operands=(in_qubit, angle),
+            properties=properties,
+            result_types=(IntegerType(1), QubitType()),
+        )
+
+    def verify_(self):
+        """Verify operation when rewriting."""
+        if self.postselect is None:
+            return
+
+        if self.postselect.value.data not in [0, 1]:
+            raise VerifyException("'postselect' must be 0 or 1.")
+
+
+MBQCDialect = Dialect(
+    "mbqc",
+    [
+        MeasureInBasisOp,
+    ],
+    [
+        MeasurementPlaneAttr,
+    ],
+)
diff --git a/pennylane/ops/functions/assert_valid.py b/pennylane/ops/functions/assert_valid.py
index 455697cc4..2b652624d 100644
--- a/pennylane/ops/functions/assert_valid.py
+++ b/pennylane/ops/functions/assert_valid.py
@@ -436,7 +436,17 @@ def _check_wires(op, skip_wire_mapping):
     assert mapped_op.wires == new_wires, "wires must be mappable with map_wires"
 
 
-def assert_valid(op: qml.operation.Operator, **kwargs) -> None:
+# pylint: disable=too-many-arguments
+def assert_valid(
+    op: qml.operation.Operator,
+    *,
+    skip_deepcopy=False,
+    skip_pickle=False,
+    skip_wire_mapping=False,
+    skip_differentiation=False,
+    skip_new_decomp=False,
+    heuristic_resources=False,
+) -> None:
     """Runs basic validation checks on an :class:`~.operation.Operator` to make
     sure it has been correctly defined.
 
@@ -496,20 +506,20 @@ def assert_valid(op: qml.operation.Operator, **kwargs) -> None:
         assert qml.math.allclose(d, p), "data and parameters must match."
 
     if len(op.wires) <= 26:
-        _check_wires(op, kwargs.get("skip_wire_mapping", False))
-    _check_copy(op, kwargs.get("skip_deepcopy", False))
+        _check_wires(op, skip_wire_mapping=skip_wire_mapping)
+    _check_copy(op, skip_deepcopy=skip_deepcopy)
     _check_pytree(op)
-    if not kwargs.get("skip_pickle", False):
+    if not skip_pickle:
         _check_pickle(op)
     _check_bind_new_parameters(op)
-    _check_decomposition(op, kwargs.get("skip_wire_mapping", False))
-    if not kwargs.get("skip_new_decomp", False):
-        _check_decomposition_new(op, heuristic_resources=kwargs.get("heuristic_resources", False))
+    _check_decomposition(op, skip_wire_mapping=skip_wire_mapping)
+    if not skip_new_decomp:
+        _check_decomposition_new(op, heuristic_resources=heuristic_resources)
     _check_matrix(op)
     _check_matrix_matches_decomp(op)
     _check_sparse_matrix(op)
     _check_eigendecomposition(op)
     _check_generator(op)
-    if not kwargs.get("skip_differentiation", False):
+    if not skip_differentiation:
         _check_differentiation(op)
     _check_capture(op)
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index dc85d2fcd..8d4ac3423 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -217,7 +217,7 @@ add_decomps(BasisState, _basis_state_decomp)
 
 
 class StatePrep(StatePrepBase):
-    r"""StatePrep(state, wires, pad_with = None, normalize = False, validate_norm = True)
+    r"""StatePrep(state, wires, pad_with = None, normalize = False, validate_norm = False)
     Prepare subsystems using a state vector in the computational basis.
 
     **Details:**
@@ -503,7 +503,7 @@ class StatePrep(StatePrepBase):
                 padding = math.convert_like(padding, state)
                 state = math.hstack([state, padding])
 
-        if not validate_norm:
+        if not (validate_norm or normalize):
             return state
 
         # normalize


Test patch:
diff --git a/tests/ops/qubit/test_state_prep.py b/tests/ops/qubit/test_state_prep.py
index d20ecf68c..7dda42db1 100644
--- a/tests/ops/qubit/test_state_prep.py
+++ b/tests/ops/qubit/test_state_prep.py
@@ -162,14 +162,15 @@ class TestDecomposition:
             (np.array([1, 1j, 1j, 1])),
         ],
     )
-    def test_StatePrep_normalize(self, state):
+    @pytest.mark.parametrize("validate_norm", [True, False])
+    def test_StatePrep_normalize(self, state, validate_norm):
         """Test that StatePrep normalizes the input state correctly."""
 
         wires = (0, 1)
 
         @qml.qnode(qml.device("default.qubit", wires=2))
         def circuit():
-            qml.StatePrep(state, normalize=True, wires=wires, validate_norm=True)
+            qml.StatePrep(state, normalize=True, wires=wires, validate_norm=validate_norm)
             return qml.state()
 
         assert np.allclose(circuit(), state / 2)
diff --git a/tests/python_compiler/test_mbqc_dialect.py b/tests/python_compiler/test_mbqc_dialect.py
new file mode 100644
index 000000000..947a7f3d0
--- /dev/null
+++ b/tests/python_compiler/test_mbqc_dialect.py
@@ -0,0 +1,163 @@
+# 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.
+
+"""Unit test module for pennylane/compiler/python_compiler/mbqc_dialect.py."""
+
+import pytest
+
+# pylint: disable=wrong-import-position
+
+xdsl = pytest.importorskip("xdsl")
+filecheck = pytest.importorskip("filecheck")
+
+pytestmark = pytest.mark.external
+
+from xdsl.dialects import builtin, test
+from xdsl.utils.exceptions import VerifyException
+
+from pennylane.compiler.python_compiler import mbqc_dialect as mbqc
+from pennylane.compiler.python_compiler.mbqc_dialect import MBQCDialect
+from pennylane.compiler.python_compiler.quantum_dialect import QuantumDialect
+
+all_ops = list(MBQCDialect.operations)
+all_attrs = list(MBQCDialect.attributes)
+
+expected_ops_names = {
+    "MeasureInBasisOp": "mbqc.measure_in_basis",
+}
+
+expected_attrs_names = {
+    "MeasurementPlaneAttr": "mbqc.measurement_plane",
+}
+
+
+def test_mbqc_dialect_name():
+    """Test that the MBQCDialect name is correct."""
+    assert MBQCDialect.name == "mbqc"
+
+
+@pytest.mark.parametrize("op", all_ops)
+def test_all_operations_names(op):
+    """Test that all operations have the expected name."""
+    op_class_name = op.__name__
+    expected_name = expected_ops_names.get(op_class_name)
+    assert expected_name is not None, f"Unexpected operation {op_class_name} found in MBQCDialect"
+    assert op.name == expected_name
+
+
+@pytest.mark.parametrize("attr", all_attrs)
+def test_all_attributes_names(attr):
+    """Test that all attributes have the expected name."""
+    attr_class_name = attr.__name__
+    expected_name = expected_attrs_names.get(attr_class_name)
+    assert expected_name is not None, f"Unexpected attribute {attr_class_name} found in MBQCDialect"
+    assert attr.name == expected_name
+
+
+def test_assembly_format(run_filecheck):
+    """Test the assembly format of the mbqc ops."""
+    program = r"""
+    // CHECK: [[angle:%.+]] = "test.op"() : () -> f64
+    %angle = "test.op"() : () -> f64
+
+    // CHECK: [[qubit:%.+]] = "test.op"() : () -> !quantum.bit
+    %qubit = "test.op"() : () -> !quantum.bit
+
+    // CHECK: [[mres0:%.+]], [[out_qubit0:%.+]] = mbqc.measure_in_basis{{\s*}}[XY, [[angle]]] [[qubit]] : i1, !quantum.bit
+    %mres0, %out_qubit0 = mbqc.measure_in_basis [XY, %angle] %qubit : i1, !quantum.bit
+
+    // CHECK: [[mres1:%.+]], [[out_qubit1:%.+]] = mbqc.measure_in_basis{{\s*}}[YZ, [[angle]]] [[qubit]] : i1, !quantum.bit
+    %mres1, %out_qubit1 = mbqc.measure_in_basis [YZ, %angle] %qubit : i1, !quantum.bit
+
+    // CHECK: [[mres2:%.+]], [[out_qubit2:%.+]] = mbqc.measure_in_basis{{\s*}}[ZX, [[angle]]] [[qubit]] : i1, !quantum.bit
+    %mres2, %out_qubit2 = mbqc.measure_in_basis [ZX, %angle] %qubit : i1, !quantum.bit
+
+    // CHECK: [[mres3:%.+]], [[out_qubit3:%.+]] = mbqc.measure_in_basis{{\s*}}[XY, [[angle]]] [[qubit]] postselect 0 : i1, !quantum.bit
+    %mres3, %out_qubit3 = mbqc.measure_in_basis [XY, %angle] %qubit postselect 0 : i1, !quantum.bit
+
+    // CHECK: [[mres4:%.+]], [[out_qubit4:%.+]] = mbqc.measure_in_basis{{\s*}}[XY, [[angle]]] [[qubit]] postselect 1 : i1, !quantum.bit
+    %mres4, %out_qubit4 = mbqc.measure_in_basis [XY, %angle] %qubit postselect 1 : i1, !quantum.bit
+    """
+
+    ctx = xdsl.context.Context()
+
+    ctx.load_dialect(builtin.Builtin)
+    ctx.load_dialect(test.Test)
+    ctx.load_dialect(QuantumDialect)
+    ctx.load_dialect(MBQCDialect)
+
+    module = xdsl.parser.Parser(ctx, program).parse_module()
+
+    run_filecheck(program, module)
+
+
+class TestMeasureInBasisOp:
+    """Unit tests for the mbqc.measure_in_basis op."""
+
+    @pytest.mark.parametrize("plane,", ["XY", "YZ", "ZX"])
+    @pytest.mark.parametrize("postselect", ["", "postselect 0", "postselect 1"])
+    def test_measure_in_basis_properties(self, plane, postselect):
+        """Test the parsing of the mbqc.measure_in_basis op's properties."""
+        program = rf"""
+        %angle = "test.op"() : () -> f64
+        %qubit = "test.op"() : () -> !quantum.bit
+
+        %mres, %out_qubit = mbqc.measure_in_basis [{plane}, %angle] %qubit {postselect} : i1, !quantum.bit
+        """
+
+        ctx = xdsl.context.Context()
+
+        ctx.load_dialect(builtin.Builtin)
+        ctx.load_dialect(test.Test)
+        ctx.load_dialect(QuantumDialect)
+        ctx.load_dialect(MBQCDialect)
+
+        module = xdsl.parser.Parser(ctx, program).parse_module()
+
+        measure_in_basis_op: mbqc.MeasureInBasisOp = module.ops.last
+        assert isinstance(measure_in_basis_op, mbqc.MeasureInBasisOp)
+
+        assert measure_in_basis_op.properties["plane"].data == plane
+
+        if postselect:
+            assert measure_in_basis_op.properties["postselect"].value.data == int(postselect[-1])
+        else:
+            assert measure_in_basis_op.properties.get("postselect") is None
+
+    @pytest.mark.parametrize("postselect", [-1, 2])
+    def test_invalid_postselect_raises_on_verify(self, postselect):
+        """Test that using an invalid postselect value (a value other than 0 or 1) raises a
+        VerifyException during verification."""
+
+        program = rf"""
+        %angle = "test.op"() : () -> f64
+        %qubit = "test.op"() : () -> !quantum.bit
+
+        %mres, %out_qubit = mbqc.measure_in_basis [XY, %angle] %qubit postselect {postselect} : i1, !quantum.bit
+        """
+
+        ctx = xdsl.context.Context()
+
+        ctx.load_dialect(builtin.Builtin)
+        ctx.load_dialect(test.Test)
+        ctx.load_dialect(QuantumDialect)
+        ctx.load_dialect(MBQCDialect)
+
+        module = xdsl.parser.Parser(ctx, program).parse_module()
+
+        measure_in_basis_op: mbqc.MeasureInBasisOp = module.ops.last
+        assert isinstance(measure_in_basis_op, mbqc.MeasureInBasisOp)
+
+        with pytest.raises(VerifyException, match="'postselect' must be 0 or 1"):
+            measure_in_basis_op.verify_()