Repository: PennyLaneAI/pennylane
Pull request #1889 (2021-11-11 09:19:06)

[OpRefactor] Delete DiagonalOperation class and replace with attribute
**Context:**

As part of the operation refactor we decided to not tag properties of operations by class inheritence due to problems with scalability. Rather, we use attributes (basically, sets) that list all gates with a specific property.  This PR switches from inheritance to attributes for the property of an operation being diagonal in the computational basis.

**Description of the Change:**

Deletes `DiagonalOperation` class and adds a `diagonal_in_z_basis` attribute.

**Benefits:**

More scalable and clean approach. Only very few changes needed.

**Possible Drawbacks:**

I don't know if external devices rely on `DiagonalOperation`, but should be easy to rewrite..

**Related GitHub Issues:**


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1889](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3622a33) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/9de11822b679b9e4e2edd44b443ac6304a4b5cc0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9de1182) will **decrease** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1889      +/-   ##
==========================================
- Coverage   98.98%   98.97%   -0.01%     
==========================================
  Files         218      218              
  Lines       16457    16454       -3     
==========================================
- Hits        16290    16286       -4     
- Misses        167      168       +1     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `98.88% <ø> (-0.03%)` | :arrow_down: |
| [pennylane/devices/default\_mixed.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/default\_qubit\_torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF90b3JjaC5weQ==) | `90.21% <100.00%> (ø)` | |
| [pennylane/ops/qubit/attributes.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9hdHRyaWJ1dGVzLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/qubit/matrix\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9tYXRyaXhfb3BzLnB5) | `99.14% <100.00%> (+0.02%)` | :arrow_up: |
| [pennylane/ops/qubit/non\_parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9ub25fcGFyYW1ldHJpY19vcHMucHk=) | `97.20% <100.00%> (ø)` | |
| [pennylane/ops/qubit/parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9wYXJhbWV0cmljX29wcy5weQ==) | `99.82% <100.00%> (-0.18%)` | :arrow_down: |
| ... and [1 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [9de1182...3622a33](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1889?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

mariaschuld: @glassnotes I am so sorry, I did not see your two comments earlier! The scratch file was from me trying VS Code as editor and trying to reproduce PyCharm's scratch feature :) Thanks for catching those.

Patch:
diff --git a/doc/development/plugins.rst b/doc/development/plugins.rst
index 6d420a9ff..cec4bd83d 100644
--- a/doc/development/plugins.rst
+++ b/doc/development/plugins.rst
@@ -550,27 +550,9 @@ The user can then import this operation directly from your plugin, and use it wh
     that the plugin unit tests **do** provide tests to ensure that PennyLane returns the correct
     gradient for the custom operations.
 
-If the custom operation is diagonal in the computational basis, it is recommended to subclass
-:class:`~.DiagonalOperation`. For diagonal operations, :attr:`_eigvals` has to be overridden
-instead of :attr:`_matrix`. 
-
-.. code-block:: python
-
-    class CustomDiagonalGate(DiagonalOperation):
-        """Custom gate"""
-        num_params = 0
-        num_wires = 2
-
-        @classmethod
-        def _eigvals(cls, *params):
-            """Returns the eigenvalues of the operation in the computational basis."""
-            return np.array([1j, 1j, -1j, -1j])
-
-        @staticmethod
-        def decomposition(*params, wires):
-            """(Optional) Returns a list of PennyLane operations that decompose
-            the custom gate."""
-            return [qml.DiagonalQubitUnitary([1j, 1j, -1j, -1j], wires=wires)]
+If the custom operation is diagonal in the computational basis, it can be added to the
+``diagonal_in_z_basis`` attribute in ``pennylane.ops.qubit.attributes``. Devices can use this
+information to implement faster simulations.
 
 Supporting new observables
 --------------------------
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 78ee1ffd0..5a4414908 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -20,6 +20,19 @@
 
 <h3>Breaking changes</h3>
 
+* `DiagonalOperation` is not a separate subclass any more. 
+  [(#1889)](https://github.com/PennyLaneAI/pennylane/pull/1889) 
+
+  Instead, devices can check for the diagonal 
+  property using attributes:
+
+  ``` python
+  from pennylane.ops.qubit.attributes import diagonal_in_z_basis
+
+  if op in diagonal_in_z_basis:
+      # do something
+  ``` 
+
 <h3>Deprecations</h3>
 
 <h3>Bug fixes</h3>
@@ -38,4 +51,4 @@
 
 This release contains contributions from (in alphabetical order): 
 
-Jalani Kanem, Christina Lee, Guillermo Alonso-Linaje, Alejandro Montanez, Jay Soni
+Jalani Kanem, Christina Lee, Guillermo Alonso-Linaje, Alejandro Montanez, Maria Schuld, Jay Soni
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index 8a965e853..53570d3d1 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -25,8 +25,9 @@ from string import ascii_letters as ABC
 
 import numpy as np
 from pennylane import QubitDevice, QubitStateVector, BasisState, DeviceError
-from pennylane.operation import DiagonalOperation, Channel
+from pennylane.operation import Channel
 from pennylane.wires import Wires
+from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from .._version import __version__
 
 ABC_ARRAY = np.array(list(ABC))
@@ -209,7 +210,7 @@ class DefaultMixed(QubitDevice):
             the operation is unitary, returns a single Kraus operator. In the case of a diagonal
             unitary, returns a 1D array representing the matrix diagonal.
         """
-        if isinstance(operation, DiagonalOperation):
+        if operation in diagonal_in_z_basis:
             return operation.eigvals
 
         if isinstance(operation, Channel):
@@ -408,7 +409,7 @@ class DefaultMixed(QubitDevice):
 
         matrices = self._get_kraus(operation)
 
-        if isinstance(operation, DiagonalOperation):
+        if operation in diagonal_in_z_basis:
             self._apply_diagonal_unitary(matrices, wires)
         else:
             self._apply_channel(matrices, wires)
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 6312b5e06..6c8c04a15 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -27,7 +27,7 @@ from scipy.sparse import coo_matrix
 
 import pennylane as qml
 from pennylane import QubitDevice, DeviceError, QubitStateVector, BasisState
-from pennylane.operation import DiagonalOperation
+from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.wires import WireError
 from .._version import __version__
 
@@ -242,7 +242,7 @@ class DefaultQubit(QubitDevice):
 
         matrix = self._get_unitary_matrix(operation)
 
-        if isinstance(operation, DiagonalOperation):
+        if operation in diagonal_in_z_basis:
             return self._apply_diagonal_unitary(state, matrix, wires)
         if len(wires) <= 2:
             # Einsum is faster for small gates
@@ -544,7 +544,7 @@ class DefaultQubit(QubitDevice):
             the unitary in the computational basis, or, in the case of a diagonal unitary,
             a 1D array representing the matrix diagonal.
         """
-        if isinstance(unitary, DiagonalOperation):
+        if unitary in diagonal_in_z_basis:
             return unitary.eigvals
 
         return unitary.matrix
diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py
index 179f21e9e..2e5816932 100644
--- a/pennylane/devices/default_qubit_torch.py
+++ b/pennylane/devices/default_qubit_torch.py
@@ -27,7 +27,7 @@ except ImportError as e:
     raise ImportError("default.qubit.torch device requires Torch>=1.8.1") from e
 
 import numpy as np
-from pennylane.operation import DiagonalOperation
+from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from . import DefaultQubit
 
 
@@ -237,7 +237,7 @@ class DefaultQubitTorch(DefaultQubit):
             the unitary in the computational basis, or, in the case of a diagonal unitary,
             a 1D array representing the matrix diagonal.
         """
-        if isinstance(unitary, DiagonalOperation):
+        if unitary in diagonal_in_z_basis:
             return self._asarray(unitary.eigvals, dtype=self.C_DTYPE)
         return self._asarray(unitary.matrix, dtype=self.C_DTYPE)
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index f57fee938..ba430d39b 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -821,91 +821,6 @@ class Operation(Operator):
             assert self.grad_recipe is None, "Gradient recipe is only used by the A method!"
 
 
-class DiagonalOperation(Operation):
-    r"""Base class for diagonal quantum operations supported by a device.
-
-    As with :class:`~.Operation`, the following class attributes must be
-    defined for all operations:
-
-    * :attr:`~.Operator.num_params`
-    * :attr:`~.Operator.num_wires`
-    * :attr:`~.Operator.par_domain`
-
-    The following two class attributes are optional, but in most cases
-    should be clearly defined to avoid unexpected behavior during
-    differentiation.
-
-    * :attr:`~.Operation.grad_method`
-    * :attr:`~.Operation.grad_recipe`
-
-    Finally, there are some additional optional class attributes
-    that may be set, and used by certain quantum optimizers:
-
-    * :attr:`~.Operation.generator`
-
-    Args:
-        params (tuple[float, int, array]): operation parameters
-
-    Keyword Args:
-        wires (Sequence[int]): Subsystems it acts on. If not given, args[-1]
-            is interpreted as wires.
-        do_queue (bool): Indicates whether the operation should be
-            immediately pushed into a :class:`BaseQNode` circuit queue.
-            This flag is useful if there is some reason to run an Operation
-            outside of a BaseQNode context.
-    """
-    # pylint: disable=abstract-method
-
-    @classmethod
-    def _eigvals(cls, *params):
-        """Eigenvalues of the operator.
-
-        The order of the eigenvalues needs to match the order of
-        the computational basis vectors.
-
-        This is a *class method* that must be defined for all
-        new diagonal operations, that returns the eigenvalues
-        of the operator in the computational basis.
-
-        This private method allows eigenvalues to be computed
-        directly without instantiating the operators first.
-
-        To return the eigenvalues of *instantiated* operators,
-        please use the :attr:`~.Operator.eigvals` property instead.
-
-        **Example:**
-
-        >>> qml.RZ._eigvals(0.5)
-        >>> array([0.96891242-0.24740396j, 0.96891242+0.24740396j])
-
-        Returns:
-            array: eigenvalue representation
-        """
-        raise NotImplementedError
-
-    @property
-    def eigvals(self):
-        r"""Eigenvalues of an instantiated diagonal operation.
-
-        The order of the eigenvalues needs to match the order of
-        the computational basis vectors.
-
-        **Example:**
-
-        >>> U = qml.RZ(0.5, wires=1)
-        >>> U.eigvals
-        >>> array([0.96891242-0.24740396j, 0.96891242+0.24740396j])
-
-        Returns:
-            array: eigvals representation
-        """
-        return super().eigvals
-
-    @classmethod
-    def _matrix(cls, *params):
-        return np.diag(cls._eigvals(*params))
-
-
 class Channel(Operation, abc.ABC):
     r"""Base class for quantum channels.
 
diff --git a/pennylane/ops/qubit/attributes.py b/pennylane/ops/qubit/attributes.py
index 1f1078202..2c143a2c2 100644
--- a/pennylane/ops/qubit/attributes.py
+++ b/pennylane/ops/qubit/attributes.py
@@ -16,7 +16,7 @@ This file contains a number of attributes that may be held by operators,
 and lists all operators satisfying those criteria.
 """
 from inspect import isclass
-from pennylane.operation import Operator
+from pennylane.operation import Operator, Tensor
 
 
 class Attribute(set):
@@ -75,6 +75,13 @@ class Attribute(set):
             return super().__contains__(obj)
 
         try:
+
+            # Hotfix: return False for all tensors.
+            # Can be removed or updated when tensor class is
+            # improved.
+            if isinstance(obj, Tensor):
+                return False
+
             if isinstance(obj, Operator):
                 return super().__contains__(obj.name)
 
@@ -158,3 +165,23 @@ For example, ``qml.Toffoli(wires=[0, 1, 2])`` has the same effect as
 ``qml.Toffoli(wires=[1, 0, 2])``, but neither are the same as
 ``qml.Toffoli(wires=[0, 2, 1])``.
 """
+
+diagonal_in_z_basis = Attribute(
+    [
+        "PauliZ",
+        "S",
+        "T",
+        "CZ",
+        "DiagonalQubitUnitary",
+        "RZ",
+        "PhaseShift",
+        "ControlledPhaseShift",
+        "MultiRZ",
+        "CRZ",
+    ]
+)
+"""Attribute: Operations that are diagonal in the computational basis.
+
+For such operations, the eigenvalues provide all necessary information to 
+construct the matrix representation in the computational basis.
+"""
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 8de44d2ed..a586cfd0c 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -20,7 +20,7 @@ import warnings
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import AnyWires, DiagonalOperation, Operation
+from pennylane.operation import AnyWires, Operation
 from pennylane.wires import Wires
 
 
@@ -242,7 +242,7 @@ class ControlledQubitUnitary(QubitUnitary):
         ControlledQubitUnitary(*self.parameters, control_wires=ctrl_wires, wires=self._target_wires)
 
 
-class DiagonalQubitUnitary(DiagonalOperation):
+class DiagonalQubitUnitary(Operation):
     r"""DiagonalQubitUnitary(D, wires)
     Apply an arbitrary fixed diagonal unitary matrix.
 
@@ -261,6 +261,10 @@ class DiagonalQubitUnitary(DiagonalOperation):
     par_domain = "A"
     grad_method = None
 
+    @classmethod
+    def _matrix(cls, *params):
+        return qml.math.diag(cls._eigvals(*params))
+
     @classmethod
     def _eigvals(cls, *params):
         D = qml.math.asarray(params[0])
diff --git a/pennylane/ops/qubit/non_parametric_ops.py b/pennylane/ops/qubit/non_parametric_ops.py
index 7699de2c2..22d21e28d 100644
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -21,7 +21,7 @@ import numpy as np
 from scipy.linalg import block_diag
 
 import pennylane as qml
-from pennylane.operation import AnyWires, DiagonalOperation, Observable, Operation
+from pennylane.operation import AnyWires, Observable, Operation
 from pennylane.utils import pauli_eigs
 from pennylane.wires import Wires
 
@@ -230,7 +230,7 @@ class PauliY(Observable, Operation):
         return [0.0, np.pi, 0.0]
 
 
-class PauliZ(Observable, DiagonalOperation):
+class PauliZ(Observable, Operation):
     r"""PauliZ(wires)
     The Pauli Z operator
 
@@ -281,7 +281,7 @@ class PauliZ(Observable, DiagonalOperation):
         return [np.pi, 0.0, 0.0]
 
 
-class S(DiagonalOperation):
+class S(Operation):
     r"""S(wires)
     The single-qubit phase gate
 
@@ -326,7 +326,7 @@ class S(DiagonalOperation):
         return [np.pi / 2, 0.0, 0.0]
 
 
-class T(DiagonalOperation):
+class T(Operation):
     r"""T(wires)
     The single-qubit T gate
 
@@ -466,7 +466,7 @@ class CNOT(Operation):
         return Wires(self.wires[0])
 
 
-class CZ(DiagonalOperation):
+class CZ(Operation):
     r"""CZ(wires)
     The controlled-Z operator
 
diff --git a/pennylane/ops/qubit/parametric_ops.py b/pennylane/ops/qubit/parametric_ops.py
index b8a5ed97d..e1f15c394 100644
--- a/pennylane/ops/qubit/parametric_ops.py
+++ b/pennylane/ops/qubit/parametric_ops.py
@@ -21,7 +21,7 @@ import math
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import AnyWires, DiagonalOperation, Operation
+from pennylane.operation import AnyWires, Operation
 from pennylane.ops.qubit.non_parametric_ops import PauliX, PauliY, PauliZ, Hadamard
 from pennylane.utils import expand, pauli_eigs
 from pennylane.wires import Wires
@@ -134,7 +134,7 @@ class RY(Operation):
         return [0.0, self.data[0], 0.0]
 
 
-class RZ(DiagonalOperation):
+class RZ(Operation):
     r"""RZ(phi, wires)
     The single qubit Z rotation
 
@@ -194,7 +194,7 @@ class RZ(DiagonalOperation):
         return [self.data[0], 0.0, 0.0]
 
 
-class PhaseShift(DiagonalOperation):
+class PhaseShift(Operation):
     r"""PhaseShift(phi, wires)
     Arbitrary single qubit local phase shift
 
@@ -262,7 +262,7 @@ class PhaseShift(DiagonalOperation):
         return [self.data[0], 0.0, 0.0]
 
 
-class ControlledPhaseShift(DiagonalOperation):
+class ControlledPhaseShift(Operation):
     r"""ControlledPhaseShift(phi, wires)
     A qubit controlled phase shift.
 
@@ -428,7 +428,7 @@ class Rot(Operation):
         return self.data
 
 
-class MultiRZ(DiagonalOperation):
+class MultiRZ(Operation):
     r"""MultiRZ(theta, wires)
     Arbitrary multi Z rotation.
 
@@ -947,7 +947,7 @@ class CRY(Operation):
         return Wires(self.wires[0])
 
 
-class CRZ(DiagonalOperation):
+class CRZ(Operation):
     r"""CRZ(phi, wires)
     The controlled-RZ operator
 


Test patch:
diff --git a/tests/devices/test_default_mixed.py b/tests/devices/test_default_mixed.py
index 5958c2dda..b9c019bd6 100644
--- a/tests/devices/test_default_mixed.py
+++ b/tests/devices/test_default_mixed.py
@@ -448,9 +448,9 @@ class TestApplyDiagonal:
         dev = qml.device("default.mixed", wires=nr_wires)
         kraus = dev._get_kraus(op)
         if op == CZ:
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
+            dev._apply_diagonal_unitary(kraus, wires=Wires([0, 1]))
         else:
-            dev._apply_channel(kraus, wires=Wires(0))
+            dev._apply_diagonal_unitary(kraus, wires=Wires(0))
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
@@ -467,9 +467,9 @@ class TestApplyDiagonal:
         dev._state = max_mixed
         kraus = dev._get_kraus(op)
         if op == CZ:
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
+            dev._apply_diagonal_unitary(kraus, wires=Wires([0, 1]))
         else:
-            dev._apply_channel(kraus, wires=Wires(0))
+            dev._apply_diagonal_unitary(kraus, wires=Wires(0))
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
@@ -500,9 +500,9 @@ class TestApplyDiagonal:
         dev._state = root
         kraus = dev._get_kraus(op)
         if op == CZ:
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
+            dev._apply_diagonal_unitary(kraus, wires=Wires([0, 1]))
         else:
-            dev._apply_channel(kraus, wires=Wires(0))
+            dev._apply_diagonal_unitary(kraus, wires=Wires(0))
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
diff --git a/tests/devices/test_default_qubit.py b/tests/devices/test_default_qubit.py
index 0a6cd7960..56d1cf589 100644
--- a/tests/devices/test_default_qubit.py
+++ b/tests/devices/test_default_qubit.py
@@ -2216,10 +2216,10 @@ class TestInverseDecomposition:
         assert np.allclose(dev.state, expected, atol=tol, rtol=0)
 
 
-@pytest.mark.parametrize("inverse", [True, False])
 class TestApplyOperationUnit:
     """Unit tests for the internal _apply_operation method."""
 
+    @pytest.mark.parametrize("inverse", [True, False])
     def test_internal_apply_ops_case(self, inverse, monkeypatch):
         """Tests that if we provide an operation that has an internal
         implementation, then we use that specific implementation.
@@ -2243,19 +2243,15 @@ class TestApplyOperationUnit:
             res = dev._apply_operation(test_state, op)
             assert np.allclose(res, expected_test_output)
 
-    def test_diagonal_operation_case(self, inverse, mocker, monkeypatch):
-        """Tests the case when the operation to be applied is a
-        DiagonalOperation and the _apply_diagonal_unitary method is used."""
+    def test_diagonal_operation_case(self, mocker, monkeypatch):
+        """Tests the case when the operation to be applied is
+        diagonal in the computational basis and the _apply_diagonal_unitary method is used."""
         dev = qml.device("default.qubit", wires=1)
         par = 0.3
 
         test_state = np.array([1, 0])
         wires = 0
-        op = (
-            qml.PhaseShift(par, wires=wires)
-            if not inverse
-            else qml.PhaseShift(par, wires=wires).inv()
-        )
+        op = qml.PhaseShift(par, wires=wires)
         assert op.name not in dev._apply_ops
 
         # Set the internal _apply_diagonal_unitary
@@ -2273,7 +2269,7 @@ class TestApplyOperationUnit:
             assert np.allclose(res_mat, np.diag(op.matrix))
             assert np.allclose(res_wires, wires)
 
-    def test_apply_einsum_case(self, inverse, mocker, monkeypatch):
+    def test_apply_einsum_case(self, mocker, monkeypatch):
         """Tests the case when np.einsum is used to apply an operation in
         default.qubit."""
         dev = qml.device("default.qubit", wires=1)
@@ -2282,7 +2278,7 @@ class TestApplyOperationUnit:
         wires = 0
 
         # Redefine the S gate so that it is an example for a one-qubit gate
-        # that does not inherit from DiagonalOperation
+        # that is not registered in the diagonal_in_z_basis attribute
         class TestSGate(qml.operation.Operation):
             matrix = np.array([[0, 1], [1, 0]])
             num_params = 0
@@ -2299,9 +2295,6 @@ class TestApplyOperationUnit:
         assert op.name in dev.operations
         assert op.name not in dev._apply_ops
 
-        if inverse:
-            op = op.inv()
-
         # Set the internal _apply_unitary_einsum
         history = []
         mock_apply_einsum = lambda state, matrix, wires: history.append((state, matrix, wires))
@@ -2316,6 +2309,7 @@ class TestApplyOperationUnit:
             assert np.allclose(res_mat, op.matrix)
             assert np.allclose(res_wires, wires)
 
+    @pytest.mark.parametrize("inverse", [True, False])
     def test_apply_tensordot_case(self, inverse, mocker, monkeypatch):
         """Tests the case when np.tensordot is used to apply an operation in
         default.qubit."""
diff --git a/tests/ops/qubit/test_attributes.py b/tests/ops/qubit/test_attributes.py
index b9a8b344b..095e2e091 100644
--- a/tests/ops/qubit/test_attributes.py
+++ b/tests/ops/qubit/test_attributes.py
@@ -75,3 +75,7 @@ class TestAttribute:
         assert "PhaseShift" in new_attribute
         assert "RY" in new_attribute
         assert len(new_attribute) == 8
+
+    def test_tensor_check(self):
+        """Test that we can ask if a tensor is in the attribute."""
+        assert not qml.PauliX(wires=0) @ qml.PauliZ(wires=1) in new_attribute