Adds shots to experimental device interface and integrate with QNode

doc/releases/changelog-dev.md

@@ -76,6 +76,13 @@
 * When given a callable, `qml.ctrl` now does its custom pre-processing on all queued operators from the callable.
   [(#4370)](https://github.com/PennyLaneAI/pennylane/pull/4370)
 
+* `qml.interfaces.set_shots` now has null behaviour for the experimental Device and accepts
+  a `Shots` object as well as `int`'s and tuples of `int`'s.
+
+* `pennylane.devices.experimental.Device` now accepts a shots keyword argument and has a `shots`
+  property. This property is merely used to set defaults for a workflow, and does not directly influence the number of shots used in executions or derivatives.
+  [(#4388)](https://github.com/PennyLaneAI/pennylane/pull/4388)
+
 <h3>Breaking changes 💔</h3>
 
 * The `do_queue` keyword argument in `qml.operation.Operator` has been removed. Instead of

pennylane/devices/experimental/default_qubit_2.py

@@ -44,6 +44,8 @@ class DefaultQubit2(Device):
     """A PennyLane device written in Python and capable of backpropagation derivatives.
 
     Args:
+        shots (int, Tuple[int], Tuple[Tuple[int]]): The default number of shots to use in executions involving
+            this device.
         seed (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
             seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
             If no value is provided, a default RNG will be used.
@@ -134,8 +136,8 @@ def name(self):
         """The name of the device."""
         return "default.qubit.2"
 
-    def __init__(self, seed=None, max_workers=None) -> None:
-        super().__init__()
+    def __init__(self, shots=None, seed=None, max_workers=None) -> None:
+        super().__init__(shots=shots)
         self._max_workers = max_workers
         self._rng = np.random.default_rng(seed)
         self._debugger = None

pennylane/devices/experimental/device_api.py

@@ -20,6 +20,7 @@
 from numbers import Number
 from typing import Callable, Union, Sequence, Tuple, Optional
 
+from pennylane.measurements import Shots
 from pennylane.tape import QuantumTape, QuantumScript
 from pennylane.typing import Result, ResultBatch
 from pennylane import Tracker
@@ -134,9 +135,21 @@ def name(self) -> str:
             self.tracker.record()
     """
 
-    def __init__(self) -> None:
+    def __init__(self, shots=None) -> None:
         # each instance should have its own Tracker.
         self.tracker = Tracker()
+        self._shots = Shots(shots)
+
+    @property
+    def shots(self) -> Shots:
+        """Default shots for execution workflows containing this device.
+
+        Note that the device itself should **always** pull shots from the provided :class:`QuantumTape` at
+        :property:`~.QuantumTape.shots`, not from this property. This property is merely to provide consistency
+        with the old interface.
+
+        """
+        return self._shots
 
     def preprocess(
         self,

pennylane/interfaces/execution.py

@@ -23,9 +23,7 @@
 # pylint: disable=unused-argument,unnecessary-lambda-assignment,inconsistent-return-statements,
 # pylint: disable=too-many-statements, invalid-unary-operand-type, function-redefined
 
-import inspect
 import warnings
-from contextlib import _GeneratorContextManager
 from functools import wraps, partial
 from typing import Callable, Sequence, Optional, Union, Tuple
 
@@ -244,23 +242,7 @@ def wrapper(tapes: Sequence[QuantumTape], **kwargs):
                 # Tape exists within the cache, store the cached result
                 cached_results[i] = cache[hashes[i]]
 
-                # Introspect the set_shots decorator of the input function:
-                #   warn the user in case of finite shots with cached results
-                finite_shots = False
-
-                closure = inspect.getclosurevars(fn).nonlocals
-                if "original_fn" in closure:  # deal with expand_fn wrapper above
-                    closure = inspect.getclosurevars(closure["original_fn"]).nonlocals
-
-                # retrieve the captured context manager instance (for set_shots)
-                if "self" in closure and isinstance(closure["self"], _GeneratorContextManager):
-                    # retrieve the shots from the arguments or device instance
-                    if closure["self"].func.__name__ == "set_shots":
-                        dev, shots = closure["self"].args
-                        shots = dev.shots if shots is False else shots
-                        finite_shots = isinstance(shots, int)
-
-                if finite_shots and getattr(cache, "_persistent_cache", True):
+                if tape.shots and getattr(cache, "_persistent_cache", True):
                     warnings.warn(
                         "Cached execution with finite shots detected!\n"
                         "Note that samples as well as all noisy quantities computed via sampling "

pennylane/interfaces/set_shots.py

@@ -18,6 +18,9 @@
 # pylint: disable=protected-access
 import contextlib
 
+import pennylane as qml
+from pennylane.measurements import Shots
+
 
 @contextlib.contextmanager
 def set_shots(device, shots):
@@ -40,6 +43,11 @@
     >>> set_shots(dev, shots=100)(lambda: dev.shots)()
     100
     """
+    if isinstance(device, qml.devices.experimental.Device):
+        yield
+        return
+    if isinstance(shots, Shots):
+        shots = shots.shot_vector if shots.has_partitioned_shots else shots.total_shots
     if shots == device.shots:
         yield
         return

pennylane/measurements/shots.py

@@ -190,7 +190,11 @@ def __repr__(self):
 
     def __eq__(self, other):
         """Equality between Shot instances."""
-        return self.total_shots == other.total_shots and self.shot_vector == other.shot_vector
+        return (
+            isinstance(other, Shots)
+            and self.total_shots == other.total_shots
+            and self.shot_vector == other.shot_vector
+        )
 
     def __hash__(self):
         """Hash for a given Shot instance."""

pennylane/qnode.py

@@ -27,7 +27,7 @@
 import pennylane as qml
 from pennylane import Device
 from pennylane.interfaces import INTERFACE_MAP, SUPPORTED_INTERFACES, set_shots
-from pennylane.measurements import ClassicalShadowMP, CountsMP, MidMeasureMP
+from pennylane.measurements import ClassicalShadowMP, CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumTape, make_qscript
 
 
@@ -49,6 +49,15 @@ def _convert_to_interface(res, interface):
     return qml.math.asarray(res, like=interface if interface != "tf" else "tensorflow")
 
 
+# pylint: disable=protected-access
+def _get_device_shots(device) -> Shots:
+    if isinstance(device, Device):
+        if device._shot_vector:
+            return Shots(device._raw_shot_sequence)
+        return Shots(device.shots)
+    return device.shots
+
+
 class QNode:
     """Represents a quantum node in the hybrid computational graph.
 
@@ -693,7 +702,8 @@ def best_method_str(device, interface):
 
     @staticmethod
     def _validate_backprop_method(device, interface, shots=None):
-        if shots is not None or getattr(device, "shots", None) is not None:
+        if shots is not None or _get_device_shots(device):
+            print(shots, _get_device_shots(device))
             raise qml.QuantumFunctionError("Backpropagation is only supported when shots=None.")
 
         if isinstance(device, qml.devices.experimental.Device):
@@ -828,14 +838,10 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         """Call the quantum function with a tape context, ensuring the operations get queued."""
         old_interface = self.interface
 
-        if not self._qfunc_uses_shots_arg:
-            shots = kwargs.pop("shots", None)
+        if self._qfunc_uses_shots_arg:
+            shots = _get_device_shots(self._original_device)
         else:
-            shots = (
-                self._original_device._raw_shot_sequence
-                if self._original_device._shot_vector
-                else self._original_device.shots
-            )
+            shots = kwargs.pop("shots", _get_device_shots(self._original_device))
 
         if old_interface == "auto":
             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
@@ -933,20 +939,12 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
 
                 # pylint: disable=not-callable
                 # update the gradient function
-                if isinstance(self._original_device, Device):
-                    set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
-                else:
-                    self._update_gradient_fn(shots=override_shots)
+                set_shots(self._original_device, override_shots)(self._update_gradient_fn)(
+                    shots=override_shots
+                )
 
             else:
-                if isinstance(self._original_device, Device):
-                    kwargs["shots"] = (
-                        self._original_device._raw_shot_sequence
-                        if self._original_device._shot_vector
-                        else self._original_device.shots
-                    )
-                else:
-                    kwargs["shots"] = None
+                kwargs["shots"] = _get_device_shots(self._original_device)
 
         # construct the tape
         self.construct(args, kwargs)

tests/devices/experimental/test_default_qubit_2.py

@@ -29,6 +29,15 @@ def test_name():
     assert DefaultQubit2().name == "default.qubit.2"
 
 
+def test_shots():
+    """Test the shots property of DefaultQubit2."""
+    assert DefaultQubit2().shots == qml.measurements.Shots(None)
+    assert DefaultQubit2(shots=100).shots == qml.measurements.Shots(100)
+
+    with pytest.raises(AttributeError):
+        DefaultQubit2().shots = 10
+
+
 def test_no_jvp_functionality():
     """Test that jvp is not supported on DefaultQubit2."""
     dev = DefaultQubit2()

tests/devices/experimental/test_device_api.py

@@ -49,6 +49,17 @@ def test_device_name(self):
         """Test the default name is the name of the class"""
         assert self.dev.name == "MinimalDevice"
 
+    def test_shots(self):
+        """Test default behavior for shots."""
+
+        assert self.dev.shots == qml.measurements.Shots(None)
+
+        shots_dev = self.MinimalDevice(shots=100)
+        assert shots_dev.shots == qml.measurements.Shots(100)
+
+        with pytest.raises(AttributeError):
+            self.dev.shots = 100
+
     def test_tracker_set_on_initialization(self):
         """Test that a new tracker instance is initialized with the class."""
         assert isinstance(self.dev.tracker, qml.Tracker)

tests/test_qnode.py

@@ -1426,7 +1426,7 @@ def test_warning_finite_shots_tape(self):
             qml.RZ(0.3, wires=0)
             qml.expval(qml.PauliZ(0))
 
-        tape = QuantumScript.from_queue(q)
+        tape = QuantumScript.from_queue(q, shots=5)
         # no warning on the first execution
         cache = {}
         qml.execute([tape], dev, None, cache=cache)
@@ -1696,16 +1696,6 @@ def supports_derivatives(self, execution_config=None, circuit=None) -> bool:
         assert not kwargs
         assert new_dev is dev
 
-    def test_shots_not_set_on_device(self):
-        """Test that shots are not set on the device when override shots are passed on a call."""
-
-        def f():
-            return qml.expval(qml.PauliZ(0))
-
-        qn = QNode(f, self.dev)
-        qn(shots=10)
-        assert getattr(self.dev, "shots", "not here") == "not here"
-
     def test_shots_integration(self):
         """Test that shots provided at call time are passed through the workflow."""
 
@@ -1744,11 +1734,8 @@ class UnsupportedOp(qml.operation.Operation):
 
             num_wires = 1
 
-            def expand(self):
-                with qml.queuing.AnnotatedQueue() as q:
-                    qml.RX(3 * self.data[0], wires=self.wires)
-                tape = QuantumScript.from_queue(q)
-                return tape
+            def decomposition(self):
+                return [qml.RX(3 * self.data[0], wires=self.wires)]
 
         @qnode(dev, diff_method=diff_method, mode=mode)
         def circuit(x):
@@ -1783,11 +1770,8 @@ class UnsupportedOp(qml.operation.Operation):
             grad_method = "A"
             grad_recipe = ([[3 / 2, 1, np.pi / 6], [-3 / 2, 1, -np.pi / 6]],)
 
-            def expand(self):
-                with qml.queuing.AnnotatedQueue() as q:
-                    qml.RX(3 * self.data[0], wires=self.wires)
-                tape = QuantumScript.from_queue(q)
-                return tape
+            def decomposition(self):
+                return [qml.RX(3 * self.data[0], wires=self.wires)]
 
         @qnode(dev, interface="autograd", diff_method="parameter-shift", max_diff=2)
         def circuit(x):
@@ -1828,11 +1812,8 @@ class PhaseShift(qml.PhaseShift):
 
             grad_method = None
 
-            def expand(self):
-                with qml.queuing.AnnotatedQueue() as q:
-                    qml.RY(3 * self.data[0], wires=self.wires)
-                tape = QuantumScript.from_queue(q)
-                return tape
+            def decomposition(self):
+                return [qml.RY(3 * self.data[0], wires=self.wires)]
 
         @qnode(dev, diff_method="parameter-shift", max_diff=2)
         def circuit(x):