fix adjoint and ctrl of qubit unitary and gphase (#2158)

**Context:**

The special handling for qubit unitary and gphase bypassed the handling
for operator modifiers.

**Description of the Change:**

Handle qubit unitary and gphase inside `interpret_operation` instead of
registering custom handling for the primitive.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

Fixes #2151 
[sc-102523]

Author: albi3ro

doc/releases/changelog-dev.md

@@ -35,6 +35,10 @@
   [(#2128)](https://github.com/PennyLaneAI/catalyst/pull/2128)
   [(#2133)](https://github.com/PennyLaneAI/catalyst/pull/2133)
 
+* Fixes the translation of `QubitUnitary` and `GlobalPhase` ops
+  when they are modified by adjoint or control.
+  [(##2158)](https://github.com/PennyLaneAI/catalyst/pull/2158)
+
 * Fixes the translation of a workflow with different transforms applied to different qnodes.
   [(#2167)](https://github.com/PennyLaneAI/catalyst/pull/2167)
 

frontend/catalyst/from_plxpr/control_flow.py

@@ -296,10 +296,10 @@ def handle_while_loop(
     jaxpr = ClosedJaxpr(jaxpr_body_fn, consts_body)
 
     f = partial(_calling_convention, self, jaxpr)
-    converted_body_jaxpr_branch = jax.make_jaxpr(f)(*args_plus_qreg).jaxpr
-
+    converted_body_jaxpr_branch = jax.make_jaxpr(f)(*args_plus_qreg)
+    new_consts_body = converted_body_jaxpr_branch.consts
     converted_body_closed_jaxpr_branch = ClosedJaxpr(
-        convert_constvars_jaxpr(converted_body_jaxpr_branch), ()
+        convert_constvars_jaxpr(converted_body_jaxpr_branch.jaxpr), ()
     )
 
     # Convert for condition from plxpr to Catalyst jaxpr
@@ -324,21 +324,22 @@ def remove_qreg(*args_plus_qreg):
 
         return converter(jaxpr, *args)
 
-    converted_cond_jaxpr_branch = jax.make_jaxpr(remove_qreg)(*args_plus_qreg).jaxpr
+    converted_cond_jaxpr_branch = jax.make_jaxpr(remove_qreg)(*args_plus_qreg)
     converted_cond_closed_jaxpr_branch = ClosedJaxpr(
-        convert_constvars_jaxpr(converted_cond_jaxpr_branch), ()
+        convert_constvars_jaxpr(converted_cond_jaxpr_branch.jaxpr), ()
     )
 
+    new_consts_cond = converted_cond_jaxpr_branch.consts
     # Build Catalyst compatible input values
-    while_loop_invals = [*consts_cond, *consts_body, *args_plus_qreg]
+    while_loop_invals = [*new_consts_cond, *new_consts_body, *args_plus_qreg]
 
     # Perform the binding
     outvals = while_p.bind(
         *while_loop_invals,
         cond_jaxpr=converted_cond_closed_jaxpr_branch,
         body_jaxpr=converted_body_closed_jaxpr_branch,
-        cond_nconsts=len(consts_cond),
-        body_nconsts=len(consts_body),
+        cond_nconsts=len(new_consts_cond),
+        body_nconsts=len(new_consts_body),
         nimplicit=0,
         preserve_dimensions=True,
     )

frontend/catalyst/from_plxpr/qfunc_interpreter.py

@@ -156,15 +156,16 @@ def interpret_operation(self, op, is_adjoint=False, control_values=(), control_w
         if any(not qreg.is_qubit_mode() and qreg.expired for qreg in in_qregs + in_ctrl_qregs):
             raise CompileError(f"Deallocated qubits cannot be used, but used in {op.name}.")
 
-        out_qubits = qinst_p.bind(
+        bind_fn = _special_op_bind_call.get(type(op), qinst_p.bind)
+
+        out_qubits = bind_fn(
             *[*in_qubits, *op.data, *in_ctrl_qubits, *control_values],
             op=op.name,
             qubits_len=len(op.wires),
             params_len=len(op.data),
             ctrl_len=len(control_wires),
             adjoint=is_adjoint,
         )
-
         out_non_ctrl_qubits = out_qubits[: len(out_qubits) - len(control_wires)]
         out_ctrl_qubits = out_qubits[-len(control_wires) :]
 
@@ -275,6 +276,27 @@ def __call__(self, jaxpr, *args):
         return self.eval(jaxpr.jaxpr, jaxpr.consts, *args)
 
 
+# pylint: disable=unused-argument
+def _qubit_unitary_bind_call(*invals, op, qubits_len, params_len, ctrl_len, adjoint):
+    wires = invals[:qubits_len]
+    mat = invals[qubits_len]
+    ctrl_inputs = invals[qubits_len + 1 :]
+    return unitary_p.bind(
+        mat, *wires, *ctrl_inputs, qubits_len=qubits_len, ctrl_len=ctrl_len, adjoint=adjoint
+    )
+
+
+# pylint: disable=unused-argument
+def _gphase_bind_call(*invals, op, qubits_len, params_len, ctrl_len, adjoint):
+    return gphase_p.bind(*invals[qubits_len:], ctrl_len=ctrl_len, adjoint=adjoint)
+
+
+_special_op_bind_call = {
+    qml.QubitUnitary: _qubit_unitary_bind_call,
+    qml.GlobalPhase: _gphase_bind_call,
+}
+
+
 # pylint: disable=unused-argument
 @PLxPRToQuantumJaxprInterpreter.register_primitive(qml.allocation.allocate_prim)
 def handle_qml_alloc(self, *, num_wires, state=None, restored=False):
@@ -441,22 +463,6 @@ def wrapper(*args):
     return ()
 
 
-@PLxPRToQuantumJaxprInterpreter.register_primitive(qml.QubitUnitary._primitive)
-def handle_qubit_unitary(self, *invals, n_wires):
-    """Handle the conversion from plxpr to Catalyst jaxpr for the QubitUnitary primitive"""
-    in_qregs, in_qubits = get_in_qubit_values(invals[1:], self.qubit_index_recorder, self.init_qreg)
-    outvals = unitary_p.bind(invals[0], *in_qubits, qubits_len=n_wires, ctrl_len=0, adjoint=False)
-    for in_qreg, w, new_wire in zip(in_qregs, invals[1:], outvals):
-        in_qreg[in_qreg.global_index_to_local_index(w)] = new_wire
-
-
-# pylint: disable=unused-argument
-@PLxPRToQuantumJaxprInterpreter.register_primitive(qml.GlobalPhase._primitive)
-def handle_global_phase(self, phase, *wires, n_wires):
-    """Handle the conversion from plxpr to Catalyst jaxpr for the GlobalPhase primitive"""
-    gphase_p.bind(phase, ctrl_len=0, adjoint=False)
-
-
 @PLxPRToQuantumJaxprInterpreter.register_primitive(qml.BasisState._primitive)
 def handle_basis_state(self, *invals, n_wires):
     """Handle the conversion from plxpr to Catalyst jaxpr for the BasisState primitive"""
@@ -590,13 +596,16 @@ def calling_convention(*args_plus_qreg):
         init_qreg.insert_all_dangling_qubits()
         return *retvals, converter.init_qreg.get()
 
-    _, args_tree = tree_flatten((consts, args, [qreg]))
-    converted_jaxpr_branch = jax.make_jaxpr(calling_convention)(*consts, *args, qreg).jaxpr
+    converted_jaxpr_branch = jax.make_jaxpr(calling_convention)(*args, qreg)
 
-    converted_closed_jaxpr_branch = ClosedJaxpr(convert_constvars_jaxpr(converted_jaxpr_branch), ())
+    converted_closed_jaxpr_branch = ClosedJaxpr(
+        convert_constvars_jaxpr(converted_jaxpr_branch.jaxpr), ()
+    )
+    new_consts = converted_jaxpr_branch.consts
+    _, args_tree = tree_flatten((new_consts, args, [qreg]))
     # Perform the binding
     outvals = adjoint_p.bind(
-        *consts,
+        *new_consts,
         *args,
         qreg,
         jaxpr=converted_closed_jaxpr_branch,

frontend/test/pytest/test_adjoint.py

@@ -22,9 +22,11 @@
 import pennylane.numpy as pnp
 import pytest
 from numpy.testing import assert_allclose
+from pennylane import adjoint, cond, for_loop, qjit, while_loop
 from pennylane.ops.op_math.adjoint import Adjoint, AdjointOperation
 
-from catalyst import adjoint, cond, debug, for_loop, measure, qjit, while_loop
+import catalyst
+from catalyst import debug, measure, qjit
 
 # pylint: disable=too-many-lines,missing-class-docstring,missing-function-docstring,too-many-public-methods
 
@@ -50,8 +52,17 @@ def pennylane_workflow(*args):
             qml.adjoint(quantum_func)(*args)
             return qml.state()
 
-        assert_allclose(catalyst_workflow(*args), pennylane_workflow(*args))
+        capture_enabled = qml.capture.enabled()
+        qml.capture.disable()
+        try:
+            pl_res = pennylane_workflow(*args)
+        finally:
+            if capture_enabled:
+                qml.capture.enable()
 
+        assert_allclose(catalyst_workflow(*args), pl_res)
+
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_func(self, backend):
         """Ensures that catalyst.adjoint accepts simple Python functions as argument. Makes sure
         that simple quantum gates are adjointed correctly."""
@@ -82,6 +93,7 @@ def PL_workflow():
         desired = PL_workflow()
         assert_allclose(actual, desired)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("theta, val", [(jnp.pi, 0), (-100.0, 1)])
     def test_adjoint_op(self, theta, val, backend):
         """Ensures that catalyst.adjoint accepts single PennyLane operators classes as argument."""
@@ -91,19 +103,20 @@ def test_adjoint_op(self, theta, val, backend):
         @qml.qnode(device)
         def C_workflow(theta, val):
             adjoint(qml.RY)(jnp.pi, val)
-            adjoint(qml.RZ)(theta, wires=val)
+            adjoint(qml.RZ)(theta, val)
             return qml.state()
 
         @qml.qnode(device)
         def PL_workflow(theta, val):
             qml.adjoint(qml.RY)(jnp.pi, val)
-            qml.adjoint(qml.RZ)(theta, wires=val)
+            qml.adjoint(qml.RZ)(theta, val)
             return qml.state()
 
         actual = C_workflow(theta, val)
         desired = PL_workflow(theta, val)
         assert_allclose(actual, desired)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("theta, val", [(np.pi, 0), (-100.0, 2)])
     def test_adjoint_bound_op(self, theta, val, backend):
         """Ensures that catalyst.adjoint accepts single PennyLane operators objects as argument."""
@@ -129,6 +142,7 @@ def PL_workflow(theta, val):
         desired = PL_workflow(theta, val)
         assert_allclose(actual, desired, atol=1e-6, rtol=1e-6)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("w, p", [(0, 0.5), (0, -100.0), (1, 123.22)])
     def test_adjoint_param_fun(self, w, p, backend):
         """Ensures that catalyst.adjoint accepts parameterized Python functions as arguments."""
@@ -144,21 +158,22 @@ def func(w, theta1, theta2, theta3=1):
         @qml.qnode(device)
         def C_workflow(w, theta):
             qml.PauliX(wires=0)
-            adjoint(func)(w, theta, theta2=theta)
+            adjoint(func)(w, theta, theta)
             qml.PauliY(wires=0)
             return qml.state()
 
         @qml.qnode(device)
         def PL_workflow(w, theta):
             qml.PauliX(wires=0)
-            qml.adjoint(func)(w, theta, theta2=theta)
+            qml.adjoint(func)(w, theta, theta)
             qml.PauliY(wires=0)
             return qml.state()
 
         actual = C_workflow(w, p)
         desired = PL_workflow(w, p)
         assert_allclose(actual, desired)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_nested_fun(self, backend):
         """Ensures that catalyst.adjoint allows arbitrary nesting."""
 
@@ -186,6 +201,7 @@ def PL_workflow():
 
         assert_allclose(C_workflow(), PL_workflow())
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_qubitunitary(self, backend):
         """Ensures that catalyst.adjoint supports QubitUnitary oprtations."""
 
@@ -204,6 +220,7 @@ def func():
 
         self.verify_catalyst_adjoint_against_pennylane(func, qml.device(backend, wires=2))
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_qubitunitary_dynamic_variable_loop(self, backend):
         """Ensures that catalyst.adjoint supports QubitUnitary oprtations."""
 
@@ -228,6 +245,7 @@ def loop_body(_i, s):
 
         self.verify_catalyst_adjoint_against_pennylane(func, qml.device(backend, wires=2), _input)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_multirz(self, backend):
         """Ensures that catalyst.adjoint supports MultiRZ operations."""
 
@@ -275,6 +293,7 @@ def C_workflow():
 
             C_workflow()
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_classical_loop(self, backend):
         """Checks that catalyst.adjoint supports purely-classical Control-flows."""
 
@@ -288,6 +307,7 @@ def loop(_i, s):
 
         self.verify_catalyst_adjoint_against_pennylane(func, qml.device(backend, wires=3), 0)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("pred", [True, False])
     def test_adjoint_cond(self, backend, pred):
         """Tests that the correct gates are applied in reverse in a conditional branch"""
@@ -302,6 +322,7 @@ def cond_fn():
         dev = qml.device(backend, wires=1)
         self.verify_catalyst_adjoint_against_pennylane(func, dev, pred, jnp.pi)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_while_loop(self, backend):
         """
         Tests that the correct gates are applied in reverse in a while loop with a statically
@@ -322,6 +343,7 @@ def loop_body(carried):
         dev = qml.device(backend, wires=1)
         self.verify_catalyst_adjoint_against_pennylane(func, dev, 10)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_for_loop(self, backend):
         """Tests the correct application of gates (with dynamic wires)"""
 
@@ -335,6 +357,7 @@ def loop_body(i):
         dev = qml.device(backend, wires=5)
         self.verify_catalyst_adjoint_against_pennylane(func, dev, 4)
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_while_nested(self, backend):
         """Tests the correct handling of nested while loops."""
 
@@ -367,6 +390,7 @@ def cond_otherwise():
             func, dev, 10, jnp.array([2, 4, 3, 5, 1, 7, 4, 6, 9, 10])
         )
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_nested_with_control_flow(self, backend):
         """
         Tests that nested adjoint ops produce correct results in the presence of nested control
@@ -420,6 +444,7 @@ def pennylane_workflow(*args):
 
         assert_allclose(catalyst_workflow(jnp.pi), pennylane_workflow(jnp.pi))
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_for_nested(self, backend):
         """
         Tests the adjoint op with nested and interspersed for/while loops that produce classical
@@ -542,6 +567,7 @@ def cond_fn():
         # It returns `-1` instead of `0`
         assert circuit() == qml.wires.Wires([0])
 
+    @pytest.mark.usefixtures("use_both_frontend")
     def test_adjoint_ctrl_ctrl_subroutine(self, backend):
         """https://github.com/PennyLaneAI/catalyst/issues/589"""
 
@@ -585,7 +611,7 @@ def qfunc(x):
             qml.RY(x, wires=0)
             qml.Hadamard(0)
 
-        adj_op = adjoint(qfunc)(0.7)
+        adj_op = catalyst.adjoint(qfunc)(0.7)
         decomp = adj_op.decomposition()
 
         assert len(decomp) == 2
@@ -602,7 +628,7 @@ def qfunc(x, w):
             qml.CNOT(wires=[1, w])
 
         with pytest.raises(ValueError, match="Eagerly computing the adjoint"):
-            adjoint(qfunc, lazy=False)(0.1, 0)
+            catalyst.adjoint(qfunc, lazy=False)(0.1, 0)
 
 
 #####################################################################################

frontend/test/pytest/test_global_phase.py

@@ -17,10 +17,10 @@
 import numpy as np
 import pennylane as qml
 import pytest
-
-from catalyst import cond, qjit
+from pennylane import cond, qjit
 
 
+@pytest.mark.usefixtures("use_both_frontend")
 def test_global_phase(backend):
     """Test vanilla global phase"""
     dev = qml.device(backend, wires=1)
@@ -36,6 +36,7 @@ def qnn():
     assert np.allclose(expected, observed)
 
 
+@pytest.mark.usefixtures("use_both_frontend")
 @pytest.mark.parametrize("inp", [True, False])
 def test_global_phase_in_region(backend, inp):
     """Test global phase in region"""
@@ -52,11 +53,13 @@ def cir():
         cir()
         return qml.state()
 
-    expected = qnn(inp)
     observed = qjit(qnn)(inp)
+    qml.capture.disable()
+    expected = qnn(inp)
     assert np.allclose(expected, observed)
 
 
+@pytest.mark.usefixtures("use_both_frontend")
 def test_global_phase_control(backend):
     """Test global phase controlled"""