Repository: PennyLaneAI/pennylane
Pull request #4795 (2023-11-06 21:35:57)

Fix gradient transforms with default qubit and overridden shot vectors
[sc-44395]

With this change, we can now do:
```
dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

x = qml.numpy.array(0.543, requires_grad=True)

qml.gradients.param_shift(circuit)(x, shots=(10000, 10000, 1000) )
```
```
(array(-0.5209), array(-0.5258), array(-0.522))
```

albi3ro: @mudit2812 @timmysilv More or less how I figured this out:


First I saw that shots weren't on the tape getting sent to `param_shift`.  So where was that tape getting created? 

Instead of getting created by the call to `QNode.construct` inside of `QNode.__call__`, it was being created by the call to `QNode.construct` inside `TransformProgram._set_all_classical_jacobians`.  But by the time the kwargs got to `_set_all_classical_jacobians`, shots had been popped out by the first call to `QNode.construct`.

If a bug is weird, it's usually in-place mutation causing problems in a separate part of the code base.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e7cd3acc3..90af0bf7c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -69,6 +69,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Gradient transforms now work with overridden shot vectors and default qubit.
+  [(#4795)](https://github.com/PennyLaneAI/pennylane/pull/4795)
+
 * `qml.defer_measurements` now correctly transforms circuits when terminal measurements include wires
   used in mid-circuit measurements.
   [(#4787)](https://github.com/PennyLaneAI/pennylane/pull/4787)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index f20556166..0a4ac2584 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -15,6 +15,7 @@
 This module contains the QNode class and qnode decorator.
 """
 # pylint: disable=too-many-instance-attributes,too-many-arguments,protected-access,unnecessary-lambda-assignment, too-many-branches, too-many-statements
+import copy
 import functools
 import inspect
 import warnings
@@ -842,6 +843,7 @@ class QNode:
 
     def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         """Call the quantum function with a tape context, ensuring the operations get queued."""
+        kwargs = copy.copy(kwargs)
         old_interface = self.interface
 
         if self._qfunc_uses_shots_arg:


Test patch:
diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py
index 596c0060e..4eb735fad 100644
--- a/tests/gradients/core/test_gradient_transform.py
+++ b/tests/gradients/core/test_gradient_transform.py
@@ -605,7 +605,7 @@ class TestGradientTransformIntegration:
         """Test that setting the number of shots works correctly for
         a gradient transform"""
 
-        dev = qml.device("default.qubit.legacy", wires=1, shots=1000)
+        dev = qml.device("default.qubit", wires=1, shots=1000)
 
         @qml.qnode(dev)
         def circuit(x):
@@ -617,7 +617,7 @@ class TestGradientTransformIntegration:
         # the gradient function can be called with different shot values
         grad_fn = qml.gradients.param_shift(circuit)
         assert grad_fn(x).shape == ()
-        assert grad_fn(x, shots=[(1, 1000)]).shape == (1000,)
+        assert len(grad_fn(x, shots=[(1, 1000)])) == 1000
 
         # the original QNode is unaffected
         assert circuit(x).shape == tuple()