Complete legacy tracing for dynamic magic wires

doc/releases/changelog-dev.md

@@ -2,13 +2,13 @@
 
 <h3>New features since last release</h3>
 
-* You can now dynamically prepare magic T states inside captured Catalyst workflows using
-  ``qp.allocate(state="magic-T")`` and ``qp.allocate(state="magic-T-adj")``, which makes it
-  easier to compile FTQC-style routines that need T-state ancillas on the fly (for example
-  TemporaryAND) with ``qjit(capture=True)``.
+* You can now dynamically allocate magic T states (and their conjugates) with Catalyst on
+  both the capture and legacy tracing paths. This lets you build workflows that combine
+  dynamic wire allocation, mid-circuit measurements, and control flow — for example preparing
+  magic-T ancillas, applying a Pauli-product measurement, and branching with ``qp.cond``.
 
   ```python
-  @qjit(capture=True)
+  @qjit  # works with capture=True or capture=False
   @qnode(dev)
   def circuit():
       qb = qp.allocate(state="magic-T")
@@ -19,6 +19,7 @@
   Requires a PennyLane release that supports the new ``AllocateState`` values
   (see [PennyLane #9846](https://github.com/PennyLaneAI/pennylane/pull/9846)).
   [(#3029)](https://github.com/PennyLaneAI/catalyst/pull/3029)
+  [(#3027)](https://github.com/PennyLaneAI/catalyst/pull/3027)
 
 * The `local-random` unitary folding option for :func:`~.mitigate_with_zne` is now implemented,
   reproducing Mitiq's ``fold_gates_at_random``: every gate is folded ``floor((scale_factor-1)/2)``

frontend/catalyst/api_extensions/control_flow.py

@@ -1631,7 +1631,7 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
                 partial(inner_trace.new_arg, source_info=current_source_info()), [new_qreg]
             )[0]
             qrp_out = trace_quantum_operations(
-                inner_tape, device, qreg_in, ctx, inner_trace, **kwargs
+                inner_tape, device, qreg_in, ctx, inner_trace, parent_qrp=qrp, **kwargs
             )
             qreg_out = qrp_out.actualize()
 
@@ -1683,8 +1683,10 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
                 apply_reverse_transform=self.apply_reverse_transform,
                 num_implicit_inputs=num_implicit_inputs,
                 preserve_dimensions=not expansion_strategy.input_unshare_variables,
-            )
+            ),
+            num_device_wires=qrp.num_device_wires,
         )
+        qrp2.inherit_dynamic_wire_state(qrp)
         return qrp2
 
 
@@ -1754,7 +1756,7 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
                 partial(body_trace.new_arg, source_info=current_source_info()), [AbstractQreg()]
             )[0]
             qrp_out = trace_quantum_operations(
-                body_tape, device, qreg_in, ctx, body_trace, **kwargs
+                body_tape, device, qreg_in, ctx, body_trace, parent_qrp=qrp, **kwargs
             )
             qreg_out = qrp_out.actualize()
             arg_expanded_tracers = expand_args(
@@ -1804,8 +1806,10 @@ def trace_quantum(self, ctx, device, trace, qrp, **kwargs) -> QRegPromise:
                 body_nconsts=len(body_consts),
                 num_implicit_inputs=num_implicit_inputs,
                 preserve_dimensions=not expansion_strategy.input_unshare_variables,
-            )
+            ),
+            num_device_wires=qrp.num_device_wires,
         )
+        qrp2.inherit_dynamic_wire_state(qrp)
         return qrp2
 
 
@@ -1888,7 +1892,13 @@ def trace_quantum_branches(op, ctx, device, trace, qrp, **kwargs) -> QRegPromise
                 partial(region.trace.new_arg, source_info=current_source_info()), [new_qreg]
             )[0]
             qreg_out = trace_quantum_operations(
-                region.quantum_tape, device, qreg_in, ctx, region.trace, **kwargs
+                region.quantum_tape,
+                device,
+                qreg_in,
+                ctx,
+                region.trace,
+                parent_qrp=qrp,
+                **kwargs,
             ).actualize()
 
             constants = []
@@ -1928,6 +1938,8 @@ def trace_quantum_branches(op, ctx, device, trace, qrp, **kwargs) -> QRegPromise
             out_expanded_tracers=out_expanded_classical_tracers,
             branch_jaxprs=branch_jaxprs,
             num_implicit_outputs=num_implicit_outputs[0],
-        )
+        ),
+        num_device_wires=qrp.num_device_wires,
     )
+    qrp2.inherit_dynamic_wire_state(qrp)
     return qrp2

frontend/catalyst/api_extensions/quantum_operators.py

@@ -418,6 +418,29 @@ def __init__(
     # pylint: disable=too-many-arguments
     def trace_quantum(self, ctx, device, trace, qrp, postselect_mode=None) -> QRegPromise:
         qubit = qrp.extract(self.wires)[0]
+        if qrp.qref_mode:
+            from catalyst.from_plxpr.qref_jax_primitives import (  # pylint: disable=import-outside-toplevel
+                qref_measure_p,
+            )
+
+            kwargs = {}
+            if postselect_mode != "hw-like":
+                kwargs["postselect"] = self.postselect
+
+            original_binder = self.binder
+            self.binder = qref_measure_p.bind
+            try:
+                self.bind_overwrite_classical_tracers(
+                    trace,
+                    in_expanded_tracers=[qubit],
+                    out_expanded_tracers=self.out_classical_tracers,
+                    num_quantum_outs=0,
+                    **kwargs,
+                )
+            finally:
+                self.binder = original_binder
+            return qrp
+
         if postselect_mode == "hw-like":
             qubit2 = self.bind_overwrite_classical_tracers(
                 trace,
@@ -460,6 +483,26 @@ def __init__(
 
     def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
         qubits = qrp.extract(self.wires)
+        if qrp.qref_mode:
+            from catalyst.from_plxpr.qref_jax_primitives import (  # pylint: disable=import-outside-toplevel
+                qref_pauli_measure_p,
+            )
+
+            original_binder = self.binder
+            self.binder = qref_pauli_measure_p.bind
+            try:
+                self.bind_overwrite_classical_tracers(
+                    trace,
+                    in_expanded_tracers=qubits,
+                    out_expanded_tracers=self.out_classical_tracers,
+                    num_quantum_outs=0,
+                    pauli_word=self.pauli_word,
+                    qubits_len=len(qubits),
+                )
+            finally:
+                self.binder = original_binder
+            return qrp
+
         out_qubits = self.bind_overwrite_classical_tracers(
             trace,
             in_expanded_tracers=qubits,
@@ -594,7 +637,7 @@ def trace_quantum(self, ctx, device, _trace, qrp, **kwargs) -> QRegPromise:
                 partial(body_trace.new_arg, source_info=current_source_info()), [AbstractQreg()]
             )[0]
             qrp_out = trace_quantum_operations(
-                body_tape, device, qreg_in, ctx, body_trace, **kwargs
+                body_tape, device, qreg_in, ctx, body_trace, parent_qrp=qrp, **kwargs
             )
             qreg_out = qrp_out.actualize()
             body_jaxpr, _, body_consts = body_trace.frame.to_jaxpr2(
@@ -608,7 +651,8 @@ def trace_quantum(self, ctx, device, _trace, qrp, **kwargs) -> QRegPromise:
             args_tree=args_tree,
             jaxpr=ClosedJaxpr(convert_constvars_jaxpr(body_jaxpr), ()),
         )
-        qrp2 = QRegPromise(op_results[-1])
+        qrp2 = QRegPromise(op_results[-1], num_device_wires=qrp.num_device_wires)
+        qrp2.inherit_dynamic_wire_state(qrp)
         return qrp2
 
     @property

frontend/catalyst/jax_extras/tracing.py

@@ -208,9 +208,24 @@ def __lt__(self, other):
         return self.id < other.id
 
 
+def _eqn_invar_counts(eqn: JaxprEqn) -> Set[int]:
+    """Return the set of JAX variable ids used as inputs to an equation."""
+    counts: Set[int] = set()
+    for v in eqn.invars:
+        if hasattr(v, "val") and isinstance(v.val, jax._src.core.Var):
+            actual_var = v.val
+        elif isinstance(v, jax._src.core.Var):
+            actual_var = v
+        else:
+            continue
+        counts.add(actual_var.count)
+    return counts
+
+
 def sort_eqns(
     eqns: List[JaxprEqn | Callable[[], JaxprEqn]],
     forced_order_primitives: Set[JaxprPrimitive],
+    dealloc_after_use_primitives: Set[JaxprPrimitive] | None = None,
 ) -> List[JaxprEqn | Callable[[], JaxprEqn]]:
     """Topologically sort TracingEqns in a unsorted list of equations, based on their
     input/output variables and additional criterias."""
@@ -253,6 +268,17 @@ def sort_eqns(
         for b in boxes[i + 1 :]:
             b.parents.append(q)  # [3]
 
+    if dealloc_after_use_primitives:
+        for b in boxes:
+            if b.e.primitive not in dealloc_after_use_primitives:
+                continue
+            dealloc_invars = _eqn_invar_counts(b.e)
+            for other in boxes:
+                if other is b:
+                    continue
+                if dealloc_invars & _eqn_invar_counts(other.e):
+                    b.parents.append(other)
+
     sorted_boxes = stable_toposort(boxes)
 
     # Restore the original callables/wrappers for the sorted equations