Improvements to Jitting: broadcasted measurements and counts on all wires

doc/releases/changelog-dev.md

@@ -28,6 +28,11 @@
 
 <h3>Improvements 🛠</h3>
 
+* Counts measurements with `all_outcomes=True` can now be used with jax jitting. Measurements
+  broadcasted across all available wires (`qml.probs()`) can now be used with jit and devices that
+  allow variable numbers of wires (`qml.device('default.qubit')`).
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
 <h4>A Prep-Select-Prep template</h4>
 
 * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear
@@ -230,6 +235,15 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `MeasurementProcess.shape(shots: Shots, device:Device)` is now
+  `MeasurementProcess.shape(shots: Optional[int], num_device_wires:int = 0)`. This has been done to allow
+  jitting when a measurement is broadcasted across all available wires, but the device does not specify wires.
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
+* If the shape of a probability measurement is affected by a `Device.cutoff` property, it will no longer work with
+  jitting.
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
 * `GlobalPhase` is considered non-differentiable with tape transforms.
   As a consequence, `qml.gradients.finite_diff` and `qml.gradients.spsa_grad` no longer
   support differentiation of `GlobalPhase` with state-based outputs.

pennylane/devices/null_qubit.py

@@ -22,7 +22,7 @@
 from dataclasses import replace
 from functools import singledispatch
 from numbers import Number
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -37,7 +37,6 @@
     MeasurementProcess,
     MeasurementValue,
     ProbabilityMP,
-    Shots,
     StateMP,
 )
 from pennylane.tape import QuantumTape, QuantumTapeBatch
@@ -56,71 +55,71 @@
 
 @singledispatch
 def zero_measurement(
-    mp: MeasurementProcess, obj_with_wires, shots: Shots, batch_size: int, interface: str
+    mp: MeasurementProcess, num_device_wires, shots: Optional[int], batch_size: int, interface: str
 ):
     """Create all-zero results for various measurement processes."""
-    return _zero_measurement(mp, obj_with_wires, shots, batch_size, interface)
+    return _zero_measurement(mp, num_device_wires, shots, batch_size, interface)
 
 
-def _zero_measurement(mp, obj_with_wires, shots, batch_size, interface):
-    shape = mp.shape(obj_with_wires, shots)
-    if all(isinstance(s, int) for s in shape):
-        if batch_size is not None:
-            shape = (batch_size,) + shape
-        return math.zeros(shape, like=interface, dtype=mp.numeric_type)
+def _zero_measurement(
+    mp: MeasurementProcess, num_device_wires: int, shots: Optional[int], batch_size, interface
+):
+    shape = mp.shape(shots, num_device_wires)
     if batch_size is not None:
-        shape = ((batch_size,) + s for s in shape)
-    return tuple(math.zeros(s, like=interface, dtype=mp.numeric_type) for s in shape)
+        shape = (batch_size,) + shape
+    return math.zeros(shape, like=interface, dtype=mp.numeric_type)
 
 
 @zero_measurement.register
-def _(mp: ClassicalShadowMP, obj_with_wires, shots, batch_size, interface):
-    shapes = [mp.shape(obj_with_wires, Shots(s)) for s in shots]
+def _(mp: ClassicalShadowMP, num_device_wires, shots: Optional[int], batch_size, interface):
     if batch_size is not None:
         # shapes = [(batch_size,) + shape for shape in shapes]
         raise ValueError(
             "Parameter broadcasting is not supported with null.qubit and qml.classical_shadow"
         )
-    results = tuple(math.zeros(shape, like=interface, dtype=np.int8) for shape in shapes)
-    return results if shots.has_partitioned_shots else results[0]
+    shape = mp.shape(shots, num_device_wires)
+    return math.zeros(shape, like=interface, dtype=np.int8)
 
 
 @zero_measurement.register
-def _(mp: CountsMP, obj_with_wires, shots, batch_size, interface):
+def _(mp: CountsMP, num_device_wires, shots, batch_size, interface):
     outcomes = []
     if mp.obs is None and not isinstance(mp.mv, MeasurementValue):
-        num_wires = len(obj_with_wires.wires)
-        state = "0" * num_wires
-        results = tuple({state: math.asarray(s, like=interface)} for s in shots)
+        state = "0" * num_device_wires
+        results = {state: math.asarray(shots, like=interface)}
         if mp.all_outcomes:
-            outcomes = [f"{x:0{num_wires}b}" for x in range(1, 2**num_wires)]
+            outcomes = [f"{x:0{num_device_wires}b}" for x in range(1, 2**num_device_wires)]
     else:
         outcomes = sorted(mp.eigvals())  # always assign shots to the smallest
-        results = tuple({outcomes[0]: math.asarray(s, like=interface)} for s in shots)
+        results = {outcomes[0]: math.asarray(shots, like=interface)}
         outcomes = outcomes[1:] if mp.all_outcomes else []
 
     if outcomes:
         zero = math.asarray(0, like=interface)
-        for res in results:
-            for val in outcomes:
-                res[val] = zero
+        for val in outcomes:
+            results[val] = zero
     if batch_size is not None:
-        results = tuple([r] * batch_size for r in results)
-    return results[0] if len(results) == 1 else results
+        results = tuple(results for _ in range(batch_size))
+    return results
 
 
 zero_measurement.register(DensityMatrixMP)(_zero_measurement)
 
 
 @zero_measurement.register(StateMP)
 @zero_measurement.register(ProbabilityMP)
-def _(mp: Union[StateMP, ProbabilityMP], obj_with_wires, shots, batch_size, interface):
-    wires = mp.wires or obj_with_wires.wires
-    state = [1.0] + [0.0] * (2 ** len(wires) - 1)
+def _(
+    mp: Union[StateMP, ProbabilityMP],
+    num_device_wires: int,
+    shots: Optional[int],
+    batch_size,
+    interface,
+):
+    num_wires = len(mp.wires) or num_device_wires
+    state = [1.0] + [0.0] * (2**num_wires - 1)
     if batch_size is not None:
         state = [state] * batch_size
-    result = math.asarray(state, like=interface)
-    return (result,) * shots.num_copies if shots.has_partitioned_shots else result
+    return math.asarray(state, like=interface)
 
 
 @simulator_tracking
@@ -224,26 +223,26 @@ def __init__(self, wires=None, shots=None) -> None:
         self._debugger = None
 
     def _simulate(self, circuit, interface):
-        shots = circuit.shots
-        obj_with_wires = self if self.wires else circuit
-        results = tuple(
-            zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface)
-            for mp in circuit.measurements
-        )
-        if len(results) == 1:
-            return results[0]
-        if shots.has_partitioned_shots:
-            return tuple(zip(*results))
-        return results
+        num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
+        results = []
+        for s in circuit.shots or [None]:
+            r = tuple(
+                zero_measurement(mp, num_device_wires, s, circuit.batch_size, interface)
+                for mp in circuit.measurements
+            )
+            results.append(r[0] if len(circuit.measurements) == 1 else r)
+        if circuit.shots.has_partitioned_shots:
+            return tuple(results)
+        return results[0]
 
     def _derivatives(self, circuit, interface):
         shots = circuit.shots
-        obj_with_wires = self if self.wires else circuit
+        num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
         n = len(circuit.trainable_params)
         derivatives = tuple(
             (
                 math.zeros_like(
-                    zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface)
+                    zero_measurement(mp, num_device_wires, shots, circuit.batch_size, interface)
                 ),
             )
             * n

pennylane/measurements/classical_shadow.py

@@ -450,17 +450,17 @@ def _abstract_eval(
     ) -> tuple:
         return (2, shots, n_wires), np.int8
 
-    def shape(self, device, shots):  # pylint: disable=unused-argument
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, int, int]:
         # otherwise, the return type requires a device
-        if not shots:
+        if shots is None:
             raise MeasurementShapeError(
                 "Shots must be specified to obtain the shape of a classical "
                 "shadow measurement process."
             )
 
         # the first entry of the tensor represents the measured bits,
         # and the second indicate the indices of the unitaries used
-        return (2, shots.total_shots, len(self.wires))
+        return (2, shots, len(self.wires))
 
     def __copy__(self):
         return self.__class__(
@@ -559,12 +559,8 @@ def numeric_type(self):
     def return_type(self):
         return ShadowExpval
 
-    def shape(self, device, shots):
-        is_single_op = isinstance(self.H, Operator)
-        if not shots.has_partitioned_shots:
-            return () if is_single_op else (len(self.H),)
-        base = () if is_single_op else (len(self.H),)
-        return (base,) * sum(s.copies for s in shots.shot_vector)
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return () if isinstance(self.H, Operator) else (len(self.H),)
 
     @property
     def wires(self):

pennylane/measurements/expval.py

@@ -98,11 +98,8 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_samples(
         self,

pennylane/measurements/measurements.py

@@ -17,7 +17,6 @@
 and measurement samples using AnnotatedQueues.
 """
 import copy
-import functools
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from enum import Enum
@@ -30,8 +29,6 @@
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
-from .shots import Shots
-
 # =============================================================================
 # ObservableReturnTypes types
 # =============================================================================
@@ -289,58 +286,35 @@ def numeric_type(self) -> type:
             f"The numeric type of the measurement {self.__class__.__name__} is not defined."
         )
 
-    def shape(self, device, shots: Shots) -> tuple:
-        """The expected output shape of the MeasurementProcess.
-
-        Note that the output shape is dependent on the shots or device when:
-
-        * The measurement type is either ``_Probability``, ``_State`` (from :func:`.state`) or
-          ``_Sample``;
-        * The shot vector was defined.
-
-        For example, assuming a device with ``shots=None``, expectation values
-        and variances define ``shape=(,)``, whereas probabilities in the qubit
-        model define ``shape=(2**num_wires)`` where ``num_wires`` is the
-        number of wires the measurement acts on.
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, ...]:
+        """Calculate the shape of the result object tensor.
 
         Args:
-            device (pennylane.Device): a PennyLane device to use for determining the shape
-            shots (~.Shots): object defining the number and batches of shots
+            shots (Optional[int]) = None: the number of shots used execute the circuit. ``None``
+               indicates an analytic simulation.  Shot vectors are handled by calling this method
+               multiple times.
+            num_device_wires (int)=0 : The number of wires that will be used if the measurement is
+               broadcasted across all available wires (``len(mp.wires) == 0``). If the device
+               itself doesn't provide a number of wires, the number of tape wires will be provided
+               here instead:
 
         Returns:
-            tuple: the output shape
+            tuple[int,...]: An arbitrary length tuple of ints.  May be an empty tuple.
+
+        >>> qml.probs(wires=(0,1)).shape()
+        (4,)
+        >>> qml.sample(wires=(0,1)).shape(shots=50)
+        (50, 2)
+        >>> qml.state().shape(num_device_wires=4)
+        (16,)
+        >>> qml.expval(qml.Z(0)).shape()
+        ()
 
-        Raises:
-            QuantumFunctionError: the return type of the measurement process is
-                unrecognized and cannot deduce the numeric type
         """
         raise qml.QuantumFunctionError(
             f"The shape of the measurement {self.__class__.__name__} is not defined"
         )
 
-    @staticmethod
-    @functools.lru_cache()
-    def _get_num_basis_states(num_wires, device):
-        """Auxiliary function to determine the number of basis states given the
-        number of systems and a quantum device.
-
-        This function is meant to be used with the Probability measurement to
-        determine how many outcomes there will be. With qubit based devices
-        we'll have two outcomes for each subsystem. With continuous variable
-        devices that impose a Fock cutoff the number of basis states per
-        subsystem equals the cutoff value.
-
-        Args:
-            num_wires (int): the number of qubits/qumodes
-            device (pennylane.Device): a PennyLane device
-
-        Returns:
-            int: the number of basis states
-        """
-        cutoff = getattr(device, "cutoff", None)
-        base = 2 if cutoff is None else cutoff
-        return base**num_wires
-
     @qml.QueuingManager.stop_recording()
     def diagonalizing_gates(self):
         """Returns the gates that diagonalize the measured wires such that they

pennylane/measurements/mutual_info.py

@@ -154,11 +154,8 @@ def map_wires(self, wire_map: dict):
         ]
         return new_measurement
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         state = qml.math.dm_from_state_vector(state)

pennylane/measurements/probs.py

@@ -163,16 +163,9 @@ def _abstract_eval(cls, n_wires=None, has_eigvals=False, shots=None, num_device_
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        num_shot_elements = (
-            sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
-        )
-        len_wires = len(self.wires)
-        if len_wires == 0:
-            len_wires = len(device.wires) if device.wires else 0
-        dim = self._get_num_basis_states(len_wires, device)
-
-        return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int]:
+        len_wires = len(self.wires) if self.wires else num_device_wires
+        return (2**len_wires,)
 
     def process_samples(
         self,

pennylane/measurements/purity.py

@@ -85,11 +85,8 @@ def return_type(self):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         wire_map = dict(zip(wire_order, list(range(len(wire_order)))))