Repository: PennyLaneAI/pennylane
Pull request #7710 (2025-06-18 16:47:34)

Capture dynamic qubit allocation v2
**Context:**

This PR Replaces #7673 

**Description of the Change:**

Adds allocate_prim and deallocate_prim that generate and consume dynamic wires (abstract ints) respectively.

**Benefits:**

Can capture qubit allocation for integration with catalyst.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-93191]


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7710?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`aed2c5f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/aed2c5faf7db30fd2f950b69edbb8c64f1ac8090?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`b50d7af`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b50d7af2826a36f1560b0f835aadbb223fca2203?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 13 commits behind head on allocate-option.

<details><summary>Additional details and impacted files</summary>


```diff
@@               Coverage Diff                @@
##           allocate-option    #7710   +/-   ##
================================================
  Coverage            99.67%   99.67%           
================================================
  Files                  534      534           
  Lines                52672    52765   +93     
================================================
+ Hits                 52502    52595   +93     
  Misses                 170      170           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7710?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>

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ff9320418..e212e7aea 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -415,6 +415,7 @@
 * Adds a new `allocation` module containing `allocate` and `deallocate` instructions for requesting dynamic wires. This is currently
   experimental and not integrated.
   [(#7704)](https://github.com/PennyLaneAI/pennylane/pull/7704)
+  [(#7710)](https://github.com/PennyLaneAI/pennylane/pull/7710)
 
 * Caching with finite shots now always warns about the lack of expected noise.
   [(#7644)](https://github.com/PennyLaneAI/pennylane/pull/7644)
diff --git a/pennylane/allocation.py b/pennylane/allocation.py
index 9bb65f251..b11e007fb 100644
--- a/pennylane/allocation.py
+++ b/pennylane/allocation.py
@@ -17,9 +17,48 @@ This module contains the commands for allocating and deallocating wires dynamica
 import uuid
 from typing import Optional, Sequence
 
+from pennylane.capture import enabled as capture_enabled
 from pennylane.operation import Operator
 from pennylane.wires import Wires
 
+has_jax = True
+try:
+    import jax
+except ImportError:
+    jax = None
+    has_jax = False
+
+
+if not has_jax:
+    allocate_prim = None
+    deallocate_prim = None
+else:
+    allocate_prim = jax.extend.core.Primitive("allocate")
+    allocate_prim.multiple_results = True
+
+    # pylint: disable=unused-argument
+    @allocate_prim.def_impl
+    def _(*, num_wires, require_zeros=True, restored=False):
+        raise NotImplementedError("jaxpr containing qubit allocation cannot be executed.")
+
+    # pylint: disable=unused-argument
+    @allocate_prim.def_abstract_eval
+    def _(*, num_wires, require_zeros=True, restored=False):
+        return [jax.core.ShapedArray((), dtype=int) for _ in range(num_wires)]
+
+    deallocate_prim = jax.extend.core.Primitive("deallocate")
+    deallocate_prim.multiple_results = True
+
+    # pylint: disable=unused-argument
+    @deallocate_prim.def_impl
+    def _(*wires):
+        raise NotImplementedError("jaxpr containing qubit deallocation cannot be executed.")
+
+    # pylint: disable=unused-argument
+    @deallocate_prim.def_abstract_eval
+    def _(*wires):
+        return []
+
 
 class DynamicWire:
     """A wire whose concrete value will be determined later during a compilation step or execution.
@@ -142,6 +181,10 @@ def deallocate(wires: DynamicWire | Wires | Sequence[DynamicWire]) -> Deallocate
     gate as well.
 
     """
+    if capture_enabled():
+        if not isinstance(wires, Sequence):
+            wires = (wires,)
+        return deallocate_prim.bind(*wires)
     wires = Wires(wires)
     if not_dynamic_wires := [w for w in wires if not isinstance(w, DynamicWire)]:
         raise ValueError(f"deallocate only accepts DynamicWire wires. Got {not_dynamic_wires}")
@@ -159,7 +202,7 @@ class DynamicRegister(Wires):
         return self
 
     def __exit__(self, *_, **__):
-        Deallocate(self)
+        deallocate(self)
 
 
 def allocate(num_wires: int, require_zeros: bool = True, restored: bool = False) -> DynamicRegister:
@@ -208,7 +251,13 @@ def allocate(num_wires: int, require_zeros: bool = True, restored: bool = False)
     before being re-used again in the second block.
 
     """
-    wires = [DynamicWire() for _ in range(num_wires)]
+    if capture_enabled():
+        wires = allocate_prim.bind(
+            num_wires=num_wires, require_zeros=require_zeros, restored=restored
+        )
+    else:
+        wires = [DynamicWire() for _ in range(num_wires)]
     reg = DynamicRegister(wires)
-    Allocate(reg, require_zeros=require_zeros, restored=restored)
+    if not capture_enabled():
+        Allocate(reg, require_zeros=require_zeros, restored=restored)
     return reg


Test patch:
diff --git a/tests/test_allocation.py b/tests/test_allocation.py
index ce3e92032..46d8fd6f6 100644
--- a/tests/test_allocation.py
+++ b/tests/test_allocation.py
@@ -25,7 +25,9 @@ from pennylane.allocation import (
     DynamicRegister,
     DynamicWire,
     allocate,
+    allocate_prim,
     deallocate,
+    deallocate_prim,
 )
 
 
@@ -193,3 +195,84 @@ def test_allocate_context_manager():
     qml.assert_equal(q.queue[0], Allocate(wires, require_zeros=False, restored=True))
     qml.assert_equal(q.queue[1], qml.I(wires))
     qml.assert_equal(q.queue[2], Deallocate(wires))
+
+
+@pytest.mark.jax
+@pytest.mark.usefixtures("enable_disable_plxpr")
+class TestCaptureIntegration:
+
+    @pytest.mark.parametrize("use_context", (True, False))
+    def test_capturing_allocate_and_deallocate(self, use_context):
+        """Test that allocate and deallcoate can be captured."""
+
+        import jax
+
+        def f():
+            if use_context:
+                with allocate(2, require_zeros=True, restored=True) as wires:
+                    qml.H(wires[0])
+                    qml.Z(wires[1])
+            else:
+                w, w2 = allocate(2, require_zeros=True, restored=True)
+                qml.H(w)
+                qml.Z(w2)
+                deallocate((w, w2))
+
+        jaxpr = jax.make_jaxpr(f)()
+        assert len(jaxpr.eqns) == 4
+        assert jaxpr.eqns[0].primitive == allocate_prim
+        assert len(jaxpr.eqns[0].invars) == 0
+        assert jaxpr.eqns[0].params == {"num_wires": 2, "require_zeros": True, "restored": True}
+        assert len(jaxpr.eqns[0].outvars) == 2
+        assert all(v.aval.shape == () for v in jaxpr.eqns[0].outvars)
+        for v in jaxpr.eqns[0].outvars:
+            assert v.aval.dtype == jax.numpy.int64
+
+        assert jaxpr.eqns[1].invars[0] is jaxpr.eqns[0].outvars[0]
+        assert jaxpr.eqns[2].invars[0] is jaxpr.eqns[0].outvars[1]
+
+        assert jaxpr.eqns[3].primitive == deallocate_prim
+        assert jaxpr.eqns[3].params == {}
+        assert jaxpr.eqns[3].invars == jaxpr.eqns[0].outvars
+
+        with pytest.raises(NotImplementedError):
+            jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+    def test_deallocate_single_wire(self):
+        """Test deallocate can accept a single wire."""
+
+        import jax
+
+        def f():
+            [w] = allocate(1)
+            qml.X(w)
+            deallocate(w)
+
+        jaxpr = jax.make_jaxpr(f)()
+
+        assert len(jaxpr.eqns) == 3
+        assert jaxpr.eqns[0].primitive == allocate_prim
+        assert len(jaxpr.eqns[0].invars) == 0
+        assert jaxpr.eqns[0].params == {"num_wires": 1, "require_zeros": True, "restored": False}
+        assert len(jaxpr.eqns[0].outvars) == 1
+        assert all(v.aval.shape == () for v in jaxpr.eqns[0].outvars)
+        for v in jaxpr.eqns[0].outvars:
+            assert v.aval.dtype == jax.numpy.int64
+
+        assert jaxpr.eqns[1].invars[0] is jaxpr.eqns[0].outvars[0]
+
+        assert jaxpr.eqns[2].primitive == deallocate_prim
+        assert jaxpr.eqns[2].params == {}
+        assert jaxpr.eqns[2].invars == jaxpr.eqns[0].outvars
+
+    def test_no_implementation(self):
+        """Test that (de)allocation has no concrete implementation."""
+        with pytest.raises(NotImplementedError):
+            allocate(2)
+
+        with pytest.raises(NotImplementedError):
+            with allocate(2) as wires:
+                qml.X(wires)
+
+        with pytest.raises(NotImplementedError):
+            deallocate(2)