fix translation of transforms with multiple qnodes (#2167)

**Context:**

Transforms were global for the entire workflow. We needed to make them
specific to the code they are actually applied to.

**Description of the Change:**

Create new interpreters for transforms and only append transforms to the
new interpreter.

**Benefits:**

Proper application of transforms.

**Possible Drawbacks:**

**Related GitHub Issues:**


[sc-102818]

---------

Co-authored-by: Paul <[email protected]>

Author: albi3ro

doc/releases/changelog-dev.md

@@ -35,6 +35,9 @@
   [(#2128)](https://github.com/PennyLaneAI/catalyst/pull/2128)
   [(#2133)](https://github.com/PennyLaneAI/catalyst/pull/2133)
 
+* Fixes the translation of a workflow with different transforms applied to different qnodes.
+  [(#2167)](https://github.com/PennyLaneAI/catalyst/pull/2167)
+
 <h3>Internal changes ⚙️</h3>
 
 * Refactor Catalyst pass registering so that it's no longer necessary to manually add new

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -18,6 +18,7 @@
 
 
 import warnings
+from copy import copy
 from functools import partial
 from typing import Callable
 
@@ -137,6 +138,14 @@ def f(x):
 class WorkflowInterpreter(PlxprInterpreter):
     """An interpreter that converts a qnode primitive from a plxpr variant to a catalyst jaxpr variant."""
 
+    def __copy__(self):
+        new_version = WorkflowInterpreter()
+        new_version._pass_pipeline = copy(self._pass_pipeline)
+        new_version.init_qreg = self.init_qreg
+        new_version.requires_decompose_lowering = self.requires_decompose_lowering
+        new_version.decompose_tkwargs = copy(self.decompose_tkwargs)
+        return new_version
+
     def __init__(self):
         self._pass_pipeline = []
         self.init_qreg = None
@@ -284,12 +293,13 @@ def handle_transform(
                     "Multiple decomposition transforms are not yet supported."
                 )
 
+            next_eval = copy(self)
             # Update the decompose_gateset to be used by the quantum kernel primitive
             # TODO: we originally wanted to treat decompose_gateset as a queue of
             # gatesets to be used by the decompose-lowering pass at MLIR
             # but this requires a C++ implementation of the graph-based decomposition
             # which doesn't exist yet.
-            self.decompose_tkwargs = tkwargs
+            next_eval.decompose_tkwargs = tkwargs
 
             # Note. We don't perform the compiler-specific decomposition here
             # to be able to support multiple decomposition transforms
@@ -300,7 +310,7 @@ def handle_transform(
             # in the qnode handler.
 
             # Add the decompose-lowering pass to the start of the pipeline
-            self._pass_pipeline.insert(0, Pass("decompose-lowering"))
+            next_eval._pass_pipeline.insert(0, Pass("decompose-lowering"))
 
             # We still need to construct and solve the graph based on
             # the current jaxpr based on the current gateset
@@ -313,7 +323,7 @@ def handle_transform(
 
             # final_jaxpr = jax.make_jaxpr(gds_wrapper)(*args)
             # return self.eval(final_jaxpr.jaxpr, consts, *non_const_args)
-            return self.eval(inner_jaxpr, consts, *non_const_args)
+            return next_eval.eval(inner_jaxpr, consts, *non_const_args)
 
         if catalyst_pass_name is None:
             # Use PL's ExpandTransformsInterpreter to expand this and any embedded
@@ -333,11 +343,12 @@ def wrapper(*args):
                     final_jaxpr.jaxpr, final_jaxpr.consts, targs, tkwargs, *non_const_args
                 )
 
-            return self.eval(final_jaxpr.jaxpr, final_jaxpr.consts, *non_const_args)
+            return copy(self).eval(final_jaxpr.jaxpr, final_jaxpr.consts, *non_const_args)
 
         # Apply the corresponding Catalyst pass counterpart
-        self._pass_pipeline.insert(0, Pass(catalyst_pass_name, *targs, **tkwargs))
-        return self.eval(inner_jaxpr, consts, *non_const_args)
+        next_eval = copy(self)
+        next_eval._pass_pipeline.insert(0, Pass(catalyst_pass_name, *targs, **tkwargs))
+        return next_eval.eval(inner_jaxpr, consts, *non_const_args)
 
 
 # This is our registration factory for PL transforms. The loop below iterates

frontend/test/lit/test_from_plxpr.py

@@ -421,3 +421,36 @@ def circuit3():
 
 
 test_pass_decomposition()
+
+
+def test_two_qnodes_with_different_passes_in_one_workflow():
+    """Two qnodes with different passes in one workflow."""
+
+    dev = qml.device("null.qubit", wires=1)
+
+    qml.capture.enable()
+
+    @qml.qjit(target="mlir")
+    def workflow():
+        @qml.transforms.merge_rotations
+        @qml.qnode(dev)
+        def circuit1():
+            return qml.probs()
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        def circuit2():
+            return qml.probs()
+
+        return circuit1() + circuit2()
+
+    # CHECK: module @module_circuit1 {
+    # CHECK: transform.apply_registered_pass "merge-rotations"
+    # CHECK: module @module_circuit2 {
+    # CHECK: transform.apply_registered_pass "remove-chained-self-inverse"
+
+    print(workflow.mlir)
+    qml.capture.disable()
+
+
+test_two_qnodes_with_different_passes_in_one_workflow()