Better error message for dynamic allocation size when capture is enabled

doc/releases/changelog-dev.md

@@ -71,6 +71,7 @@
   [(#8163)](https://github.com/PennyLaneAI/pennylane/pull/8163)
   [(#8179)](https://github.com/PennyLaneAI/pennylane/pull/8179)
   [(#8198)](https://github.com/PennyLaneAI/pennylane/pull/8198)
+  [(#8381)](https://github.com/PennyLaneAI/pennylane/pull/8381)
 
   The :func:`~.allocate` function can accept three arguments that dictate how dynamically allocated
   wires are handled:
@@ -1572,5 +1573,6 @@ Justin Pickering,
 Alex Preciado,
 Shuli Shu,
 Jay Soni,
+Paul Haochen Wang,
 David Wierichs,
 Jake Zaia

pennylane/allocation.py

@@ -19,6 +19,7 @@
 from typing import Literal
 
 from pennylane.capture import enabled as capture_enabled
+from pennylane.math import is_abstract
 from pennylane.operation import Operator
 from pennylane.wires import DynamicWire, Wires
 
@@ -227,6 +228,9 @@ def allocate(
         The ``allocate`` function can be used as a context manager with automatic deallocation
         (recommended for most cases) or with manual deallocation via :func:`~.deallocate`.
 
+    .. note::
+        The ``num_wires`` argument must be static when capture is enabled.
+
     .. seealso::
         :func:`~.deallocate`
 
@@ -357,6 +361,10 @@ def circuit():
     """
     state = AllocateState(state)
     if capture_enabled():
+        if is_abstract(num_wires):
+            raise NotImplementedError(
+                "Number of allocated wires must be static when capture is enabled."
+            )
         wires = allocate_prim.bind(num_wires=num_wires, state=state, restored=restored)
     else:
         wires = [DynamicWire() for _ in range(num_wires)]

tests/test_allocation.py

@@ -304,6 +304,19 @@ def test_no_implementation(self):
         with pytest.raises(NotImplementedError):
             deallocate(2)
 
+    def test_no_dynamic_allocation_size(self):
+        """Test that allocation size must be static with capture."""
+        with pytest.raises(
+            NotImplementedError,
+            match="Number of allocated wires must be static when capture is enabled.",
+        ):
+
+            @qml.qnode(qml.device("default.qubit", wires=2))
+            def c(n: int):
+                allocate(n)
+
+            c(1)
+
 
 @pytest.mark.integration
 class TestDeviceIntegration: