Make the pipeline compatible with shots decoupling

doc/releases/changelog-dev.md

@@ -6,6 +6,7 @@
 
 * A new QNode transform called :func:`~.transforms.set_shots` has been added to set or update the number of shots to be performed, overriding shots specified in the device.
   [(#7337)](https://github.com/PennyLaneAI/pennylane/pull/7337)
+  [(#7358)](https://github.com/PennyLaneAI/pennylane/pull/7358)
 
   The :func:`~.transforms.set_shots` transform can be used as a decorator:
 
@@ -261,6 +262,9 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Internal changes ⚙️</h3>
 
+* Ajusted execution pipeline to be compatible with the new `set_shots` transform.
+  [(#7358)](https://github.com/PennyLaneAI/pennylane/pull/7358)
+
 * Wheel releases for PennyLane now follow the `PyPA binary-distribution format <https://packaging.python.org/en/latest/specifications/binary-distribution-format/>_` guidelines more closely.
   [(#7382)](https://github.com/PennyLaneAI/pennylane/pull/7382)
 

pennylane/workflow/_setup_transform_program.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """Contains a function for setting up the inner and outer transform programs for execution of a QNode."""
 
-import warnings
 
 from cachetools import LRUCache
 
@@ -49,21 +48,10 @@ def _prune_dynamic_transform(outer_transform, inner_transform):
     inner_contains_one_shot = inner_transform.prune_dynamic_transform(type_to_keep)
     if inner_contains_one_shot:
         type_to_keep = 0
-    original_len = len(outer_transform)
     outer_transform.prune_dynamic_transform(type_to_keep)
-    outer_contained_one_shot = len(outer_transform) < original_len
-    if inner_contains_one_shot and outer_contained_one_shot:
-        warnings.warn(
-            "A dynamic_one_shot transform already exists in the preprocessing program of the "
-            "device. Therefore, the dynamic_one_shot applied on the qnode will be ignored. "
-            "See https://docs.pennylane.ai/en/stable/code/api/pennylane.dynamic_one_shot.html "
-            "for more information on the recommended way to use dynamic_one_shot.",
-            UserWarning,
-        )
 
 
 def _setup_transform_program(
-    user_transform_program: TransformProgram,
     device: "qml.devices.Device",
     resolved_execution_config: "qml.devices.ExecutionConfig",
     cache=None,
@@ -85,24 +73,22 @@ def _setup_transform_program(
 
     device_transform_program = device.preprocess_transforms(resolved_execution_config)
 
-    full_transform_program = qml.transforms.core.TransformProgram(
-        user_transform_program, cotransform_cache=user_transform_program.cotransform_cache
-    )
+    outer_transform_program = qml.transforms.core.TransformProgram()
     inner_transform_program = qml.transforms.core.TransformProgram()
 
     # Add the gradient expand to the program if necessary
     if getattr(resolved_execution_config.gradient_method, "expand_transform", False):
-        full_transform_program.add_transform(
+        outer_transform_program.add_transform(
             qml.transform(resolved_execution_config.gradient_method.expand_transform),
             **resolved_execution_config.gradient_keyword_arguments,
         )
     if resolved_execution_config.use_device_gradient:
-        full_transform_program += device_transform_program
+        outer_transform_program += device_transform_program
     else:
         inner_transform_program += device_transform_program
 
     # Making sure dynamic_one_shot occurs at most once between the inner and outer transform programs
-    _prune_dynamic_transform(full_transform_program, inner_transform_program)
+    _prune_dynamic_transform(outer_transform_program, inner_transform_program)
 
     # If caching is desired but an explicit cache is not provided, use an ``LRUCache``.
     if cache is True:
@@ -124,4 +110,4 @@ def _setup_transform_program(
     if cache is not None:
         inner_transform_program.add_transform(_cache_transform, cache=cache)
 
-    return full_transform_program, inner_transform_program
+    return outer_transform_program, inner_transform_program

pennylane/workflow/construct_execution_config.py

@@ -106,9 +106,7 @@ def wrapper(*args, **kwargs):
         if resolve:
             tape = construct_tape(qnode, level=0)(*args, **kwargs)
             # pylint:disable=protected-access
-            config = _resolve_execution_config(
-                config, qnode.device, (tape,), qnode._transform_program
-            )
+            config = _resolve_execution_config(config, qnode.device, (tape,))
 
         return config
 

pennylane/workflow/execution.py

@@ -186,8 +186,17 @@ def cost_fn(params, x):
     if not tapes:
         return ()
 
-    ### Specifying and preprocessing variables ####
+    ### Apply the user transforms ####
+    transform_program = transform_program or TransformProgram()
+    tapes, user_post_processing = transform_program(tapes)
+    if transform_program.is_informative:
+        return user_post_processing(tapes)
+    transform_program = TransformProgram()
+
+    if not tapes:
+        return user_post_processing(())
 
+    ### Specifying and preprocessing variables ####
     interface = _resolve_interface(interface, tapes)
 
     config = qml.devices.ExecutionConfig(
@@ -200,18 +209,14 @@ def cost_fn(params, x):
         derivative_order=max_diff,
         executor_backend=executor_backend,
     )
-    config = _resolve_execution_config(config, device, tapes, transform_program=transform_program)
+    config = _resolve_execution_config(config, device, tapes)
 
-    transform_program = transform_program or qml.transforms.core.TransformProgram()
-    transform_program, inner_transform = _setup_transform_program(
-        transform_program, device, config, cache, cachesize
-    )
+    outer_transform, inner_transform = _setup_transform_program(device, config, cache, cachesize)
 
     #### Executing the configured setup #####
-    tapes, post_processing = transform_program(tapes)
+    tapes, outer_post_processing = outer_transform(tapes)
 
-    if transform_program.is_informative:
-        return post_processing(tapes)
+    assert not outer_transform.is_informative, "should only contain device preprocessing"
 
     results = run(tapes, device, config, inner_transform)
-    return post_processing(results)
+    return user_post_processing(outer_post_processing(results))

pennylane/workflow/resolution.py

@@ -17,7 +17,7 @@
 from dataclasses import replace
 from importlib.metadata import version
 from importlib.util import find_spec
-from typing import Literal, Optional, Union, get_args
+from typing import Literal, Union, get_args
 from warnings import warn
 
 from packaging.version import Version
@@ -26,7 +26,7 @@
 from pennylane.logging import debug_logger
 from pennylane.math import Interface, get_canonical_interface_name, get_interface
 from pennylane.tape import QuantumScriptBatch
-from pennylane.transforms.core import TransformDispatcher, TransformProgram
+from pennylane.transforms.core import TransformDispatcher
 
 SupportedDiffMethods = Literal[
     None,
@@ -260,15 +260,13 @@ def _resolve_execution_config(
     execution_config: "qml.devices.ExecutionConfig",
     device: "qml.devices.Device",
     tapes: QuantumScriptBatch,
-    transform_program: Optional[TransformProgram] = None,
 ) -> "qml.devices.ExecutionConfig":
     """Resolves the execution configuration for non-device specific properties.
 
     Args:
         execution_config (qml.devices.ExecutionConfig): an execution config to be executed on the device
         device (qml.devices.Device): a Pennylane device
         tapes (QuantumScriptBatch): a batch of tapes
-        transform_program (TransformProgram): a program of transformations to be applied to the tapes
 
     Returns:
         qml.devices.ExecutionConfig: resolved execution configuration
@@ -279,14 +277,6 @@ def _resolve_execution_config(
         execution_config.gradient_method, Callable
     ):
         updated_values["grad_on_execution"] = False
-
-    if (
-        "lightning" in device.name
-        and transform_program
-        and qml.metric_tensor in transform_program
-        and execution_config.gradient_method == "best"
-    ):
-        execution_config = replace(execution_config, gradient_method=qml.gradients.param_shift)
     execution_config = _resolve_diff_method(execution_config, device, tape=tapes[0])
 
     if execution_config.use_device_jacobian_product and not device.supports_vjp(

tests/gradients/core/test_hadamard_gradient.py

@@ -870,7 +870,7 @@ def test_output_shape_matches_qnode_expval_array(self, cost, exp_shape, exp_type
         (cost4, [4, 3], np.ndarray),
         (cost5, [4, 3], list),
         (cost6, [2, 4, 3], tuple),
-        (cost10, [32, 3], np.ndarray),  # Note that the shape here depends on the device
+        (cost10, [8, 3], np.ndarray),  # Note that the shape here depends on the device
         (cost11, [2, 32, 3], tuple),
         (cost12, [8, 3], np.ndarray),
     ]
@@ -881,7 +881,7 @@ def test_output_shape_matches_qnode_expval_array(self, cost, exp_shape, exp_type
     @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_probs)
     def test_output_shape_matches_qnode_probs(self, cost, exp_shape, exp_type, mode):
         """Test that the transform output shape matches that of the QNode."""
-        dev = qml.device("default.qubit", wires=5)
+        dev = qml.device("default.qubit", wires=6)
 
         x = np.random.rand(3)
         circuit = qml.QNode(cost, dev)

tests/gradients/core/test_metric_tensor.py

@@ -1560,7 +1560,7 @@ def circuit(x):
         return qml.expval(qml.PauliZ(0))
 
     x = np.array([1.3, 0.2])
-    with pytest.raises(qml.wires.WireError, match=r"contain wires not found on the device: \{1\}"):
+    with pytest.raises(qml.wires.WireError, match=r"no free wire"):
         qml.metric_tensor(circuit)(x)
 
 

tests/transforms/test_set_shots.py

@@ -80,24 +80,48 @@ def test_returns_tuple(self):
         assert isinstance(result[0], (list, tuple))
         assert callable(result[1])
 
-    @pytest.mark.xfail(
-        reason="This test is expected to fail until the pipeline is updated to use the new set_shots transform"
-    )
     @pytest.mark.integration
-    def test_error_finite_shots_with_backprop(self):
-        """Test that DeviceError is raised if finite shots are used with backprop + default.qubit."""
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    @pytest.mark.parametrize("interface", qml.math.SUPPORTED_INTERFACE_NAMES)
+    def test_compatibility_finite_shots(self, interface, shots):
+        """Test that setting shots works with default.qubit with default setting up"""
         dev = qml.device("default.qubit", wires=2)
 
-        @partial(set_shots, shots=1000)
-        @qml.qnode(dev)
+        @partial(set_shots, shots=shots)
+        @qml.qnode(dev, interface=interface)
         def circuit(x):
             qml.RX(x, wires=0)
             qml.RY(x, wires=1)
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
         x = qml.numpy.array(0.5)
-        circuit(x)
+        assert len(circuit(x)) == 2
+
+    @pytest.mark.integration
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("shots", [(10, 10, 10), ((10, 3),)])
+    @pytest.mark.parametrize("interface", qml.math.SUPPORTED_INTERFACE_NAMES)
+    def test_shot_vector(self, interface, shots):
+        """Test that setting shots with shot vectors works with default setting up."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @partial(set_shots, shots=shots)
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(x, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
+
+        x = qml.numpy.array(0.5)
+        assert len(circuit(x)) == 3
+
+    @pytest.mark.integration
+    def test_toplevel_accessible(self):
+        """Test that qml.set_shots is available at top-level."""
+        assert hasattr(qml, "set_shots")
 
     @pytest.mark.integration
     def test_circuit_specification(self):

tests/workflow/interfaces/run/test_autograd_run.py

@@ -25,7 +25,6 @@
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.transforms.core import TransformProgram
 from pennylane.workflow import _resolve_execution_config, _setup_transform_program, run
 
 
@@ -46,10 +45,8 @@ def cost(a, b):
             ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")]
             tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots)
 
-            resolved_config = _resolve_execution_config(
-                config, device, [tape1, tape2], TransformProgram()
-            )
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape1, tape2])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape1, tape2], device, resolved_config, inner_tp)
 
         a = pnp.array(0.1, requires_grad=True)
@@ -80,8 +77,8 @@ def test_scalar_jacobian(self, device, config, shots, seed):
 
         def cost(a):
             tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots)
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape], device, resolved_config, inner_tp)[0]
 
         a = pnp.array(0.1, requires_grad=True)
@@ -106,8 +103,8 @@ def cost(a, b):
             ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])]
             m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
             tape = qml.tape.QuantumScript(ops, m, shots=shots)
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return autograd.numpy.hstack(run([tape], device, resolved_config, inner_tp)[0])
 
         a = pnp.array(0.1, requires_grad=True)

tests/workflow/interfaces/run/test_jax_jit_run.py

@@ -23,7 +23,6 @@
 from conftest import atol_for_shots, get_device, test_matrix
 
 import pennylane as qml
-from pennylane.transforms.core import TransformProgram
 from pennylane.workflow import _resolve_execution_config, _setup_transform_program, run
 
 jax = pytest.importorskip("jax")
@@ -53,10 +52,8 @@ def cost(a, b):
             ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")]
             tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots)
 
-            resolved_config = _resolve_execution_config(
-                config, device, [tape1, tape2], TransformProgram()
-            )
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape1, tape2])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape1, tape2], device, resolved_config, inner_tp)
 
         a = jnp.array(0.1)
@@ -87,8 +84,8 @@ def test_scalar_jacobian(self, device, config, shots, seed):
 
         def cost(a):
             tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots)
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape], device, resolved_config, inner_tp)[0]
 
         a = jnp.array(0.1)
@@ -112,8 +109,8 @@ def cost(a, b):
             m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
             tape = qml.tape.QuantumScript(ops, m, shots=shots)
 
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape], device, resolved_config, inner_tp)[0]
 
         a = jnp.array(0.1)

tests/workflow/interfaces/run/test_jax_run.py

@@ -23,7 +23,6 @@
 from conftest import atol_for_shots, get_device, test_matrix
 
 import pennylane as qml
-from pennylane.transforms.core import TransformProgram
 from pennylane.workflow import _resolve_execution_config, _setup_transform_program, run
 
 jax = pytest.importorskip("jax")
@@ -48,10 +47,8 @@ def cost(a, b):
             ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")]
             tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots)
 
-            resolved_config = _resolve_execution_config(
-                config, device, [tape1, tape2], TransformProgram()
-            )
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape1, tape2])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape1, tape2], device, resolved_config, inner_tp)
 
         a = jnp.array(0.1)
@@ -82,8 +79,8 @@ def test_scalar_jacobian(self, device, config, shots, seed):
 
         def cost(a):
             tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots)
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape], device, resolved_config, inner_tp)[0]
 
         a = jnp.array(0.1)
@@ -107,8 +104,8 @@ def cost(a, b):
             m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
             tape = qml.tape.QuantumScript(ops, m, shots=shots)
 
-            resolved_config = _resolve_execution_config(config, device, [tape], TransformProgram())
-            inner_tp = _setup_transform_program(TransformProgram(), device, resolved_config)[1]
+            resolved_config = _resolve_execution_config(config, device, [tape])
+            inner_tp = _setup_transform_program(device, resolved_config)[1]
             return run([tape], device, resolved_config, inner_tp)[0]
 
         a = jnp.array(0.1)