Repository: PennyLaneAI/pennylane
Pull request #3360 (2022-11-15 15:50:28)

CMP #13: Stop using `return_type` attribute in `QubitDevice`
Change:
```python
if m.return_type is Probability:
    # do sth
```
to
```python
if isinstance(m, Probability):
    # do sth
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3360](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (b2f30b3) into [remove-return-type](https://codecov.io/gh/PennyLaneAI/pennylane/commit/a6d4702b9a14b244750b26a42e3fc33810995051?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a6d4702) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@                 Coverage Diff                 @@
##           remove-return-type    #3360   +/-   ##
===================================================
  Coverage               99.70%   99.70%           
===================================================
  Files                     298      298           
  Lines                   26429    26430    +1     
===================================================
+ Hits                    26350    26351    +1     
  Misses                     79       79           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/measurements/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `99.33% <ø> (ø)` | |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.66% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/tape/qscript.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3360/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvcXNjcmlwdC5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index b8f7d3926..baa52b964 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -34,6 +34,7 @@ from pennylane.math import multiply as qmlmul
 from pennylane.math import sum as qmlsum
 from pennylane.measurements import (
     AllCounts,
+    ClassicalShadow,
     Counts,
     Expectation,
     MeasurementProcess,
@@ -48,6 +49,15 @@ from pennylane.measurements import (
     StateMeasurement,
     Variance,
     VnEntropy,
+    _Counts,
+    _Expectation,
+    _MutualInfo,
+    _Probability,
+    _Sample,
+    _ShadowExpval,
+    _State,
+    _Variance,
+    _VnEntropy,
 )
 from pennylane.operation import operation_derivative
 from pennylane.tape import QuantumScript
@@ -343,7 +353,7 @@ class QubitDevice(Device):
 
         return results
 
-    def execute(self, circuit, **kwargs):
+    def execute(self, circuit: QuantumScript, **kwargs):
         """Execute a queue of quantum operations on the device and then
         measure the given observables.
 
@@ -381,10 +391,8 @@ class QubitDevice(Device):
         if self.shots is not None or circuit.is_sampled:
             self._samples = self.generate_samples()
 
-        ret_types = [m.return_type for m in circuit.measurements]
-        counts_exist = any(
-            ret in (qml.measurements.Counts, qml.measurements.AllCounts) for ret in ret_types
-        )
+        measurements = circuit.measurements
+        counts_exist = any(isinstance(m, _Counts) for m in measurements)
 
         # compute the required statistics
         if not self.analytic and self._shot_vector is not None:
@@ -394,8 +402,8 @@ class QubitDevice(Device):
 
         if not circuit.is_sampled:
 
-            if len(circuit.measurements) == 1:
-                if ret_types[0] is qml.measurements.State:
+            if len(measurements) == 1:
+                if isinstance(measurements[0], _State):
                     # State: assumed to only be allowed if it's the only measurement
                     results = self._asarray(results, dtype=self.C_DTYPE)
                 else:
@@ -403,10 +411,7 @@ class QubitDevice(Device):
                     with contextlib.suppress(TypeError):
                         # Feature for returning custom objects: if the type cannot be cast to float then we can still allow it as an output
                         results = self._asarray(results, dtype=self.R_DTYPE)
-            elif all(
-                ret in (qml.measurements.Expectation, qml.measurements.Variance)
-                for ret in ret_types
-            ):
+            elif all(isinstance(m, (_Expectation, _Variance)) for m in measurements):
                 # Measurements with expval or var
                 results = self._asarray(results, dtype=self.R_DTYPE)
             elif not counts_exist:
@@ -429,7 +434,7 @@ class QubitDevice(Device):
 
         return results
 
-    def shot_vec_statistics(self, circuit):
+    def shot_vec_statistics(self, circuit: QuantumScript):
         """Process measurement results from circuit execution using a device
         with a shot vector and return statistics.
 
@@ -450,11 +455,9 @@ class QubitDevice(Device):
         results = []
         s1 = 0
 
-        ret_types = [m.return_type for m in circuit.measurements]
-        counts_exist = any(
-            ret in (qml.measurements.Counts, qml.measurements.AllCounts) for ret in ret_types
-        )
-        single_measurement = len(circuit.measurements) == 1
+        measurements = circuit.measurements
+        counts_exist = any(isinstance(m, _Counts) for m in measurements)
+        single_measurement = len(measurements) == 1
 
         for shot_tuple in self._shot_vector:
             s2 = s1 + np.prod(shot_tuple)
@@ -509,7 +512,7 @@ class QubitDevice(Device):
 
         return tuple(results)
 
-    def _multi_meas_with_counts_shot_vec(self, circuit, shot_tuple, r):
+    def _multi_meas_with_counts_shot_vec(self, circuit: QuantumScript, shot_tuple, r):
         """Auxiliary function of the shot_vec_statistics and execute_new
         functions for post-processing the results of multiple measurements at
         least one of which was a counts measurement.
@@ -527,17 +530,16 @@ class QubitDevice(Device):
 
             for idx2, r_ in enumerate(r):
                 measurement_proc = circuit.measurements[idx2]
-                if measurement_proc.return_type is Probability or (
-                    measurement_proc.return_type is Sample and measurement_proc.obs
+                if isinstance(measurement_proc, _Probability) or (
+                    isinstance(measurement_proc, _Sample) and measurement_proc.obs
                 ):
-
                     # Here, the result has a shape of (num_basis_states, shot_tuple.copies)
                     # Extract a single row -> shape (num_basis_states,)
                     result = r_[:, idx]
                 else:
                     result = r_[idx]
 
-                if circuit.observables[idx2].return_type not in (Counts, AllCounts):
+                if not isinstance(measurement_proc, _Counts):
                     result = self._asarray(result.T)
 
                 result_group.append(result)
@@ -746,6 +748,7 @@ class QubitDevice(Device):
                 obs.return_type = m.return_type
             else:
                 obs = m
+            # TODO: Remove return_type when `observables` argument is removed from this method
             # Pass instances directly
             if obs.return_type is Expectation:
                 # Appends a result of shape (num_bins,) if bin_size is not None, else a scalar
@@ -970,25 +973,25 @@ class QubitDevice(Device):
                 obs.return_type = m.return_type
             else:
                 obs = m
-            # 1. Based on the return_type, compute statistics
+            # 1. Based on the measurement type, compute statistics
             # Pass instances directly
-            if obs.return_type is Expectation:
+            if isinstance(m, _Expectation):
                 result = self.expval(obs, shot_range=shot_range, bin_size=bin_size)
 
-            elif obs.return_type is Variance:
+            elif isinstance(m, _Variance):
                 result = self.var(obs, shot_range=shot_range, bin_size=bin_size)
 
-            elif obs.return_type is Sample:
+            elif isinstance(m, _Sample):
                 samples = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=False)
                 result = self._asarray(qml.math.squeeze(samples))
 
-            elif obs.return_type in (Counts, AllCounts):
+            elif isinstance(m, _Counts):
                 result = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
 
-            elif obs.return_type is Probability:
+            elif isinstance(m, _Probability):
                 result = self.probability(wires=obs.wires, shot_range=shot_range, bin_size=bin_size)
 
-            elif obs.return_type is State:
+            elif isinstance(m, _State):
                 if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "The state or density matrix cannot be returned in combination"
@@ -1008,7 +1011,7 @@ class QubitDevice(Device):
                 state = self.access_state(wires=obs.wires)
                 result = self._asarray(state, dtype=self.C_DTYPE)
 
-            elif obs.return_type is VnEntropy:
+            elif isinstance(m, _VnEntropy):
                 if self.wires.labels != tuple(range(self.num_wires)):
                     raise qml.QuantumFunctionError(
                         "Returning the Von Neumann entropy is not supported when using custom wire labels"
@@ -1029,7 +1032,7 @@ class QubitDevice(Device):
                     )
                 result = self.vn_entropy(wires=obs.wires, log_base=obs.log_base)
 
-            elif obs.return_type is MutualInfo:
+            elif isinstance(m, _MutualInfo):
                 if self.wires.labels != tuple(range(self.num_wires)):
                     raise qml.QuantumFunctionError(
                         "Returning the mutual information is not supported when using custom wire labels"
@@ -1051,22 +1054,20 @@ class QubitDevice(Device):
                 wires0, wires1 = obs.raw_wires
                 result = self.mutual_info(wires0=wires0, wires1=wires1, log_base=obs.log_base)
 
-            elif obs.return_type is Shadow:
+            elif isinstance(m, ClassicalShadow):
                 if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
                     )
+                result = self.classical_shadow(obs, circuit)
 
-                result = self.classical_shadow(obs, circuit=circuit)
-
-            elif obs.return_type is ShadowExpval:
+            elif isinstance(m, _ShadowExpval):
                 if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
                     )
-
                 result = self.shadow_expval(obs, circuit=circuit)
 
             elif isinstance(m, MeasurementTransform):
@@ -1080,22 +1081,17 @@ class QubitDevice(Device):
                     result = method(obs, shot_range=shot_range, bin_size=bin_size)
                 else:
                     result = self._measure(m, shot_range=shot_range, bin_size=bin_size)
+
             elif obs.return_type is not None:
                 raise qml.QuantumFunctionError(
                     f"Unsupported return type specified for observable {obs.name}"
                 )
 
             # 2. Post-process statistics results (if need be)
-            float_return_types = {
-                Expectation,
-                Variance,
-                Probability,
-                VnEntropy,
-                MutualInfo,
-                ShadowExpval,
-            }
-
-            if obs.return_type in float_return_types:
+            if isinstance(
+                m, (_Expectation, _Variance, _Probability, _VnEntropy, _MutualInfo, _ShadowExpval)
+            ):
+                # Result is a float
                 result = self._asarray(result, dtype=self.R_DTYPE)
 
             if self._shot_vector is not None and isinstance(result, np.ndarray):
@@ -1656,6 +1652,7 @@ class QubitDevice(Device):
 
         powers_of_two = 2 ** np.arange(len(device_wires))[::-1]
         perm = basis_states @ powers_of_two
+        # The permutation happens on the last axis both with and without broadcasting
         return self._gather(prob, perm, axis=1 if batch_size is not None else 0)
 
     def expval(self, observable, shot_range=None, bin_size=None):
@@ -1880,7 +1877,7 @@ class QubitDevice(Device):
         )
 
     def adjoint_jacobian(
-        self, tape, starting_state=None, use_device_state=False
+        self, tape: QuantumScript, starting_state=None, use_device_state=False
     ):  # pylint: disable=too-many-statements
         """Implements the adjoint method outlined in
         `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
@@ -1900,7 +1897,7 @@ class QubitDevice(Device):
               :class:`~.Hamiltonian` or :class:`~.Hermitian`.
 
         Args:
-            tape (.QuantumTape): circuit that the function takes the gradient of
+            tape (.QuantumScript): circuit that the function takes the gradient of
 
         Keyword Args:
             starting_state (tensor_like): post-forward pass state to start execution with. It should be
@@ -1923,10 +1920,10 @@ class QubitDevice(Device):
         dot_product_real = lambda b, k: self._real(qmlsum(self._conj(b) * k, axis=sum_axes))
 
         for m in tape.measurements:
-            if m.return_type is not Expectation:
+            if not isinstance(m, _Expectation):
                 raise qml.QuantumFunctionError(
                     "Adjoint differentiation method does not support"
-                    f" measurement {m.return_type.value}"
+                    f" measurement {m.__class__.__name__}"
                 )
 
             if m.obs.name == "Hamiltonian":
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index a736123c2..0d69f0a40 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -205,8 +205,7 @@ class MeasurementProcess(ABC):
                 unrecognized and cannot deduce the numeric type
         """
         raise qml.QuantumFunctionError(
-            "Cannot deduce the numeric type of the measurement process with unrecognized "
-            + f"return_type {self.return_type}."
+            f"The numeric type of the measurement {self.__class__.__name__} is not defined."
         )
 
     def shape(self, device=None):
@@ -214,8 +213,8 @@ class MeasurementProcess(ABC):
 
         Note that the output shape is dependent on the device when:
 
-        * The ``return_type`` is either ``Probability``, ``State`` (from :func:`.state`) or
-          ``Sample``;
+        * The measurement type is either ``_Probability``, ``_State`` (from :func:`.state`) or
+          ``_Sample``;
         * The shot vector was defined in the device.
 
         For example, assuming a device with ``shots=None``, expectation values
@@ -223,14 +222,13 @@ class MeasurementProcess(ABC):
         model define ``shape=(1, 2**num_wires)`` where ``num_wires`` is the
         number of wires the measurement acts on.
 
-        Note that the shapes for vector-valued return types such as
-        ``Probability`` and ``State`` are adjusted to the output of
+        Note that the shapes for vector-valued measurements such as
+        ``_Probability`` and ``_State`` are adjusted to the output of
         ``qml.execute`` and may have an extra first element that is squeezed
         when using QNodes.
 
         Args:
-            device (.Device): a PennyLane device to use for determining the
-                shape
+            device (.Device): a PennyLane device to use for determining the shape
 
         Returns:
             tuple: the output shape
@@ -242,8 +240,7 @@ class MeasurementProcess(ABC):
         if qml.active_return():
             return self._shape_new(device=device)
         raise qml.QuantumFunctionError(
-            "Cannot deduce the shape of the measurement process with unrecognized return_type "
-            + f"{self.return_type}."
+            f"The shape of the measurement {self.__class__.__name__} is not defined"
         )
 
     def _shape_new(self, device=None):
@@ -251,8 +248,8 @@ class MeasurementProcess(ABC):
 
         Note that the output shape is dependent on the device when:
 
-        * The ``return_type`` is either ``Probability``, ``State`` (from :func:`.state`) or
-          ``Sample``;
+        * The measurement type is either ``_Probability``, ``_State`` (from :func:`.state`) or
+          ``_Sample``;
         * The shot vector was defined in the device.
 
         For example, assuming a device with ``shots=None``, expectation values
@@ -261,8 +258,7 @@ class MeasurementProcess(ABC):
         number of wires the measurement acts on.
 
         Args:
-            device (.Device): a PennyLane device to use for determining the
-                shape
+            device (.Device): a PennyLane device to use for determining the shape
 
         Returns:
             tuple: the output shape
@@ -272,8 +268,7 @@ class MeasurementProcess(ABC):
                 unrecognized and cannot deduce the numeric type
         """
         raise qml.QuantumFunctionError(
-            "Cannot deduce the shape of the measurement process with unrecognized return_type "
-            + f"{self.return_type}."
+            f"The shape of the measurement {self.__class__.__name__} is not defined"
         )
 
     @staticmethod
@@ -468,14 +463,14 @@ class MeasurementProcess(ABC):
                 str(self.name),
                 tuple(self.wires.tolist()),
                 str(self.data),
-                self.return_type,
+                self.__class__.__name__,
             )
         else:
             fingerprint = (
                 str(self.obs.name),
                 tuple(self.wires.tolist()),
                 str(self.obs.data),
-                self.return_type,
+                self.__class__.__name__,
             )
 
         return hash(fingerprint)
@@ -620,6 +615,20 @@ class MeasurementTransform(MeasurementProcess):
             '''
     """
 
+    method_name = ""
+    """Devices can override the logic of a measurement process by defining a method with the
+    name ``method_name`` of the corresponding class. The method should have the following signature:
+
+    .. code-block:: python
+
+        def method_name(self, qscript: QuantumScript):
+            '''Device's custom measurement implementation.
+
+            Args:
+                qscript: quantum script to transform
+            '''
+    """
+
     @abstractmethod
     def process(self, qscript, device):
         """Process the given quantum script."""


Test patch:
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 09d521c81..386c19650 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -91,7 +91,10 @@ def test_numeric_type_unrecognized_error():
     unrecognized return type raises an error."""
 
     mp = NotValidMeasurement()
-    with pytest.raises(qml.QuantumFunctionError, match="Cannot deduce the numeric type"):
+    with pytest.raises(
+        qml.QuantumFunctionError,
+        match="The numeric type of the measurement NotValidMeasurement is not defined",
+    ):
         mp.numeric_type()
 
 
@@ -100,8 +103,11 @@ def test_shape_unrecognized_error():
     unrecognized return type raises an error."""
     dev = qml.device("default.qubit", wires=2)
     mp = NotValidMeasurement()
-    with pytest.raises(qml.QuantumFunctionError, match="Cannot deduce the shape"):
-        mp.shape()
+    with pytest.raises(
+        qml.QuantumFunctionError,
+        match="The shape of the measurement NotValidMeasurement is not defined",
+    ):
+        mp.shape(dev)
 
 
 @pytest.mark.parametrize(
diff --git a/tests/returntypes/test_new_qscript.py b/tests/returntypes/test_new_qscript.py
index f99961557..94b717db7 100644
--- a/tests/returntypes/test_new_qscript.py
+++ b/tests/returntypes/test_new_qscript.py
@@ -125,7 +125,7 @@ class TestMeasurementProcess:
     def test_undefined_shape_error(self):
         """Test that an error is raised for a measurement with an undefined shape"""
         measurement = qml.counts(wires=[0, 1])
-        msg = "Cannot deduce the shape of the measurement process with unrecognized return_type"
+        msg = "The shape of the measurement _Counts is not defined"
 
         with pytest.raises(qml.QuantumFunctionError, match=msg):
             measurement.shape()
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 8f23803c7..fe92aa818 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -29,6 +29,12 @@ from pennylane.measurements import (
     Sample,
     State,
     Variance,
+    _Expectation,
+    _Probability,
+    _Sample,
+    _State,
+    _Variance,
+    state,
 )
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
@@ -340,16 +346,13 @@ class TestExtractStatistics:
             ):
                 dev.statistics()
 
-    @pytest.mark.parametrize("returntype", [Expectation, Variance, Sample, Probability, State])
-    def test_results_created(self, mock_qubit_device_extract_stats, monkeypatch, returntype):
+    @pytest.mark.parametrize(
+        "measurement", [_Expectation, _Variance, _Sample, _Probability, _State]
+    )
+    def test_results_created(self, mock_qubit_device_extract_stats, monkeypatch, measurement):
         """Tests that the statistics method simply builds a results list without any side-effects"""
 
-        class UnsupportedMeasurement(MeasurementProcess):
-            @property
-            def return_type(self):
-                return returntype
-
-        qscript = QuantumScript(measurements=[UnsupportedMeasurement()])
+        qscript = QuantumScript(measurements=[measurement(obs=qml.PauliX(0))])
 
         with monkeypatch.context() as m:
             dev = mock_qubit_device_extract_stats()
diff --git a/tests/test_qutrit_device.py b/tests/test_qutrit_device.py
index 49714b525..702fede05 100644
--- a/tests/test_qutrit_device.py
+++ b/tests/test_qutrit_device.py
@@ -31,6 +31,12 @@ from pennylane.measurements import (
     Sample,
     State,
     Variance,
+    _Counts,
+    _Expectation,
+    _Probability,
+    _Sample,
+    _State,
+    _Variance,
 )
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
@@ -296,17 +302,12 @@ class TestExtractStatistics:
     """Test the statistics method"""
 
     @pytest.mark.parametrize(
-        "returntype", [Expectation, Variance, Sample, Probability, State, Counts]
+        "measurement", [_Expectation, _Variance, _Sample, _Probability, _State, _Counts]
     )
-    def test_results_created(self, mock_qutrit_device_extract_stats, monkeypatch, returntype):
+    def test_results_created(self, mock_qutrit_device_extract_stats, monkeypatch, measurement):
         """Tests that the statistics method simply builds a results list without any side-effects"""
 
-        class UnsupportedMeasurement(MeasurementProcess):
-            @property
-            def return_type(self):
-                return returntype
-
-        qscript = QuantumScript(measurements=[UnsupportedMeasurement()])
+        qscript = QuantumScript(measurements=[measurement(obs=qml.PauliX(0))])
 
         with monkeypatch.context() as m:
             dev = mock_qutrit_device_extract_stats()