Repository: PennyLaneAI/pennylane
Pull request #4744 (2023-10-27 02:51:34)

Daily rc sync to master
Automatic sync from the release candidate to master during a feature freeze.


Patch:
diff --git a/doc/releases/changelog-0.33.0.md b/doc/releases/changelog-0.33.0.md
index 33f3de25d..e16c38b81 100644
--- a/doc/releases/changelog-0.33.0.md
+++ b/doc/releases/changelog-0.33.0.md
@@ -382,6 +382,11 @@
   decomposition.
   [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
 
+* Shallow copies of the `QNode` now also copy the `execute_kwargs` and transform program. When applying
+  a transform to a `QNode`, the new qnode is only a shallow copy of the original and thus keeps the same
+  device.
+  [(#4736)](https://github.com/PennyLaneAI/pennylane/pull/4736)
+
 <h3>Breaking changes 💔</h3>
 
 * ``qml.defer_measurements`` now raises an error if a transformed circuit measures ``qml.probs``,
diff --git a/docker/pennylane.dockerfile b/docker/pennylane.dockerfile
index a63d596b2..83f00a7c2 100644
--- a/docker/pennylane.dockerfile
+++ b/docker/pennylane.dockerfile
@@ -40,7 +40,7 @@ ENV PATH="/opt/venv/bin:$PATH"
 WORKDIR /opt/pennylane
 COPY  . .
 RUN git submodule update --init --recursive
-RUN pip install wheel && pip install "numpy<1.25" -r requirements.txt
+RUN pip install wheel && pip install -r requirements.txt
 RUN python3 setup.py install
 RUN pip install pytest pytest-cov pytest-mock flaky
 RUN pip install -i https://test.pypi.org/simple/ pennylane-lightning --pre --upgrade
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 08379a72f..5b3432233 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -259,7 +259,6 @@ def _draw_qnode(
             qnode.construct(args, kwargs)
             program, _ = qnode.device.preprocess()
             tapes = program([qnode.tape])
-            _wire_order = wire_order or qnode.tape.wires
         else:
             original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
             try:
@@ -280,7 +279,7 @@ def _draw_qnode(
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 
-            _wire_order = wire_order or qnode.device.wires
+        _wire_order = wire_order or qnode.device.wires or qnode.tape.wires
 
         if tapes is not None:
             cache = {"tape_offset": 0, "matrices": []}
@@ -572,18 +571,30 @@ def _draw_mpl_qnode(
             program, _ = qnode.device.preprocess()
             tapes, _ = program([qnode.tape])
             tape = tapes[0]
-            _wire_order = wire_order or qnode.tape.wires
         else:
             original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
 
             try:
                 qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
                 qnode.construct(args, kwargs_qnode)
+                if isinstance(qnode.device, qml.devices.Device):
+                    program = qnode.transform_program
+                    if any(
+                        isinstance(op, qml.measurements.MidMeasureMP)
+                        for op in qnode.tape.operations
+                    ):
+                        tapes, _ = qml.defer_measurements(qnode.tape, device=qnode.device)
+                    else:
+                        tapes = [qnode.tape]
+
+                    tapes, _ = program(tapes)
+                    tape = tapes[0]
+                else:
+                    tape = qnode.tape
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 
-            tape = qnode.tape
-            _wire_order = wire_order or qnode.device.wires
+        _wire_order = wire_order or qnode.device.wires or tape.wires
 
         return tape_mpl(
             tape,
diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py
index 6bddef769..45b94da65 100644
--- a/pennylane/gradients/spsa_gradient.py
+++ b/pennylane/gradients/spsa_gradient.py
@@ -330,8 +330,10 @@ def spsa_grad(
     all_coeffs = []
     for idx_rep in range(num_directions):
         direction = sampler(indices, num_trainable_params, idx_rep, rng=sampler_rng)
+        # the `where` arg is being cast to list to avoid unexpected side effects from types that
+        # override __array_ufunc__. See https://github.com/numpy/numpy/pull/23240 for more details
         inv_direction = qml.math.divide(
-            1, direction, where=(direction != 0), out=qml.math.zeros_like(direction)
+            1, direction, where=list(direction != 0), out=qml.math.zeros_like(direction)
         )
         # Use only the non-zero part of `direction` for the shifts, to skip redundant zero shifts
         _shifts = qml.math.tensordot(h * shifts, direction[indices], axes=0)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 1a9362f14..63c0d2227 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -482,6 +482,19 @@ class QNode:
         functools.update_wrapper(self, func)
         self._transform_program = qml.transforms.core.TransformProgram()
 
+    def __copy__(self):
+        copied_qnode = QNode.__new__(QNode)
+        for attr, value in vars(self).items():
+            if attr not in {"execute_kwargs", "_transform_program", "gradient_kwargs"}:
+                setattr(copied_qnode, attr, value)
+
+        copied_qnode.execute_kwargs = dict(self.execute_kwargs)
+        copied_qnode._transform_program = qml.transforms.core.TransformProgram(
+            self.transform_program
+        )  # pylint: disable=protected-access
+        copied_qnode.gradient_kwargs = dict(self.gradient_kwargs)
+        return copied_qnode
+
     def __repr__(self):
         """String representation."""
         if isinstance(self.device, qml.devices.Device):
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 0c1004b1f..6b95e0770 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -196,16 +196,8 @@ class TransformDispatcher:
         The default method that takes in a QNode and returns another QNode
         with the transform applied.
         """
-        # pylint: disable=protected-access
-        if (
-            isinstance(qnode._original_device, qml.Device)
-            and hasattr(qnode._original_device, "_state")
-            and qml.math.is_abstract(qnode._original_device._state)
-        ):
-            qnode._original_device.reset()
-            qnode.device.reset()
-
-        qnode = copy.deepcopy(qnode)
+
+        qnode = copy.copy(qnode)
 
         if self.expand_transform:
             qnode.add_transform(TransformContainer(self._expand_transform, targs, tkwargs))


Test patch:
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index b5551a2bb..d6ec97bf4 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -339,3 +339,38 @@ def test_draw_mpl_with_qfunc_warns_with_expansion_strategy():
 
     with pytest.warns(UserWarning, match="the expansion_strategy argument is ignored"):
         _ = qml.draw_mpl(qfunc, expansion_strategy="gradient")
+
+
+def test_mid_circuit_measurement_device_api(mocker):
+    """Test that a circuit containing mid-circuit measurements is transformed by the drawer
+    to use deferred measurements if the device uses the new device API."""
+
+    @qml.qnode(qml.device("default.qubit"))
+    def circ():
+        qml.PauliX(0)
+        qml.measure(0)
+        return qml.probs(wires=0)
+
+    draw_qnode = qml.draw_mpl(circ)
+    spy = mocker.spy(qml.defer_measurements, "_transform")
+
+    _ = draw_qnode()
+    spy.assert_called_once()
+
+
+def test_qnode_transform_program(mocker):
+    """Test that qnode transforms are applied before drawing a circuit."""
+
+    @qml.compile
+    @qml.qnode(qml.device("default.qubit"))
+    def circuit():
+        qml.RX(1.1, 0)
+        qml.RX(2.2, 0)
+        return qml.state()
+
+    draw_qnode = qml.draw_mpl(circuit, decimals=2)
+    qnode_transform = circuit.transform_program[0]
+    spy = mocker.spy(qnode_transform, "_transform")
+
+    _ = draw_qnode()
+    spy.assert_called_once()
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index fc4ca54bc..08ed2b79a 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -56,6 +56,27 @@ class CustomDeviceWithDiffMethod(qml.devices.Device):
         return 0
 
 
+def test_copy():
+    """Test that a shallow copy also copies the execute kwargs, gradient kwargs, and transform program."""
+    dev = CustomDevice()
+
+    qn = qml.QNode(dummyfunc, dev)
+    copied_qn = copy.copy(qn)
+    assert copied_qn is not qn
+    assert copied_qn.execute_kwargs == qn.execute_kwargs
+    assert copied_qn.execute_kwargs is not qn.execute_kwargs
+    assert list(copied_qn.transform_program) == list(qn.transform_program)
+    assert copied_qn.transform_program is not qn.transform_program
+    assert copied_qn.gradient_kwargs == qn.gradient_kwargs
+    assert copied_qn.gradient_kwargs is not qn.gradient_kwargs
+
+    assert copied_qn.func is qn.func
+    assert copied_qn.device is qn.device
+    assert copied_qn.interface is qn.interface
+    assert copied_qn.diff_method == qn.diff_method
+    assert copied_qn.expansion_strategy == qn.expansion_strategy
+
+
 # pylint: disable=too-many-public-methods
 class TestValidation:
     """Tests for QNode creation and validation"""
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
index 5dbb1d6bc..2ae36088c 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -182,6 +182,12 @@ class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
         qnode_transformed = dispatched_transform(qnode_circuit, 0)
         assert not qnode_circuit.transform_program
 
+        assert qnode_transformed.device is qnode_circuit.device
+
+        with dev.tracker:
+            qnode_circuit(0.1)
+        assert dev.tracker.totals["executions"] == 1
+
         assert isinstance(qnode_transformed, qml.QNode)
         assert isinstance(qnode_transformed.transform_program, qml.transforms.core.TransformProgram)
         assert isinstance(