Test supported gradient kwargs and raise error for shots kwarg (Pen…

…nyLaneAI#5748)

**Context:**
Additional kwargs passed to a QNode at initialization are interpreted as
gradient kwargs.
They are checked against the hardcoded collection
`qml.gradients.SUPPORTED_GRADIENT_KWARGS` to make sure that this
boilerplate interpretation of additional kwargs makes sense.

**Description of the Change:**
This PR adds a test that makes sure that
`qml.gradients.SUPPORTED_GRADIENT_KWARGS` actually matches all gradient
transform kwargs supported by PL.
In addition, this PR introduces an error being raised when the keyword
argument `shots` is passed at QNode initialization, because this is not
a valid gradient kwarg and even if it were, it would lead to the
confusing behaviour of `QNode(..., shots=100)` not executing with
`shots=100`.

**Benefits:**
Code quality, better user input validation.

**Possible Drawbacks:**

**Related GitHub Issues:**
PennyLaneAI#5745 

[sc-64175]

---------

Co-authored-by: Pietropaolo Frisoni <[email protected]>
Co-authored-by: Christina Lee <[email protected]>

doc/releases/changelog-dev.md

@@ -158,6 +158,9 @@
 
 <h3>Breaking changes 💔</h3>
 
+* Passing `shots` as a keyword argument to a `QNode` initialization now raises an error, instead of ignoring the input.
+  [(#5748)](https://github.com/PennyLaneAI/pennylane/pull/5748)
+
 * A custom decomposition can no longer be provided to `QDrift`. Instead, apply the operations in your custom
   operation directly with `qml.apply`.
   [(#5698)](https://github.com/PennyLaneAI/pennylane/pull/5698)

pennylane/gradients/gradient_transform.py

@@ -21,34 +21,31 @@
 import pennylane as qml
 from pennylane.measurements import MutualInfoMP, ProbabilityMP, StateMP, VarianceMP, VnEntropyMP
 
-SUPPORTED_GRADIENT_KWARGS = [
+SUPPORTED_GRADIENT_KWARGS = {
     "approx_order",
     "argnum",
+    "atol",
     "aux_wire",
     "broadcast",  # [TODO: This is in param_shift. Unify with use_broadcasting in stoch_pulse_grad
     "device_wires",
     "diagonal_shifts",
+    "fallback_fn",
     "f0",
     "force_order2",
     "gradient_recipes",
-    "gradient_kwargs",
     "h",
     "n",
-    "num",
     "num_directions",
     "num_split_times",
     "off_diagonal_shifts",
-    "order",
-    "reduction",
     "sampler",
     "sampler_rng",
     "sampler_seed",
     "shifts",
-    "shots",
     "strategy",
     "use_broadcasting",
     "validate_params",
-]
+}
 
 
 def assert_multimeasure_not_broadcasted(measurements, broadcast):

pennylane/workflow/qnode.py

@@ -487,14 +487,20 @@ def __init__(
         for kwarg in gradient_kwargs:
             if kwarg in ["gradient_fn", "grad_method"]:
                 warnings.warn(
-                    f"It appears you may be trying to set the method of differentiation via the kwarg "
-                    f"{kwarg}. This is not supported in qnode and will default to backpropogation. Use "
-                    f"diff_method instead."
+                    "It appears you may be trying to set the method of differentiation via the "
+                    f"keyword argument {kwarg}. This is not supported in qnode and will default to "
+                    "backpropogation. Use diff_method instead."
+                )
+            elif kwarg == "shots":
+                raise ValueError(
+                    "'shots' is not a valid gradient_kwarg. If your quantum function takes the "
+                    "argument 'shots' or if you want to set the number of shots with which the "
+                    "QNode is executed, pass it to the QNode call, not its definition."
                 )
             elif kwarg not in qml.gradients.SUPPORTED_GRADIENT_KWARGS:
                 warnings.warn(
-                    f"Received gradient_kwarg {kwarg}, which is not included in the list of standard qnode "
-                    f"gradient kwargs."
+                    f"Received gradient_kwarg {kwarg}, which is not included in the list of "
+                    "standard qnode gradient kwargs."
                 )
 
         # input arguments

tests/gradients/core/test_gradient_transform.py

@@ -12,15 +12,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for the gradients.gradient_transform module."""
+import inspect
+
 import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.gradients.gradient_transform import (
+    SUPPORTED_GRADIENT_KWARGS,
     _find_gradient_methods,
     _validate_gradient_methods,
     choose_trainable_params,
 )
+from pennylane.transforms.core import TransformDispatcher
+
+
+def test_supported_gradient_kwargs():
+    """Test that all keyword arguments of gradient transforms are
+    registered as supported gradient kwargs, and no others."""
+    # Collect all gradient transforms
+    grad_transforms = []
+    for attr in qml.gradients.__dir__():
+        if attr == "metric_tensor":
+            # Skip metric_tensor because it is not a diff_method
+            continue
+        obj = getattr(qml.gradients, attr)
+        if isinstance(obj, TransformDispatcher):
+            grad_transforms.append(obj)
+
+    # Collect arguments of all gradient transforms
+    grad_kwargs = set()
+    for tr in grad_transforms:
+        grad_kwargs |= set(inspect.signature(tr).parameters)
+
+    # Remove arguments that are not keyword arguments
+    grad_kwargs -= {"tape"}
+    # Remove "dev", because we decided against supporting this kwarg, although
+    # it is an argument to param_shift_cv, to avoid confusion.
+    grad_kwargs -= {"dev"}
+
+    # Check equality of required and supported gradient kwargs
+    assert grad_kwargs == SUPPORTED_GRADIENT_KWARGS
+    # Shots should never be used as a gradient kwarg
+    assert "shots" not in grad_kwargs
 
 
 def test_repr():

tests/test_debugging.py

@@ -242,12 +242,12 @@ def circuit():
 
         assert result == expected
 
-    @pytest.mark.parametrize("shots", [None, 0, 1, 100])
+    @pytest.mark.parametrize("shots", [None, 1, 1000000])
     def test_different_shots(self, shots):
         """Test that snapshots are returned correctly with different QNode shot values."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.qnode(dev, shots=shots)
+        @qml.qnode(dev)
         def circuit():
             qml.Snapshot()
             qml.Hadamard(wires=0)
@@ -265,7 +265,18 @@ def circuit():
         }
 
         assert all(k1 == k2 for k1, k2 in zip(result.keys(), expected.keys()))
-        assert all(np.allclose(v1, v2) for v1, v2 in zip(result.values(), expected.values()))
+        for v1, v2 in zip(result.values(), expected.values()):
+            if v1.shape == ():
+                # Expectation value comparison
+                assert v2.shape == ()
+                if shots != 1:
+                    # If many shots or analytic mode, we can compare the result
+                    assert np.allclose(v1, v2, atol=0.006)
+                else:
+                    # If a single shot, it must be in the eigvals of qml.Z
+                    assert v1.item() in {-1.0, 1.0}
+            else:
+                assert np.allclose(v1, v2)
 
     @pytest.mark.parametrize(
         "m,expected_result",

tests/test_qnode.py

@@ -1233,6 +1233,22 @@ def ansatz1(a, shots):
         assert len(ansatz1(0.8, shots=0)) == 10
         assert ansatz1.qtape.operations[0].wires.labels == (0,)
 
+    def test_shots_passed_as_unrecognized_kwarg(self):
+        """Test that an error is raised if shots are passed to QNode initialization."""
+        dev = qml.device("default.qubit", wires=[0, 1], shots=10)
+
+        def ansatz0():
+            return qml.expval(qml.X(0))
+
+        with pytest.raises(ValueError, match="'shots' is not a valid gradient_kwarg."):
+            qml.QNode(ansatz0, dev, shots=100)
+
+        with pytest.raises(ValueError, match="'shots' is not a valid gradient_kwarg."):
+
+            @qml.qnode(dev, shots=100)
+            def _():
+                return qml.expval(qml.X(0))
+
     # pylint: disable=unexpected-keyword-arg
     def test_shots_setting_does_not_mutate_device(self):
         """Tests that per-call shots setting does not change the number of shots in the device."""