Repository: PennyLaneAI/pennylane
Pull request #3938 (2023-03-23 19:04:47)

introduce QubitDevice._get_diagonalizing_gates; skip Hamiltonian in default.qubit
**Context:**
Can be done instead of #3929. Some devices compute diagonalizing gates, and this can be costly. If they aren't going to be used, a device should be able to skip trying.

**Description of the Change:**
1. Add `_get_diagonalizing_gates` to the `QubitDevice` class. It is a copy-paste of the `QuantumScript.diagonalizing_gates` property.
2. Replace `circuit.diagonalizing_gates` with `self._get_diagonalizing_gates(circuit)` in `QubitDevice.execute` (when it calls `self.apply`).
3. Override the method in `DefaultQubit` to filter out any measurements with `isinstance(m.obs, qml.Hamiltonian)`.

**Benefits:**
- The device can weigh in on what rotation gates should be executed/evaluated for a circuit
- We move away from `QuantumTape.diagonalizing_gates`, whose reason for existence is questionable

**Possible Drawbacks:**
- More code in `default.qubit` that needs maintenance/porting over
- Duplicate of `QuantumScript.diagonalizing_gates` code

**Additional Notes**
- I put this in QubitDevice rather than Device to minimize changes to the outward-facing Device API, because I know the scope is limited. I can move it to Device if people prefer it, but this was an intentional choice
- I prefixed the method name with `get_` on purpose, to show that it's doing something (unlike the aforementioned property which was intentionally hiding that fact), and to distinguish it from existing code. Again, I can change this if people prefer it

**Related GitHub Issues:**
#3929, blocks PennyLaneAI/pennylane-lightning#424

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/3938?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3938](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3938?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7681d49) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/2038cee0085ae6175098f982fe8b5964bc635327?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2038cee) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3938   +/-   ##
=======================================
  Coverage   99.74%   99.74%           
=======================================
  Files         344      344           
  Lines       30154    30164   +10     
=======================================
+ Hits        30076    30086   +10     
  Misses         78       78           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.38% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `99.70% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/devices/null\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvbnVsbF9xdWJpdC5weQ==) | `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/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b86cd9fda..fa93df5fa 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -112,6 +112,10 @@
   are real, providing a significant performance improvement.
   [(#3915)](https://github.com/PennyLaneAI/pennylane/pull/3915)
 
+* `QubitDevice` now defines the private `_get_diagonalizing_gates(circuit)` method and uses it when executing circuits.
+  This allows devices that inherit from `QubitDevice` to override and customize their definition of diagonalizing gates.
+  [(#3938)](https://github.com/PennyLaneAI/pennylane/pull/3938)
+
 <h3>Breaking changes</h3>
 
 * Both JIT interfaces are not compatible with Jax `>0.4.3`, we raise an error for those versions.
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index ee6e90a9d..43f586ba7 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -24,7 +24,7 @@ import contextlib
 import itertools
 import warnings
 from collections import defaultdict
-from typing import Union
+from typing import Union, List
 
 import numpy as np
 
@@ -60,7 +60,7 @@ from pennylane.measurements import (
     VnEntropy,
     VnEntropyMP,
 )
-from pennylane.operation import operation_derivative
+from pennylane.operation import operation_derivative, Operation
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.wires import Wires
 
@@ -303,7 +303,7 @@ class QubitDevice(Device):
         the ``QNode`` hash that can be used later for parametric compilation.
 
         Args:
-            circuit (~.tapes.QuantumTape): circuit to execute on the device
+            circuit (~.tape.QuantumTape): circuit to execute on the device
 
         Raises:
             QuantumFunctionError: if the value of :attr:`~.Observable.return_type` is not supported
@@ -314,7 +314,7 @@ class QubitDevice(Device):
         self.check_validity(circuit.operations, circuit.observables)
 
         # apply all circuit operations
-        self.apply(circuit.operations, rotations=circuit.diagonalizing_gates, **kwargs)
+        self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
 
         # generate computational basis samples
         if self.shots is not None:
@@ -356,7 +356,7 @@ class QubitDevice(Device):
         the ``QNode`` hash that can be used later for parametric compilation.
 
         Args:
-            circuit (~.tapes.QuantumTape): circuit to execute on the device
+            circuit (~.tape.QuantumTape): circuit to execute on the device
 
         Raises:
             QuantumFunctionError: if the value of :attr:`~.Observable.return_type` is not supported
@@ -370,7 +370,7 @@ class QubitDevice(Device):
         self.check_validity(circuit.operations, circuit.observables)
 
         # apply all circuit operations
-        self.apply(circuit.operations, rotations=circuit.diagonalizing_gates, **kwargs)
+        self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
 
         # generate computational basis samples
         if self.shots is not None or circuit.is_sampled:
@@ -428,7 +428,7 @@ class QubitDevice(Device):
         vector are contained in a tuple.
 
         Args:
-            circuit (~.tapes.QuantumTape): circuit to execute on the device
+            circuit (~.tape.QuantumTape): circuit to execute on the device
 
         Raises:
             QuantumFunctionError: if the value of :attr:`~.Observable.return_type` is not supported
@@ -543,7 +543,7 @@ class QubitDevice(Device):
         circuits on a backend, for example using parallel and/or asynchronous executions.
 
         Args:
-            circuits (list[.tapes.QuantumTape]): circuits to execute on the device
+            circuits (list[~.tape.QuantumTape]): circuits to execute on the device
 
         Returns:
             list[array[float]]: list of measured value(s)
@@ -576,7 +576,7 @@ class QubitDevice(Device):
         circuits on a backend, for example using parallel and/or asynchronous executions.
 
         Args:
-            circuits (list[.tapes.QuantumTape]): circuits to execute on the device
+            circuits (list[~.tape.QuantumTape]): circuits to execute on the device
 
         Returns:
             list[array[float]]: list of measured value(s)
@@ -1357,7 +1357,7 @@ class QubitDevice(Device):
 
         Args:
             obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow measurement process
-            circuit (~.tapes.QuantumTape): The quantum tape that is being executed
+            circuit (~.tape.QuantumTape): The quantum tape that is being executed
 
         Returns:
             tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
@@ -1394,7 +1394,9 @@ class QubitDevice(Device):
                 ]
 
                 self.reset()
-                self.apply(circuit.operations, rotations=circuit.diagonalizing_gates + rotations)
+                self.apply(
+                    circuit.operations, rotations=self._get_diagonalizing_gates(circuit) + rotations
+                )
 
                 outcomes[t] = self.generate_samples()[0][mapped_wires]
 
@@ -1408,7 +1410,7 @@ class QubitDevice(Device):
         Args:
             obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow expectation
                 value measurement process
-            circuit (~.tapes.QuantumTape): The quantum tape that is being executed
+            circuit (~.tape.QuantumTape): The quantum tape that is being executed
 
         Returns:
             float: expectation value estimate.
@@ -2011,3 +2013,20 @@ class QubitDevice(Device):
 
         # must be 2-dimensional
         return tuple(tuple(np.array(j_) for j_ in j) for j in jac)
+
+    def _get_diagonalizing_gates(self, circuit: QuantumTape) -> List[Operation]:
+        """Returns the gates that diagonalize the measured wires such that they
+        are in the eigenbasis of the circuit observables.
+
+        Note that this exists as a method of the Device class to enable child classes to
+        override the implementation if necessary (for example, to skip computing rotation
+        gates for a measurement that doesn't need them).
+
+        Args:
+            circuit (~.tape.QuantumTape): The circuit containing observables that may need diagonalizing
+
+        Returns:
+            List[~.Operation]: the operations that diagonalize the observables
+        """
+        # pylint:disable=no-self-use
+        return circuit.diagonalizing_gates
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 6e8536fe0..5362134fb 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -21,12 +21,14 @@ simulation of a qubit-based quantum circuit architecture.
 import functools
 import itertools
 from string import ascii_letters as ABC
+from typing import List
 
 import numpy as np
 from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import BasisState, DeviceError, QubitDevice, QubitStateVector, Snapshot
+from pennylane.operation import Operation
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.pulse import ParametrizedEvolution
 from pennylane.typing import TensorLike
@@ -948,7 +950,7 @@ class DefaultQubit(QubitDevice):
 
         Args:
             obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow measurement process
-            circuit (~.tapes.QuantumTape): The quantum tape that is being executed
+            circuit (~.tape.QuantumTape): The quantum tape that is being executed
 
         Returns:
             tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
@@ -1036,3 +1038,11 @@ class DefaultQubit(QubitDevice):
             stacked_state /= norms
 
         return self._cast(self._stack([outcomes, recipes]), dtype=np.int8)
+
+    def _get_diagonalizing_gates(self, circuit: qml.tape.QuantumTape) -> List[Operation]:
+        meas_filtered = [
+            m
+            for m in circuit.measurements
+            if m.obs is None or not isinstance(m.obs, qml.Hamiltonian)
+        ]
+        return super()._get_diagonalizing_gates(qml.tape.QuantumScript(measurements=meas_filtered))
diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py
index 5ee3fdf62..c5d278ee5 100644
--- a/pennylane/devices/null_qubit.py
+++ b/pennylane/devices/null_qubit.py
@@ -280,7 +280,7 @@ class NullQubit(QubitDevice):
         return self._operation_calls
 
     def execute(self, circuit, **kwargs):
-        self.apply(circuit.operations, rotations=circuit.diagonalizing_gates, **kwargs)
+        self.apply(circuit.operations, rotations=self._get_diagonalizing_gates(circuit), **kwargs)
 
         if self.tracker.active:
             self.tracker.update(executions=1, shots=self._shots)
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 51f8b8baf..7e88dab60 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -35,7 +35,7 @@ from pennylane.measurements import (
     StateMP,
     VarianceMP,
 )
-from pennylane.operation import Observable, Operator
+from pennylane.operation import Observable, Operator, Operation
 from pennylane.queuing import AnnotatedQueue, process_queue
 
 _empty_wires = qml.wires.Wires([])
@@ -338,7 +338,7 @@ class QuantumScript:
         return self._output_dim
 
     @property
-    def diagonalizing_gates(self):
+    def diagonalizing_gates(self) -> List[Operation]:
         """Returns the gates that diagonalize the measured wires such that they
         are in the eigenbasis of the circuit observables.
 


Test patch:
diff --git a/tests/devices/test_default_qubit.py b/tests/devices/test_default_qubit.py
index f7d4acaf3..26ed6a731 100644
--- a/tests/devices/test_default_qubit.py
+++ b/tests/devices/test_default_qubit.py
@@ -2335,6 +2335,24 @@ class TestHamiltonianSupport:
         ):
             dev.expval(H)
 
+    def test_Hamiltonian_filtered_from_rotations(self, mocker):
+        """Tests that the device does not attempt to get rotations for Hamiltonians."""
+        dev = qml.device("default.qubit", wires=2, shots=10)
+        H = qml.Hamiltonian([0.1, 0.2], [qml.PauliX(0), qml.PauliZ(1)])
+
+        spy = mocker.spy(qml.QubitDevice, "_get_diagonalizing_gates")
+        qs = qml.tape.QuantumScript([qml.RX(1, 0)], [qml.expval(qml.PauliX(0)), qml.expval(H)])
+        rotations = dev._get_diagonalizing_gates(qs)
+
+        assert len(rotations) == 1
+        assert qml.equal(rotations[0], qml.Hadamard(0))
+
+        call_args = spy.call_args.args[1]  # 0 is self (the device)
+        assert isinstance(call_args, qml.tape.QuantumScript)
+        assert len(call_args.operations) == 0
+        assert len(call_args.measurements) == 1
+        assert qml.equal(call_args.measurements[0], qml.expval(qml.PauliX(0)))
+
 
 class TestGetBatchSize:
     """Tests for the helper method ``_get_batch_size`` of ``QubitDevice``."""
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index e5b6a0f49..a5c47c033 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -1231,6 +1231,14 @@ class TestExecution:
             node_3(0.432, 0.12)
         assert dev_1.num_executions == num_evals_1 + num_evals_3
 
+    def test_get_diagonalizing_gates(self, mock_qubit_device):
+        """Test the private _get_diagonalizing_gates helper method."""
+        circuit = qml.tape.QuantumScript([qml.RX(1, 0)], [qml.probs(), qml.expval(qml.PauliX(0))])
+        dev = mock_qubit_device(wires=1)
+        rotations = dev._get_diagonalizing_gates(circuit)
+        assert len(rotations) == 1
+        assert qml.equal(rotations[0], qml.Hadamard(0))
+
 
 class TestExecutionBroadcasted:
     """Tests for the execute method with broadcasted parameters"""