Move to specify shots on the QNode

pennylane/__init__.py

@@ -305,6 +305,12 @@ def run_cnot():
     if name in plugin_devices:
         options = {}
 
+        if "shots" in kwargs:
+            warnings.warn(
+                "In v0.33, the shots will always be determined by the QNode. Please specify shots there.",
+                UserWarning,
+            )
+
         # load global configuration settings if available
         config = kwargs.get("config", default_config)
 

pennylane/optimize/shot_adaptive.py

@@ -242,7 +242,7 @@ def weighted_random_sampling(qnodes, coeffs, shots, argnums, *args, **kwargs):
                 continue
 
             # set the QNode device shots
-            h.device.shots = 1 if s == 1 else [(1, int(s))]
+            new_shots = 1 if s == 1 else [(1, int(s))]
 
             jacs = []
             for i in argnums:
@@ -255,7 +255,7 @@ def cost(*args, **kwargs):
                 else:
                     cost = h
 
-                j = qml.jacobian(cost, argnum=i)(*args, **kwargs)
+                j = qml.jacobian(cost, argnum=i)(*args, shots=new_shots, **kwargs)
 
                 if s == 1:
                     j = np.expand_dims(j, 0)
@@ -303,65 +303,54 @@ def _single_shot_expval_gradients(self, expval_cost, args, kwargs):
         """Compute the single shot gradients of an ExpvalCost object"""
         qnodes = expval_cost.qnodes
         coeffs = expval_cost.hamiltonian.coeffs
-        device = qnodes[0].device
 
-        self.check_device(device)
-        original_shots = device.shots
+        self.check_device(qnodes[0].device)
 
         if self.lipschitz is None:
             self.check_learning_rate(coeffs)
 
-        try:
-            if self.term_sampling == "weighted_random_sampling":
-                grads = self.weighted_random_sampling(
-                    qnodes, coeffs, self.max_shots, self.trainable_args, *args, **kwargs
-                )
-            elif self.term_sampling is None:
-                device.shots = [(1, int(self.max_shots))]
-                # We iterate over each trainable argument, rather than using
-                # qml.jacobian(expval_cost), to take into account the edge case where
-                # different arguments have different shapes and cannot be stacked.
-                grads = [
-                    qml.jacobian(expval_cost, argnum=i)(*args, **kwargs)
-                    for i in self.trainable_args
-                ]
-            else:
-                raise ValueError(
-                    f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
-                    "Only term_sampling='weighted_random_sampling' and "
-                    "term_sampling=None currently supported."
+        if self.term_sampling == "weighted_random_sampling":
+            grads = self.weighted_random_sampling(
+                qnodes, coeffs, self.max_shots, self.trainable_args, *args, **kwargs
+            )
+        elif self.term_sampling is None:
+            # We iterate over each trainable argument, rather than using
+            # qml.jacobian(expval_cost), to take into account the edge case where
+            # different arguments have different shapes and cannot be stacked.
+            grads = [
+                qml.jacobian(expval_cost, argnum=i)(
+                    *args, shots=[(1, int(self.max_shots))], **kwargs
                 )
-        finally:
-            device.shots = original_shots
-
+                for i in self.trainable_args
+            ]
+        else:
+            raise ValueError(
+                f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
+                "Only term_sampling='weighted_random_sampling' and "
+                "term_sampling=None currently supported."
+            )
         return grads
 
     def _single_shot_qnode_gradients(self, qnode, args, kwargs):
         """Compute the single shot gradients of a QNode."""
-        device = qnode.device
 
         self.check_device(qnode.device)
-        original_shots = device.shots
 
         if self.lipschitz is None:
             self.check_learning_rate(1)
 
-        try:
-            device.shots = [(1, int(self.max_shots))]
-
-            if qml.active_return():
+        if qml.active_return():
 
-                def cost(*args, **kwargs):
-                    return qml.math.stack(qnode(*args, **kwargs))
+            def cost(*args, **kwargs):
+                return qml.math.stack(qnode(*args, **kwargs))
 
-            else:
-                cost = qnode
-
-            grads = [qml.jacobian(cost, argnum=i)(*args, **kwargs) for i in self.trainable_args]
-        finally:
-            device.shots = original_shots
+        else:
+            cost = qnode
 
-        return grads
+        return [
+            qml.jacobian(cost, argnum=i)(*args, shots=[(1, int(self.max_shots))], **kwargs)
+            for i in self.trainable_args
+        ]
 
     def compute_grad(
         self, objective_fn, args, kwargs
@@ -515,17 +504,6 @@ def step_and_cost(self, objective_fn, *args, **kwargs):
         """
         new_args = self.step(objective_fn, *args, **kwargs)
 
-        if isinstance(objective_fn, qml.ExpvalCost):
-            device = objective_fn.qnodes[0].device
-        elif isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
-            device = objective_fn.device
-
-        original_shots = device.shots
-
-        try:
-            device.shots = int(self.max_shots)
-            forward = objective_fn(*args, **kwargs)
-        finally:
-            device.shots = original_shots
+        forward = objective_fn(*args, shots=int(self.max_shots), **kwargs)
 
         return new_args, forward

pennylane/qnode.py

@@ -384,6 +384,7 @@ def __init__(
         device: Union[Device, "qml.devices.experimental.Device"],
         interface="auto",
         diff_method="best",
+        shots=None,
         expansion_strategy="gradient",
         max_expansion=10,
         grad_on_execution="best",
@@ -443,13 +444,26 @@ def __init__(
                 UserWarning,
             )
 
+        if (
+            hasattr(device, "shots")
+            and device.shots != shots
+            and device.shots is not None
+            and shots is None
+        ):
+            warnings.warn(
+                "Shots should now be specified on the qnode instead of on the device."
+                "Using shots from the device. QNode specified shots will be used in v0.33."
+            )
+            shots = device.shots
+
         # input arguments
         self.func = func
         self.device = device
         self._interface = interface
         self.diff_method = diff_method
         self.expansion_strategy = expansion_strategy
         self.max_expansion = max_expansion
+        self._shots = shots
 
         # execution keyword arguments
         self.execute_kwargs = {
@@ -490,6 +504,15 @@ def __repr__(self):
             self.diff_method,
         )
 
+    @property
+    def default_shots(self) -> qml.measurements.Shots:
+        """The default shots to use for an execution.
+
+        Can be overridden on a per-call basis using the ``qnode(*args, shots=new_shots, **kwargs)`` syntax.
+
+        """
+        return qml.measurements.Shots(self._shots)
+
     @property
     def interface(self):
         """The interface used by the QNode"""
@@ -736,9 +759,7 @@ def _validate_backprop_method(device, interface, shots=None):
                 expand_fn = device.expand_fn
                 batch_transform = device.batch_transform
 
-                device = qml.device(
-                    backprop_devices[mapped_interface], wires=device.wires, shots=device.shots
-                )
+                device = qml.device(backprop_devices[mapped_interface], wires=device.wires)
                 device.expand_fn = expand_fn
                 device.batch_transform = batch_transform
 
@@ -826,23 +847,16 @@ def tape(self) -> QuantumTape:
 
     qtape = tape  # for backwards compatibility
 
-    def construct(self, args, kwargs):  # pylint: disable=too-many-branches
+    def construct(self, args, kwargs, override_shots="unset"):  # 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)
-        else:
-            shots = (
-                self._original_device._raw_shot_sequence
-                if self._original_device._shot_vector
-                else self._original_device.shots
-            )
+        override_shots = self._shots if isinstance(override_shots, str) else override_shots
 
         if old_interface == "auto":
             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
 
-        self._tape = make_qscript(self.func, shots)(*args, **kwargs)
+        self._tape = make_qscript(self.func, override_shots)(*args, **kwargs)
         self._qfunc_output = self.tape._qfunc_output
 
         params = self.tape.get_parameters(trainable_only=False)
@@ -920,38 +934,24 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
             self.device.tracker = self._original_device.tracker
 
+        original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
+        override_shots = self._shots
         if not self._qfunc_uses_shots_arg:
-            # If shots specified in call but not in qfunc signature,
-            # interpret it as device shots value for this call.
-            override_shots = kwargs.get("shots", False)
-
-            if override_shots is not False:
+            if "shots" in kwargs:
+                override_shots = kwargs.pop("shots")
                 # Since shots has changed, we need to update the preferred gradient function.
                 # This is because the gradient function chosen at initialization may
                 # no longer be applicable.
 
-                # store the initialization gradient function
-                original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
-
                 # 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)
 
-            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
-
         # construct the tape
-        self.construct(args, kwargs)
+        self.construct(args, kwargs, override_shots=override_shots)
 
         cache = self.execute_kwargs.get("cache", False)
         using_custom_cache = (
@@ -1002,9 +1002,8 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
                 else:
                     res = type(self.tape._qfunc_output)(res)
 
-            if override_shots is not False:
-                # restore the initialization gradient function
-                self.gradient_fn, self.gradient_kwargs, self.device = original_grad_fn
+            # restore the initialization gradient function
+            self.gradient_fn, self.gradient_kwargs, self.device = original_grad_fn
 
             self._update_original_device()
 
@@ -1063,9 +1062,7 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
             qfunc_output_type = type(self._qfunc_output)
             return qfunc_output_type(res)
 
-        if override_shots is not False:
-            # restore the initialization gradient function
-            self.gradient_fn, self.gradient_kwargs, self.device = original_grad_fn
+        self.gradient_fn, self.gradient_kwargs, self.device = original_grad_fn
 
         self._update_original_device()
 

tests/test_debugging.py

@@ -244,7 +244,7 @@ def circuit():
 
         assert result == expected
 
-    @pytest.mark.parametrize("shots", [None, 0, 1, 100])
+    @pytest.mark.parametrize("shots", [None, 1, 100])
     def test_different_shots(self, shots):
         """Test that snapshots are returned correctly with different QNode shot values."""
         dev = qml.device("default.qubit", wires=2)
@@ -265,9 +265,10 @@ def circuit():
             2: np.array([1 / np.sqrt(2), 0, 0, 1 / np.sqrt(2)]),
             "execution_results": np.array(0),
         }
-
         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 key in (0, "very_importand_state", 2):
+            assert np.allclose(result[key], expected[key])
+        # can't compare results as too noisy
 
     @pytest.mark.parametrize(
         "m,expected_result",