All SampleMeasurement processes include a dtype optional parameter

doc/releases/changelog-dev.md

@@ -3,9 +3,10 @@
 
 <h3>New features since last release</h3>
 
-* The `qml.sample` function can now receive an optional `dtype` parameter 
-  which sets the type and precision of the samples returned by this measurement process.
+* `qml.expval`, `qml.var`, `qml.sample`, and `qml.probs` can now accept an optional `dtype` 
+  parameter which sets the type and precision of the samples returned by these measurement processes.
   [(#8189)](https://github.com/PennyLaneAI/pennylane/pull/8189)
+  [(#8214)](https://github.com/PennyLaneAI/pennylane/pull/8214)
 
 * The Resource estimation toolkit was upgraded and has migrated from
   :mod:`~.labs` to PennyLane as the :mod:`~.estimator` module.

pennylane/devices/_qubit_device.py

@@ -1388,7 +1388,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
             # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
             return self._dot(prob, eigvals)
 
-        # estimate the ev
+        # estimate the expectation value
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size

pennylane/devices/default_gaussian.py

@@ -798,7 +798,7 @@ def expval(self, observable, wires, par):
             ev, var = self._observable_map[observable](cov, mu, par, hbar=self.hbar)
 
         if self.shots is not None:
-            # estimate the ev
+            # estimate the expectation value
             # use central limit theorem, sample normal distribution once, only ok if n_eval is large
             # (see https://en.wikipedia.org/wiki/Berry%E2%80%93Esseen_theorem)
             ev = np.random.normal(ev, math.sqrt(var / self.shots))

pennylane/measurements/expval.py

@@ -20,7 +20,7 @@
 from pennylane.operation import Operator
 from pennylane.ops import I
 from pennylane.queuing import QueuingManager
-from pennylane.wires import Wires
+from pennylane.wires import WiresLike
 
 from .measurement_value import MeasurementValue
 from .measurements import SampleMeasurement, StateMeasurement
@@ -43,12 +43,15 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
             This can only be specified if an observable was not provided.
         id (str): custom label given to a measurement instance, can be useful for some applications
             where the instance has to be identified
+        dtype: The dtype of the samples returned by this measurement process.
     """
 
     _shortname = "expval"
 
     @property
     def numeric_type(self):
+        if self._dtype is not None:
+            return self._dtype
         return float
 
     def shape(self, shots: int | None = None, num_device_wires: int = 0) -> tuple:
@@ -57,19 +60,21 @@ def shape(self, shots: int | None = None, num_device_wires: int = 0) -> tuple:
     def process_samples(
         self,
         samples: Sequence[complex],
-        wire_order: Wires,
+        wire_order: WiresLike,
         shot_range: tuple[int, ...] | None = None,
         bin_size: int | None = None,
     ):
         if not self.wires:
             return math.squeeze(self.eigvals())
-        # estimate the ev
+        # estimate the expectation value
         op = self.mv if self.mv is not None else self.obs
+
         with QueuingManager.stop_recording():
             samples = SampleMP(
                 obs=op,
                 eigvals=self._eigvals,
                 wires=self.wires if self._eigvals is not None else None,
+                dtype=self._dtype,
             ).process_samples(
                 samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
             )
@@ -80,7 +85,7 @@ def process_samples(
         # TODO: do we need to squeeze here? Maybe remove with new return types
         return math.squeeze(math.mean(samples, axis=axis))
 
-    def process_state(self, state: Sequence[complex], wire_order: Wires):
+    def process_state(self, state: Sequence[complex], wire_order: WiresLike):
         # This also covers statistics for mid-circuit measurements manipulated using
         # arithmetic operators
         # we use ``self.wires`` instead of ``self.obs`` because the observable was
@@ -94,14 +99,14 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
         return self._calculate_expectation(probabilities)
 
-    def process_counts(self, counts: dict, wire_order: Wires):
+    def process_counts(self, counts: dict, wire_order: WiresLike):
         with QueuingManager.stop_recording():
             probabilties = probs(wires=self.wires).process_counts(
                 counts=counts, wire_order=wire_order
             )
         return self._calculate_expectation(probabilties)
 
-    def process_density_matrix(self, density_matrix: Sequence[complex], wire_order: Wires):
+    def process_density_matrix(self, density_matrix: Sequence[complex], wire_order: WiresLike):
         if not self.wires:
             return math.squeeze(self.eigvals())
         with QueuingManager.stop_recording():
@@ -122,6 +127,7 @@ def _calculate_expectation(self, probabilities):
 
 def expval(
     op: Operator | MeasurementValue,
+    dtype=None,
 ) -> ExpectationMP:
     r"""Expectation value of the supplied observable.
 
@@ -143,17 +149,38 @@ def circuit(x):
     >>> circuit(0.5)
     -0.4794255386042029
 
+    The ``dtype`` argument can be used to specify the precision of the returned expectation value when
+    sampling is used and ``op`` is an operator. If sampling is not used, the ``dtype`` argument is ignored.
+
+    By default, the dtype is ``float64``.
+
+    **Example:**
+
+    .. code-block:: python3
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.set_shots(10)
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.Y(0), dtype='float32')
+
+    Executing this QNode, we see that the returned samples have the specified dtype:
+
+    >>> samples = circuit(0.5)
+    >>> samples.dtype
+    dtype('float32')
+
     Args:
         op (Union[Operator, MeasurementValue]): a quantum observable object. To
             get expectation values for mid-circuit measurements, ``op`` should be
             a ``MeasurementValue``.
+        dtype: The dtype of the samples returned by this measurement process.
 
     Returns:
         ExpectationMP: measurement process instance
     """
-    if isinstance(op, MeasurementValue):
-        return ExpectationMP(obs=op)
-
     if isinstance(op, Sequence):
         raise ValueError(
             "qml.expval does not support measuring sequences of measurements or observables"
@@ -166,4 +193,4 @@ def circuit(x):
             "Expectation values of qml.Identity() without wires are currently not allowed."
         )
 
-    return ExpectationMP(obs=op)
+    return ExpectationMP(obs=op, dtype=dtype)

pennylane/measurements/measurements.py

@@ -503,6 +503,7 @@ class SampleMeasurement(MeasurementProcess):
     * bin_size (int): Divides the shot range into bins of size ``bin_size``, and
         returns the measurement statistic separately over each bin. If not
         provided, the entire shot range is treated as a single bin.
+    * dtype: The dtype of the samples returned by this measurement process.
 
     **Example:**
 
@@ -528,6 +529,18 @@ class SampleMeasurement(MeasurementProcess):
 
     _shortname = "sample"
 
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        obs: None | (Operator | MeasurementValue | Sequence[MeasurementValue]) = None,
+        wires: Wires | None = None,
+        eigvals: TensorLike | None = None,
+        dtype=None,
+        id: str | None = None,
+    ):
+        self._dtype = dtype
+        super().__init__(obs=obs, wires=wires, eigvals=eigvals, id=id)
+
     @abstractmethod
     def process_samples(
         self,

pennylane/measurements/probs.py

@@ -47,6 +47,7 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
             This can only be specified if an observable was not provided.
         id (str): custom label given to a measurement instance, can be useful for some applications
             where the instance has to be identified
+        dtype: The dtype of the samples returned by this measurement process.
     """
 
     _shortname = "probs"
@@ -59,6 +60,8 @@ def _abstract_eval(cls, n_wires=None, has_eigvals=False, shots=None, num_device_
 
     @property
     def numeric_type(self):
+        if self._dtype is not None:
+            return self._dtype
         return float
 
     def shape(self, shots: int | None = None, num_device_wires: int = 0) -> tuple[int]:
@@ -97,7 +100,7 @@ def process_samples(
         new_bin_size = bin_size or samples.shape[-2]
         new_shape = (-1, new_bin_size) if batch_size is None else (batch_size, -1, new_bin_size)
         indices = indices.reshape(new_shape)
-        prob = self._count_samples(indices, batch_size, dim)
+        prob = self._count_samples(indices, batch_size, dim, dtype=self._dtype)
         return math.squeeze(prob) if bin_size is None else prob
 
     def process_state(self, state: TensorLike, wire_order: Wires):
@@ -147,7 +150,7 @@ def process_counts(self, counts: dict, wire_order: Wires) -> np.ndarray:
 
         # constructs the probability vector
         # converts outcomes from binary strings to integers (base 10 representation)
-        prob_vector = math.zeros((dim), dtype="float64")
+        prob_vector = math.zeros((dim), dtype=self._dtype if self._dtype is not None else "float64")
         for outcome, occurrence in mapped_counts.items():
             prob_vector[int(outcome, base=2)] = occurrence / num_shots
 
@@ -167,7 +170,7 @@ def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
         return self.process_state(p_state, wire_order)
 
     @staticmethod
-    def _count_samples(indices, batch_size, dim):
+    def _count_samples(indices, batch_size, dim, dtype="float64"):
         """Count the occurrences of sampled indices and convert them to relative
         counts in order to estimate their occurrence probability."""
         num_bins, bin_size = indices.shape[-2:]
@@ -184,7 +187,7 @@ def _count_samples_core(indices, dim, interface):
         else:
 
             def _count_samples_core(indices, dim, *_):
-                probabilities = math.zeros((dim, num_bins), dtype="float64")
+                probabilities = math.zeros((dim, num_bins), dtype=dtype)
                 for b, idx in enumerate(indices):
                     basis_states, counts = math.unique(idx, return_counts=True)
                     probabilities[basis_states, b] = counts
@@ -203,7 +206,7 @@ def _count_samples_core(indices, dim, *_):
         return probabilities / bin_size
 
 
-def probs(wires=None, op=None) -> ProbabilityMP:
+def probs(wires=None, op=None, dtype=None) -> ProbabilityMP:
     r"""Probability of each computational basis state.
 
     This measurement function accepts either a wire specification or
@@ -223,8 +226,9 @@ def probs(wires=None, op=None) -> ProbabilityMP:
     Args:
         wires (Sequence[int] or int): the wire the operation acts on
         op (Operator or MeasurementValue or Sequence[MeasurementValue]): Observable (with a ``diagonalizing_gates``
-            attribute) that rotates the computational basis, or a  ``MeasurementValue``
+            attribute) that rotates the computational basis, or a ``MeasurementValue``
             corresponding to mid-circuit measurements.
+        dtype: The dtype of the samples returned by this measurement process.
 
     Returns:
         ProbabilityMP: Measurement process instance
@@ -294,6 +298,29 @@ def circuit():
     Both outputs are in the eigenbasis of the observable, but at different locations in a degenerate subspace.  Both
     correspond to half in the ``-1`` eigenvalue and half in the ``+1`` eigenvalue.
 
+    Note that the output shape of this measurement process depends on whether
+    the device simulates qubit or continuous variable quantum systems.
+
+    The ``dtype`` argument can be used to specify the precision of the returned probabilities when
+    sampling is used and the ``op`` argument is an operator. If sampling is not used, the ``dtype`` argument is ignored.
+
+    By default, the dtype is ``float64``.
+
+    **Example:**
+
+    .. code-block:: python3
+
+        @qml.set_shots(10)
+        @qml.qnode(qml.device("default.qubit", wires=2), interface="torch")
+        def circuit():
+            qml.Hadamard(0)
+            return qml.probs(dtype="float32")
+
+    Executing this QNode, we see that the returned samples have the specified dtype:
+
+    >>> samples = circuit()
+    >>> samples.dtype
+    torch.float32
     """
     if isinstance(op, MeasurementValue):
         if len(op.measurements) > 1:
@@ -339,4 +366,4 @@ def circuit():
             UserWarning,
         )
 
-    return ProbabilityMP(obs=op, wires=wires)
+    return ProbabilityMP(obs=op, wires=wires, dtype=dtype)

pennylane/measurements/sample.py

@@ -55,8 +55,6 @@
     # pylint: disable=too-many-arguments
     def __init__(self, obs=None, wires=None, eigvals=None, id=None, dtype=None):
 
-        self._dtype = dtype
-
         if isinstance(obs, MeasurementValue):
             super().__init__(obs=obs)
             return
@@ -71,7 +69,7 @@
                     "used when collecting statistics for a sequence of mid-circuit measurements."
                 )
 
-            super().__init__(obs=obs)
+            super().__init__(obs=obs, dtype=dtype)
             return
 
         if wires is not None:
@@ -82,7 +80,7 @@
                 )
             wires = Wires(wires)
 
-        super().__init__(obs=obs, wires=wires, eigvals=eigvals, id=id)
+        super().__init__(obs=obs, wires=wires, eigvals=eigvals, id=id, dtype=dtype)
 
     @classmethod
     def _abstract_eval(
@@ -127,7 +125,7 @@
 
         return (shots, num_values_per_shot)
 
     def process_samples(
         self,
         samples: Sequence[complex],
         wire_order: WiresLike,