Repository: PennyLaneAI/catalyst
Pull request #1576: Fix a bug when chained passes are used on a qnode
Author: login: mehrdad2m
Created: 2025-03-18T21:05:50Z | Merged: 2025-03-21T19:39:46Z

**Context:**
In the current status of the Catalyst pass api, pass decorators recieve a qnode and return a classical wrapper function. This make stacking multiple passes/pipelines not possible. For more details please refer to Issue #1547.

In the current status of the Catalyst pass api, there is significant code duplication across different pass decorators (cancel_inverses, merge_rotations, ions_decomposition, to_ppr, etc.) and the main pass API functions (pipeline, apply_pass, apply_pass_plugin). Each of these implements similar QNode wrapping logic, type checking, and pass pipeline management. This duplication makes the code harder to maintain and increases the risk of inconsistencies.

**Description of the Change:**
Added a new QNodeWrapper class which can be used to return in the decorators. 
Added a PassPipelineWrapper class that extends QNodeWrapper to handle pass pipeline management
refactored the pass-api to eliminate duplicate codes. 

**Benefits:**
Catalyst's support for chained pass decorators is restored.
Reduced code duplicates.

**Possible Drawbacks:**
The new class hierarchy adds a layer of abstraction that might make the code slightly harder to understand. Ideally we liked for the decorators to receive and return a qnode. However, this is not possible at this point.

**Related GitHub Issues:**
closes issue #1547 
Note that as a side benefit, this PR resolves issue #1548 
[sc-85468]
[sc-85470]


Diff:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 63efcd2670..069db8a70d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -190,6 +190,14 @@
   autograph in conjunction with catalyst passes causes a crash.
   [(#1541)](https://github.com/PennyLaneAI/catalyst/pull/1541)
 
+* Fixes an issue ([(#1548)](https://github.com/PennyLaneAI/catalyst/issues/1548)) where using
+  autograph in conjunction with catalyst pipeline causes a crash.
+  [(#1576)](https://github.com/PennyLaneAI/catalyst/pull/1576)
+
+* Fixes an issue ([(#1547)](https://github.com/PennyLaneAI/catalyst/issues/1547)) where using
+  chained catalyst passe decorators causes a crash.
+  [(#1576)](https://github.com/PennyLaneAI/catalyst/pull/1576)
+
 <h3>Internal changes ⚙️</h3>
 
 * Updated the call signature for the PLXPR `qnode_prim` primitive.
diff --git a/frontend/catalyst/passes/builtin_passes.py b/frontend/catalyst/passes/builtin_passes.py
index ddcc4b29ee..434cb48f59 100644
--- a/frontend/catalyst/passes/builtin_passes.py
+++ b/frontend/catalyst/passes/builtin_passes.py
@@ -14,12 +14,7 @@
 
 """This module exposes built-in Catalyst MLIR passes to the frontend."""
 
-import copy
-import functools
-
-import pennylane as qml
-
-from catalyst.passes.pass_api import Pass
+from catalyst.passes.pass_api import PassPipelineWrapper
 
 
 ## API ##
@@ -133,20 +128,7 @@ def circuit(x: float):
         %2 = quantum.namedobs %out_qubits[ PauliZ] : !quantum.obs
         %3 = quantum.expval %2 : f64
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_cancel_inverses"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("remove-chained-self-inverse"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "remove-chained-self-inverse")
 
 
 def merge_rotations(qnode):
@@ -210,20 +192,7 @@ def circuit(x: float):
     >>> circuit(0.54)
     Array(0.5965506257017892, dtype=float64)
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_merge_rotations"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("merge-rotations"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "merge-rotations")
 
 
 def ions_decomposition(qnode):  # pragma: nocover
@@ -345,18 +314,7 @@ def circuit():
         %out_qubits_8 = quantum.custom "RY"(%cst_2) %out_qubits_6#1 : !quantum.bit
         %out_qubits_9 = quantum.custom "RY"(%cst_2) %out_qubits_7 : !quantum.bit
     """
-
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    @functools.wraps(qnode)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("ions-decomposition"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return qnode(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "ions-decomposition")
 
 
 def to_ppr(qnode):
@@ -431,20 +389,7 @@ def circuit():
         . . .
 
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_to_ppr"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("to_ppr"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "to_ppr")
 
 
 def commute_ppr(qnode):
@@ -502,17 +447,4 @@ def circuit():
             . . .
 
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_commute_ppr"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("commute_ppr"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "commute_ppr")
diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py
index defcd9eaf7..adb05ba0ca 100644
--- a/frontend/catalyst/passes/pass_api.py
+++ b/frontend/catalyst/passes/pass_api.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import functools
 from importlib.metadata import entry_points
 from pathlib import Path
@@ -26,6 +25,60 @@
 
 
 ## API ##
+class QNodeWrapper:
+    """A wrapper class for QNodes that preserves the pass pipeline when called.
+
+    This class is used internally by pass decorators to ensure that compiler passes
+    are properly applied when the QNode is called.
+
+    Args:
+        qnode (QNode or callable): The QNode or wrapped QNode function to call
+    """
+
+    def __init__(self, qnode):
+        self.qnode = qnode
+        if isinstance(qnode, QNodeWrapper):
+            self.original_qnode = qnode.original_qnode
+            self.func = qnode.func
+        elif isinstance(qnode, qml.QNode):
+            self.original_qnode = qnode
+            self.func = qnode.func
+        else:
+            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
+
+        # Update the wrapper to preserve the original function's metadata (signature, annotations,
+        # etc.) This is crucial for qjit to properly inspect the function signature and make
+        # compilation decisions.
+        functools.update_wrapper(self, self.func)
+
+    def __call__(self, *args, **kwargs):
+        return self.qnode(*args, **kwargs)
+
+
+class PassPipelineWrapper(QNodeWrapper):
+    """A QNodeWrapper subclass that adds passes to the pipeline when called.
+
+    This class can handle both single passes and multiple passes in a pipeline.
+    For single passes, it takes a pass name and optional flags/options.
+    For multiple passes, it takes a dictionary of pass configurations.
+    """
+
+    def __init__(self, qnode, pass_name_or_pipeline, *flags, **valued_options):
+        super().__init__(qnode)
+        self.pass_name_or_pipeline = pass_name_or_pipeline
+        self.flags = flags
+        self.valued_options = valued_options
+
+    def __call__(self, *args, **kwargs):
+        if EvaluationContext.is_tracing():
+            pass_pipeline = kwargs.pop("pass_pipeline", [])
+            pass_pipeline += dictionary_to_list_of_passes(
+                self.pass_name_or_pipeline, *self.flags, **self.valued_options
+            )
+            kwargs["pass_pipeline"] = pass_pipeline
+        return super().__call__(*args, **kwargs)
+
+
 def pipeline(pass_pipeline: PipelineDict):
     """Configures the Catalyst MLIR pass pipeline for quantum circuit transformations for a QNode
     within a qjit-compiled program.
@@ -114,22 +167,8 @@ def fn(x):
     will always take precedence over global pass pipelines.
     """
 
-    def _decorator(qnode=None):
-        if not isinstance(qnode, qml.QNode):
-            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-        clone = copy.copy(qnode)
-        clone.__name__ += "_transformed"
-
-        @functools.wraps(clone)
-        def wrapper(*args, **kwargs):
-            if EvaluationContext.is_tracing():
-                passes = kwargs.pop("pass_pipeline", [])
-                passes += dictionary_to_list_of_passes(pass_pipeline)
-                kwargs["pass_pipeline"] = passes
-            return clone(*args, **kwargs)
-
-        return wrapper
+    def _decorator(qnode):
+        return PassPipelineWrapper(qnode, pass_pipeline)
 
     return _decorator
 
@@ -161,21 +200,7 @@ def module():
     """
 
     def decorator(qnode):
-
-        if not isinstance(qnode, qml.QNode):
-            # Technically, this apply pass is general enough that it can apply to
-            # classical functions too. However, since we lack the current infrastructure
-            # to denote a function, let's limit it to qnodes
-            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-        @functools.wraps(qnode)
-        def qnode_call(*args, **kwargs):
-            pass_pipeline = kwargs.get("pass_pipeline", [])
-            pass_pipeline.append(Pass(pass_name, *flags, **valued_options))
-            kwargs["pass_pipeline"] = pass_pipeline
-            return qnode(*args, **kwargs)
-
-        return qnode_call
+        return PassPipelineWrapper(qnode, pass_name, *flags, **valued_options)
 
     return decorator
 
@@ -215,20 +240,7 @@ def module():
         raise FileNotFoundError(f"File '{path_to_plugin}' does not exist.")
 
     def decorator(qnode):
-        if not isinstance(qnode, qml.QNode):
-            # Technically, this apply pass is general enough that it can apply to
-            # classical functions too. However, since we lack the current infrastructure
-            # to denote a function, let's limit it to qnodes
-            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-        @functools.wraps(qnode)
-        def qnode_call(*args, **kwargs):
-            pass_pipeline = kwargs.get("pass_pipeline", [])
-            pass_pipeline.append(PassPlugin(path_to_plugin, pass_name, *flags, **valued_options))
-            kwargs["pass_pipeline"] = pass_pipeline
-            return qnode(*args, **kwargs)
-
-        return qnode_call
+        return PassPipelineWrapper(qnode, pass_name, *flags, **valued_options)
 
     return decorator
 
@@ -332,21 +344,29 @@ def __init__(
 
 
 ## PRIVATE ##
-def dictionary_to_list_of_passes(pass_pipeline: PipelineDict):
-    """Convert dictionary of passes into list of passes."""
+def dictionary_to_list_of_passes(pass_pipeline: PipelineDict | str, *flags, **valued_options):
+    """Convert dictionary of passes or single pass name into list of passes.
 
-    if pass_pipeline == None:
+    Args:
+        pass_pipeline: Either a dictionary of pass configurations or a single pass name.
+        *flags: Optional flags for single pass
+        **valued_options: Optional valued options for single pass
+    """
+    if pass_pipeline is None:
         return []
 
-    if type(pass_pipeline) != dict:
-        return pass_pipeline
+    if isinstance(pass_pipeline, str):
+        return [Pass(pass_pipeline, *flags, **valued_options)]
+
+    if isinstance(pass_pipeline, dict):
+        passes = []
+        pass_names = _API_name_to_pass_name()
+        for API_name, pass_options in pass_pipeline.items():
+            name = pass_names.get(API_name, API_name)
+            passes.append(Pass(name, **pass_options))
+        return passes
 
-    passes = []
-    pass_names = _API_name_to_pass_name()
-    for API_name, pass_options in pass_pipeline.items():
-        name = pass_names.get(API_name, API_name)
-        passes.append(Pass(name, **pass_options))
-    return passes
+    return pass_pipeline
 
 
 def _API_name_to_pass_name():
diff --git a/frontend/test/lit/test_peephole_optimizations.py b/frontend/test/lit/test_peephole_optimizations.py
index ab0520f298..b8ff96cde1 100644
--- a/frontend/test/lit/test_peephole_optimizations.py
+++ b/frontend/test/lit/test_peephole_optimizations.py
@@ -30,7 +30,7 @@
 
 from catalyst import pipeline
 from catalyst.debug import get_compilation_stage
-from catalyst.passes import cancel_inverses, merge_rotations
+from catalyst.passes import apply_pass, cancel_inverses, merge_rotations
 
 
 def flush_peephole_opted_mlir_to_iostream(QJIT):
@@ -76,8 +76,8 @@ def test_pipeline_lowering_workflow(x):
     # CHECK-NEXT: transform.yield
     print_mlir(test_pipeline_lowering_workflow, 1.2)
 
-    # CHECK: {{%.+}} = call @test_pipeline_lowering_workflow_transformed_0(
-    # CHECK: func.func public @test_pipeline_lowering_workflow_transformed_0(
+    # CHECK: {{%.+}} = call @test_pipeline_lowering_workflow_0(
+    # CHECK: func.func public @test_pipeline_lowering_workflow_0(
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
@@ -130,12 +130,12 @@ def test_pipeline_lowering_keep_original_workflow(x):
 
     # CHECK: func.func public @jit_test_pipeline_lowering_keep_original_workflow
     # CHECK: {{%.+}} = call @f_0(
-    # CHECK: {{%.+}} = call @f_transformed_0(
+    # CHECK: {{%.+}} = call @f_1_0(
     # CHECK: func.func public @f_0(
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
-    # CHECK: func.func public @f_transformed_0(
+    # CHECK: func.func public @f_1_0(
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
@@ -256,13 +256,13 @@ def h(x):
     print_mlir(global_wf)
 
     # CHECK: func.func public @jit_global_wf()
-    # CHECK {{%.+}} = call @g(
-    # CHECK {{%.+}} = call @h_transformed(
-    # CHECK: func.func public @g_0
+    # CHECK {{%.+}} = call @g_0(
+    # CHECK {{%.+}} = call @h_0(
+    # CHECK: func.func public @g_0(
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
-    # CHECK: func.func public @h_transformed_0
+    # CHECK: func.func public @h_0
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
@@ -272,6 +272,213 @@ def h(x):
 
 test_pipeline_lowering_globloc_override()
 
+
+def test_chained_pipeline_lowering():
+    """
+    Test that chained pipelines are correctly lowered.
+    """
+    pipeline1 = {
+        "cancel_inverses": {},
+        "merge_rotations": {},
+    }
+    pipeline2 = {
+        "cancel_inverses": {},
+    }
+
+    @qjit()
+    @pipeline(pipeline1)
+    @pipeline(pipeline2)
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def test_chained_pipeline_lowering_workflow(x: float):
+        qml.Hadamard(wires=[1])
+        qml.RX(x, wires=[0])
+        qml.RX(-x, wires=[0])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    print(test_chained_pipeline_lowering_workflow.mlir)
+
+    # # CHECK: transform.named_sequence @__transform_main
+    # # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # # CHECK-NEXT: transform.yield
+
+
+test_chained_pipeline_lowering()
+
+
+def test_chained_pipeline_lowering_keep_original():
+    """
+    Test when the chained pipeline and the original qnode are both used,
+    and the original is correctly kept and untransformed.
+    """
+    pipeline1 = {
+        "cancel_inverses": {},
+        "merge_rotations": {},
+    }
+    pipeline2 = {
+        "cancel_inverses": {},
+    }
+
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def f(x: float):
+        qml.RX(x, wires=[0])
+        qml.Hadamard(wires=[1])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    f_pipeline1 = pipeline(pipeline1)(f)
+    f_pipeline2 = pipeline(pipeline2)(f_pipeline1)
+
+    @qjit()
+    def test_chained_pipeline_lowering_keep_original_workflow(x: float):
+        return f(x), f_pipeline1(x), f_pipeline2(x)
+
+    # COM: The unchanged qnode
+    # CHECK: transform.named_sequence @__transform_main
+    # COM: The qnode after pipeline1
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # COM: The qnode after pipeline2
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print(test_chained_pipeline_lowering_keep_original_workflow.mlir)
+
+
+test_chained_pipeline_lowering_keep_original()
+
+
+def test_chained_apply_passes():
+    """
+    Test that chained passes are correctly applied in sequence using apply_pass.
+    """
+
+    @qjit()
+    @apply_pass("remove-chained-self-inverse")
+    @apply_pass("merge-rotations")
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def test_chained_apply_passes_workflow(x: float):
+        qml.Hadamard(wires=[1])
+        qml.RX(x, wires=[0])
+        qml.RX(-x, wires=[0])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    print(test_chained_apply_passes_workflow.mlir)
+
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+
+
+test_chained_apply_passes()
+
+
+def test_chained_apply_passes_keep_original():
+    """
+    Test when the chained apply passes and the original qnode are both used,
+    and the original is correctly kept and untransformed.
+    """
+
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def f(x: float):
+        qml.RX(x, wires=[0])
+        qml.Hadamard(wires=[1])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    f_pass1 = apply_pass("remove-chained-self-inverse")(f)
+    f_pass2 = apply_pass("merge-rotations")(f_pass1)
+
+    @qjit()
+    def test_chained_apply_passes_keep_original_workflow(x: float):
+        return f(x), f_pass1(x), f_pass2(x)
+
+    # COM: The unchanged qnode
+    # CHECK: transform.named_sequence @__transform_main
+    # COM: The qnode after pipeline1
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # COM: The qnode after merge_rotations
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print(test_chained_apply_passes_keep_original_workflow.mlir)
+
+
+test_chained_apply_passes_keep_original()
+
+
+def test_chained_peephole_passes():
+    """
+    Test that chained peephole passes are correctly applied in sequence.
+    """
+
+    @qjit()
+    @cancel_inverses
+    @merge_rotations
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def test_chained_peephole_passes_workflow(x: float):
+        qml.Hadamard(wires=[1])
+        qml.RX(x, wires=[0])
+        qml.RX(-x, wires=[0])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    print(test_chained_peephole_passes_workflow.mlir)
+
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+
+
+test_chained_peephole_passes()
+
+
+def test_chained_peephole_passes_keep_original():
+    """
+    Test when the chained peephole passes and the original qnode are both used,
+    and the original is correctly kept and untransformed.
+    """
+
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def f(x: float):
+        qml.RX(x, wires=[0])
+        qml.Hadamard(wires=[1])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    f_pass1 = cancel_inverses(f)
+    f_pass2 = merge_rotations(f_pass1)
+
+    @qjit()
+    def test_chained_peephole_passes_keep_original_workflow(x: float):
+        return f(x), f_pass1(x), f_pass2(x)
+
+    # COM: The unchanged qnode
+    # CHECK: transform.named_sequence @__transform_main
+    # COM: The qnode after cancel_inverses
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # COM: The qnode after merge_rotations
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print(test_chained_peephole_passes_keep_original_workflow.mlir)
+
+
+test_chained_peephole_passes_keep_original()
+
+
 #
 # autograph
 #
@@ -279,7 +486,7 @@ def h(x):
 
 def test_single_pass_with_autograph():
     """
-    Test tha peephole optimization work with autograph
+    Test that peephole optimization works with autograph
     """
 
     @qjit(autograph=True, target="mlir")
@@ -299,6 +506,35 @@ def f(x: float):
 
 test_single_pass_with_autograph()
 
+
+def test_pipeline_with_autograph():
+    """
+    Test that pipeline works with autograph
+    """
+
+    my_pipeline = {
+        "cancel_inverses": {},
+        "merge_rotations": {},
+    }
+
+    @qjit(autograph=True, target="mlir")
+    @pipeline(my_pipeline)
+    @qml.qnode(qml.device("lightning.qubit", wires=1))
+    def f(x: float):
+        qml.RX(x, wires=0)
+        qml.RX(x, wires=0)
+        qml.Hadamard(wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    # CHECK: transform.named_sequence @__transform_main(
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}}
+    # CHECK-NEXT: transform.yield
+    print(f.mlir)
+
+
+test_pipeline_with_autograph()
+
 #
 # cancel_inverses
 #
@@ -469,12 +705,12 @@ def f(x: float):
 
     # CHECK-LABEL: public @jit_test_cancel_inverses_keep_original_workflow0
     # CHECK: {{%.+}} = call @f_0({{%.+}})
-    # CHECK-NOT: {{%.+}} = call @f_cancel_inverses_0({{%.+}})
+    # CHECK-NOT: {{%.+}} = call @f_0({{%.+}})
     # CHECK-LABEL: public @f_0({{%.+}})
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NEXT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NEXT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
-    # CHECK-NOT: public @f_cancel_inverses_0
+    # CHECK-NOT: public @f_0
     @qjit(keep_intermediate=True)
     def test_cancel_inverses_keep_original_workflow0():
         return f(1.0)
@@ -483,9 +719,9 @@ def test_cancel_inverses_keep_original_workflow0():
     flush_peephole_opted_mlir_to_iostream(test_cancel_inverses_keep_original_workflow0)
 
     # CHECK-LABEL: public @jit_test_cancel_inverses_keep_original_workflow1
-    # CHECK: {{%.+}} = call @f_cancel_inverses_0({{%.+}})
+    # CHECK: {{%.+}} = call @f_0({{%.+}})
     # CHECK-NOT: {{%.+}} = call @f_0({{%.+}})
-    # CHECK-LABEL: public @f_cancel_inverses_0({{%.+}})
+    # CHECK-LABEL: public @f_0({{%.+}})
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
@@ -499,12 +735,12 @@ def test_cancel_inverses_keep_original_workflow1():
 
     # CHECK-LABEL: public @jit_test_cancel_inverses_keep_original_workflow2
     # CHECK: {{%.+}} = call @f_0({{%.+}})
-    # CHECK: {{%.+}} = call @f_cancel_inverses_0({{%.+}})
+    # CHECK: {{%.+}} = call @f_1_0({{%.+}})
     # CHECK-LABEL: public @f_0({{%.+}})
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NEXT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NEXT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
-    # CHECK-LABEL: public @f_cancel_inverses_0({{%.+}})
+    # CHECK-LABEL: public @f_1_0({{%.+}})
     # CHECK: {{%.+}} = quantum.custom "RX"({{%.+}}) {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
     # CHECK-NOT: {{%.+}} = quantum.custom "Hadamard"() {{%.+}} : !quantum.bit
diff --git a/frontend/test/pytest/test_autograph.py b/frontend/test/pytest/test_autograph.py
index 7e70af7d7d..e4cb13e3ba 100644
--- a/frontend/test/pytest/test_autograph.py
+++ b/frontend/test/pytest/test_autograph.py
@@ -27,7 +27,6 @@
 
 from catalyst import *
 from catalyst.autograph.transformer import TRANSFORMER
-from catalyst.passes import merge_rotations
 from catalyst.utils.dummy import dummy_func
 from catalyst.utils.exceptions import CompileError
 
diff --git a/frontend/test/pytest/test_peephole_optimizations.py b/frontend/test/pytest/test_peephole_optimizations.py
index 029cf26657..7e10d69ecb 100644
--- a/frontend/test/pytest/test_peephole_optimizations.py
+++ b/frontend/test/pytest/test_peephole_optimizations.py
@@ -169,5 +169,28 @@ def classical_func():
     test_cancel_inverses_not_on_qnode()
 
 
+def test_chained_passes():
+    """
+    Test that chained passes are present in the transform passes.
+    """
+
+    @qjit()
+    @cancel_inverses
+    @merge_rotations
+    @qml.qnode(qml.device("lightning.qubit", wires=2))
+    def test_chained_apply_passes_workflow(x: float):
+        qml.Hadamard(wires=[1])
+        qml.RX(x, wires=[0])
+        qml.RX(-x, wires=[0])
+        qml.Hadamard(wires=[1])
+        return qml.expval(qml.PauliY(wires=0))
+
+    assert "remove-chained-self-inverse" in test_chained_apply_passes_workflow.mlir
+    assert "merge-rotations" in test_chained_apply_passes_workflow.mlir
+
+
+test_chained_passes()
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])


Reviews:
user: login: erick-xanadu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:24:43Z
user: login: mehrdad2m
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:42:32Z
user: login: erick-xanadu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:45:02Z
user: login: paul0403
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T15:50:34Z
user: login: paul0403
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T17:43:29Z
user: login: erick-xanadu
body: :100: :partying_face: 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T18:14:21Z

Review comments:
diff_hunk: @@ -26,6 +25,61 @@
 
 
 ## API ##
+class QNodeWrapper:
+    """A wrapper class for QNodes that preserves the pass pipeline when called.
+
+    This class is used internally by pass decorators to ensure that compiler passes
+    are properly applied when the QNode is called.
+
+    Args:
+        qnode (QNode or callable): The QNode or wrapped QNode function to call
+    """
+
+    def __init__(self, qnode):
+        self.qnode = qnode
+        if isinstance(qnode, QNodeWrapper):
+            self.original_qnode = qnode.original_qnode
+            self.func = qnode.func
+        elif isinstance(qnode, qml.QNode):
+            self.original_qnode = qnode
+            self.func = qnode.func
+        else:
+            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
+
+        functools.update_wrapper(self, self.func)
path: frontend/catalyst/passes/pass_api.py
user: login: erick-xanadu
body: This line is a bit confusing. :sweat_smile: 
created_at: 2025-03-19T19:24:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -26,6 +25,61 @@
 
 
 ## API ##
+class QNodeWrapper:
+    """A wrapper class for QNodes that preserves the pass pipeline when called.
+
+    This class is used internally by pass decorators to ensure that compiler passes
+    are properly applied when the QNode is called.
+
+    Args:
+        qnode (QNode or callable): The QNode or wrapped QNode function to call
+    """
+
+    def __init__(self, qnode):
+        self.qnode = qnode
+        if isinstance(qnode, QNodeWrapper):
+            self.original_qnode = qnode.original_qnode
+            self.func = qnode.func
+        elif isinstance(qnode, qml.QNode):
+            self.original_qnode = qnode
+            self.func = qnode.func
+        else:
+            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
+
+        functools.update_wrapper(self, self.func)
path: frontend/catalyst/passes/pass_api.py
user: login: mehrdad2m
body: I will add a comment. This is so that the original function metadata (signature, annotations, etc.) is passed to the wrapper. Otherwise, `qjit` would see a `QnodeWrapper` object instead of the original `qnode` which can cause issue e.g. `qjit` inspects the signature of the `user_function` to judge whether aot compilation should be triggered.

I can also be more explicit and set whatever is needed like `__signiture__`, etc., but I thought that this would be more general and can avoid issue in future.
created_at: 2025-03-19T19:42:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -26,6 +25,61 @@
 
 
 ## API ##
+class QNodeWrapper:
+    """A wrapper class for QNodes that preserves the pass pipeline when called.
+
+    This class is used internally by pass decorators to ensure that compiler passes
+    are properly applied when the QNode is called.
+
+    Args:
+        qnode (QNode or callable): The QNode or wrapped QNode function to call
+    """
+
+    def __init__(self, qnode):
+        self.qnode = qnode
+        if isinstance(qnode, QNodeWrapper):
+            self.original_qnode = qnode.original_qnode
+            self.func = qnode.func
+        elif isinstance(qnode, qml.QNode):
+            self.original_qnode = qnode
+            self.func = qnode.func
+        else:
+            raise TypeError(f"A QNode is expected, got the classical function {qnode}")
+
+        functools.update_wrapper(self, self.func)
path: frontend/catalyst/passes/pass_api.py
user: login: erick-xanadu
body: The comment is good enough :) Thanks!
created_at: 2025-03-19T19:45:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -418,17 +376,4 @@ def circuit():
         %mres_0, %out_qubits_1 = qec.ppm ["Z"] %10 : !quantum.bit
         . . .
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_to_ppr"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("convert-clifford-t-to-ppr"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "convert-clifford-t-to-ppr")
path: frontend/catalyst/passes/builtin_passes.py
user: login: paul0403
body: Can we add a python lit test for this pass? Nothing too crazy, just checking the pass name in the transform dialect's "apply pass" operation is fine.
created_at: 2025-03-20T15:50:34Z
author_association: MEMBER
diff_hunk: @@ -418,17 +376,4 @@ def circuit():
         %mres_0, %out_qubits_1 = qec.ppm ["Z"] %10 : !quantum.bit
         . . .
     """
-    if not isinstance(qnode, qml.QNode):
-        raise TypeError(f"A QNode is expected, got the classical function {qnode}")
-
-    clone = copy.copy(qnode)
-    clone.__name__ += "_to_ppr"
-
-    @functools.wraps(clone)
-    def wrapper(*args, **kwargs):
-        pass_pipeline = kwargs.pop("pass_pipeline", [])
-        pass_pipeline.append(Pass("convert-clifford-t-to-ppr"))
-        kwargs["pass_pipeline"] = pass_pipeline
-        return clone(*args, **kwargs)
-
-    return wrapper
+    return PassPipelineWrapper(qnode, "convert-clifford-t-to-ppr")
path: frontend/catalyst/passes/builtin_passes.py
user: login: paul0403
body: Discussed, the ppr python decorator is special and does not go through the usual catalyst dialect in its entirety.
created_at: 2025-03-20T17:43:28Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-03-19T20:42:08Z
body: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?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 `91.30435%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 96.09%. Comparing base [(`a3d78dd`)](https://app.codecov.io/gh/PennyLaneAI/catalyst/commit/a3d78dd8bcc407638807cd18d3bc9dd85bf9fe42?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`d8364bd`)](https://app.codecov.io/gh/PennyLaneAI/catalyst/commit/d8364bd875cb49ac6776d82e0a672df81c7fc032?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 190 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [frontend/catalyst/passes/builtin\_passes.py](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?src=pr&el=tree&filepath=frontend%2Fcatalyst%2Fpasses%2Fbuiltin_passes.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-ZnJvbnRlbmQvY2F0YWx5c3QvcGFzc2VzL2J1aWx0aW5fcGFzc2VzLnB5) | 60.00% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [frontend/catalyst/passes/pass\_api.py](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?src=pr&el=tree&filepath=frontend%2Fcatalyst%2Fpasses%2Fpass_api.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-ZnJvbnRlbmQvY2F0YWx5c3QvcGFzc2VzL3Bhc3NfYXBpLnB5) | 95.12% | [1 Missing and 1 partial :warning: ](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?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             @@
##             main    #1576      +/-   ##
==========================================
+ Coverage   95.76%   96.09%   +0.32%     
==========================================
  Files          81       81              
  Lines        8652     8608      -44     
  Branches      819      815       -4     
==========================================
- Hits         8286     8272      -14     
+ Misses        308      279      -29     
+ Partials       58       57       -1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/1576?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>
author_association: NONE
user: login: erick-xanadu
created_at: 2025-03-20T18:15:40Z
body: I think there is a bit of coverage that is missing. Are you able to handle those cases? @mehrdad2m 
author_association: CONTRIBUTOR
user: login: mehrdad2m
created_at: 2025-03-20T20:45:59Z
body: > I think there is a bit of coverage that is missing. Are you able to handle those cases? @mehrdad2m

The codecov report does not make a lot of sense to me. I think the cases should be covered by the tests. 🤔 
author_association: CONTRIBUTOR
user: login: mehrdad2m
created_at: 2025-03-20T21:09:02Z
body: > > I think there is a bit of coverage that is missing. Are you able to handle those cases? @mehrdad2m
> 
> The codecov report does not make a lot of sense to me. I think the cases should be covered by the tests. 🤔

Discussed offline, apparently lit tests are not considered by codecov. Will add a single pytest to check the jaxpr.
author_association: CONTRIBUTOR
user: login: mehrdad2m
created_at: 2025-03-21T15:33:15Z
body: There are two more code cov issues remaining. one of them will be resolved by #1579. The other is related to the pass plugin passes, but I think it should be covered by [this test](https://github.com/PennyLaneAI/catalyst/blob/f3900f895bd28035e36122874f687792a600404d/standalone_plugin_wheel/standalone_plugin/test/test_mlir_plugin.py#L60). I think code coverage does not see plugin tests either
author_association: CONTRIBUTOR
user: login: paul0403
created_at: 2025-03-21T18:11:34Z
body: > I think code coverage does not see plugin tests either

I would imagine so. I think we saw that codecov just sees `/frontend/pytest`.

Edit: just checked `Makefile`, codecov sees frontend pytest, oqc/oqd tests and async tests. But yes it doesn't see plugin tests.
author_association: MEMBER
user: login: erick-xanadu
created_at: 2025-03-21T18:13:25Z
body: @mehrdad2m I'm happy with this. You can wait until the other PR is merged or merge this one before.
author_association: CONTRIBUTOR