Repository: PennyLaneAI/pennylane
Pull request #1707 (2021-09-30 11:48:51)

Allows shots to be passed to transforms at runtime
**Context:** The number of shots is currently a reserved runtime argument when calling QNodes. This PR enables this for batch transforms as well.

**Description of the Change:**

- When a batch transform is called with `shots` as keyword argument, this is intercepted and used as a runtime argument when calling `qml.execute()`.

- In addition, a small bug was found and fixed at the same time, whereby the number of shots was not returned to the original value after execution.

**Benefits:** Transforms can also be dynamically called with different shot values.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** n/a


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1707](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (aa43219) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/8a14365e206771874dcfb48ef6d68d51de9118ba?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8a14365) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1707   +/-   ##
=======================================
  Coverage   99.20%   99.20%           
=======================================
  Files         201      201           
  Lines       15127    15132    +5     
=======================================
+ Hits        15007    15012    +5     
  Misses        120      120           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/batch\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfdHJhbnNmb3JtLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [8a14365...aa43219](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1707?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0584a6794..c060de1b5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -280,6 +280,27 @@
           return ([[c, 0.0, 2 * x], [-c, 0.0, 0.0]],)
   ```
 
+* Shots can now be passed as a runtime argument to transforms that execute circuits in batches, similarly
+  to QNodes.
+  [(#1707)](https://github.com/PennyLaneAI/pennylane/pull/1707)
+  
+  An example of such a transform are the gradient transforms in the
+  `qml.gradients` module. As a result, we can now call gradient transforms
+  (such as `qml.gradients.param_shift`) and set the number of shots at runtime.
+
+  ```pycon
+  >>> dev = qml.device("default.qubit", wires=1, shots=1000)
+  >>> @qml.beta.qnode(dev)
+  ... def circuit(x):
+  ...     qml.RX(x, wires=0)
+  ...     return qml.expval(qml.PauliZ(0))
+  >>> grad_fn = qml.gradients.param_shift(circuit)
+  >>> grad_fn(0.564, shots=[(1, 10)]).T
+  array([[-1., -1., -1., -1., -1.,  0., -1.,  0., -1.,  0.]])
+  >>> grad_fn(0.1233, shots=None)
+  array([[-0.53457096]])
+  ```
+
 <h3>Breaking changes</h3>
 
 - The `QNode.metric_tensor` method has been deprecated, and will be removed in an upcoming release.
diff --git a/pennylane/interfaces/batch/__init__.py b/pennylane/interfaces/batch/__init__.py
index d5750b47f..72cfea03e 100644
--- a/pennylane/interfaces/batch/__init__.py
+++ b/pennylane/interfaces/batch/__init__.py
@@ -59,6 +59,10 @@ def set_shots(device, shots):
     >>> set_shots(dev, shots=100)(lambda: dev.shots)()
     100
     """
+    if shots == device.shots:
+        yield
+        return
+
     original_shots = device.shots
 
     try:
@@ -66,8 +70,7 @@ def set_shots(device, shots):
             device.shots = shots
         yield
     finally:
-        if device.shots != original_shots:
-            device.shots = original_shots
+        device.shots = original_shots
 
 
 def cache_execute(fn, cache, pass_kwargs=False, return_tuple=True, expand_fn=None):
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 731ab5404..45ec66850 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -14,6 +14,7 @@
 """Contains tools and decorators for registering batch transforms."""
 # pylint: disable=too-few-public-methods
 import functools
+import inspect
 import types
 
 import pennylane as qml
@@ -221,8 +222,15 @@ class batch_transform:
         the QNode, and return the output of the applying the tape transform
         to the QNode's constructed tape.
         """
+        if "shots" in inspect.signature(qnode.func).parameters:
+            raise ValueError(
+                "Detected 'shots' as an argument of the quantum function to transform. "
+                "The 'shots' argument name is reserved for overriding the number of shots "
+                "taken by the device."
+            )
 
         def _wrapper(*args, **kwargs):
+            shots = kwargs.pop("shots", False)
             qnode.construct(args, kwargs)
             tapes, processing_fn = self.construct(qnode.qtape, *targs, **tkwargs)
 
@@ -242,7 +250,12 @@ class batch_transform:
                 gradient_fn = qml.gradients.finite_diff
 
             res = qml.execute(
-                tapes, device=qnode.device, gradient_fn=gradient_fn, interface=interface, max_diff=2
+                tapes,
+                device=qnode.device,
+                gradient_fn=gradient_fn,
+                interface=interface,
+                max_diff=2,
+                override_shots=shots,
             )
 
             return processing_fn(res)
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index 44f30da46..a96a1dcce 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -138,6 +138,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None):
     def classical_preprocessing(*args, **kwargs):
         """Returns the trainable gate parameters for a given QNode input."""
         trainable_only = kwargs.pop("_trainable_only", True)
+        kwargs.pop("shots", None)
         qnode.construct(args, kwargs)
         tape = qnode.qtape
 


Test patch:
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
index a1882d4aa..452b05368 100644
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -347,6 +347,42 @@ class TestGradientTransformIntegration:
         expected = qml.jacobian(circuit)(weights)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_setting_shots(self):
+        """Test that setting the number of shots works correctly for
+        a gradient transform"""
+
+        dev = qml.device("default.qubit", wires=1, shots=1000)
+
+        @qml.beta.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        x = 0.543
+
+        # the gradient function can be called with different shot values
+        grad_fn = qml.gradients.param_shift(circuit)
+        assert grad_fn(x).shape == (1, 1)
+        assert grad_fn(x, shots=[(1, 1000)]).shape == (1000, 1)
+
+        # the original QNode is unaffected
+        assert circuit(x).shape == tuple()
+        assert circuit(x, shots=1000).shape == tuple()
+
+    def test_shots_error(self):
+        """Raise an exception if shots is used within the QNode"""
+        dev = qml.device("default.qubit", wires=1, shots=1000)
+
+        @qml.beta.qnode(dev)
+        def circuit(x, shots):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        with pytest.raises(
+            ValueError, match="'shots' argument name is reserved for overriding the number of shots"
+        ):
+            qml.gradients.param_shift(circuit)(0.2, shots=100)
+
 
 class TestInterfaceIntegration:
     """Test that the gradient transforms are differentiable