remove excess copies in QuantumScript.copy

doc/releases/changelog-0.43.0.md

@@ -801,6 +801,7 @@
   qubit allocation.
   [(#7678)](https://github.com/PennyLaneAI/pennylane/pull/7678)
   [(#8184)](https://github.com/PennyLaneAI/pennylane/pull/8184)
+  [(#8406)](https://github.com/PennyLaneAI/pennylane/pull/8406)
 
 * The :func:`qml.workflow.set_shots` transform can now be directly applied to a QNode without the need for `functools.partial`, providing a more user-friendly syntax and negating having to import the `functools` package.
   [(#7876)](https://github.com/PennyLaneAI/pennylane/pull/7876)
@@ -1178,6 +1179,9 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* Removes excess copies in `QuantumScript.copy`, and some other performance improvements to `resolve_dynamic_wires`.
+  [(#8406)](https://github.com/PennyLaneAI/pennylane/pull/8406)
+
 * GitHub actions and workflows (`interface-unit-tests.yml`, `tests-labs.yml`, `unit-test.yml`, `upload-nightly-release.yml` and `upload.yml`) have been updated to 
   use `ubuntu-24.04` runners. 
   [(#8371)](https://github.com/PennyLaneAI/pennylane/pull/8371)

pennylane/tape/qscript.py

@@ -878,8 +878,12 @@ def copy(self, copy_operations: bool = False, **update) -> "QuantumScript":
             # Perform a shallow copy of all operations in the operation and measurement
             # queues. The operations will continue to share data with the original script operations
             # unless modified.
-            _ops = update.get("operations", [copy.copy(op) for op in self.operations])
-            _measurements = update.get("measurements", [copy.copy(op) for op in self.measurements])
+            _ops = update.get("operations")
+            _measurements = update.get("measurements")
+            if _ops is None:
+                _ops = (copy.copy(op) for op in self.operations)
+            if _measurements is None:
+                _measurements = (copy.copy(mp) for mp in self.measurements)
         else:
             # Perform a shallow copy of the operation and measurement queues. The
             # operations within the queues will be references to the original script operations;
@@ -891,7 +895,7 @@ def copy(self, copy_operations: bool = False, **update) -> "QuantumScript":
 
         update_trainable_params = "operations" in update or "measurements" in update
         # passing trainable_params=None will re-calculate trainable_params
-        default_trainable_params = None if update_trainable_params else self.trainable_params
+        default_trainable_params = None if update_trainable_params else self._trainable_params
 
         new_qscript = self.__class__(
             ops=_ops,

pennylane/transforms/resolve_dynamic_wires.py

@@ -16,7 +16,7 @@
 """
 from collections.abc import Hashable, Sequence
 
-from pennylane.allocation import Allocate, AllocateState, Deallocate
+from pennylane.allocation import AllocateState
 from pennylane.exceptions import AllocationError
 from pennylane.measurements import measure
 from pennylane.tape import QuantumScript, QuantumScriptBatch
@@ -91,6 +91,28 @@ def null_postprocessing(results: ResultBatch) -> Result:
     return results[0]
 
 
+def _new_ops(operations, manager, wire_map, deallocated):
+    for op in operations:
+        # check name faster than isinstance
+        if op.name == "Allocate":
+            for w in op.wires:
+                wire, ops = manager.get_wire(**op.hyperparameters)
+                yield from ops
+                wire_map[w] = wire
+        elif op.name == "Deallocate":
+            for w in op.wires:
+                deallocated.add(w)
+                manager.return_wire(wire_map.pop(w))
+        else:
+            if wire_map:
+                op = op.map_wires(wire_map)
+            if deallocated and (intersection := deallocated.intersection(set(op.wires))):
+                raise AllocationError(
+                    f"Encountered deallocated wires {intersection} in {op}. Dynamic wires cannot be used after deallocation."
+                )
+            yield op
+
+
 @transform
 def resolve_dynamic_wires(
     tape: QuantumScript,
@@ -234,25 +256,8 @@ def multiple_allocations():
     wire_map = {}
     deallocated = set()
 
-    new_ops = []
-    for op in tape.operations:
-        if isinstance(op, Allocate):
-            for w in op.wires:
-                wire, ops = manager.get_wire(**op.hyperparameters)
-                new_ops += ops
-                wire_map[w] = wire
-        elif isinstance(op, Deallocate):
-            for w in op.wires:
-                deallocated.add(w)
-                manager.return_wire(wire_map.pop(w))
-        else:
-            if wire_map:
-                op = op.map_wires(wire_map)
-            if intersection := deallocated.intersection(set(op.wires)):
-                raise AllocationError(
-                    f"Encountered deallocated wires {intersection} in {op}. Dynamic wires cannot be used after deallocation."
-                )
-            new_ops.append(op)
+    # note that manager, wire_map, and deallocated updated in place
+    new_ops = list(_new_ops(tape.operations, manager, wire_map, deallocated))
 
     if wire_map:
         mps = [mp.map_wires(wire_map) for mp in tape.measurements]
@@ -263,6 +268,11 @@ def multiple_allocations():
             raise AllocationError(
                 f"Encountered deallocated wires {intersection} in {mp}. Dynamic wires cannot be used after deallocation."
             )
+
+    if not wire_map and not deallocated:
+        return (tape,), null_postprocessing
+    # use private trainable params to avoid calculating them if they haven't already been set
+    # pylint: disable=protected-access
     return (
-        tape.copy(ops=new_ops, measurements=mps, trainable_params=tape.trainable_params),
+        tape.copy(ops=new_ops, measurements=mps, trainable_params=tape._trainable_params),
     ), null_postprocessing