Controlled regions keeps work_wire_type

.dep-versions

@@ -8,7 +8,7 @@ enzyme=v0.0.238
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt'
-pennylane=0.46.0-dev130
+pennylane=0.46.0-dev131
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -811,6 +811,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Controlled quantum regions now keeps track of the `work_wire_type` passed from `qp.ctrl`.
+  [(#3366)](https://github.com/PennyLaneAI/catalyst/pull/3366)
+
 * MLIR passes that recreate `scf.for`, `scf.while`, `scf.if`, and `scf.index_switch` ops
   (`ctrl-lowering`, `adjoint-lowering`, `--convert-to-value-semantics`) now preserve discardable
   attributes such as `catalyst.estimated_iterations`, `catalyst.estimated_probability`, and

doc/requirements.txt

@@ -34,4 +34,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.46.0-dev27
 pennylane-lightning==0.46.0-dev27
-pennylane==0.46.0-dev130
+pennylane==0.46.0-dev131

frontend/catalyst/from_plxpr/qfunc_interpreter.py

@@ -888,9 +888,18 @@ def handle_adjoint_transform(self, *invals, jaxpr, lazy, n_consts):
     )
 
 
-# pylint: disable=unused-argument
+# pylint: disable=unused-argument,too-many-arguments
 @PLxPRToQuantumJaxprInterpreter.register_primitive(plxpr_ctrl_transform_prim)
-def handle_ctrl_transform(self, *invals, jaxpr, n_control, control_values, work_wires, n_consts):
+def handle_ctrl_transform(
+    self,
+    *invals,
+    jaxpr,
+    n_control,
+    control_values,
+    work_wires,
+    n_consts,
+    work_wire_type="borrowed",
+):
     """Interpret a control transform, then re-bind it for lowering to a `qref.ctrl` region op."""
     consts = invals[:n_consts]
     args = invals[n_consts:-n_control]
@@ -912,6 +921,7 @@ def handle_ctrl_transform(self, *invals, jaxpr, n_control, control_values, work_
         jaxpr=body.jaxpr,
         control_values=control_values,
         work_wires=work_wires,
+        work_wire_type=work_wire_type,
         n_consts=len(body.consts),
     )
 

frontend/catalyst/from_plxpr/qref_jax_primitives.py

@@ -745,6 +745,7 @@ def _pl_ctrl_lowering(
     control_values,
     work_wires,
     n_consts,
+    work_wire_type="borrowed",
 ):
     """Lower a `qp.ctrl` transform to a `qref.ctrl` region op."""
     body_invals = plxpr_invals[: len(plxpr_invals) - n_control]

frontend/test/pytest/from_plxpr/test_from_plxpr.py

@@ -520,6 +520,24 @@ def circuit():
         assert s_eqn.primitive == qref_operator_p
         assert s_eqn.params["n_ctrls"] == 0
 
+    @pytest.mark.usefixtures("use_capture")
+    @pytest.mark.parametrize("work_wire_type", ("zeroed", "borrowed"))
+    def test_qfunc_ctrl_work_wire_type(self, work_wire_type):
+        """Test that a qfunc control with a work wire type is converted and keeps the type."""
+
+        @qp.qnode(qp.device("lightning.qubit", wires=3))
+        def circuit():
+            qp.ctrl(qp.S, control=1, work_wires=[2], work_wire_type=work_wire_type)(0)
+            return qp.state()
+
+        plxpr = jax.make_jaxpr(circuit)()
+        catalyst_xpr = from_plxpr(plxpr)()
+        qfunc_xpr = catalyst_xpr.eqns[0].params["call_jaxpr"]
+
+        eqn = next(e for e in qfunc_xpr.eqns if e.primitive == ctrl_transform_prim)
+        assert eqn.params["work_wires"] == qp.wires.Wires([2])
+        assert eqn.params["work_wire_type"] == work_wire_type
+
     @pytest.mark.parametrize("as_qfunc", (True, False))
     def test_doubly_ctrl(self, as_qfunc):
         """Test doubly controlled op."""

frontend/test/pytest/python_interface/transforms/quantum/test_xdsl_diagonalize_measurements.py

@@ -538,9 +538,6 @@ def test_with_composite_observables(self, capture_mode):
         """Test the transform works for an observable built using operator arithmetic
         (sprod, prod, sum)"""
 
-        if capture_mode:
-            pytest.xfail("Requires fixes from PennyLane #10156")
-
         dev = qp.device("lightning.qubit", wires=3)
 
         @qp.qnode(dev)
@@ -683,9 +680,6 @@ def test_with_composite_observables(self, capture_mode):
         """Test the transform works for an observable built using operator arithmetic
         (sprod, prod, sum)"""
 
-        if capture_mode:
-            pytest.xfail("Requires fixes from PennyLane #10156")
-
         dev = qp.device("lightning.qubit", wires=3)
 
         @qp.qnode(dev)