Replacing device(..., shots=...) with set_shots (#1952)

**Context:**
PennyLane is
[deprecating](PennyLaneAI/pennylane#7979) the
interface of setting shots on device level. Hopefully this PR can help
Catalyst to avoid `PennyLaneDeprecationWarning`.
We will start from tests and then carefully check if within source code
there's any places to modify

**Description of the Change:**
Source code:
- `get_device_capabilities` accepts an extra kwarg `shots=` in order to
be able to receive `tape.shots`.
- The `capabilities` of `QJITDevice` has been heavily adjusted, but only
to extend to accept `shots` information from `tape(s)`. As the typical
deprecation cycle intended all the time, the previous `QJITDevice.shots`
and `QJITDevice.capabilities` passes keep as the same as before if a
user want to use them.
- `validate_measurements` arg `shots` has been changed from positional
to keyword. Internally, this help prefers the shots of its input `tape`
- `QFunc` does not re-initialize the copied devices with shots any more.

Tests:
- A lot. But just to transit from `device(..., shots=...)` to
`set_shots`.
 - Tests related to `capture` remain the same but with warnings caught.

**Benefits:**

**Possible Drawbacks:**
Unlike a typical deprecation, where we completely remove the deprecated
usage, this one will be co-existing with the new UI.

**Related GitHub Issues:**
[sc-95431]

---------

Co-authored-by: David Ittah <[email protected]>
Co-authored-by: Paul <[email protected]>

Author: 3 people

.dep-versions

@@ -9,7 +9,8 @@ enzyme=v0.0.186
 # Always remove custom PL/LQ versions before release.
 
 # For a custom PL version, update the package version here and at
-pennylane=0.43.0.dev24
+# 'doc/requirements.txt'
+pennylane=0.43.0.dev30
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -54,6 +54,10 @@
 
 <h3>Deprecations 👋</h3>
 
+* Deprecated usages of `Device.shots` along with setting `device(..., shots=...)`.
+  Heavily adjusted frontend pipelines within qfunc, tracer, verification and QJITDevice to account for this change.
+  [(#1952)](https://github.com/PennyLaneAI/catalyst/pull/1952)
+
 <h3>Bug fixes 🐛</h3>
 
 * Fix wrong handling of partitioned shots in the decomposition pass of `measurements_from_samples`.
@@ -162,6 +166,7 @@
 This release contains contributions from (in alphabetical order):
 
 Joey Carter,
+Yushao Chen,
 Sengthai Heng,
 David Ittah,
 Christina Lee,

doc/requirements.txt

@@ -33,4 +33,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.43.0.dev12
 pennylane-lightning==0.43.0.dev12
-pennylane==0.43.0.dev24
+pennylane==0.43.0.dev30

frontend/catalyst/device/qjit_device.py

@@ -306,20 +306,10 @@ def __init__(self, original_device):
 
         check_device_wires(original_device.wires)
 
-        super().__init__(wires=original_device.wires, shots=original_device.shots)
+        super().__init__(wires=original_device.wires)
 
         # Capability loading
-        device_capabilities = get_device_capabilities(original_device)
-
-        # TODO: This is a temporary measure to ensure consistency of behaviour. Remove this
-        #       when customizable multi-pathway decomposition is implemented. (Epic 74474)
-        if hasattr(original_device, "_to_matrix_ops"):
-            _to_matrix_ops = getattr(original_device, "_to_matrix_ops")
-            setattr(device_capabilities, "to_matrix_ops", _to_matrix_ops)
-            if _to_matrix_ops and not device_capabilities.supports_operation("QubitUnitary"):
-                raise CompileError(
-                    "The device that specifies to_matrix_ops must support QubitUnitary."
-                )
+        device_capabilities = get_device_capabilities(original_device, self.original_device.shots)
 
         backend = QJITDevice.extract_backend_info(original_device)
 
@@ -333,6 +323,7 @@ def preprocess(
         self,
         ctx,
         execution_config: Optional[qml.devices.ExecutionConfig] = None,
+        shots=None,
     ):
         """This function defines the device transform program to be applied and an updated device
         configuration. The transform program will be created and applied to the tape before
@@ -357,22 +348,27 @@ def preprocess(
 
         if execution_config is None:
             execution_config = qml.devices.ExecutionConfig()
-
         _, config = self.original_device.preprocess(execution_config)
 
         program = TransformProgram()
+        if shots is None:
+            capabilities = self.capabilities
+        else:
+            # recompute device capabilities if shots were provided through set_shots
+            device_caps = get_device_capabilities(self.original_device, shots)
+            capabilities = get_qjit_device_capabilities(device_caps)
 
         # measurement transforms may change operations on the tape to accommodate
         # measurement transformations, so must occur before decomposition
-        measurement_transforms = self._measurement_transform_program()
+        measurement_transforms = self._measurement_transform_program(capabilities)
         config = replace(config, device_options=deepcopy(config.device_options))
         program = program + measurement_transforms
 
         # decomposition to supported ops/measurements
         program.add_transform(
             catalyst_decompose,
             ctx=ctx,
-            capabilities=self.capabilities,
+            capabilities=capabilities,
             grad_method=config.gradient_method,
         )
 
@@ -382,9 +378,9 @@ def preprocess(
         )
         program.add_transform(
             validate_measurements,
-            self.capabilities,
+            capabilities,
             self.original_device.name,
-            self.original_device.shots,
+            shots,
         )
 
         if config.gradient_method is not None:
@@ -396,47 +392,47 @@ def preprocess(
 
         return program, config
 
-    def _measurement_transform_program(self):
-
+    def _measurement_transform_program(self, capabilities=None):
+        capabilities = capabilities or self.capabilities
         measurement_program = TransformProgram()
         if isinstance(self.original_device, SoftwareQQPP):
             return measurement_program
 
-        supports_sum_observables = "Sum" in self.capabilities.observables
+        supports_sum_observables = "Sum" in capabilities.observables
 
-        if self.capabilities.non_commuting_observables is False:
+        if capabilities.non_commuting_observables is False:
             measurement_program.add_transform(split_non_commuting)
         elif not supports_sum_observables:
             measurement_program.add_transform(split_to_single_terms)
 
         # if no observables are supported, we apply a transform to convert *everything* to the
         # readout basis, using either sample or counts based on device specification
-        if not self.capabilities.observables:
+        if not capabilities.observables:
             if not split_non_commuting in measurement_program:
                 # this *should* be redundant, a TOML that doesn't have observables should have
                 # a False non_commuting_observables flag, but we aren't enforcing that
                 measurement_program.add_transform(split_non_commuting)
-            if "SampleMP" in self.capabilities.measurement_processes:
+            if "SampleMP" in capabilities.measurement_processes:
                 measurement_program.add_transform(measurements_from_samples, self.wires)
-            elif "CountsMP" in self.capabilities.measurement_processes:
+            elif "CountsMP" in capabilities.measurement_processes:
                 measurement_program.add_transform(measurements_from_counts, self.wires)
             else:
                 raise RuntimeError("The device does not support observables or sample/counts")
 
-        elif not self.capabilities.measurement_processes.keys() - {"CountsMP", "SampleMP"}:
+        elif not capabilities.measurement_processes.keys() - {"CountsMP", "SampleMP"}:
             # ToDo: this branch should become unnecessary when selective conversion of
             # unsupported MPs is finished, see ToDo below
             if not split_non_commuting in measurement_program:  # pragma: no branch
                 measurement_program.add_transform(split_non_commuting)
             mp_transform = (
                 measurements_from_samples
-                if "SampleMP" in self.capabilities.measurement_processes
+                if "SampleMP" in capabilities.measurement_processes
                 else measurements_from_counts
             )
             measurement_program.add_transform(mp_transform, self.wires)
 
         # if only some observables are supported, we try to diagonalize those that aren't
-        elif not {"PauliX", "PauliY", "PauliZ", "Hadamard"}.issubset(self.capabilities.observables):
+        elif not {"PauliX", "PauliY", "PauliZ", "Hadamard"}.issubset(capabilities.observables):
             if not split_non_commuting in measurement_program:
                 # the device might support non commuting measurements but not all the
                 # Pauli + Hadamard observables, so here it is needed
@@ -449,7 +445,7 @@ def _measurement_transform_program(self):
             }
             # checking which base observables are unsupported and need to be diagonalized
             supported_observables = {"PauliX", "PauliY", "PauliZ", "Hadamard"}.intersection(
-                self.capabilities.observables
+                capabilities.observables
             )
             supported_observables = [_obs_dict[obs] for obs in supported_observables]
 
@@ -520,15 +516,23 @@ def _load_device_capabilities(device) -> DeviceCapabilities:
     return capabilities
 
 
-def get_device_capabilities(device) -> DeviceCapabilities:
+def get_device_capabilities(device, shots=None) -> DeviceCapabilities:
     """Get or load the original DeviceCapabilities from device"""
 
     assert not isinstance(device, QJITDevice)
 
-    shots_present = bool(device.shots)
-    device_capabilities = _load_device_capabilities(device)
+    shots_present = bool(shots)
+    device_capabilities = _load_device_capabilities(device).filter(finite_shots=shots_present)
+
+    # TODO: This is a temporary measure to ensure consistency of behaviour. Remove this
+    #       when customizable multi-pathway decomposition is implemented. (Epic 74474)
+    if hasattr(device, "_to_matrix_ops"):
+        _to_matrix_ops = getattr(device, "_to_matrix_ops")
+        setattr(device_capabilities, "to_matrix_ops", _to_matrix_ops)
+        if _to_matrix_ops and not device_capabilities.supports_operation("QubitUnitary"):
+            raise CompileError("The device that specifies to_matrix_ops must support QubitUnitary.")
 
-    return device_capabilities.filter(finite_shots=shots_present)
+    return device_capabilities
 
 
 def is_dynamic_wires(wires: qml.wires.Wires):

frontend/catalyst/device/verification.py

@@ -17,7 +17,7 @@
 with the compiler and device.
 """
 
-from typing import Any, Callable, List, Sequence, Union
+from typing import Any, Callable, List, Sequence
 
 from pennylane import transform
 from pennylane.devices.capabilities import DeviceCapabilities, OperatorProperties
@@ -30,7 +30,7 @@
     VarianceMP,
     VnEntropyMP,
 )
-from pennylane.measurements.shots import Shots
+from pennylane.measurements.shots import Shots, ShotsLike
 from pennylane.operation import Operation
 from pennylane.ops import (
     Adjoint,
@@ -268,7 +268,7 @@ def _obs_checker(obs):
 
 @transform
 def validate_measurements(
-    tape: QuantumTape, capabilities: DeviceCapabilities, name: str, shots: Union[int, Shots]
+    tape: QuantumTape, capabilities: DeviceCapabilities, name: str, shots: ShotsLike = None
 ) -> (Sequence[QuantumTape], Callable):
     """Validates the observables and measurements for a circuit against the capabilites
     from the TOML file.
@@ -287,6 +287,7 @@ def validate_measurements(
         CompileError: if a measurement is not supported by the given device with Catalyst
 
     """
+    shots = tape.shots if shots is None else Shots(shots)
 
     def _obs_checker(obs):
         if not obs.name in capabilities.observables:

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -43,7 +43,7 @@
 from pennylane.transforms import single_qubit_fusion as pl_single_qubit_fusion
 from pennylane.transforms import unitary_to_rot as pl_unitary_to_rot
 
-from catalyst.device import extract_backend_info, get_device_capabilities
+from catalyst.device import extract_backend_info
 from catalyst.from_plxpr.qreg_manager import QregManager
 from catalyst.jax_extras import jaxpr_pad_consts, make_jaxpr2, transient_jax_config
 from catalyst.jax_primitives import (

frontend/catalyst/jax_tracer.py

@@ -1382,7 +1382,9 @@ def is_leaf(obj):
             )
             if isinstance(device, qml.devices.Device):
                 config = _make_execution_config(qnode)
-                device_program, config = device.preprocess(ctx, config)
+                device_program, config = device.preprocess(
+                    ctx, execution_config=config, shots=shots
+                )
             else:
                 device_program = TransformProgram()
 

frontend/catalyst/qfunc.py

@@ -141,7 +141,6 @@ def __call__(self, *args, **kwargs):
                 return Function(dynamic_one_shot(self, mcm_config=mcm_config))(*args, **kwargs)
 
         new_device = copy(self.device)
-        new_device._shots = self._shots  # pylint: disable=protected-access
         qjit_device = QJITDevice(new_device)
 
         static_argnums = kwargs.pop("static_argnums", ())
@@ -264,18 +263,13 @@ def processing_fn(results):
         return dynamic_one_shot_partial(qnode)
 
     single_shot_qnode = transform_to_single_shot(qnode)
+    single_shot_qnode = qml.set_shots(single_shot_qnode, shots=1)
     if mcm_config is not None:
         single_shot_qnode.execute_kwargs["postselect_mode"] = mcm_config.postselect_mode
         single_shot_qnode.execute_kwargs["mcm_method"] = mcm_config.mcm_method
     single_shot_qnode._dynamic_one_shot_called = True
-    dev = qnode.device
     total_shots = _get_total_shots(qnode)
 
-    new_dev = copy(dev)
-    new_dev._shots = qml.measurements.Shots(1)
-    single_shot_qnode.device = new_dev
-    single_shot_qnode._set_shots(qml.measurements.Shots(1))  # pylint: disable=protected-access
-
     def one_shot_wrapper(*args, **kwargs):
         def wrap_single_shot_qnode(*_):
             return single_shot_qnode(*args, **kwargs)

frontend/test/lit/test_mitigation.py

@@ -26,13 +26,13 @@
 @qjit(target="mlir")
 def mcm_method_with_zne():
     """Test that the dynamic_one_shot works with ZNE."""
-    dev = qml.device("lightning.qubit", wires=1, shots=5)
+    dev = qml.device("lightning.qubit", wires=1)
 
     def circuit():
         return qml.expval(qml.PauliZ(0))
 
     s = [1, 3]
-    g = qml.QNode(circuit, dev, mcm_method="one-shot")
+    g = qml.set_shots(qml.QNode(circuit, dev, mcm_method="one-shot"), shots=5)
     return mitigate_with_zne(g, scale_factors=s)()
 
 

frontend/test/pytest/device/test_decomposition.py

@@ -156,8 +156,9 @@ def matrix(self):
                     dtype=np.complex128,
                 )
 
-        dev = NoUnitaryDevice(4, wires=4)
+        dev = NoUnitaryDevice(wires=4)
 
+        @qml.set_shots(4)
         @qml.qnode(dev)
         def f():
             ctrl(UnknownOp(wires=[0, 1]), control=[2, 3])