Repository: PennyLaneAI/pennylane
Pull request #7489 (2025-05-26 20:34:09)

[Decomposition] Implement decomposition rules for Exp
**Context:**

**Description of the Change:**

Decomposition rules compatible with the new graph-based decomposition system has been implemented for `Exp`, specifically the following decompositions:

- Use Suzuki-Trotter decomposition when `num_steps` is specified.
- Decompose to a `PauliRot` when the base is a single-term pauli word.

Also removes a decomposition rule that is incorrect.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-88546]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7489?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 `98.46154%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 99.69%. Comparing base [(`68e6a22`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/68e6a22b3d81a7f02f676344ba5e2d5c45c71c75?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`67abfa8`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/67abfa8b4ac2d7259d619bd8532d46f7bc6a7673?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/7489?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/ops/op\_math/exp.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7489?src=pr&el=tree&filepath=pennylane%2Fops%2Fop_math%2Fexp.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2V4cC5weQ==) | 98.46% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7489?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    #7489      +/-   ##
==========================================
- Coverage   99.69%   99.69%   -0.01%     
==========================================
  Files         529      529              
  Lines       51289    51353      +64     
==========================================
+ Hits        51133    51196      +63     
- Misses        156      157       +1     
```

</details>

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

Patch:
diff --git a/doc/code/qml_ftqc.rst b/doc/code/qml_ftqc.rst
index 4fb293dd1..64d0ad47f 100644
--- a/doc/code/qml_ftqc.rst
+++ b/doc/code/qml_ftqc.rst
@@ -27,4 +27,5 @@ This module contains functions for tracking, commuting Pauli operations in a Cli
     ~pauli_to_xz
     ~xz_to_pauli
     ~pauli_prod
-    ~commute_clifford_op
\ No newline at end of file
+    ~commute_clifford_op
+    ~get_byproduct_corrections
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8c391fa29..84fa8187c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -7,6 +7,7 @@
 * A new function called :func:`qml.from_qasm3` has been added, which converts OpenQASM 3.0 circuits into quantum functions
   that can be subsequently loaded into QNodes and executed. 
   [(#7432)](https://github.com/PennyLaneAI/pennylane/pull/7432)
+  [(#7486)](https://github.com/PennyLaneAI/pennylane/pull/7486)
 
   ```python
   import pennylane as qml
@@ -140,6 +141,7 @@
   expressed in a limited, compatible gate-set into the MBQC formalism. Circuits can be converted to the 
   relevant gate-set with the `convert_to_mbqc_gateset` transform.
   [(#7355)](https://github.com/PennyLaneAI/pennylane/pull/7355)
+  [(#7586)](https://github.com/PennyLaneAI/pennylane/pull/7586)
 
 <h4>Resource-efficient Decompositions 🔎</h4>
 
@@ -322,9 +324,18 @@
   0: ──RX(-0.50)─╭●────────────┤  <Z>
   1: ────────────╰X──RY(-0.50)─┤
   ```
+  
+* Decomposition rules compatible with the new graph-based decomposition system have been implemented
+  for :class:`~pennylane.ops.Exp`. Specifically, the following decompositions have been added:
+  * Suzuki-Trotter decomposition when the `num_steps` keyword argument is specified.
+  * Decomposition to a :class:`~pennylane.PauliRot` when the base is a single-term Pauli word.
+  [(#7489)](https://github.com/PennyLaneAI/pennylane/pull/7489)
 
 <h3>Improvements 🛠</h3>
 
+* `qml.grad` and `qml.jacobian` can now handle inputs with dynamic shapes being captured into plxpr.
+  [(#7544)](https://github.com/PennyLaneAI/pennylane/pull/7544/)
+
 * Improved the drawing of `GlobalPhase`, `ctrl(GlobalPhase)`, `Identity` and `ctrl(Identity)` operations.
   The labels are grouped together like for other multi-qubit operations, and the drawing
   no longer depends on the wires of `GlobalPhase` or `Identity`. Control nodes of controlled global phases
@@ -363,6 +374,7 @@
   [(#7367)](https://github.com/PennyLaneAI/pennylane/pull/7367)
   [(#7462)](https://github.com/PennyLaneAI/pennylane/pull/7462)
   [(#7470)](https://github.com/PennyLaneAI/pennylane/pull/7470)
+  [(#7590)](https://github.com/PennyLaneAI/pennylane/pull/7590)
 
 * PennyLane supports `JAX` version 0.6.0.
   [(#7299)](https://github.com/PennyLaneAI/pennylane/pull/7299)
@@ -387,6 +399,9 @@
   accessible via the `commute_clifford_op` function.
   [(#7444)](https://github.com/PennyLaneAI/pennylane/pull/7444)
 
+* Add offline byproduct correction support to the `ftqc` module.
+  [(#7447)](https://github.com/PennyLaneAI/pennylane/pull/7447)
+
 * 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)
@@ -524,6 +539,7 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
   [(#7292)](https://github.com/PennyLaneAI/pennylane/pull/7292)
   [(#7477)](https://github.com/PennyLaneAI/pennylane/pull/7477)
   [(#7508)](https://github.com/PennyLaneAI/pennylane/pull/7508)
+  [(#7603)](https://github.com/PennyLaneAI/pennylane/pull/7603)
 
 * `qml.operation.Observable` and the corresponding `Observable.compare` have been deprecated, as
   pennylane now depends on the more general `Operator` interface instead. The
@@ -636,6 +652,9 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.metric_tensor` can now be calculated with catalyst.
+  [(#7528)](https://github.com/PennyLaneAI/pennylane/pull/7528)
+
 * The mapping to standard wires (consecutive integers) of `qml.tape.QuantumScript` has been fixed
   to correctly consider work wires that are not used otherwise in the circuit.
   [(#7581)](https://github.com/PennyLaneAI/pennylane/pull/7581)
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index d5b3c437a..5dfacf99b 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -199,7 +199,7 @@ def __getattr__(name):
         warnings.warn(
             f"pennylane.{name} is no longer accessible at top-level \
                 and must be imported as pennylane.exceptions.{name}. \
-                    Support for top-level access will be removed in v0.42.",
+                    Support for top-level access will be removed in v0.43.",
             pennylane.exceptions.PennyLaneDeprecationWarning,
         )
         return getattr(pennylane.exceptions, name)
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index 754622426..af30ba8ef 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -23,7 +23,7 @@ from autograd.extend import vspace
 from autograd.numpy.numpy_boxes import ArrayBox
 from autograd.wrap_util import unary_to_nary
 
-from pennylane.capture import enabled
+from pennylane.capture import determine_abstracted_axes, enabled
 from pennylane.capture.capture_diff import _get_grad_prim, _get_jacobian_prim
 from pennylane.capture.flatfn import FlatFn
 from pennylane.compiler import compiler
@@ -72,9 +72,21 @@ def _capture_diff(func, argnum=None, diff_prim=None, method=None, h=None):
         # Create fully flattened function (flat inputs & outputs)
         flat_fn = FlatFn(partial(func, **kwargs) if kwargs else func, full_in_tree)
         flat_args = sum(flat_args, start=[])
-        jaxpr = jax.make_jaxpr(flat_fn)(*flat_args)
-        prim_kwargs = {"argnum": flat_argnum, "jaxpr": jaxpr.jaxpr, "n_consts": len(jaxpr.consts)}
-        out_flat = diff_prim.bind(*jaxpr.consts, *flat_args, **prim_kwargs, method=method, h=h)
+        abstracted_axes, abstract_shapes = determine_abstracted_axes(tuple(flat_args))
+        jaxpr = jax.make_jaxpr(flat_fn, abstracted_axes=abstracted_axes)(*flat_args)
+
+        num_abstract_shapes = len(abstract_shapes)
+        shifted_argnum = [a + num_abstract_shapes for a in flat_argnum]
+
+        prim_kwargs = {
+            "argnum": shifted_argnum,
+            "jaxpr": jaxpr.jaxpr,
+            "n_consts": len(jaxpr.consts),
+        }
+        out_flat = diff_prim.bind(
+            *jaxpr.consts, *abstract_shapes, *flat_args, **prim_kwargs, method=method, h=h
+        )
+
         # flatten once more to go from 2D derivative structure (outputs, args) to flat structure
         out_flat = tree_leaves(out_flat)
         assert flat_fn.out_tree is not None, "out_tree should be set after executing flat_fn"
@@ -535,7 +547,7 @@ def jacobian(func, argnum=None, method=None, h=None):
     return _jacobian_function
 
 
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments, too-many-positional-arguments
 def vjp(f, params, cotangents, method=None, h=None, argnum=None):
     """A :func:`~.qjit` compatible Vector-Jacobian product of PennyLane programs.
 
@@ -601,7 +613,7 @@ def vjp(f, params, cotangents, method=None, h=None, argnum=None):
     raise CompileError("Pennylane does not support the VJP function without QJIT.")
 
 
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments, too-many-positional-arguments
 def jvp(f, params, tangents, method=None, h=None, argnum=None):
     """A :func:`~.qjit` compatible Jacobian-vector product of PennyLane programs.
 
diff --git a/pennylane/_version.py b/pennylane/_version.py
index dd6fee12d..6848d97ca 100644
--- a/pennylane/_version.py
+++ b/pennylane/_version.py
@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.42.0-dev44"
+__version__ = "0.42.0-dev45"
diff --git a/pennylane/capture/capture_diff.py b/pennylane/capture/capture_diff.py
index 2bd128ce4..c41b25aaa 100644
--- a/pennylane/capture/capture_diff.py
+++ b/pennylane/capture/capture_diff.py
@@ -55,11 +55,17 @@ def _get_grad_prim():
     def _(*args, argnum, jaxpr, n_consts, method, h):
         if len(jaxpr.outvars) != 1 or jaxpr.outvars[0].aval.shape != ():
             raise TypeError("Grad only applies to scalar-output functions. Try jacobian.")
-        return tuple(jaxpr.invars[i].aval for i in argnum)
+        return tuple(args[i + n_consts] for i in argnum)
 
     return grad_prim
 
 
+def _shape(shape, dtype):
+    if jax.config.jax_dynamic_shapes and any(not isinstance(s, int) for s in shape):
+        return jax.core.DShapedArray(shape, dtype)
+    return jax.core.ShapedArray(shape, dtype)
+
+
 @lru_cache
 def _get_jacobian_prim():
     """Create a primitive for Jacobian computations.
@@ -89,10 +95,10 @@ def _get_jacobian_prim():
     # pylint: disable=unused-argument
     @jacobian_prim.def_abstract_eval
     def _(*args, argnum, jaxpr, n_consts, method, h):
-        in_avals = [jaxpr.invars[i].aval for i in argnum]
-        out_shapes = (outvar.aval.shape for outvar in jaxpr.outvars)
+        in_avals = tuple(args[i + n_consts] for i in argnum)
+        out_shapes = tuple(outvar.aval.shape for outvar in jaxpr.outvars)
         return [
-            jax.core.ShapedArray(out_shape + in_aval.shape, in_aval.dtype)
+            _shape(out_shape + in_aval.shape, in_aval.dtype)
             for out_shape in out_shapes
             for in_aval in in_avals
         ]
diff --git a/pennylane/compiler/python_compiler/quantum_dialect.py b/pennylane/compiler/python_compiler/quantum_dialect.py
index 7eeb183a0..b1c9b8698 100644
--- a/pennylane/compiler/python_compiler/quantum_dialect.py
+++ b/pennylane/compiler/python_compiler/quantum_dialect.py
@@ -245,9 +245,9 @@ class CustomOp(IRDLOperation):
 
     name = "quantum.custom"
 
-    # assembly_format = """
-    #        $gate_name `(` $params `)` $in_qubits attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
-    #    """
+    assembly_format = """
+           $gate_name `(` $params `)` $in_qubits (`adj` $adjoint^)? attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
+       """
 
     irdl_options = [
         AttrSizedOperandSegments(as_property=True),
@@ -342,7 +342,7 @@ class DeviceInitOp(IRDLOperation):
 
     # assembly_format = """
     #      (`shots` `(` $shots^ `)`)? `[` $lib `,` $name `,` $kwargs `]` attr-dict
-    #    """
+    #   """
 
     shots = opt_operand_def(EqAttrConstraint(IntegerType(64)))
 
@@ -547,9 +547,9 @@ class MeasureOp(IRDLOperation):
 
     name = "quantum.measure"
 
-    # assembly_format = """
-    #        $in_qubit attr-dict `:` type(results)
-    #    """
+    assembly_format = """
+           $in_qubit (`postselect` $postselect^)? attr-dict `:` type(results)
+       """
 
     in_qubit = operand_def(BaseAttr(QubitType))
 
@@ -589,9 +589,9 @@ class MultiRZOp(IRDLOperation):
 
     name = "quantum.multirz"
 
-    # assembly_format = """
-    #        `(` $theta `)` $in_qubits attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
-    #    """
+    assembly_format = """
+              `(` $theta `)` $in_qubits (`adj` $adjoint^)? attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
+       """
 
     irdl_options = [
         AttrSizedOperandSegments(as_property=True),
@@ -667,9 +667,9 @@ class QubitUnitaryOp(IRDLOperation):
 
     name = "quantum.unitary"
 
-    # assembly_format = """
-    #        `(` $matrix `:` type($matrix) `)` $in_qubits attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
-    #    """
+    assembly_format = """
+           `(` $matrix `:` type($matrix) `)` $in_qubits (`adj` $adjoint^)? attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
+       """
 
     irdl_options = [
         AttrSizedOperandSegments(as_property=True),
diff --git a/pennylane/decomposition/resources.py b/pennylane/decomposition/resources.py
index 87c913a0c..c2aff6e37 100644
--- a/pennylane/decomposition/resources.py
+++ b/pennylane/decomposition/resources.py
@@ -288,12 +288,13 @@ def resource_rep(op_type: Type[Operator], **params) -> CompressedResourceOp:
     return CompressedResourceOp(op_type, params)
 
 
-def controlled_resource_rep(
+def controlled_resource_rep(  # pylint: disable=too-many-arguments
     base_class: Type[Operator],
     base_params: dict,
     num_control_wires: int,
     num_zero_control_values: int = 0,
     num_work_wires: int = 0,
+    work_wire_type="clean",
 ):
     """Creates a ``CompressedResourceOp`` representation of a controlled operator.
 
@@ -308,6 +309,7 @@ def controlled_resource_rep(
         num_control_wires (int): the number of control wires
         num_zero_control_values (int): the number of control values that are 0
         num_work_wires (int): the number of work wires
+        work_wire_type (str): the type of work wire
 
     """
 
@@ -355,6 +357,7 @@ def controlled_resource_rep(
             num_control_wires,
             num_zero_control_values,
             num_work_wires,
+            work_wire_type,
         )
 
     return CompressedResourceOp(
@@ -448,8 +451,13 @@ def _controlled_qubit_unitary_rep(
     )
 
 
-def _controlled_x_rep(
-    base_class, base_params, num_control_wires, num_zero_control_values, num_work_wires
+def _controlled_x_rep(  # pylint: disable=too-many-arguments
+    base_class,
+    base_params,
+    num_control_wires,
+    num_zero_control_values,
+    num_work_wires,
+    work_wire_type="clean",
 ) -> Optional[CompressedResourceOp]:
     """Helper function that handles custom logic for controlled X gates."""
 
@@ -463,6 +471,7 @@ def _controlled_x_rep(
             num_control_wires=num_control_wires,
             num_zero_control_values=num_zero_control_values,
             num_work_wires=num_work_wires,
+            work_wire_type=work_wire_type,
         )
 
     # base_class is qml.MultiControlledX:
@@ -474,4 +483,5 @@ def _controlled_x_rep(
         num_control_wires=num_control_wires,
         num_zero_control_values=num_zero_control_values,
         num_work_wires=num_work_wires,
+        work_wire_type=work_wire_type,
     )
diff --git a/pennylane/decomposition/symbolic_decomposition.py b/pennylane/decomposition/symbolic_decomposition.py
index 302791fac..c896ef510 100644
--- a/pennylane/decomposition/symbolic_decomposition.py
+++ b/pennylane/decomposition/symbolic_decomposition.py
@@ -226,7 +226,7 @@ def make_controlled_decomp(base_decomposition):
     return _impl
 
 
-def flip_zero_control(inner_decomp):
+def flip_zero_control(inner_decomp: DecompositionRule) -> DecompositionRule:
     """Wraps a decomposition for a controlled operator with X gates to flip zero control wires."""
 
     def _condition_fn(**resource_params):
diff --git a/pennylane/ftqc/__init__.py b/pennylane/ftqc/__init__.py
index 11bfe0be7..eda44913c 100644
--- a/pennylane/ftqc/__init__.py
+++ b/pennylane/ftqc/__init__.py
@@ -42,6 +42,7 @@ from .lattice import Lattice, generate_lattice
 from .graph_state_preparation import GraphStatePrep, make_graph_state
 from .qubit_graph import QubitGraph
 from .utils import QubitMgr
+from .pauli_tracker import get_byproduct_corrections
 
 __all__ = [
     "GraphStatePrep",
@@ -62,4 +63,5 @@ __all__ = [
     "convert_to_mbqc_gateset",
     "generate_lattice",
     "make_graph_state",
+    "get_byproduct_corrections",
 ]
diff --git a/pennylane/ftqc/decomposition.py b/pennylane/ftqc/decomposition.py
index 280c219b7..4150934b1 100644
--- a/pennylane/ftqc/decomposition.py
+++ b/pennylane/ftqc/decomposition.py
@@ -204,55 +204,50 @@ def _s_measurements(op: S, wires):
     return [m1, m2, m3, m4]
 
 
-@singledispatch
 def queue_corrections(op, measurements):
     """Queue the byproduct corrections associated with the operation in the
     MBQC formalism, based on the operation and the measurement results"""
-    raise NotImplementedError(f"Received unsupported gate of type {op}")
-
-
-@queue_corrections.register(RotXZX)
-@queue_corrections.register(RZ)
-def _rotation_corrections(op, measurements):
-    """Queue the byproduct corrections associated with the rotation
-    gate RotXZX in the MBQC formalism, based on the operation and the
-    measurement results. Note that these corrections also apply in the
-    more specific rotation case, RZ = RotXZX(0, Z, 0)"""
-
-    m1, m2, m3, m4 = measurements
+    x_corr, z_corr = _single_xz_corrections(op, *measurements)
 
-    def correction_func(wire):
-        cond(m1 ^ m3, Z)(wire)
-        cond(m2 ^ m4, X)(wire)
-
-    return correction_func
+    def corrections_func(wire):
+        cond(z_corr, Z)(wire)
+        cond(x_corr, X)(wire)
 
+    return corrections_func
 
-@queue_corrections.register(H)
-def _hadamard_corrections(op, measurements):
-    """Queue the byproduct corrections associated with the Hadamard gate in
-    the MBQC formalism, based on the operation and the measurement results"""
-    m1, m2, m3, m4 = measurements
 
-    def correction_func(wire):
-        cond(m2 ^ m3, Z)(wire)
-        cond(parity(m1, m3, m4), X)(wire)
+@singledispatch
+def _single_xz_corrections(op, m1, m2, m3, m4):
+    """Get the xz corrections based on the measurements. Returns a tuple with
+    two boolean elements, indicating the need for PauliX and PauliZ
+    corrections respectively."""
+    raise NotImplementedError(f"Received unsupported gate of type {op}")
 
-    return correction_func
 
+@_single_xz_corrections.register(RotXZX)
+@_single_xz_corrections.register(RZ)
+def _rotation_corrections(op, m1, m2, m3, m4):
+    """Get the xz corrections based on the measurements. Returns a tuple with
+    two boolean elements, indicating the need for PauliX and PauliZ
+    corrections respectively. Note that these corrections also apply in the
+    more specific rotation case, RZ = RotXZX(0, Z, 0)"""
+    return m2 ^ m4, m1 ^ m3
 
-@queue_corrections.register(S)
-def _s_corrections(op, measurements):
-    """Queue the byproduct corrections associated with the S gate in
-    the MBQC formalism, based on the operation and the measurement results"""
 
-    m1, m2, m3, m4 = measurements
+@_single_xz_corrections.register(H)
+def _hadamard_corrections(op, m1, m2, m3, m4):
+    """Get the xz corrections based on the measurements. Returns a tuple with
+    two boolean elements, indicating the need for PauliX and PauliZ
+    corrections respectively."""
+    return parity(m1, m3, m4), m2 ^ m3
 
-    def correction_func(wire):
-        cond(parity(m1, m2, m3, 1), Z)(wire)
-        cond(m2 ^ m4, X)(wire)
 
-    return correction_func
+@_single_xz_corrections.register(S)
+def _s_corrections(op, m1, m2, m3, m4):
+    """Get the xz corrections based on the measurements. Returns a tuple with
+    two boolean elements, indicating the need for PauliX and PauliZ
+    corrections respectively."""
+    return m2 ^ m4, parity(m1, m2, m3, 1)
 
 
 def queue_cnot(q_mgr, ctrl_idx, target_idx):
@@ -312,6 +307,24 @@ def cnot_corrections(measurements):
     """Queue the byproduct corrections associated with the CNOT gate in
     the MBQC formalism, based on measurement results"""
 
+    (x_cor_ctrl, z_cor_ctrl), (x_cor_tgt, z_cor_tgt) = _cnot_xz_corrections(measurements)
+
+    def correction_func(ctrl_wire, target_wire):
+        cond(z_cor_ctrl, Z)(ctrl_wire)
+        cond(x_cor_ctrl, X)(ctrl_wire)
+        cond(z_cor_tgt, Z)(target_wire)
+        cond(x_cor_tgt, X)(target_wire)
+
+    return correction_func
+
+
+def _cnot_xz_corrections(measurements):
+    """Get the xz corrections for the control and target wire based on the measurements.
+    Returns a list of two tuples indicating corrections for the control and target wires
+    respectively. For each tuple, the first element is a boolean indicating whether an
+    PauliX correction is needed, and the second element indicates whether a PauliZ
+    correction is needed."""
+
     # Numbering convention follows the procedure in Raussendorf et al. 2003,
     # https://doi.org/10.1103/PhysRevA.68.022312, Fig 2
     m1, m2, m3, m4, m5, m6, m8, m9, m10, m11, m12, m13, m14 = measurements
@@ -324,13 +337,7 @@ def cnot_corrections(measurements):
     x_cor_tgt = parity(m2, m3, m8, m10, m12, m14)
     z_cor_tgt = parity(m9, m11, m13)
 
-    def correction_func(ctrl_wire, target_wire):
-        cond(z_cor_ctrl, Z)(ctrl_wire)
-        cond(x_cor_ctrl, X)(ctrl_wire)
-        cond(z_cor_tgt, Z)(target_wire)
-        cond(x_cor_tgt, X)(target_wire)
-
-    return correction_func
+    return [(x_cor_ctrl, z_cor_ctrl), (x_cor_tgt, z_cor_tgt)]
 
 
 def _generate_cnot_graph():
diff --git a/pennylane/ftqc/pauli_tracker.py b/pennylane/ftqc/pauli_tracker.py
index 7676fd16d..ce4f92c9b 100644
--- a/pennylane/ftqc/pauli_tracker.py
+++ b/pennylane/ftqc/pauli_tracker.py
@@ -15,12 +15,18 @@
 r"""
 This module contains Pauli Tracking functions.
 """
-
 import itertools
 from typing import List, Tuple
 
-from pennylane import CNOT, H, I, S, X, Y, Z
+import numpy as np
+
+from pennylane import math
 from pennylane.operation import Operator
+from pennylane.ops import CNOT, RZ, H, I, S, X, Y, Z
+from pennylane.tape import QuantumScript
+
+from .decomposition import _cnot_xz_corrections, _single_xz_corrections
+from .operations import RotXZX
 
 _OPS_TO_XZ = {
     I: (0, 0),
@@ -38,6 +44,10 @@ _XZ_TO_OPS = {
 
 _PAULIS = (X, Y, Z, I)
 
+_CLIFFORD_GATES_SUPPORTED = (H, S, CNOT)
+
+_NON_CLIFFORD_GATES_SUPPORTED = (RZ, RotXZX)
+
 
 def pauli_to_xz(op: Operator) -> Tuple[int, int]:
     r"""
@@ -235,3 +245,250 @@ def commute_clifford_op(clifford_op: Operator, xz: List[Tuple[int, int]]) -> Lis
         return _commute_cnot(_xc, _zc, _xt, _zt)
 
     raise NotImplementedError("Only qml.H, qml.S and qml.CNOT are supported.")
+
+
+def _parse_mid_measurements(tape: QuantumScript, mid_meas: List):
+    r"""Parse a series of mid-measurement results of a quantum tape where
+    all gates are from the set {X, Y, Z, I, H, S, CNOT, RZ, RotXZX}, and
+    where only the first gate on each wire is permitted to be a non-Clifford
+    gates (RZ, RotXZX). Both non-Clifford and Clifford gates mentioned above
+    are measured in the way defined in Raussendorf et al. <https://arxiv.org/abs/quant-ph/0301052>`__.
+
+    Args:
+        tape (QuantumScript): The quantum tape in the standard circuit mode (Gates are not transformed into the MBQC formalism).
+        mid_meas (list): Mid-measurements results.
+
+    Returns:
+        A list of `byproduct` ops in a reversed manner.
+    """
+    ops = tape.operations
+
+    by_ops = []
+
+    mid_meas_offset = 0
+    _wires_used = [0] * (max(tape.wires) + 1)
+    for op in ops:
+        for wire in op.wires:
+            _wires_used[wire] += 1
+
+        gate_offset = 4 if op.num_wires == 1 else 13
+        ms = mid_meas[mid_meas_offset : mid_meas_offset + gate_offset]
+        by_op = []
+        if isinstance(op, (S, H)):
+            by_op = [_single_xz_corrections(op, *ms)]
+        elif isinstance(op, CNOT):
+            by_op = _cnot_xz_corrections(ms)
+        elif isinstance(op, (RZ, RotXZX)):
+            if _wires_used[op.wires[0]] > 1:
+                raise NotImplementedError(
+                    "The current implementation requires that only a single non-Clifford gate comes before Clifford and Pauli gates for each wire."
+                )
+            by_op = [_single_xz_corrections(op, *ms)]
+        elif isinstance(op, _PAULIS):
+            continue
+        else:
+            raise NotImplementedError(f"{op.name} is not supported.")
+
+        by_ops.append(by_op)
+
+        mid_meas_offset += gate_offset
+
+    # To use list in a stack manner
+    by_ops.reverse()
+    return by_ops
+
+
+def _get_xz_record(tape: QuantumScript, by_ops: List[Tuple[int, int]]):
+    """Commutate/merge the Pauli/byproduct ops of a Clifford circuit.
+
+    Args:
+        tape (QuantumScript): A quantum tape.
+        by_ops (list): List of byproduct operators for Clifford gates
+
+    Return:
+        The final recorded x and z for each wire.
+    """
+
+    num_wires = max(tape.wires) + 1
+
+    x_record = math.zeros(num_wires, dtype=np.uint8)
+    z_record = math.zeros(num_wires, dtype=np.uint8)
+
+    for op in tape.operations:
+        wires = list(op.wires)
+
+        # Get the recorded xz
+        xz = [(x_record[wire], z_record[wire]) for wire in wires]
+
+        # Updated xz
+        new_xz = []
+        if isinstance(op, _NON_CLIFFORD_GATES_SUPPORTED):
+            # Branch for non-Clifford gates
+            new_xz.append(by_ops.pop()[0])
+        elif isinstance(op, _CLIFFORD_GATES_SUPPORTED):
+            # Branch for Clifford gates
+            # Step 1: Commutate the recorded xz with the Clifford gate to a new xz.
+            xz_commutated = commute_clifford_op(op, xz)
+
+            # Step 2: Merge the new xz with the byproduct by_op
+            by_op = by_ops.pop()
+            for _by_op, _xz_comm in zip(by_op, xz_commutated):
+                new_xz.append(math.bitwise_xor(_by_op, _xz_comm))
+        else:  # branch for Paulis
+            # Commutate step is skipped.
+            # Get the new xz by merging the recorded xz with the Pauli ops directly.
+            new_xz.append(math.bitwise_xor(pauli_to_xz(op), xz[0]))
+        # Assign the updated the xz to the x, z record
+        for idx, wire in enumerate(wires):
+            x_record[wire], z_record[wire] = new_xz[idx]
+
+    return x_record, z_record
+
+
+def _correct_samples(tape: QuantumScript, x_record: math.array, measurement_vals: List):
+    """Correct sample measurements in a tape. The samples are corrected based on the `samples`
+    at `wires` with the corresponding recorded x.
+
+        Args:
+            tape (tape: qml.tape.QuantumScript): A quantum tape.
+            measurement_vals (List) : A list of measurement values.
+            x_record (math.array): The array of recorded x for each wire.
+            measurement_vals (List): Measurement values.
+
+        Return:
+            A list of corrected measurement values.
+    """
+    measured_wires = tape.measurements[0].wires
+    x_at_target = np.array([x_record[w] for w in measured_wires])
+    correct_meas = np.bitwise_xor(np.array(measurement_vals), x_at_target)
+
+    return correct_meas
+
+
+def get_byproduct_corrections(tape: QuantumScript, mid_meas: List, measurement_vals: List):
+    r"""Correct sample results offline based on the executed quantum script and the mid-circuit measurement results for each shot.
+    The mid measurement results are first parsed with the quantum script to get the byproduct operations for each Clifford
+    and non-Clifford gates. Note that byproduct operations are stored as a list and accessed in a stack manner. The calculation iteratively
+    pops out the first operation in the tape and applies commutation rules for the first byproduct ops in the byproduct stack and
+    then the results are commutated to the byproduct of the current operations in the tape if it is a Clifford gate. The calculation
+    starts from applying commutate rules for :class:`qml.I` gate or :math:`encode\_xz(x,z)=(0,0)` to the first gate in the tape. The
+    measurement corrections are returned based on the observable operators and the xz recorded.
+
+    Args:
+        tape (tape: qml.tape.QuantumScript): A Clifford quantum tape with :class:`~pennylane.X`, :class:`~pennylane.Y`, :class:`~pennylane.Z`,
+            :class:`~pennylane.I`, :class:`~pennylane.H`, :class:`~pennylane.S`, :class:`~pennylane.CNOT` and non-Clifford gates (:class:`~pennylane.RZ`
+            and :class:`~pennylane.ftqc.RotXZX`) at the beginning of circuit in the standard circuit formalism. Note that one non-Clifford gate per wire
+            at most is supported.
+        mid_meas (list): MidMeasurement results per shot.
+        measurement_vals (list): Raw measurement results.
+
+    Return:
+        A list of corrected measurement results.
+
+
+    **Note**
+    This work is to be integrated into the MBQC transform pipeline.
+
+    **Example:**
+
+        .. code-block:: python3
+
+            import pennylane as qml
+
+            from pennylane.ftqc import diagonalize_mcms, generate_lattice, measure_x, measure_y
+            from pennylane.ftqc import GraphStatePrep
+
+            from pennylane.ftqc.pauli_tracker import get_byproduct_corrections
+            import numpy as np
+
+
+            def generate_random_state(n):
+                seed_value = 42  # You can use any integer as the seed
+                np.random.seed(seed_value)
+                input_state = np.random.random(2**n) + 1j * np.random.random(2**n)
+                return input_state / np.linalg.norm(input_state)
+
+
+            def generate_rot_gate_graph():
+                lattice = generate_lattice([4], "chain")
+                return lattice.graph
+
+
+            num_shots = 1000
+            dev = qml.device("lightning.qubit", shots=num_shots)
+
+            @diagonalize_mcms
+            @qml.qnode(dev, mcm_method="one-shot")
+            def circ(start_state):
+                qml.StatePrep(start_state, wires=[0])
+                GraphStatePrep(generate_rot_gate_graph(), wires=[1, 2, 3, 4])
+                qml.CZ(wires=[0, 1])
+                m0 = measure_x(0, reset=True)
+                m1 = measure_y(1, reset=True)
+                m2 = measure_y(2, reset=True)
+                m3 = measure_y(3, reset=True)
+
+                GraphStatePrep(generate_rot_gate_graph(), wires=[3, 2, 1, 0])
+                qml.CZ(wires=[3, 4])
+                m4 = measure_x(4, reset=True)
+                m5 = measure_y(3, reset=True)
+                m6 = measure_y(2, reset=True)
+                m7 = measure_y(1, reset=True)
+
+                GraphStatePrep(generate_rot_gate_graph(), wires=[1, 2, 3, 4])
+                qml.CZ(wires=[0, 1])
+                m8 = measure_x(0, reset=True)
+                m9 = measure_y(1, reset=True)
+                m10 = measure_y(2, reset=True)
+                m11 = measure_y(3, reset=True)
+
+                return (
+                    qml.sample(wires=[4]),
+                    qml.sample(m0),
+                    qml.sample(m1),
+                    qml.sample(m2),
+                    qml.sample(m3),
+                    qml.sample(m4),
+                    qml.sample(m5),
+                    qml.sample(m6),
+                    qml.sample(m7),
+                    qml.sample(m8), qml.sample(m9), qml.sample(m10), qml.sample(m11)
+                )
+
+            init_state = generate_random_state(1)
+            res = circ(init_state)
+
+            ops = [qml.H(wires=[0]), qml.H(wires=[0]), qml.H(wires=[0])]
+            measurements = [qml.sample(qml.Z(0))]
+
+            meas_res = res[0]
+            mid_meas_res = res[1:]
+            corrected_meas_res = []
+
+            script = qml.tape.QuantumScript(ops, measurements, shots=num_shots)
+
+            for i in range(num_shots):
+                mid_meas = [row[i] for row in mid_meas_res]
+                corrected_meas_res.extend(get_byproduct_corrections(script, mid_meas, [meas_res[i]]))
+
+            res_corrected = 1 - 2*np.sum(corrected_meas_res) / num_shots
+
+            dev_ref = qml.device("default.qubit")
+
+            @diagonalize_mcms
+            @qml.qnode(dev)
+            def circ_ref(start_state):
+                qml.StatePrep(start_state, wires=[0])
+                qml.H(0)
+                qml.H(0)
+                qml.H(0)
+                return qml.expval(qml.Z(0))
+
+            np.allclose(res_corrected, circ_ref(init_state))
+
+    """
+    by_ops = _parse_mid_measurements(tape, mid_meas)
+
+    x_record, _ = _get_xz_record(tape, by_ops)
+
+    return _correct_samples(tape, x_record, measurement_vals)
diff --git a/pennylane/gradients/metric_tensor.py b/pennylane/gradients/metric_tensor.py
index 5746f1d4d..bc0bfc801 100644
--- a/pennylane/gradients/metric_tensor.py
+++ b/pennylane/gradients/metric_tensor.py
@@ -190,6 +190,34 @@ def metric_tensor(  # pylint:disable=too-many-arguments, too-many-positional-arg
         a warning is raised and the block-diagonal approximation is computed instead.
         It is significantly cheaper in this case to explicitly set ``approx="block-diag"`` .
 
+    .. note::
+
+        When used with Catalyst, the classical component of the circuit is not included.
+        This matches the results of setting ``hybrid=False``.
+
+        For example,
+
+        >>> from jax import numpy as jnp
+        >>> @qml.qnode(qml.device('lightning.qubit', wires=4))
+        ... def c(x, y):
+        ...     qml.RX(2*x, 0)
+        ...     qml.RY(y, 0)
+        ...     return qml.expval(qml.Z(0))
+        ...
+        >>> qml.qjit(qml.metric_tensor(c))(jnp.array(0.5), jnp.array(0.6))
+        Array([[0.25      , 0.        ],
+                [0.        , 0.07298165]], dtype=float64)
+        >>> qml.metric_tensor(c, argnums=(0,1))(jnp.array(0.5), jnp.array(0.6))
+        (Array(1., dtype=float64), Array(0.07298165, dtype=float64))
+        >>> qml.metric_tensor(c, hybrid=False)(qml.numpy.array(0.5), qml.numpy.array(0.6))
+        array([[0.25      , 0.        ],
+                [0.        , 0.07298165]])
+
+        Here you can see that the ``qjit`` and ``hybrid=False`` options did not postprocess
+        the metric tensor to match the shape of the arguments, and they do not include the factor
+        of ``4`` from the derivative of ``2*x``.
+
+
     The flag ``allow_nonunitary`` should be set to ``True`` whenever the device with
     which the metric tensor is computed supports non-unitary operations.
     This will avoid additional decompositions of gates, in turn avoiding a potentially
@@ -737,9 +765,12 @@ def _metric_tensor_hadamard(
         first_term = math.zeros_like(diag_mt)
         if ids:
             off_diag_res = math.stack(off_diag_res, 1)[0]
-            inv_ids = [_id[::-1] for _id in ids]
-            first_term = math.scatter_element_add(first_term, list(zip(*ids)), off_diag_res)
-            first_term = math.scatter_element_add(first_term, list(zip(*inv_ids)), off_diag_res)
+
+            for loc, r in zip(ids, off_diag_res):
+                # not sure if we can promise ordering of locations
+                # so need to loop over indices for compatibility with catalyst
+                first_term = math.scatter_element_add(first_term, loc, r)
+                first_term = math.scatter_element_add(first_term, (loc[1], loc[0]), r)
 
         # Second terms of off block-diagonal metric tensor
         expvals = math.zeros_like(first_term[0])
diff --git a/pennylane/io/io.py b/pennylane/io/io.py
index d732677e6..5ff744597 100644
--- a/pennylane/io/io.py
+++ b/pennylane/io/io.py
@@ -848,8 +848,8 @@ def from_qasm3(quantum_circuit: str, wire_map: dict = None):
     Converts an OpenQASM 3.0 circuit into a quantum function that can be used within a QNode.
 
     .. note::
-        The following OpenQASM 3.0 gates are not supported: sdg, tdg, cu. The remaining standard library gates are
-        all supported. End statements, robust variables, subroutines, control flow, measurements, built-in mathematical
+        The following OpenQASM 3.0 gates are not supported: sdg, tdg, cu. The remaining standard library gates as well
+        as ``end`` statements are supported. Robust variables, subroutines, control flow, measurements, built-in mathematical
         functions and constants, custom gates, and pulses are not yet supported.
 
         In order to use this function, ``openqasm3`` and ``'openqasm3[parser]'`` must be installed in the user's
diff --git a/pennylane/io/qasm_interpreter.py b/pennylane/io/qasm_interpreter.py
index 15da4ecea..8289b9984 100644
--- a/pennylane/io/qasm_interpreter.py
+++ b/pennylane/io/qasm_interpreter.py
@@ -7,6 +7,7 @@ import re
 
 from openqasm3.ast import (
     ClassicalDeclaration,
+    EndStatement,
     Identifier,
     QuantumGate,
     QuantumGateModifier,
@@ -52,6 +53,10 @@ PARAMETERIZED_GATES = {
 }
 
 
+class EndProgram(Exception):
+    """Exception raised when it encounters an end statement in the QASM circuit."""
+
+
 class QasmInterpreter:
     """
     Takes the top level node of the AST as a parameter and recursively descends the AST, calling the
@@ -94,14 +99,30 @@ class QasmInterpreter:
         context.update({"wires": [], "vars": {}, "gates": [], "callable": None})
 
         # begin recursive descent traversal
-        for value in node.__dict__.values():
-            if not isinstance(value, list):
-                value = [value]
-            for item in value:
-                if isinstance(item, QASMNode):
-                    self.visit(item, context)
+        try:
+            for value in node.__dict__.values():
+                if not isinstance(value, list):
+                    value = [value]
+                for item in value:
+                    if isinstance(item, QASMNode):
+                        self.visit(item, context)
+        except EndProgram:
+            pass
         return context
 
+    @visit.register(EndStatement)
+    def visit_end_statement(self, node: QASMNode, context: dict):  # pylint: disable=no-self-use
+        """
+        Ends the program.
+        Args:
+            node (QASMNode): The end statement QASMNode.
+            context (dict): the current context.
+        """
+        raise EndProgram(
+            f"The QASM program was terminated om line {node.span.start_line}."
+            f"There may be unprocessed QASM code."
+        )
+
     # needs to have same signature as visit()
     @visit.register(QubitDeclaration)
     def visit_qubit_declaration(
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 3b58d2a09..c3560342f 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -709,8 +709,11 @@ def scatter(indices, array, new_dims, like=None):
     return np.scatter(indices, array, new_dims, like=like)
 
 
+# pylint: disable=too-many-arguments
 @multi_dispatch(argnum=[0, 2])
-def scatter_element_add(tensor, index, value, like=None):
+def scatter_element_add(
+    tensor, index, value, like=None, *, indices_are_sorted=False, unique_indices=False
+):
     """In-place addition of a multidimensional value over various
     indices of a tensor.
 
@@ -719,6 +722,13 @@ def scatter_element_add(tensor, index, value, like=None):
         index (tuple or list[tuple]): Indices to which to add the value
         value (float or tensor_like[float]): Value to add to ``tensor``
         like (str): Manually chosen interface to dispatch to.
+
+    Keyword Args:
+        indices_are_sorted=False (bool): If ``True``, jax will assume that the indices are in
+            ascending order. Required to be ``True`` with catalyst.
+        unique_indices=False (bool): If ``True``, jax will assume each index is unique.
+            Required to be ``True`` with catalyst.
+
     Returns:
         tensor_like[float]: The tensor with the value added at the given indices.
 
@@ -743,7 +753,14 @@ def scatter_element_add(tensor, index, value, like=None):
     if len(np.shape(tensor)) == 0 and index == ():
         return tensor + value
 
-    return np.scatter_element_add(tensor, index, value, like=like)
+    return np.scatter_element_add(
+        tensor,
+        index,
+        value,
+        like=like,
+        indices_are_sorted=indices_are_sorted,
+        unique_indices=unique_indices,
+    )
 
 
 def unwrap(values, max_depth=None):
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 4d97ab4b9..9393ae259 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -163,7 +163,7 @@ def _scatter_numpy(indices, array, shape):
     return new_array
 
 
-def _scatter_element_add_numpy(tensor, index, value):
+def _scatter_element_add_numpy(tensor, index, value, **_):
     """In-place addition of a multidimensional value over various
     indices of a tensor."""
     new_tensor = tensor.copy()
@@ -262,7 +262,7 @@ def _to_numpy_autograd(x, max_depth=None, _n=0):
 ar.register_function("autograd", "to_numpy", _to_numpy_autograd)
 
 
-def _scatter_element_add_autograd(tensor, index, value):
+def _scatter_element_add_autograd(tensor, index, value, **_):
     """In-place addition of a multidimensional value over various
     indices of a tensor. Since Autograd doesn't support indexing
     assignment, we have to be clever and use ravel_multi_index."""
@@ -526,7 +526,7 @@ def _scatter_tf(indices, array, new_dims):
     return tf.scatter_nd(indices, array, new_dims)
 
 
-def _scatter_element_add_tf(tensor, index, value):
+def _scatter_element_add_tf(tensor, index, value, **_):
     """In-place addition of a multidimensional value over various
     indices of a tensor."""
     import tensorflow as tf
@@ -798,7 +798,7 @@ def _scatter_torch(indices, tensor, new_dimensions):
     return new_tensor
 
 
-def _scatter_element_add_torch(tensor, index, value):
+def _scatter_element_add_torch(tensor, index, value, **_):
     """In-place addition of a multidimensional value over various
     indices of a tensor. Note that Torch only supports index assignments
     on non-leaf nodes; if the node is a leaf, we must clone it first."""
@@ -924,7 +924,7 @@ ar.register_function("jax", "scatter", _scatter_jax)
 ar.register_function(
     "jax",
     "scatter_element_add",
-    lambda x, index, value: x.at[tuple(index)].add(value),
+    lambda x, index, value, **kwargs: x.at[tuple(index)].add(value, **kwargs),
 )
 ar.register_function("jax", "unstack", list)
 # pylint: disable=unnecessary-lambda
diff --git a/pennylane/ops/op_math/controlled_decompositions.py b/pennylane/ops/op_math/controlled_decompositions.py
index 2a1d14f9b..73611c8bd 100644
--- a/pennylane/ops/op_math/controlled_decompositions.py
+++ b/pennylane/ops/op_math/controlled_decompositions.py
@@ -20,10 +20,17 @@ from typing import Literal
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import Operation, Operator
+from pennylane.operation import Operator
 from pennylane.wires import Wires, WiresLike
 
-from .decompositions.controlled_decompositions import ctrl_decomp_bisect, ctrl_decomp_zyz
+# pylint: disable=protected-access
+from .decompositions.controlled_decompositions import (
+    _decompose_mcx_with_many_workers,
+    _decompose_mcx_with_one_worker,
+    _decompose_mcx_with_two_workers,
+    ctrl_decomp_bisect,
+    ctrl_decomp_zyz,
+)
 
 
 def _is_single_qubit_special_unitary(op):
@@ -45,9 +52,9 @@ def decompose_mcx(
 
     if n_work_wires >= n_ctrl_wires - 2:
         # Lemma 7.2 of `Barenco et al. (1995) <https://arxiv.org/abs/quant-ph/9503016>`_
-        return _decompose_mcx_with_many_workers(control_wires, target_wire, work_wires)
+        return _decompose_mcx_with_many_workers_old(control_wires, target_wire, work_wires)
     if n_work_wires >= 2:
-        return _decompose_mcx_with_two_workers(
+        return _decompose_mcx_with_two_workers_old(
             control_wires, target_wire, work_wires[0:2], work_wire_type
         )
     if n_work_wires == 1:
@@ -127,47 +134,20 @@ def _decompose_recursive(op, power, control_wires, target_wire, work_wires):
     return decomposition
 
 
-def _decompose_mcx_with_many_workers(control_wires, target_wire, work_wires):
+def _decompose_mcx_with_many_workers_old(control_wires, target_wire, work_wires):
     """Decomposes the multi-controlled PauliX gate using the approach in Lemma 7.2 of
     https://arxiv.org/abs/quant-ph/9503016, which requires a suitably large register of
     work wires"""
-    num_work_wires_needed = len(control_wires) - 2
-    work_wires = work_wires[:num_work_wires_needed]
 
-    work_wires_reversed = list(reversed(work_wires))
-    control_wires_reversed = list(reversed(control_wires))
+    with qml.queuing.AnnotatedQueue() as q:
+        wires = list(control_wires) + [target_wire]
+        _decompose_mcx_with_many_workers(wires=wires, work_wires=work_wires)
 
-    gates = []
-
-    for i in range(len(work_wires)):
-        ctrl1 = control_wires_reversed[i]
-        ctrl2 = work_wires_reversed[i]
-        t = target_wire if i == 0 else work_wires_reversed[i - 1]
-        gates.append(qml.Toffoli(wires=[ctrl1, ctrl2, t]))
-
-    gates.append(qml.Toffoli(wires=[*control_wires[:2], work_wires[0]]))
-
-    for i in reversed(range(len(work_wires))):
-        ctrl1 = control_wires_reversed[i]
-        ctrl2 = work_wires_reversed[i]
-        t = target_wire if i == 0 else work_wires_reversed[i - 1]
-        gates.append(qml.Toffoli(wires=[ctrl1, ctrl2, t]))
-
-    for i in range(len(work_wires) - 1):
-        ctrl1 = control_wires_reversed[i + 1]
-        ctrl2 = work_wires_reversed[i + 1]
-        t = work_wires_reversed[i]
-        gates.append(qml.Toffoli(wires=[ctrl1, ctrl2, t]))
-
-    gates.append(qml.Toffoli(wires=[*control_wires[:2], work_wires[0]]))
-
-    for i in reversed(range(len(work_wires) - 1)):
-        ctrl1 = control_wires_reversed[i + 1]
-        ctrl2 = work_wires_reversed[i + 1]
-        t = work_wires_reversed[i]
-        gates.append(qml.Toffoli(wires=[ctrl1, ctrl2, t]))
+    if qml.QueuingManager.recording():
+        for op in q.queue:  # pragma: no cover
+            qml.apply(op)
 
-    return gates
+    return q.queue
 
 
 def _decompose_mcx_with_one_worker_b95(control_wires, target_wire, work_wire):
@@ -192,55 +172,6 @@ def _decompose_mcx_with_one_worker_b95(control_wires, target_wire, work_wire):
     return gates
 
 
-def _linear_depth_ladder_ops(wires: WiresLike) -> tuple[list[Operator], int]:
-    r"""
-    Helper function to create linear-depth ladder operations used in Khattar and Gidney's MCX synthesis.
-    In particular, this implements Step-1 and Step-2 on Fig. 3 of [1] except for the first and last
-    CCX gates.
-
-    Preconditions:
-        - The number of wires must be greater than 2.
-
-    Args:
-        wires (Wires): Wires to apply the ladder operations on.
-
-    Returns:
-        tuple[list[Operator], int]: Linear-depth ladder circuit and the index of control qubit to
-        apply the final CCX gate.
-
-    References:
-        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
-        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
-    """
-
-    n = len(wires)
-    assert n > 2, "n_ctrls > 2 to use MCX ladder. Otherwise, use CCX"
-
-    gates = []
-    # up-ladder
-    for i in range(1, n - 2, 2):
-        gates.append(qml.Toffoli(wires=[wires[i + 1], wires[i + 2], wires[i]]))
-        gates.append(qml.PauliX(wires=wires[i]))
-
-    # down-ladder
-    if n % 2 == 0:
-        ctrl_1, ctrl_2, target = n - 3, n - 5, n - 6
-    else:
-        ctrl_1, ctrl_2, target = n - 1, n - 4, n - 5
-
-    if target >= 0:
-        gates.append(qml.Toffoli(wires=[wires[ctrl_1], wires[ctrl_2], wires[target]]))
-        gates.append(qml.PauliX(wires=wires[target]))
-
-    for i in range(target, 1, -2):
-        gates.append(qml.Toffoli(wires=[wires[i], wires[i - 1], wires[i - 2]]))
-        gates.append(qml.PauliX(wires=wires[i - 2]))
-
-    final_ctrl = max(0, 5 - n)
-
-    return gates, final_ctrl
-
-
 def _decompose_mcx_with_one_worker_kg24(
     control_wires: WiresLike,
     target_wire: int,
@@ -267,109 +198,20 @@ def _decompose_mcx_with_one_worker_kg24(
         `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
     """
 
-    gates = []
-    gates.append(qml.Toffoli(wires=[control_wires[0], control_wires[1], work_wire]))
-    ladder_ops, final_ctrl = _linear_depth_ladder_ops(control_wires)
-    gates += ladder_ops
-    gates.append(qml.Toffoli(wires=[work_wire, control_wires[final_ctrl], target_wire]))
-    gates += ladder_ops[::-1]
-    gates.append(qml.Toffoli(wires=[control_wires[0], control_wires[1], work_wire]))
-
-    if work_wire_type == "dirty":
-        # perform toggle-detection if ancilla is dirty
-        gates += ladder_ops
-        gates.append(qml.Toffoli(wires=[work_wire, control_wires[final_ctrl], target_wire]))
-        gates += ladder_ops[::-1]
-
-    return gates
-
-
-def _n_parallel_ccx_x(
-    control_wires_x: WiresLike, control_wires_y: WiresLike, target_wires: WiresLike
-) -> list[Operation]:
-    r"""
-    Construct a quantum circuit for creating n-condionally clean ancillae using 3n qubits. This
-    implements Fig. 4a of [1]. Each wire is of the same size :math:`n`.
-
-    Args:
-        control_wires_x (Wires): The control wires for register 1.
-        control_wires_y (Wires): The control wires for register 2.
-        target_wires (Wires): The wires for target register.
-
-    Returns:
-        list[Operation]: The quantum circuit for creating n-condionally clean ancillae.
-
-    References:
-        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
-        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
-    """
-
-    assert (
-        len(control_wires_x) == len(control_wires_y) == len(target_wires)
-    ), "The number of wires must be the same for x, y, and target."
-
-    gates = []
-    for i, ctrl_x in enumerate(control_wires_x):
-        gates.append(qml.X(wires=target_wires[i]))
-        gates.append(qml.Toffoli(wires=[ctrl_x, control_wires_y[i], target_wires[i]]))
-
-    return gates
-
+    with qml.queuing.AnnotatedQueue() as q:
+        wires = list(control_wires) + [target_wire]
+        _decompose_mcx_with_one_worker(
+            wires=wires, work_wires=[work_wire], work_wire_type=work_wire_type
+        )
 
-def _build_logn_depth_ccx_ladder(
-    work_wire: int, control_wires: WiresLike
-) -> tuple[list[Operator], list[Operator]]:
-    r"""
-    Helper function to build a log-depth ladder compose of CCX and X gates as shown in Fig. 4b of [1].
+    if qml.QueuingManager.recording():
+        for op in q.queue:  # pragma: no cover
+            qml.apply(op)
 
-    Args:
-        work_wire (int): The work wire.
-        control_wires (list[Wire]): The control wires.
+    return q.queue
 
-    Returns:
-        tuple[list[Operator], WiresLike: log-depth ladder circuit of cond. clean ancillae and
-        control_wires to apply the linear-depth MCX gate on.
 
-    References:
-        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
-        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
-    """
-
-    gates = []
-    anc = [work_wire]
-    final_ctrls = []
-
-    while len(control_wires) > 1:
-        next_batch_len = min(len(anc) + 1, len(control_wires))
-        control_wires, nxt_batch = control_wires[next_batch_len:], control_wires[:next_batch_len]
-        new_anc = []
-        while len(nxt_batch) > 1:
-            ccx_n = len(nxt_batch) // 2
-            st = int(len(nxt_batch) % 2)
-            ccx_x, ccx_y, ccx_t = (
-                nxt_batch[st : st + ccx_n],
-                nxt_batch[st + ccx_n :],
-                anc[-ccx_n:],
-            )
-            assert len(ccx_x) == len(ccx_y) == len(ccx_t) == ccx_n >= 1
-            if ccx_t != [work_wire]:
-                gates += _n_parallel_ccx_x(ccx_x, ccx_y, ccx_t)
-            else:
-                gates.append(qml.Toffoli(wires=[ccx_x[0], ccx_y[0], ccx_t[0]]))
-            new_anc += nxt_batch[st:]  # newly created cond. clean ancilla
-            nxt_batch = ccx_t + nxt_batch[:st]
-            anc = anc[:-ccx_n]
-
-        anc = sorted(anc + new_anc)
-        final_ctrls += nxt_batch
-
-    final_ctrls += control_wires
-    final_ctrls = sorted(final_ctrls)
-    final_ctrls.remove(work_wire)  #                        # exclude ancilla
-    return gates, final_ctrls
-
-
-def _decompose_mcx_with_two_workers(
+def _decompose_mcx_with_two_workers_old(
     control_wires: WiresLike,
     target_wire: int,
     work_wires: WiresLike,
@@ -398,25 +240,14 @@ def _decompose_mcx_with_two_workers(
     if len(work_wires) < 2:
         raise ValueError("At least 2 work wires are needed for this decomposition.")
 
-    gates = []
-    ladder_ops, final_ctrls = _build_logn_depth_ccx_ladder(work_wires[0], control_wires)
-    gates += ladder_ops
-    if len(final_ctrls) == 1:  # Already a toffoli
-        gates.append(qml.Toffoli(wires=[work_wires[0], final_ctrls[0], target_wire]))
-    else:
-        mid_mcx = _decompose_mcx_with_one_worker_kg24(
-            work_wires[0:1] + final_ctrls, target_wire, work_wires[1], work_wire_type="clean"
+    with qml.queuing.AnnotatedQueue() as q:
+        wires = list(control_wires) + [target_wire]
+        _decompose_mcx_with_two_workers(
+            wires=wires, work_wires=work_wires, work_wire_type=work_wire_type
         )
-        gates += mid_mcx
-    gates += ladder_ops[::-1]
-
-    if work_wire_type == "dirty":
-        # perform toggle-detection if ancilla is dirty
-        gates += ladder_ops[1:]
-        if len(final_ctrls) == 1:
-            gates.append(qml.Toffoli(wires=[work_wires[0], final_ctrls[0], target_wire]))
-        else:
-            gates += mid_mcx
-        gates += ladder_ops[1:][::-1]
 
-    return gates
+    if qml.QueuingManager.recording():
+        for op in q.queue:  # pragma: no cover
+            qml.apply(op)
+
+    return q.queue
diff --git a/pennylane/ops/op_math/controlled_ops.py b/pennylane/ops/op_math/controlled_ops.py
index 8e62daf90..f5473c137 100644
--- a/pennylane/ops/op_math/controlled_ops.py
+++ b/pennylane/ops/op_math/controlled_ops.py
@@ -23,7 +23,7 @@ import numpy as np
 from scipy.linalg import block_diag
 
 import pennylane as qml
-from pennylane.decomposition import add_decomps, register_resources
+from pennylane.decomposition import add_decomps, register_condition, register_resources
 from pennylane.decomposition.symbolic_decomposition import (
     adjoint_rotation,
     flip_zero_control,
@@ -40,6 +40,10 @@ from .controlled_decompositions import decompose_mcx
 from .decompositions.controlled_decompositions import (
     controlled_two_qubit_unitary_rule,
     ctrl_decomp_bisect_rule,
+    decompose_mcx_with_many_workers,
+    decompose_mcx_with_no_worker,
+    decompose_mcx_with_one_worker,
+    decompose_mcx_with_two_workers,
     multi_control_decomp_zyz_rule,
     single_ctrl_decomp_zyz_rule,
 )
@@ -1361,7 +1365,12 @@ class MultiControlledX(ControlledOp):
     ndim_params = ()
     """tuple[int]: Number of dimensions per trainable parameter that the operator depends on."""
 
-    resource_keys = {"num_control_wires", "num_zero_control_values", "num_work_wires"}
+    resource_keys = {
+        "num_control_wires",
+        "num_zero_control_values",
+        "num_work_wires",
+        "work_wire_type",
+    }
 
     name = "MultiControlledX"
 
@@ -1464,6 +1473,7 @@ class MultiControlledX(ControlledOp):
             "num_control_wires": len(self.control_wires),
             "num_zero_control_values": len([val for val in self.control_values if not val]),
             "num_work_wires": len(self.work_wires),
+            "work_wire_type": self.work_wire_type,
         }
 
     def adjoint(self):
@@ -1582,6 +1592,37 @@ class MultiControlledX(ControlledOp):
         )
 
 
+def _mcx_to_cnot_or_toffoli_resource(num_control_wires, num_zero_control_values, **__):
+    if num_control_wires == 1:
+        return {qml.CNOT: 1, qml.X: num_zero_control_values}
+    return {qml.Toffoli: 1, qml.X: num_zero_control_values * 2}
+
+
+@register_condition(lambda num_control_wires, **_: num_control_wires < 3)
+@register_resources(_mcx_to_cnot_or_toffoli_resource)
+def _mcx_to_cnot_or_toffoli(wires, control_wires, control_values, **__):
+    if len(wires) == 2 and not control_values[0]:
+        qml.CNOT(wires=wires)
+        qml.X(wires[1])
+    elif len(wires) == 2:
+        qml.CNOT(wires=wires)
+    elif len(wires) == 3:
+        zero_control_wires = [w for w, val in zip(control_wires, control_values) if not val]
+        for w in zero_control_wires:
+            qml.PauliX(w)
+        qml.Toffoli(wires=wires)
+        for w in zero_control_wires:
+            qml.PauliX(w)
+
+
+add_decomps(
+    MultiControlledX,
+    _mcx_to_cnot_or_toffoli,
+    decompose_mcx_with_many_workers,
+    decompose_mcx_with_two_workers,
+    decompose_mcx_with_one_worker,
+    decompose_mcx_with_no_worker,
+)
 add_decomps("Adjoint(MultiControlledX)", self_adjoint)
 add_decomps("Pow(MultiControlledX)", pow_involutory)
 
diff --git a/pennylane/ops/op_math/decompositions/controlled_decompositions.py b/pennylane/ops/op_math/decompositions/controlled_decompositions.py
index 0a4c99066..2bf72ecbd 100644
--- a/pennylane/ops/op_math/decompositions/controlled_decompositions.py
+++ b/pennylane/ops/op_math/decompositions/controlled_decompositions.py
@@ -16,7 +16,9 @@
 
 from typing import Optional
 
-from pennylane import math, ops, queuing
+import numpy as np
+
+from pennylane import control_flow, math, ops, queuing
 from pennylane.decomposition import (
     adjoint_resource_rep,
     controlled_resource_rep,
@@ -24,6 +26,7 @@ from pennylane.decomposition import (
     register_resources,
     resource_rep,
 )
+from pennylane.decomposition.symbolic_decomposition import flip_zero_control
 from pennylane.operation import Operation, Operator
 from pennylane.ops.op_math.decompositions.unitary_decompositions import two_qubit_decomp_rule
 from pennylane.wires import Wires
@@ -98,10 +101,9 @@ def ctrl_decomp_bisect(target_operation: Operator, control_wires: Wires):
         ctrl_decomp_bisect_rule(target_operation.matrix(), control_wires + target_operation.wires)
 
     # If there is an active queuing context, queue the decomposition so that expand works
-    current_queue = queuing.QueuingManager.active_context()
-    if current_queue is not None:
+    if queuing.QueuingManager.recording():
         for op in q.queue:  # pragma: no cover
-            queuing.apply(op, context=current_queue)
+            queuing.apply(op)
 
     return q.queue
 
@@ -193,10 +195,9 @@ def ctrl_decomp_zyz(
         ops.cond(_not_zero(global_phase), _ctrl_global_phase)(global_phase, control_wires)
 
     # If there is an active queuing context, queue the decomposition so that expand works
-    current_queue = queuing.QueuingManager.active_context()
-    if current_queue is not None:
+    if queuing.QueuingManager.recording():
         for op in q.queue:  # pragma: no cover
-            queuing.apply(op, context=current_queue)
+            queuing.apply(op)
 
     return q.queue
 
@@ -217,6 +218,21 @@ def _ctrl_decomp_bisect_resources(num_target_wires, num_control_wires, **__):
     len_k1 = (num_control_wires + 1) // 2
     len_k2 = num_control_wires - len_k1
     # this is a general overestimate based on the resource requirement of the general case.
+    if len_k1 == len_k2:
+        return {
+            resource_rep(ops.QubitUnitary, num_wires=num_target_wires): 4,
+            adjoint_resource_rep(ops.QubitUnitary, {"num_wires": num_target_wires}): 4,
+            controlled_resource_rep(
+                ops.X,
+                {},
+                num_control_wires=len_k2,
+                num_work_wires=len_k1,
+                work_wire_type="dirty",
+            ): 6,
+            # we only need Hadamard for the main diagonal case (see _ctrl_decomp_bisect_md), but it still needs to be accounted for.
+            ops.Hadamard: 2,
+            controlled_resource_rep(ops.GlobalPhase, {}, num_control_wires=num_control_wires): 1,
+        }
     return {
         resource_rep(ops.QubitUnitary, num_wires=num_target_wires): 4,
         adjoint_resource_rep(ops.QubitUnitary, {"num_wires": num_target_wires}): 4,
@@ -225,14 +241,16 @@ def _ctrl_decomp_bisect_resources(num_target_wires, num_control_wires, **__):
             {},
             num_control_wires=len_k2,
             num_work_wires=len_k1,
+            work_wire_type="dirty",
         ): 4,
         controlled_resource_rep(
             ops.X,
             {},
             num_control_wires=len_k1,
             num_work_wires=len_k2,
+            work_wire_type="dirty",
         ): 2,
-        # we only need Hadamard for the md case, but it still needs to be accounted for.
+        # we only need Hadamard for the main diagonal case (see _ctrl_decomp_bisect_md), but it still needs to be accounted for.
         ops.Hadamard: 2,
         controlled_resource_rep(ops.GlobalPhase, {}, num_control_wires=num_control_wires): 1,
     }
@@ -349,6 +367,200 @@ def controlled_two_qubit_unitary_rule(U, wires, control_values, work_wires, **__
         ops.PauliX(w)
 
 
+def _decompose_mcx_with_many_workers_condition(num_control_wires, num_work_wires, **__):
+    return num_control_wires > 2 and num_work_wires >= num_control_wires - 2
+
+
+def _decompose_mcx_with_many_workers_resource(num_control_wires, **__):
+    return {ops.Toffoli: 4 * (num_control_wires - 2)}
+
+
+# pylint: disable=no-value-for-parameter
+@register_condition(_decompose_mcx_with_many_workers_condition)
+@register_resources(_decompose_mcx_with_many_workers_resource)
+def _decompose_mcx_with_many_workers(wires, work_wires, **__):
+    """Decomposes the multi-controlled PauliX gate using the approach in Lemma 7.2 of
+    https://arxiv.org/abs/quant-ph/9503016, which requires a suitably large register of
+    work wires"""
+
+    target_wire, control_wires = wires[-1], wires[:-1]
+    work_wires = work_wires[: len(control_wires) - 2]
+
+    @control_flow.for_loop(1, len(work_wires), 1)
+    def loop_up(i):
+        ops.Toffoli(wires=[control_wires[i], work_wires[i], work_wires[i - 1]])
+
+    @control_flow.for_loop(len(work_wires) - 1, 0, -1)
+    def loop_down(i):
+        ops.Toffoli(wires=[control_wires[i], work_wires[i], work_wires[i - 1]])
+
+    ops.Toffoli(wires=[control_wires[0], work_wires[0], target_wire])
+    loop_up()
+    ops.Toffoli(wires=[control_wires[-1], control_wires[-2], work_wires[-1]])
+    loop_down()
+    ops.Toffoli(wires=[control_wires[0], work_wires[0], target_wire])
+    loop_up()
+    ops.Toffoli(wires=[control_wires[-1], control_wires[-2], work_wires[-1]])
+    loop_down()
+
+
+decompose_mcx_with_many_workers = flip_zero_control(_decompose_mcx_with_many_workers)
+
+
+def _two_workers_condition(num_control_wires, num_work_wires, **__):
+    return num_control_wires > 2 and num_work_wires >= 2
+
+
+def _two_workers_resource(num_control_wires, work_wire_type, **__):
+    if work_wire_type == "clean":
+        n_ccx = 2 * num_control_wires - 3
+        return {ops.Toffoli: n_ccx, ops.X: n_ccx - 3 if num_control_wires < 6 else n_ccx - 5}
+    # Otherwise, we assume the work wires are dirty
+    n_ccx = 4 * num_control_wires - 8
+    return {ops.Toffoli: n_ccx, ops.X: n_ccx - 4 if num_control_wires < 6 else n_ccx - 8}
+
+
+@register_condition(_two_workers_condition)
+@register_resources(_two_workers_resource)
+def _decompose_mcx_with_two_workers(wires, work_wires, work_wire_type, **__):
+    r"""
+    Synthesise a multi-controlled X gate with :math:`k` controls using :math:`2` ancillary qubits.
+    It produces a circuit with :math:`2k-3` Toffoli gates and depth :math:`O(\log(k))` if using
+    clean ancillae, and :math:`4k-8` Toffoli gates and depth :math:`O(\log(k))` if using dirty
+    ancillae as described in Sec. 5 of [1].
+
+    References:
+        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
+        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
+
+    """
+
+    # First use the work wire to prepare the first two control wires as conditionally clean.
+    ops.Toffoli([wires[0], wires[1], work_wires[0]])
+    middle_ctrl_indices = _build_log_n_depth_ccx_ladder(wires[:-1])
+
+    # Apply the MCX in the middle
+    if len(middle_ctrl_indices) == 1:
+        ops.Toffoli([work_wires[0], wires[middle_ctrl_indices[0]], wires[-1]])
+    else:
+        middle_wires = [wires[i] for i in middle_ctrl_indices]
+        _decompose_mcx_with_one_worker(work_wires[:1] + middle_wires + wires[-1:], work_wires[1:])
+
+    # Uncompute the first ladder
+    ops.adjoint(_build_log_n_depth_ccx_ladder, lazy=False)(wires[:-1])
+    ops.Toffoli([wires[0], wires[1], work_wires[0]])
+
+    if work_wire_type == "dirty":
+        # Perform toggle-detection of the work wire is dirty
+        middle_ctrl_indices = _build_log_n_depth_ccx_ladder(wires[:-1])
+        if len(middle_ctrl_indices) == 1:
+            ops.Toffoli([work_wires[0], wires[middle_ctrl_indices[0]], wires[-1]])
+        else:
+            middle_wires = [wires[i] for i in middle_ctrl_indices]
+            _decompose_mcx_with_one_worker(
+                work_wires[:1] + middle_wires + wires[-1:], work_wires[1:]
+            )
+        ops.adjoint(_build_log_n_depth_ccx_ladder, lazy=False)(wires[:-1])
+
+
+decompose_mcx_with_two_workers = flip_zero_control(_decompose_mcx_with_two_workers)
+
+
+def _decompose_mcx_one_worker_condition(num_control_wires, num_work_wires, **__):
+    return num_control_wires > 2 and num_work_wires == 1
+
+
+def _decompose_mcx_one_worker_resource(num_control_wires, work_wire_type, **__):
+    if work_wire_type == "clean":
+        n_ccx = 2 * num_control_wires - 3
+        return {ops.Toffoli: n_ccx, ops.X: n_ccx - 3}
+    # Otherwise, we assume the work wire is dirty
+    n_ccx = 4 * num_control_wires - 8
+    return {ops.Toffoli: n_ccx, ops.X: n_ccx - 4}
+
+
+@register_condition(_decompose_mcx_one_worker_condition)
+@register_resources(_decompose_mcx_one_worker_resource)
+def _decompose_mcx_with_one_worker(wires, work_wires, work_wire_type="clean", **__):
+    r"""
+    Synthesise a multi-controlled X gate with :math:`k` controls using :math:`1` ancillary qubit. It
+    produces a circuit with :math:`2k-3` Toffoli gates and depth :math:`O(k)` if the ancilla is clean
+    and :math:`4k-3` Toffoli gates and depth :math:`O(k)` if the ancilla is dirty as described in
+    Sec. 5.1 of [1].
+
+    References:
+        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
+        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
+
+    """
+
+    ops.Toffoli([wires[0], wires[1], work_wires[0]])
+
+    final_ctrl_index = _build_linear_depth_ladder(wires[:-1])
+    ops.Toffoli([work_wires[0], wires[final_ctrl_index], wires[-1]])
+    ops.adjoint(_build_linear_depth_ladder, lazy=False)(wires[:-1])
+    ops.Toffoli([wires[0], wires[1], work_wires[0]])
+
+    if work_wire_type == "dirty":
+        # Perform toggle-detection of the work wire is dirty
+        _build_linear_depth_ladder(wires[:-1])
+        ops.Toffoli([work_wires[0], wires[final_ctrl_index], wires[-1]])
+        ops.adjoint(_build_linear_depth_ladder, lazy=False)(wires[:-1])
+
+
+decompose_mcx_with_one_worker = flip_zero_control(_decompose_mcx_with_one_worker)
+
+
+def _decompose_mcx_no_worker_resource(num_control_wires, **__):
+    len_k1 = (num_control_wires + 1) // 2
+    len_k2 = num_control_wires - len_k1
+    if len_k1 == len_k2:
+        return {
+            ops.Hadamard: 2,
+            resource_rep(ops.QubitUnitary, num_wires=1): 2,
+            controlled_resource_rep(
+                ops.X,
+                {},
+                num_control_wires=len_k2,
+                num_work_wires=len_k1,
+                work_wire_type="dirty",
+            ): 4,
+            adjoint_resource_rep(ops.QubitUnitary, {"num_wires": 1}): 2,
+            controlled_resource_rep(ops.GlobalPhase, {}, num_control_wires=num_control_wires): 1,
+        }
+    return {
+        ops.Hadamard: 2,
+        resource_rep(ops.QubitUnitary, num_wires=1): 2,
+        controlled_resource_rep(
+            ops.X,
+            {},
+            num_control_wires=len_k2,
+            num_work_wires=len_k1,
+            work_wire_type="dirty",
+        ): 2,
+        controlled_resource_rep(
+            ops.X,
+            {},
+            num_control_wires=len_k1,
+            num_work_wires=len_k2,
+            work_wire_type="dirty",
+        ): 2,
+        adjoint_resource_rep(ops.QubitUnitary, {"num_wires": 1}): 2,
+        controlled_resource_rep(ops.GlobalPhase, {}, num_control_wires=num_control_wires): 1,
+    }
+
+
+@register_condition(lambda num_control_wires, **_: num_control_wires > 2)
+@register_resources(_decompose_mcx_no_worker_resource)
+def _decompose_mcx_with_no_worker(wires, **_):
+    """Use ctrl_decomp_bisect_md to decompose a multi-controlled X gate with no work wires."""
+    U = ops.RX.compute_matrix(np.pi)
+    _ctrl_decomp_bisect_md(U, wires)
+    ops.ctrl(ops.GlobalPhase(-np.pi / 2), control=wires[:-1])
+
+
+decompose_mcx_with_no_worker = flip_zero_control(_decompose_mcx_with_no_worker)
+
 ####################
 # Helper Functions #
 ####################
@@ -383,21 +595,21 @@ def _ctrl_decomp_bisect_general(U, wires):
     c2t = math.matmul(b, h_matrix)
 
     mid = len(wires) // 2  # for odd n, make control_k1 bigger
-    control_k1 = wires[:mid]
-    control_k2 = wires[mid:-1]
+    ctrl_k1 = wires[:mid]
+    ctrl_k2 = wires[mid:-1]
 
     # The component
     ops.QubitUnitary(c2t, wires[-1])
-    ops.ctrl(ops.X(wires[-1]), control=control_k2, work_wires=control_k1)
+    _controlled_x(wires[-1], control=ctrl_k2, work_wires=ctrl_k1, work_wire_type="dirty")
     ops.adjoint(ops.QubitUnitary(c1, wires[-1]))
 
     # Cancel the two identity controlled X gates
     _ctrl_decomp_bisect_od(d, wires, skip_initial_cx=True)
 
     # Adjoint of the component
-    ops.ctrl(ops.X(wires[-1]), control=control_k1, work_wires=control_k2)
+    _controlled_x(wires[-1], control=ctrl_k1, work_wires=ctrl_k2, work_wire_type="dirty")
     ops.QubitUnitary(c1, wires[-1])
-    ops.ctrl(ops.X(wires[-1]), control=control_k2, work_wires=control_k1)
+    _controlled_x(wires[-1], control=ctrl_k2, work_wires=ctrl_k1, work_wire_type="dirty")
     ops.adjoint(ops.QubitUnitary(c2t, wires[-1]))
 
 
@@ -418,18 +630,18 @@ def _ctrl_decomp_bisect_od(U, wires, skip_initial_cx=False):
     a = _bisect_compute_a(U)
 
     mid = len(wires) // 2  # for odd n, make control_k1 bigger
-    control_k1 = wires[:mid]
-    control_k2 = wires[mid:-1]
+    ctrl_k1 = wires[:mid]
+    ctrl_k2 = wires[mid:-1]
 
     if not skip_initial_cx:
-        ops.ctrl(ops.X(wires[-1]), control=control_k1, work_wires=control_k2)
+        _controlled_x(wires[-1], control=ctrl_k1, work_wires=ctrl_k2, work_wire_type="dirty")
 
     ops.QubitUnitary(a, wires[-1])
-    ops.ctrl(ops.X(wires[-1]), control=control_k2, work_wires=control_k1)
+    _controlled_x(wires[-1], control=ctrl_k2, work_wires=ctrl_k1, work_wire_type="dirty")
     ops.adjoint(ops.QubitUnitary(a, wires[-1]))
-    ops.ctrl(ops.X(wires[-1]), control=control_k1, work_wires=control_k2)
+    _controlled_x(wires[-1], control=ctrl_k1, work_wires=ctrl_k2, work_wire_type="dirty")
     ops.QubitUnitary(a, wires[-1])
-    ops.ctrl(ops.X(wires[-1]), control=control_k2, work_wires=control_k1)
+    _controlled_x(wires[-1], control=ctrl_k2, work_wires=ctrl_k1, work_wire_type="dirty")
     ops.adjoint(ops.QubitUnitary(a, wires[-1]))
 
 
@@ -617,3 +829,141 @@ def _CRY(phi, wires):
 
 def _ctrl_global_phase(phase, control_wires):
     ops.ctrl(ops.GlobalPhase(-phase), control=control_wires)
+
+
+def _controlled_x(target_wire, control, work_wires, work_wire_type):
+    if len(control) == 1:
+        ops.CNOT([control[0], target_wire])
+    elif len(control) == 2:
+        ops.Toffoli([control[0], control[1], target_wire])
+    else:
+        ops.MultiControlledX(
+            control + [target_wire], work_wires=work_wires, work_wire_type=work_wire_type
+        )
+
+
+# pylint: disable=no-value-for-parameter
+def _n_parallel_ccx_x(control_wires_x, control_wires_y, target_wires):
+    r"""
+    Construct a quantum circuit for creating n-condionally clean ancillae using 3n qubits. This
+    implements Fig. 4a of [1]. Each wire is of the same size :math:`n`.
+
+    Args:
+        control_wires_x: The control wires for register 1.
+        control_wires_y: The control wires for register 2.
+        target_wires: The wires for target register.
+
+    References:
+        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
+        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
+    """
+
+    @control_flow.for_loop(0, len(control_wires_x), 1)
+    def loop(i):
+        ops.X(target_wires[i])
+        ops.Toffoli([control_wires_x[i], control_wires_y[i], target_wires[i]])
+
+    loop()
+
+
+def _build_linear_depth_ladder(wires) -> int:
+    r"""
+    Helper function to create linear-depth ladder operations used in Khattar and Gidney's MCX synthesis.
+    In particular, this implements Step-1 and Step-2 on Fig. 3 of [1] except for the first and last
+    CCX gates.
+
+    Preconditions:
+        - The number of wires must be greater than 2.
+
+    Args:
+        wires: the list of wires.
+
+    Returns:
+        int: the index of the last unmarked wire.
+
+    References:
+        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
+        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
+
+    """
+
+    if len(wires) == 3:
+        return 2
+
+    if len(wires) == 4:
+        ops.Toffoli([wires[2], wires[3], wires[1]])
+        ops.X(wires[1])
+        return 1
+
+    i = -1
+    while i + 2 < len(wires) - 2:
+        i += 2
+        ops.Toffoli([wires[i + 1], wires[i + 2], wires[i]])
+        ops.X(wires[i])
+
+    x, y = (i - 2, i) if i + 2 == len(wires) - 1 else (i, i + 3)
+    k = x - 1
+
+    ops.Toffoli([wires[x], wires[y], wires[k]])
+    ops.X(wires[k])
+
+    for i in range(k, 1, -2):
+        ops.Toffoli([wires[i - 1], wires[i], wires[i - 2]])
+        ops.X(wires[i - 2])
+
+    return 0
+
+
+def _build_log_n_depth_ccx_ladder(control_wires) -> list:
+    r"""
+    Helper function to build a log-depth ladder compose of CCX and X gates as shown in Fig. 4b of [1].
+
+    Args:
+        control_wires: The control wires.
+
+    Returns:
+        list: The list of unmarked wires to use as control wires.
+
+    References:
+        1. Khattar and Gidney, Rise of conditionally clean ancillae for optimizing quantum circuits
+        `arXiv:2407.17966 <https://arxiv.org/abs/2407.17966>`__
+    """
+
+    final_control_wires = []
+
+    # See Section 5.2 of [1] for what the following variables mean
+    rightmost_marked = 1
+    timestep = 1
+
+    while rightmost_marked < len(control_wires) - 1:
+
+        # At every time step, we aim to flip the next 2^i + 1 unmarked wires, but if
+        # there are not enough wires available, we just flip all the remaining wires.
+        n_to_flip = min(2**timestep + 1, len(control_wires) - rightmost_marked - 1)
+        rightmost_ctrl = rightmost_marked + n_to_flip
+        new_rightmost_marked = rightmost_ctrl
+        leftmost_unmarked = rightmost_marked + 1
+
+        while n_to_flip > 1:
+
+            ccx_n = n_to_flip // 2
+            ccx_t = control_wires[rightmost_marked + 1 - ccx_n : rightmost_marked + 1]
+            ccx_y = control_wires[rightmost_ctrl + 1 - ccx_n : rightmost_ctrl + 1]
+            ccx_x = control_wires[rightmost_ctrl + 1 - ccx_n * 2 : rightmost_ctrl + 1 - ccx_n]
+
+            leftmost_unmarked = rightmost_marked + 1 - ccx_n
+            _n_parallel_ccx_x(ccx_x, ccx_y, ccx_t)
+
+            # The primitive used ccx_n target wires to flip the 2 * ccx_n control wires. The
+            # total number of remaining wires to flip in this timestep is given by the original
+            # number of wires minus the 2 * ccx_n wires that were flipped, plus the ccx_n
+            # target wires that were unmarked as a result of this primitive.
+            n_to_flip = n_to_flip - ccx_n
+            rightmost_marked -= ccx_n
+            rightmost_ctrl -= ccx_n * 2
+
+        final_control_wires.append(leftmost_unmarked)
+        rightmost_marked = new_rightmost_marked
+        timestep += 1
+
+    return final_control_wires
diff --git a/pennylane/ops/op_math/decompositions/unitary_decompositions.py b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
index 5a898968f..b383867cb 100644
--- a/pennylane/ops/op_math/decompositions/unitary_decompositions.py
+++ b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
@@ -92,10 +92,9 @@ def one_qubit_decomposition(U, wire, rotations="ZYZ", return_global_phase=False)
             ops.GlobalPhase(-global_phase)
 
     # If there is an active queuing context, queue the decomposition so that expand works
-    current_queue = queuing.QueuingManager.active_context()
-    if current_queue is not None:
+    if queuing.QueuingManager.recording():
         for op in q.queue:  # pragma: no cover
-            queuing.apply(op, context=current_queue)
+            queuing.apply(op)
 
     return q.queue
 
@@ -226,10 +225,9 @@ def two_qubit_decomposition(U, wires):
             ops.GlobalPhase(-global_phase)
 
     # If there is an active queuing context, queue the decomposition so that expand works
-    current_queue = queuing.QueuingManager.active_context()
-    if current_queue is not None:
+    if queuing.QueuingManager.recording():
         for op in q.queue:  # pragma: no cover
-            queuing.apply(op, context=current_queue)
+            queuing.apply(op)
 
     return q.queue
 
@@ -267,10 +265,9 @@ def multi_qubit_decomposition(U, wires):
         multi_qubit_decomp_rule(U, wires)
 
     # If there is an active queuing context, queue the decomposition so that expand works
-    current_queue = queuing.QueuingManager.active_context()
-    if current_queue is not None:
+    if queuing.QueuingManager.recording():
         for op in q.queue:  # pragma: no cover
-            queuing.apply(op, context=current_queue)
+            queuing.apply(op)
 
     return q.queue
 
diff --git a/pennylane/ops/op_math/exp.py b/pennylane/ops/op_math/exp.py
index 2371acbe0..afea6e19f 100644
--- a/pennylane/ops/op_math/exp.py
+++ b/pennylane/ops/op_math/exp.py
@@ -20,14 +20,20 @@ import numpy as np
 from scipy.sparse.linalg import expm as sparse_expm
 
 import pennylane as qml
-from pennylane import math
-from pennylane.math import expand_matrix
+from pennylane import math, queuing
+from pennylane.decomposition import (
+    add_decomps,
+    register_condition,
+    register_resources,
+    resource_rep,
+)
 from pennylane.operation import (
     DecompositionUndefinedError,
     GeneratorUndefinedError,
     Operation,
     Operator,
     OperatorPropertyUndefined,
+    TermsUndefinedError,
 )
 from pennylane.wires import Wires
 
@@ -164,6 +170,8 @@ class Exp(ScalarSymbolicOp, Operation):
     control_wires = Wires([])
     _name = "Exp"
 
+    resource_keys = {"base", "num_steps"}
+
     def _flatten(self):
         return (self.base, self.data[0]), (self.num_steps,)
 
@@ -207,6 +215,13 @@ class Exp(ScalarSymbolicOp, Operation):
     def _queue_category(self):
         return "_ops"
 
+    @property
+    def resource_params(self) -> dict:
+        return {
+            "base": self.base,
+            "num_steps": self.num_steps,
+        }
+
     # pylint: disable=invalid-overridden-method, arguments-renamed
     @property
     def has_decomposition(self):
@@ -235,12 +250,12 @@ class Exp(ScalarSymbolicOp, Operation):
         Returns:
             list[PauliRot]: decomposition of the operator
         """
-        with qml.QueuingManager.stop_recording():
+        with queuing.QueuingManager.stop_recording():
             d = self._recursive_decomposition(self.base, self.coeff)
 
-        if qml.QueuingManager.recording():
+        if queuing.QueuingManager.recording():
             for op in d:
-                qml.apply(op)
+                queuing.apply(op)
 
         return d
 
@@ -267,7 +282,7 @@ class Exp(ScalarSymbolicOp, Operation):
             coeffs = [c * coeff for c in coeffs]
             return self._trotter_decomposition(ops, coeffs)
 
-        if not qml.math.is_abstract(coeff) and qml.math.real(coeff):
+        if not math.is_abstract(coeff) and math.real(coeff):
 
             error_msg = f"The decomposition of the {self} operator is not defined."
 
@@ -392,14 +407,14 @@ class Exp(ScalarSymbolicOp, Operation):
                 eigvals = self.eigvals()
                 eigvals_mat = (
                     math.stack([math.diag(e) for e in eigvals])
-                    if qml.math.ndim(self.scalar) > 0
+                    if math.ndim(self.scalar) > 0
                     else math.diag(eigvals)
                 )
                 if len(self.diagonalizing_gates()) == 0:
-                    return expand_matrix(eigvals_mat, wires=self.wires, wire_order=wire_order)
+                    return math.expand_matrix(eigvals_mat, wires=self.wires, wire_order=wire_order)
                 diagonalizing_mat = qml.matrix(self.diagonalizing_gates, wire_order=self.wires)()
                 mat = diagonalizing_mat.conj().T @ eigvals_mat @ diagonalizing_mat
-                return expand_matrix(mat, wires=self.wires, wire_order=wire_order)
+                return math.expand_matrix(mat, wires=self.wires, wire_order=wire_order)
             except OperatorPropertyUndefined:
                 warn(
                     f"The autograd matrix for {self} is not differentiable. "
@@ -444,10 +459,10 @@ class Exp(ScalarSymbolicOp, Operation):
         """
         base_eigvals = math.convert_like(self.base.eigvals(), self.coeff)
         base_eigvals = math.cast_like(base_eigvals, self.coeff)
-        if qml.math.ndim(self.scalar) > 0:
+        if math.ndim(self.scalar) > 0:
             # exp coeff is broadcasted
-            return qml.math.stack([qml.math.exp(c * base_eigvals) for c in self.coeff])
-        return qml.math.exp(self.coeff * base_eigvals)
+            return math.stack([math.exp(c * base_eigvals) for c in self.coeff])
+        return math.exp(self.coeff * base_eigvals)
 
     def label(self, decimals=None, base_label=None, cache=None):
         coeff = (
@@ -467,7 +482,7 @@ class Exp(ScalarSymbolicOp, Operation):
     # pylint: disable=arguments-renamed, invalid-overridden-method
     @property
     def has_generator(self):
-        if qml.math.is_abstract(self.coeff):
+        if math.is_abstract(self.coeff):
             return self.base.is_hermitian
         return self.base.is_hermitian and not np.real(self.coeff)
 
@@ -494,3 +509,99 @@ class Exp(ScalarSymbolicOp, Operation):
             f"generator. Consider using op.simplify() to simplify before finding the generator, or define the operator "
             f"in the form exp(-ixG) through the Evolution class."
         )
+
+
+def _pauli_rot_decomp_condition(base, num_steps):
+    if num_steps:
+        # If num_steps is explicitly provided, always use the Trotter decomposition
+        return False
+    with queuing.QueuingManager.stop_recording():
+        base = base.simplify()
+    # The PauliRot decomposition is only applicable when the base is a Pauli word
+    return qml.pauli.is_pauli_word(base)
+
+
+def _pauli_rot_decomp_resource(base, **_):
+    with queuing.QueuingManager.stop_recording():
+        base = base.simplify()
+    return {resource_rep(qml.PauliRot, pauli_word=qml.pauli.pauli_word_to_string(base)): 1}
+
+
+def _trotter_decomp_condition(base, num_steps):
+    if not num_steps:
+        return False
+    with queuing.QueuingManager.stop_recording():
+        base = base.simplify()
+    if qml.pauli.is_pauli_word(base):
+        return True
+    try:
+        _, ops = base.terms()
+        return all(qml.pauli.is_pauli_word(ob) for ob in ops)
+    except TermsUndefinedError:
+        return False
+
+
+def _trotter_decomp_resource(base, num_steps):
+
+    with queuing.QueuingManager.stop_recording():
+        base = base.simplify()
+
+    try:
+        _, ops = base.terms()
+    except TermsUndefinedError:
+        ops = [base]  # The condition should've already verified that this is a valid pauli word.
+
+    gate_count = {}
+    for op in ops:
+        pauli_word = qml.pauli.pauli_word_to_string(op)
+        if not all(p == "I" for p in pauli_word):
+            op_rep = resource_rep(qml.PauliRot, pauli_word=pauli_word)
+            gate_count[op_rep] = gate_count.get(op_rep, 0) + num_steps
+    return gate_count
+
+
+@register_condition(_pauli_rot_decomp_condition)
+@register_resources(_pauli_rot_decomp_resource)
+def pauli_rot_decomp(*params, wires, base, **_):  # pylint: disable=unused-argument
+    """Decompose the operator into a single PauliRot operator."""
+    with queuing.QueuingManager.stop_recording():
+        base = base.simplify()
+    coeff = params[0]
+    if isinstance(base, qml.ops.SProd):
+        coeff, base = params[0] * base.scalar, base.base
+    coeff = 2j * coeff  # The 2j cancels the coefficient added by PauliRot
+    pauli_word = qml.pauli.pauli_word_to_string(base)
+    if not all(p == "I" for p in pauli_word):
+        qml.PauliRot(coeff, pauli_word, base.wires)
+
+
+@register_condition(_trotter_decomp_condition)
+@register_resources(_trotter_decomp_resource)
+def trotter_decomp(*params, wires, base, num_steps):  # pylint: disable=unused-argument
+    """Uses the Suzuki-Trotter approximation to decompose the operator exponential."""
+
+    with queuing.QueuingManager.stop_recording():
+        simplified_base = base.simplify()
+
+    try:
+        coeffs, ops = simplified_base.terms()
+    except TermsUndefinedError:
+        # The condition should've already verified that this is a valid pauli word.
+        coeffs, ops = [1.0], [simplified_base]
+
+    new_coeffs, pauli_words, new_wires = [], [], []
+    for c, op in zip(coeffs, ops):
+        # The 2j cancels the coefficient added by PauliRot
+        c = c * params[0] / num_steps * 2j
+        pauli_word = qml.pauli.pauli_word_to_string(op)
+        if pauli_word != "I" * len(op.wires):
+            new_coeffs.append(c)
+            pauli_words.append(pauli_word)
+            new_wires.append(op.wires)
+
+    for _ in range(num_steps):
+        for c, pauli_word, new_wire in zip(new_coeffs, pauli_words, new_wires):
+            qml.PauliRot(c, pauli_word, new_wire)
+
+
+add_decomps(Exp, trotter_decomp, pauli_rot_decomp)
diff --git a/pennylane/ops/qubit/non_parametric_ops.py b/pennylane/ops/qubit/non_parametric_ops.py
index dfad19c25..501df9236 100644
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -277,6 +277,7 @@ def _controlled_h_resources(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 1,
     }
 
@@ -539,6 +540,7 @@ def _controlled_x_resource(*_, num_control_wires, num_zero_control_values, num_w
             num_control_wires=num_control_wires,
             num_zero_control_values=num_zero_control_values,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 1,
     }
 
@@ -805,6 +807,7 @@ def _controlled_y_resource(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 1,
     }
 
@@ -1077,6 +1080,7 @@ def _controlled_z_resources(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 1,
     }
 
@@ -1758,6 +1762,7 @@ def _controlled_swap_resources(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires + 1,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 1,
     }
 
diff --git a/pennylane/ops/qubit/parametric_ops_single_qubit.py b/pennylane/ops/qubit/parametric_ops_single_qubit.py
index 5a0784e46..80759c156 100644
--- a/pennylane/ops/qubit/parametric_ops_single_qubit.py
+++ b/pennylane/ops/qubit/parametric_ops_single_qubit.py
@@ -200,6 +200,7 @@ def _controlled_rx_resource(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 2,
     }
 
@@ -365,6 +366,7 @@ def _controlled_ry_resource(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 2,
     }
 
@@ -567,6 +569,7 @@ def _controlled_rz_resource(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 2,
     }
 
@@ -790,6 +793,7 @@ def _controlled_phaseshift_resource(*_, num_control_wires, num_work_wires, **__)
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires - 1,
+            work_wire_type="clean",
         ): 2,
         qml.ControlledPhaseShift: 1,
     }
@@ -1035,6 +1039,7 @@ def _controlled_rot_resource(*_, num_control_wires, num_work_wires, **__):
             num_control_wires=num_control_wires,
             num_zero_control_values=0,
             num_work_wires=num_work_wires,
+            work_wire_type="clean",
         ): 2,
     }
 


Test patch:
diff --git a/.github/workflows/interface-unit-tests.yml b/.github/workflows/interface-unit-tests.yml
index e48ef7cf2..c43cdf0d1 100644
--- a/.github/workflows/interface-unit-tests.yml
+++ b/.github/workflows/interface-unit-tests.yml
@@ -492,7 +492,7 @@ jobs:
       pytest_additional_args: ${{ needs.warnings-as-errors-setup.outputs.pytest_warning_args }}
       pytest_xml_file_path: '${{ inputs.job_name_prefix }}external-libraries-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}.xml'
       additional_pip_packages: |
-        pyzx matplotlib stim quimb mitiq!=0.45.0 ply optax scipy-openblas32>=0.3.26 qualtran openqasm3 antlr4_python3_runtime xdsl==0.38
+        pyzx matplotlib stim quimb mitiq!=0.45.0 ply optax scipy-openblas32>=0.3.26 qualtran openqasm3 antlr4_python3_runtime xdsl==0.38 filecheck
         git+https://github.com/PennyLaneAI/pennylane-qiskit.git@master
         ${{ needs.default-dependency-versions.outputs.jax-version }}
         ${{ needs.default-dependency-versions.outputs.tensorflow-version }}
diff --git a/tests/capture/test_capture_diff.py b/tests/capture/test_capture_diff.py
index d92576eb4..7f18083c9 100644
--- a/tests/capture/test_capture_diff.py
+++ b/tests/capture/test_capture_diff.py
@@ -329,6 +329,35 @@ class TestGrad:
         assert manual_out_tree == jax.tree_util.tree_flatten(manual_out_flat)[1]
         assert qml.math.allclose(jax_out_flat, manual_out_flat)
 
+    @pytest.mark.usefixtures("enable_disable_dynamic_shapes")
+    @pytest.mark.parametrize("same_dynamic_shape", (True, False))
+    def test_grad_dynamic_shape_inputs(self, same_dynamic_shape):
+        """Test that qml.grad can handle dynamic shapes"""
+
+        @qml.qnode(qml.device("default.qubit", wires=4))
+        def c(x, y):
+            qml.RX(x, 0)
+            qml.RY(y, 0)
+            return qml.expval(qml.Z(0))
+
+        def w(n):
+            x = jnp.arange(n)
+            if same_dynamic_shape:
+                y = jnp.arange(n)
+            else:
+                y = jnp.arange(n + 1)
+            return qml.grad(c, argnum=(0, 1))(x, y)
+
+        jaxpr = jax.make_jaxpr(w)(2)
+        grad_eqn = jaxpr.eqns[2] if same_dynamic_shape else jaxpr.eqns[3]
+        assert grad_eqn.primitive == grad_prim
+
+        shift = 1 if same_dynamic_shape else 2
+        assert grad_eqn.params["argnum"] == [shift, shift + 1]
+        assert len(grad_eqn.outvars) == 2
+        assert grad_eqn.outvars[0].aval.shape == grad_eqn.invars[shift].aval.shape
+        assert grad_eqn.outvars[1].aval.shape == grad_eqn.invars[shift + 1].aval.shape
+
 
 def _jac_allclose(jac1, jac2, num_axes, atol=1e-8):
     """Test that two Jacobians, given as nested sequences of arrays, are equal."""
@@ -582,3 +611,33 @@ class TestJacobian:
         # Assert that the output from the manual evaluation is flat
         assert manual_out_tree == jax.tree_util.tree_flatten(manual_out_flat)[1]
         assert qml.math.allclose(jax_out_flat, manual_out_flat)
+
+    @pytest.mark.usefixtures("enable_disable_dynamic_shapes")
+    @pytest.mark.parametrize("same_dynamic_shape", (True, False))
+    def test_jacobian_dynamic_shape_inputs(self, same_dynamic_shape):
+        """Test that qml.jacobian can handle dynamic shapes"""
+
+        @qml.qnode(qml.device("default.qubit", wires=4))
+        def c(x, y):
+            qml.RX(x, 0)
+            qml.RY(y, 0)
+            return qml.probs(wires=(0, 1))
+
+        def w(n):
+            x = jnp.arange(n)
+            if same_dynamic_shape:
+                y = jnp.arange(n)
+            else:
+                y = jnp.arange(n + 1)
+            return qml.jacobian(c, argnum=(0, 1))(x, y)
+
+        jaxpr = jax.make_jaxpr(w)(2)
+        grad_eqn = jaxpr.eqns[2] if same_dynamic_shape else jaxpr.eqns[3]
+        assert grad_eqn.primitive == jacobian_prim
+
+        shift = 1 if same_dynamic_shape else 2
+        assert grad_eqn.params["argnum"] == [shift, shift + 1]
+        assert len(grad_eqn.outvars) == 2
+
+        assert grad_eqn.outvars[0].aval.shape == (4, *grad_eqn.invars[shift].aval.shape)
+        assert grad_eqn.outvars[1].aval.shape == (4, *grad_eqn.invars[shift + 1].aval.shape)
diff --git a/tests/decomposition/test_decomposition_graph.py b/tests/decomposition/test_decomposition_graph.py
index 626e15c87..d833a7fc2 100644
--- a/tests/decomposition/test_decomposition_graph.py
+++ b/tests/decomposition/test_decomposition_graph.py
@@ -518,6 +518,7 @@ class TestControlledDecompositions:
                     num_control_wires=3,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 2,
                 qml.CRot: 1,
             }
diff --git a/tests/decomposition/test_resources.py b/tests/decomposition/test_resources.py
index 881885cca..e579b9e50 100644
--- a/tests/decomposition/test_resources.py
+++ b/tests/decomposition/test_resources.py
@@ -339,6 +339,7 @@ class TestControlledResourceRep:
                     "num_control_wires": 2,
                     "num_zero_control_values": 1,
                     "num_work_wires": 1,
+                    "work_wire_type": "clean",
                 },
                 "num_control_wires": 1,
                 "num_zero_control_values": 1,
@@ -354,6 +355,7 @@ class TestControlledResourceRep:
                 "num_control_wires": 4,
                 "num_zero_control_values": 3,
                 "num_work_wires": 3,
+                "work_wire_type": "clean",
             },
         )
 
diff --git a/tests/decomposition/test_symbolic_decomposition.py b/tests/decomposition/test_symbolic_decomposition.py
index b91c74320..35aafdbf1 100644
--- a/tests/decomposition/test_symbolic_decomposition.py
+++ b/tests/decomposition/test_symbolic_decomposition.py
@@ -327,6 +327,7 @@ def _custom_resource(num_wires):
             num_control_wires=3,
             num_zero_control_values=1,
             num_work_wires=1,
+            work_wire_type="clean",
         ): 1,
         qml.RX: 1,
         qml.Rot: 1,
@@ -416,12 +417,14 @@ class TestControlledDecomposition:
                     num_control_wires=3,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=4,
                     num_zero_control_values=1,
                     num_work_wires=2,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(qml.CRX): 1,
                 qml.resource_rep(qml.CRot): 1,
@@ -516,18 +519,21 @@ class TestControlledDecomposition:
                     num_control_wires=3,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=4,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=5,
                     num_zero_control_values=1,
                     num_work_wires=2,
+                    work_wire_type="clean",
                 ): 1,
                 qml.decomposition.controlled_resource_rep(
                     qml.RX, {}, num_control_wires=2, num_work_wires=1
@@ -629,24 +635,28 @@ class TestControlledDecomposition:
                     num_control_wires=3,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=4,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=5,
                     num_zero_control_values=0,
                     num_work_wires=1,
+                    work_wire_type="clean",
                 ): 1,
                 qml.resource_rep(
                     qml.MultiControlledX,
                     num_control_wires=6,
                     num_zero_control_values=1,
                     num_work_wires=2,
+                    work_wire_type="clean",
                 ): 1,
                 qml.decomposition.controlled_resource_rep(
                     qml.RX, {}, num_control_wires=3, num_work_wires=1
diff --git a/tests/ftqc/test_pauli_tracker.py b/tests/ftqc/test_pauli_tracker.py
index 261fd376d..80867640a 100644
--- a/tests/ftqc/test_pauli_tracker.py
+++ b/tests/ftqc/test_pauli_tracker.py
@@ -19,7 +19,14 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.ftqc.pauli_tracker import commute_clifford_op, pauli_prod, pauli_to_xz, xz_to_pauli
+from pennylane.ftqc import RotXZX, get_byproduct_corrections
+from pennylane.ftqc.pauli_tracker import (
+    _parse_mid_measurements,
+    commute_clifford_op,
+    pauli_prod,
+    pauli_to_xz,
+    xz_to_pauli,
+)
 
 _PAULIS = (qml.I, qml.X, qml.Y, qml.Z)
 
@@ -33,6 +40,12 @@ def generate_pauli_list(wire: int, num_ops: int):
     return pauli_list
 
 
+def generate_random_state(n, seed=0):
+    rng = np.random.default_rng(seed=seed)
+    input_state = rng.random(2**n) + 1j * rng.random(2**n)
+    return input_state / np.linalg.norm(input_state)
+
+
 class TestPauliTracker:
     """Test for the pauli tracker related functions."""
 
@@ -186,3 +199,120 @@ class TestPauliTracker:
             match="Please ensure xz are either 0 or 1.",
         ):
             _ = commute_clifford_op(clifford_op, xz)
+
+
+class TestOfflineCorrection:
+    """Tests for byproduct operation offline corrections."""
+
+    @pytest.mark.parametrize(
+        "ops, mid_measures, expected",
+        [
+            (
+                [qml.CNOT(wires=[0, 1]), qml.X(0), qml.X(1)],
+                [1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0],
+                [[(0, 1), (0, 1)]],
+            ),
+            ([qml.S(0), qml.Y(0)], [1, 0, 0, 1], [[(1, 0)]]),
+            ([qml.H(0), qml.Z(0)], [1, 1, 0, 0], [[(1, 1)]]),
+            ([qml.RZ(0.1, wires=[0]), qml.X(0)], [1, 0, 1, 1], [[(1, 0)]]),
+        ],
+    )
+    def test_parse_mid_measurements(self, ops, mid_measures, expected):
+        tape = qml.tape.QuantumScript(ops=ops)
+        by_ops = _parse_mid_measurements(tape, mid_measures)
+        assert by_ops == expected
+
+    @pytest.mark.parametrize("num_shots", [250])
+    @pytest.mark.parametrize(
+        "ops",
+        [
+            [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[0])],
+            [qml.RZ(0.1, wires=[0]), RotXZX(0.1, 0.2, 0.3, wires=[0])],
+            [qml.MultiControlledX(wires=[0, 1, 2, 3], control_values=[0, 1, 0])],
+        ],
+    )
+    def test_unsupported_tape_offline_correction(self, num_shots, ops):
+        measurements = [qml.sample(wires=[0])]
+
+        tape = qml.tape.QuantumScript(ops=ops, measurements=measurements, shots=num_shots)
+
+        mid_meas = [random.choice([0, 1]) for _ in range(8)]
+
+        with pytest.raises(NotImplementedError):
+            _ = get_byproduct_corrections(tape, mid_meas, [1])
+
+    @pytest.mark.parametrize(
+        "ops, mid_meas, sample_wires, measures_expected",
+        [
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[1]), qml.Z(0)],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [0],
+                [[1], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[2]), qml.Z(0)],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [0],
+                [[1], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.S(wires=[0])],
+                [0, 0, 1, 1, 1, 1, 0, 0],
+                [0],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 0, 0, 1],
+                [0],
+                [[0], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 0, 1, 0, 1, 0, 0, 1],
+                [0],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), qml.RZ(0.3, wires=[1])],
+                [0, 1, 1, 0, 1, 1, 0, 0],
+                [1],
+                [[0], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, 0.2, 0.3, wires=[1])],
+                [0, 1, 1, 0, 1, 0, 0, 0],
+                [1],
+                [[0], [0]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 0, 0, 1],
+                [1],
+                [[0], [1]],
+            ),
+            (
+                [qml.RZ(0.1, wires=[0]), RotXZX(0.1, -0.1, 0.3, wires=[1])],
+                [0, 1, 0, 1, 1, 1, 0, 1],
+                [1],
+                [[1], [1]],
+            ),
+            (
+                [qml.RZ(-0.1, wires=[0]), qml.RZ(0.1, wires=[1]), qml.X(wires=[1])],
+                [0, 1, 0, 1, 1, 1, 0, 1],
+                [0, 1],
+                [[1, 0], [1, 1]],
+            ),
+        ],
+    )
+    def test_non_clifford_offline_correction(self, ops, mid_meas, sample_wires, measures_expected):
+        measurements = [qml.sample(wires=sample_wires)]
+
+        tape = qml.tape.QuantumScript(ops=ops, measurements=measurements)
+
+        measures, expected = measures_expected
+
+        corrected_meas = get_byproduct_corrections(tape, mid_meas, measures)
+
+        assert np.allclose(corrected_meas, expected)
diff --git a/tests/gradients/core/test_metric_tensor.py b/tests/gradients/core/test_metric_tensor.py
index cf546dcf1..e4a7971be 100644
--- a/tests/gradients/core/test_metric_tensor.py
+++ b/tests/gradients/core/test_metric_tensor.py
@@ -973,6 +973,33 @@ def autodiff_metric_tensor(ansatz, num_wires):
 class TestFullMetricTensor:
     num_wires = 3
 
+    @pytest.mark.external
+    def test_catalyst_compatibility(self):
+        """Test that the metric tensor can be executed with catalyst."""
+
+        pytest.importorskip("catalyst")
+        jax = pytest.importorskip("jax")
+
+        def ansatz(params, wires=None):  # pylint: disable=unused-argument
+            qml.RX(params[0], 0)
+            qml.RX(params[1], 1)
+            qml.CNOT((0, 1))
+            qml.RX(params[2], 0)
+            qml.RY(params[3], 1)
+
+        @qml.qnode(qml.device("lightning.qubit", wires=4))
+        def circuit(params):
+            ansatz(params)
+            return qml.expval(qml.Z(0)), qml.expval(qml.Z(1))
+
+        x = jax.numpy.array([1.0, 2.0, 3.0, 4.0])
+
+        qjit_res = qml.qjit(qml.metric_tensor(circuit))(x)
+
+        cost_autograd = autodiff_metric_tensor(ansatz, num_wires=3)(qml.numpy.array(x))
+
+        assert qml.math.allclose(qjit_res, cost_autograd)
+
     @pytest.mark.autograd
     @pytest.mark.parametrize("ansatz, params", zip(fubini_ansatze, fubini_params))
     @pytest.mark.parametrize("interface", ["auto", "autograd"])
diff --git a/tests/io/qasm_interpreter/test_interpreter.py b/tests/io/qasm_interpreter/test_interpreter.py
index 8781df57f..158fa904a 100644
--- a/tests/io/qasm_interpreter/test_interpreter.py
+++ b/tests/io/qasm_interpreter/test_interpreter.py
@@ -156,6 +156,31 @@ class TestInterpreter:
 
         assert q.queue == [Identity("0q"), Hadamard("2q"), PauliX("1q"), PauliY("2q")]
 
+    def test_end_statement(self):
+        # parse the QASM program
+        ast = parse(
+            """
+            qubit q0;
+            qubit q1;
+            ch q0, q1;
+            cx q1, q0;
+            end;
+            cy q0, q1;
+            cz q1, q0;
+            swap q0, q1;
+            """,
+            permissive=True,
+        )
+
+        # execute the callable
+        with queuing.AnnotatedQueue() as q:
+            QasmInterpreter().interpret(ast, context={"wire_map": None, "name": "end-early"})
+
+        assert q.queue == [
+            CH(wires=["q0", "q1"]),
+            CNOT(wires=["q1", "q0"]),
+        ]
+
     def test_mod_with_declared_param(self):
 
         # parse the QASM program
diff --git a/tests/math/test_catalyst_math.py b/tests/math/test_catalyst_math.py
new file mode 100644
index 000000000..bf29a99f4
--- /dev/null
+++ b/tests/math/test_catalyst_math.py
@@ -0,0 +1,39 @@
+# 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.
+"""
+Tests for compatibility between the math module and catalyst.
+"""
+import pytest
+
+import pennylane as qml
+
+
+@pytest.mark.external
+def test_catalyst_integration():
+    """Test that scatter_element_add can be used with catalyst by specifying indices_are_sorted and unique_indices."""
+
+    jnp = pytest.importorskip("jax.numpy")
+
+    @qml.qjit
+    def f(x, y):
+        return qml.math.scatter_element_add(
+            x, ((0, 1), (0, 1)), y, indices_are_sorted=True, unique_indices=True
+        )
+
+    x0 = jnp.zeros((2, 2))
+    y = jnp.array([1.0, 2.0])
+
+    out = f(x0, y)
+    expected = jnp.array([[1.0, 0.0], [0.0, 2.0]])
+    assert qml.math.allclose(out, expected)
diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py
index f33a9b00e..6ee4f581d 100644
--- a/tests/measurements/test_classical_shadow.py
+++ b/tests/measurements/test_classical_shadow.py
@@ -1147,6 +1147,7 @@ def hadamard_circuit_mixed(wires, shots=10000, interface="autograd"):
     return circuit
 
 
+@pytest.mark.slow
 def test_hadamard_expval_mixed(k=1, obs=obs_hadamard, expected=expected_hadamard):
     """Test that the expval estimation is correct for a uniform
     superposition of qubits"""
diff --git a/tests/ops/op_math/test_controlled_decompositions.py b/tests/ops/op_math/test_controlled_decompositions.py
index 0793d0740..02f25501a 100644
--- a/tests/ops/op_math/test_controlled_decompositions.py
+++ b/tests/ops/op_math/test_controlled_decompositions.py
@@ -16,19 +16,22 @@ Tests for the controlled decompositions.
 """
 
 import itertools
+from collections import defaultdict
 
 import numpy as np
 import pytest
+from scipy import sparse
 
 import pennylane as qml
 from pennylane import math
 from pennylane.ops import ctrl_decomp_bisect, ctrl_decomp_zyz
+from pennylane.ops.functions.assert_valid import _test_decomposition_rule
 from pennylane.ops.op_math.controlled import _is_single_qubit_special_unitary
 from pennylane.ops.op_math.controlled_decompositions import (
-    _decompose_mcx_with_many_workers,
+    _decompose_mcx_with_many_workers_old,
     _decompose_mcx_with_one_worker_b95,
     _decompose_mcx_with_one_worker_kg24,
-    _decompose_mcx_with_two_workers,
+    _decompose_mcx_with_two_workers_old,
     _decompose_multicontrolled_unitary,
     _decompose_recursive,
     decompose_mcx,
@@ -38,6 +41,10 @@ from pennylane.ops.op_math.decompositions.controlled_decompositions import (
     _bisect_compute_b,
     _ctrl_decomp_bisect_md,
     _ctrl_decomp_bisect_od,
+    _decompose_mcx_with_no_worker,
+    _decompose_mcx_with_two_workers,
+    decompose_mcx_with_many_workers,
+    decompose_mcx_with_one_worker,
 )
 from pennylane.wires import Wires
 
@@ -48,6 +55,19 @@ def _matrix_adjoint(matrix: np.ndarray):
     return math.transpose(math.conj(matrix))
 
 
+def _tape_to_matrix(tape, wire_order):
+    """Convert a tape to a sparse matrix representation."""
+    result = sparse.eye(2 ** len(wire_order))
+    for op in tape.operations:
+        op_matrix = (
+            op.sparse_matrix(wire_order=wire_order)
+            if op.has_sparse_matrix
+            else sparse.csr_matrix(op.matrix(wire_order=wire_order))
+        )
+        result = op_matrix @ result
+    return result
+
+
 def record_from_list(func):
     """
     Decorates a function to
@@ -777,35 +797,34 @@ class TestMCXDecomposition:
                 work_wire_type="blah",
             )
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 6))
+    @pytest.mark.unit
+    @pytest.mark.parametrize("n_ctrl_wires", [3, 4, 5])
     def test_decomposition_with_many_workers(self, n_ctrl_wires):
-        """Test that the decomposed MultiControlledX gate performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there are many work wires."""
-        # pylint: disable=protected-access
-        control_wires = range(n_ctrl_wires)
-        target_wire = n_ctrl_wires
-        work_wires = range(n_ctrl_wires + 1, 2 * n_ctrl_wires + 1)
+        """Test that the decomposed MultiControlledX gate produces the correct matrix."""
 
-        dev = qml.device("default.qubit", wires=2 * n_ctrl_wires + 1)
+        control_wires = list(range(1, n_ctrl_wires + 1))
+        target_wire = 0
+        num_work_wires = n_ctrl_wires - 2
+        work_wires = list(range(n_ctrl_wires + 1, n_ctrl_wires + 1 + num_work_wires))
 
+        mcx = qml.MultiControlledX(wires=control_wires + [target_wire], work_wires=work_wires)
         with qml.queuing.AnnotatedQueue() as q:
-            _decompose_mcx_with_many_workers(control_wires, target_wire, work_wires)
-        tape = qml.tape.QuantumScript.from_queue(q)
-        assert all(isinstance(op, qml.Toffoli) for op in tape.operations)
+            decompose_mcx_with_many_workers(wires=mcx.wires, **mcx.hyperparameters)
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 1))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            for op in tape.operations:
-                op.queue()
-            return qml.probs(wires=range(n_ctrl_wires + 1))
+        # Verify that the resource estimate is correct.
+        resource = decompose_mcx_with_many_workers.compute_resources(**mcx.resource_params)
+        expected_gate_counts = {k: v for k, v in resource.gate_counts.items() if v > 0}
+        actual_gate_counts = defaultdict(int)
+        for _op in q.queue:
+            resource_rep = qml.resource_rep(type(_op), **_op.resource_params)
+            actual_gate_counts[resource_rep] += 1
+        assert actual_gate_counts == expected_gate_counts
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 1)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 1)))
+        tape = qml.tape.QuantumScript.from_queue(q)
+        matrix = _tape_to_matrix(tape, wire_order=control_wires + work_wires + [target_wire])
+        expected_matrix = mcx.sparse_matrix(wire_order=control_wires + work_wires + [target_wire])
+
+        assert qml.math.allclose(matrix, expected_matrix)
 
     @pytest.mark.parametrize("n_ctrl_wires", range(3, 6))
     def test_decomposition_with_one_worker_b95(self, n_ctrl_wires):
@@ -838,137 +857,157 @@ class TestMCXDecomposition:
         ).T
         assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 1)))
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 10))
-    def test_decomposition_with_one_clean_worker_kg24(self, n_ctrl_wires):
-        """Test that the decomposed MCX gate using 1 clean ancilla performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there is one work wire."""
+    @pytest.mark.parametrize("work_wire_type", ["clean", "dirty"])
+    @pytest.mark.parametrize("n_ctrl_wires", [3, 4, 5, 6, 7, 8, 9])
+    def test_decomposition_with_one_worker(self, n_ctrl_wires, work_wire_type):
+        """Test that the decompose_mcx_with_one_worker is correct."""
 
-        # pylint: disable=protected-access
-        control_wires = Wires(range(n_ctrl_wires))
-        target_wire = n_ctrl_wires
-        work_wires = n_ctrl_wires + 1
+        control_wires = list(range(1, n_ctrl_wires + 1))
+        target_wire = 0
+        work_wire = n_ctrl_wires + 1
 
-        dev = qml.device("default.qubit", wires=n_ctrl_wires + 2)
+        # The MultiControlledX instance to test.
+        mcx = qml.MultiControlledX(
+            wires=control_wires + [target_wire],
+            work_wires=work_wire,
+            work_wire_type=work_wire_type,
+        )
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 1))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            record_from_list(_decompose_mcx_with_one_worker_kg24)(
-                control_wires, target_wire, work_wires, work_wire_type="clean"
-            )
-            return qml.probs(wires=range(n_ctrl_wires + 1))
+        with qml.queuing.AnnotatedQueue() as q:
+            if work_wire_type == "clean":
+                qml.Projector([0], wires=work_wire)
+            decompose_mcx_with_one_worker(wires=mcx.wires, **mcx.hyperparameters)
+
+        # Verify that the resource estimate is correct.
+        resource = decompose_mcx_with_one_worker.compute_resources(**mcx.resource_params)
+        expected_gate_counts = {k: v for k, v in resource.gate_counts.items() if v > 0}
+        actual_gate_counts = defaultdict(int)
+        for _op in q.queue:
+            if isinstance(_op, qml.Projector):
+                continue
+            resource_rep = qml.resource_rep(type(_op), **_op.resource_params)
+            actual_gate_counts[resource_rep] += 1
+        assert actual_gate_counts == expected_gate_counts
+
+        # Verify that the decomposition produces an equivalent matrix.
+        tape = qml.tape.QuantumScript.from_queue(q)
+        all_wires = mcx.wires + [work_wire]
+        matrix = _tape_to_matrix(tape, wire_order=all_wires)
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 1)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 1)))
+        equivalent_op = mcx
+        if work_wire_type == "clean":
+            equivalent_op @= qml.Projector([0], wires=work_wire)
+        expected_matrix = equivalent_op.sparse_matrix(wire_order=all_wires)
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 10))
-    def test_decomposition_with_one_dirty_worker_kg24(self, n_ctrl_wires):
-        """Test that the decomposed MCX gate using 1 dirty ancilla performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there is one work wire."""
+        assert qml.math.allclose(matrix, expected_matrix)
 
-        # pylint: disable=protected-access
-        control_wires = Wires(range(n_ctrl_wires))
-        target_wire = n_ctrl_wires
-        work_wires = n_ctrl_wires + 1
+    @pytest.mark.parametrize("work_wire_type", ["clean", "dirty"])
+    @pytest.mark.parametrize("n_ctrl_wires", [3, 4, 5, 6, 7, 8, 9, 10])
+    def test_decomposition_with_two_workers(self, n_ctrl_wires, work_wire_type):
+        """Test that the decomposed MCX gate using 2 work wires produce the correct matrix."""
 
-        dev = qml.device("default.qubit", wires=n_ctrl_wires + 2)
+        control_wires = list(range(1, n_ctrl_wires + 1))
+        target_wire = 0
+        work_wires = [n_ctrl_wires + 1, n_ctrl_wires + 2]
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 2))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            record_from_list(_decompose_mcx_with_one_worker_kg24)(
-                control_wires, target_wire, work_wires, work_wire_type="dirty"
-            )
-            return qml.probs(wires=range(n_ctrl_wires + 2))
+        # The MultiControlledX instance to test.
+        mcx = qml.MultiControlledX(
+            wires=control_wires + [target_wire],
+            work_wires=work_wires,
+            work_wire_type=work_wire_type,
+        )
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 2)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 2)))
+        with qml.queuing.AnnotatedQueue() as q:
+            if work_wire_type == "clean":
+                qml.Projector([0, 0], wires=work_wires)
+            _decompose_mcx_with_two_workers(mcx.wires, work_wires, work_wire_type)
+
+        # Verify that the resource estimate is correct.
+        resource = _decompose_mcx_with_two_workers.compute_resources(**mcx.resource_params)
+        expected_gate_counts = {k: v for k, v in resource.gate_counts.items() if v > 0}
+        actual_gate_counts = defaultdict(int)
+        for _op in q.queue:
+            if isinstance(_op, qml.Projector):
+                continue
+            resource_rep = qml.resource_rep(type(_op), **_op.resource_params)
+            actual_gate_counts[resource_rep] += 1
+        assert actual_gate_counts == expected_gate_counts
+
+        # Verify that the decomposition produces an equivalent matrix.
+        tape = qml.tape.QuantumScript.from_queue(q)
+        all_wires = mcx.wires + work_wires
+        matrix = _tape_to_matrix(tape, wire_order=all_wires)
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 10))
-    def test_decomposition_with_two_clean_workers(self, n_ctrl_wires):
-        """Test that the decomposed MCX gate using 2 clean ancillae performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there is one work wire."""
+        equivalent_op = mcx
+        if work_wire_type == "clean":
+            equivalent_op @= qml.Projector([0, 0], wires=work_wires)
+        expected_matrix = equivalent_op.sparse_matrix(wire_order=all_wires)
 
-        # pylint: disable=protected-access
-        control_wires = Wires(range(n_ctrl_wires))
-        target_wire = n_ctrl_wires
-        work_wires = Wires([n_ctrl_wires + 1, n_ctrl_wires + 2])
+        assert qml.math.allclose(matrix, expected_matrix)
 
-        dev = qml.device("default.qubit", wires=n_ctrl_wires + 3)
+    @pytest.mark.parametrize("n_ctrl_wires", [3, 4, 5, 6, 7, 8, 9, 10])
+    def test_decomposition_with_no_workers(self, n_ctrl_wires):
+        """Test that the decomposed MCX gate using 2 work wires produce the correct matrix."""
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 1))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            record_from_list(_decompose_mcx_with_two_workers)(
-                control_wires, target_wire, work_wires, work_wire_type="clean"
-            )
-            return qml.probs(wires=range(n_ctrl_wires + 1))
+        control_wires = list(range(1, n_ctrl_wires + 1))
+        target_wire = 0
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 1)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 1)))
+        # The MultiControlledX instance to test.
+        mcx = qml.MultiControlledX(wires=control_wires + [target_wire])
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 10))
-    def test_decomposition_with_two_dirty_workers(self, n_ctrl_wires):
-        """Test that the decomposed MCX gate using 2 dirty ancillae performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there is one work wire."""
+        with qml.queuing.AnnotatedQueue() as q:
+            _decompose_mcx_with_no_worker(mcx.wires)
+
+        # Verify that the resource estimate is correct.
+        resource = _decompose_mcx_with_no_worker.compute_resources(**mcx.resource_params)
+        expected_gate_counts = {k: v for k, v in resource.gate_counts.items() if v > 0}
+        actual_gate_counts = defaultdict(int)
+        for _op in q.queue:
+            resource_rep = qml.resource_rep(type(_op), **_op.resource_params)
+            actual_gate_counts[resource_rep] += 1
+        assert actual_gate_counts == expected_gate_counts
+
+        # Verify that the decomposition produces an equivalent matrix.
+        tape = qml.tape.QuantumScript.from_queue(q)
+        matrix = _tape_to_matrix(tape, wire_order=mcx.wires)
 
-        # pylint: disable=protected-access
-        control_wires = Wires(range(n_ctrl_wires))
-        target_wire = n_ctrl_wires
-        work_wires = Wires([n_ctrl_wires + 1, n_ctrl_wires + 2])
+        expected_matrix = mcx.sparse_matrix()
 
-        dev = qml.device("default.qubit", wires=n_ctrl_wires + 3)
+        assert qml.math.allclose(matrix, expected_matrix)
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 3))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            record_from_list(_decompose_mcx_with_two_workers)(
-                control_wires, target_wire, work_wires, work_wire_type="dirty"
-            )
-            return qml.probs(wires=range(n_ctrl_wires + 3))
+    @pytest.mark.parametrize("n_ctrl_wires", [3, 4, 5, 6, 7])
+    def test_decompose_mcx_old(self, n_ctrl_wires):
+        """Test that the decompose_mcx produces the correct decomposition."""
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 3)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 3)))
+        control_wires = list(range(1, n_ctrl_wires + 1))
+        target_wire = 0
 
-    @pytest.mark.parametrize("n_ctrl_wires", range(3, 8))
-    def test_decomposition_with_no_workers(self, n_ctrl_wires):
-        """Test that the decomposed MultiControlledX gate performs the same unitary as the
-        matrix-based version by checking if U^dagger U applies the identity to each basis
-        state. This test focuses on the case where there is no work wires."""
+        # The MultiControlledX instance to test.
+        mcx = qml.MultiControlledX(wires=control_wires + [target_wire])
 
-        # pylint: disable=protected-access
-        control_wires = Wires(range(n_ctrl_wires))
-        target_wire = n_ctrl_wires
+        decomp = decompose_mcx(control_wires, target_wire, [])
 
-        dev = qml.device("default.qubit", wires=n_ctrl_wires + 1)
+        # Verify that the decomposition produces an equivalent matrix.
+        tape = qml.tape.QuantumScript(decomp)
+        matrix = _tape_to_matrix(tape, wire_order=mcx.wires)
+        expected_matrix = mcx.sparse_matrix()
 
-        @qml.qnode(dev)
-        def f(bitstring):
-            qml.BasisState(bitstring, wires=range(n_ctrl_wires + 1))
-            qml.MultiControlledX(wires=list(control_wires) + [target_wire])
-            record_from_list(decompose_mcx)(control_wires, Wires(target_wire), work_wires=Wires([]))
-            return qml.probs(wires=range(n_ctrl_wires + 1))
+        assert qml.math.allclose(matrix, expected_matrix)
 
-        u = np.array(
-            [f(np.array(b)) for b in itertools.product(range(2), repeat=n_ctrl_wires + 1)]
-        ).T
-        assert np.allclose(u, np.eye(2 ** (n_ctrl_wires + 1)))
+    @pytest.mark.parametrize(
+        "test_mcx",
+        [
+            qml.MultiControlledX(wires=[1, 0]),
+            qml.MultiControlledX(wires=[1, 0], control_values=[0]),
+            qml.MultiControlledX(wires=[2, 1, 0]),
+            qml.MultiControlledX(wires=[2, 1, 0], control_values=[0, 1]),
+        ],
+    )
+    def test_mcx_decompositions(self, test_mcx):
+        """Tests that MCX can be resolved into CNOT and Toffoli properly."""
+
+        for rule in qml.list_decomps(qml.MultiControlledX):
+            _test_decomposition_rule(test_mcx, rule)
 
     @pytest.mark.parametrize("work_wire_type", ["clean", "dirty"])
     @pytest.mark.parametrize("n_ctrl_wires", range(3, 10))
@@ -996,7 +1035,7 @@ class TestMCXDecomposition:
                 work_wire_type,
             )
         else:
-            expected_decomp = _decompose_mcx_with_many_workers(
+            expected_decomp = _decompose_mcx_with_many_workers_old(
                 Wires(control_wires),
                 target_wire,
                 Wires(work_wires),
@@ -1014,14 +1053,14 @@ class TestMCXDecomposition:
         computed_decomp = op.decomposition()
 
         if n_ctrl_wires > 4:
-            expected_decomp = _decompose_mcx_with_two_workers(
+            expected_decomp = _decompose_mcx_with_two_workers_old(
                 Wires(control_wires),
                 target_wire,
                 Wires(work_wires),
                 work_wire_type,
             )
         else:
-            expected_decomp = _decompose_mcx_with_many_workers(
+            expected_decomp = _decompose_mcx_with_many_workers_old(
                 Wires(control_wires),
                 target_wire,
                 Wires(work_wires),
@@ -1038,6 +1077,6 @@ class TestMCXDecomposition:
         work_wires = Wires([6])
 
         with pytest.raises(ValueError, match="At least 2 work wires are needed"):
-            _ = _decompose_mcx_with_two_workers(
+            _ = _decompose_mcx_with_two_workers_old(
                 control_wires, target_wire, work_wires, work_wire_type="clean"
             )
diff --git a/tests/ops/op_math/test_controlled_ops.py b/tests/ops/op_math/test_controlled_ops.py
index 5d615b8f1..cbd08ad53 100644
--- a/tests/ops/op_math/test_controlled_ops.py
+++ b/tests/ops/op_math/test_controlled_ops.py
@@ -63,9 +63,15 @@ class TestControlledQubitUnitary:
         with pytest.raises(qml.operation.DecompositionUndefinedError):
             op.decomposition()
 
-    def test_flatten_unflatten(self):
-        """Test that the operation can be flattened and unflattened"""
-        op = qml.ControlledQubitUnitary(np.eye(2), wires=(1, 2, 3))
+    @pytest.mark.parametrize(
+        "op",
+        [
+            qml.ControlledQubitUnitary(np.eye(2), wires=(1, 2, 3)),
+            qml.ControlledQubitUnitary(np.eye(2), wires=(1, 2, 3, 4)),
+        ],
+    )
+    def test_standard_validity(self, op):
+        """Test that the operation is valid."""
         qml.ops.functions.assert_valid(op, skip_differentiation=True, heuristic_resources=True)
 
     def test_noniterable_base(self):
diff --git a/tests/ops/op_math/test_exp.py b/tests/ops/op_math/test_exp.py
index 1b2f76ac3..4d875f44e 100644
--- a/tests/ops/op_math/test_exp.py
+++ b/tests/ops/op_math/test_exp.py
@@ -25,6 +25,7 @@ from pennylane.operation import (
     ParameterFrequenciesUndefinedError,
 )
 from pennylane.ops.op_math import Evolution, Exp
+from pennylane.ops.op_math.exp import trotter_decomp
 
 
 @pytest.mark.parametrize("constructor", (qml.exp, Exp))
@@ -609,6 +610,66 @@ class TestDecomposition:
         matrix = qml.prod(*op.decomposition()).matrix()
         assert qml.math.allclose(matrix, op.matrix())
 
+    @pytest.mark.integration
+    @pytest.mark.usefixtures("enable_graph_decomposition")
+    @pytest.mark.parametrize(
+        "coeff, hamiltonian",
+        [
+            (3j, qml.Hamiltonian([1, 2, 3], [qml.X(0), qml.Y(1), qml.Z(2)])),
+            (3, qml.Hamiltonian([1j, 2j, 3j], [qml.X(0), qml.Y(1), qml.Z(2)])),
+            (-1j, qml.Y(1) + 3 * qml.PauliX(0) @ qml.PauliZ(2)),
+            (-1, 1j * qml.PauliY(1) + 3j * qml.PauliX(0) @ qml.PauliZ(2)),
+            (0.3j, qml.Z(0) @ qml.Y(1)),
+            (0.3j, qml.Z(0)),
+        ],
+    )
+    def test_trotter_decomposition_integration_graph(self, coeff, hamiltonian):
+        """Tests that the trotter decomposition works in the new graph-based system."""
+
+        op = qml.exp(hamiltonian, coeff, num_steps=100)
+        tape = qml.tape.QuantumScript([op])
+
+        [decomp_tape], _ = qml.transforms.decompose(tape, gate_set={"PauliRot"})
+        actual_matrix = qml.matrix(decomp_tape, wire_order=op.wires)
+        expected_matrix = qml.matrix(op, wire_order=op.wires)
+        assert qml.math.allclose(actual_matrix, expected_matrix)
+
+    def test_trotter_decomposition_condition(self):
+        """Tests the condition of the trotter decomposition rule."""
+
+        hamiltonian = qml.RX(0.5, wires=0)
+        op = qml.exp(hamiltonian, coeff=1.5, num_steps=100)
+        assert not trotter_decomp.is_applicable(**op.resource_params)
+
+        hamiltonian = qml.X(0)
+        op = qml.exp(hamiltonian, coeff=1.5)
+        assert not trotter_decomp.is_applicable(**op.resource_params)
+
+        hamiltonian = qml.X(0) + qml.RX(0.5, wires=1)
+        op = qml.exp(hamiltonian, coeff=1.5, num_steps=100)
+        assert not trotter_decomp.is_applicable(**op.resource_params)
+
+    @pytest.mark.integration
+    @pytest.mark.usefixtures("enable_graph_decomposition")
+    @pytest.mark.parametrize(
+        "coeff, hamiltonian",
+        [
+            (0.3j, qml.Z(0) @ qml.Y(1)),
+            (0.5, 0.1j * qml.Y(0) @ qml.I(1) @ qml.Z(2)),
+        ],
+    )
+    def test_pauli_decomposition_integration_graph(self, coeff, hamiltonian):
+        """Tests that the pauli decomposition works in the new graph-based system."""
+
+        op = qml.exp(hamiltonian, coeff)
+        tape = qml.tape.QuantumScript([op])
+
+        [decomp_tape], _ = qml.transforms.decompose(tape, gate_set={"PauliRot"})
+        assert len(decomp_tape) == 1
+        actual_matrix = qml.matrix(decomp_tape, wire_order=op.wires)
+        expected_matrix = qml.matrix(op, wire_order=op.wires)
+        assert qml.math.allclose(actual_matrix, expected_matrix)
+
 
 class TestMiscMethods:
     """Test other representation methods."""
diff --git a/tests/python_compiler/test_quantum_dialect.py b/tests/python_compiler/test_quantum_dialect.py
index f5d33ee97..1bc19cfa8 100644
--- a/tests/python_compiler/test_quantum_dialect.py
+++ b/tests/python_compiler/test_quantum_dialect.py
@@ -14,11 +14,14 @@
 
 """Unit test module for pennylane/compiler/python_compiler/quantum_dialect.py."""
 
+import io
+
 import pytest
 
 # pylint: disable=wrong-import-position
 
 xdsl = pytest.importorskip("xdsl")
+filecheck = pytest.importorskip("filecheck")
 
 pytestmark = pytest.mark.external
 
@@ -102,3 +105,49 @@ def test_all_attributes_names(attr):
         expected_name is not None
     ), f"Unexpected attribute {attr_class_name} found in QuantumDialect"
     assert attr.name == expected_name
+
+
+def test_assembly_format():
+    program = """
+    // CHECK: [[QREG:%.+]] = "test.op"() : () -> !quantum.reg
+    %qreg = "test.op"() : () -> !quantum.reg
+
+    // CHECK: [[COMPBASIS:%.+]] = quantum.compbasis qreg [[QREG]] : !quantum.obs
+    %compbasis = quantum.compbasis qreg %qreg : !quantum.obs
+
+    // CHECK: quantum.counts [[COMPBASIS]] : tensor<2xf64>, tensor<2xi64>
+    %eigvals, %counts = quantum.counts %compbasis : tensor<2xf64>, tensor<2xi64>
+
+    // CHECK: [[PARAM:%.+]] = "test.op"() : () -> f64
+    %cst = "test.op"() : () -> f64
+
+    // CHECK: [[QUBIT:%.+]] = "quantum.extract"([[QREG]])
+    %qubit = "quantum.extract"(%qreg) <{idx_attr = 0 : i64}> : (!quantum.reg) -> !quantum.bit
+
+    // CHEKC: [[QUBIT2:%.+]] = quantum.custom "RX"([[PARAM]]) [[QUBIT]]
+    %out_qubit = quantum.custom "RX"(%cst) %qubit : !quantum.bit
+    """
+
+    ctx = xdsl.context.Context()
+    from xdsl.dialects import builtin, func, test
+
+    ctx.load_dialect(builtin.Builtin)
+    ctx.load_dialect(func.Func)
+    ctx.load_dialect(test.Test)
+    ctx.load_dialect(QuantumDialect)
+
+    module = xdsl.parser.Parser(ctx, program).parse_module()
+
+    from filecheck.finput import FInput
+    from filecheck.matcher import Matcher
+    from filecheck.options import parse_argv_options
+    from filecheck.parser import Parser, pattern_for_opts
+
+    opts = parse_argv_options(["filecheck", __file__])
+    matcher = Matcher(
+        opts,
+        FInput("no-name", str(module)),
+        Parser(opts, io.StringIO(program), *pattern_for_opts(opts)),
+    )
+
+    assert matcher.run() == 0