SampleMP.process_samples does not raise an error when the observable is multiplied by a complex coefficient (#8271)

**Context:** If an observable is multiplied by a complex coefficient,
the `indices` array created within `process_raw_samples` as `indices =
samples @ powers_of_two` might be created with complex values. This
creates issues when we try to use them as actual indices.

**Description of the Change:** We move the conversion to the requested
`dtype` selected by the user just before returning, so we avoid any
possible implicit casting unintentionally performed under the hood in
the execution workflow.

**Benefits:** `process_samples` works with complex values when the
`dtype` argument is specified

**Possible Drawbacks:** None that I can think of.

**Related GitHub Issues:** None

Author: PietropaoloFrisoni

doc/releases/changelog-dev.md

@@ -6,6 +6,7 @@
 * 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.
   [(#8189)](https://github.com/PennyLaneAI/pennylane/pull/8189)
+  [(#8271)](https://github.com/PennyLaneAI/pennylane/pull/8271)
 
 * The Resource estimation toolkit was upgraded and has migrated from
   :mod:`~.labs` to PennyLane as the :mod:`~.estimator` module.

pennylane/measurements/process_samples.py

@@ -50,7 +50,6 @@ def process_raw_samples(
 
     """
 
-    samples = samples.astype(dtype) if dtype is not None else samples
     wire_map = dict(zip(wire_order, range(len(wire_order))))
     mapped_wires = [wire_map[w] for w in mp.wires]
     # Select the samples from samples that correspond to ``shot_range`` if provided
@@ -70,6 +69,7 @@ def process_raw_samples(
     # pylint: disable=protected-access
     if mp.obs is None and not isinstance(mp.mv, MeasurementValue) and mp._eigvals is None:
         # if no observable was provided then return the raw samples
+        samples = samples.astype(dtype) if dtype is not None else samples
         return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
 
     # If we're sampling observables
@@ -97,4 +97,5 @@ def process_raw_samples(
         else:
             samples = eigvals[indices]
 
+    samples = samples.astype(dtype) if dtype is not None else samples
     return samples if bin_size is None else samples.reshape((bin_size, -1))

tests/measurements/test_sample.py

@@ -408,11 +408,27 @@ class DummyOp(Operator):  # pylint: disable=too-few-public-methods
         with pytest.raises(EigvalsUndefinedError, match="Cannot compute samples of"):
             qml.sample(op=DummyOp(0)).process_samples(samples=np.array([[1, 0]]), wire_order=[0])
 
-    def test_process_samples_dtype(self):
+    @pytest.mark.parametrize(
+        "coeffs, dtype",
+        [
+            (1, "int8"),
+            (1, "int16"),
+            (1, "int32"),
+            (1, "int64"),
+            (1, "float16"),
+            (1, "float32"),
+            (1, "float64"),
+            (1j, "complex64"),
+            (1 + 1j, "complex128"),
+        ],
+    )
+    def test_process_samples_dtype(self, coeffs, dtype):
         """Test that the dtype argument changes the dtype of the returned samples."""
         samples = np.zeros(10, dtype="int64")
-        processed_samples = qml.sample(dtype="int8").process_samples(samples, wire_order=[0])
-        assert processed_samples.dtype == np.dtype("int8")
+        processed_samples = qml.sample(coeffs * qml.X(0), dtype=dtype).process_samples(
+            samples, wire_order=[0]
+        )
+        assert processed_samples.dtype == np.dtype(dtype)
 
     def test_sample_allowed_with_parameter_shift(self):
         """Test that qml.sample doesn't raise an error with parameter-shift and autograd."""