Repository: PennyLaneAI/pennylane
Pull request #3934 (2023-03-22 18:29:30)

Update `default.qutrit` for new return types
**Context:**
Updating `default.qutrit` to work with the new return types, and validating return types correctly.

**Description of the Change:**
* Updated `QutritDevice` to raise an error if `classical_shadow` or `shadow_expval` are measured on a qutrit device.
* Added tests to validate that `DefaultQutrit` works with the new return type system.


mudit2812: [sc-35002]
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.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fa93df5fa..56e3d02e7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -112,6 +112,12 @@
   are real, providing a significant performance improvement.
   [(#3915)](https://github.com/PennyLaneAI/pennylane/pull/3915)
 
+* Added explicit errors to `QutritDevice` if `classical_shadow` or `shadow_expval` are measured.
+  [(#3934)](https://github.com/PennyLaneAI/pennylane/pull/3934)
+
+* `DefaultQutrit` supports the new return system.
+  [(#3934)](https://github.com/PennyLaneAI/pennylane/pull/3934)
+
 * `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)
diff --git a/pennylane/_qutrit_device.py b/pennylane/_qutrit_device.py
index 3787e5a3f..d9513a33d 100644
--- a/pennylane/_qutrit_device.py
+++ b/pennylane/_qutrit_device.py
@@ -206,6 +206,46 @@ class QutritDevice(QubitDevice):  # pylint: disable=too-many-public-methods
             "Unsupported return type specified for observable mutual information"
         )
 
+    def classical_shadow(self, obs, circuit):
+        """
+        Returns the measured trits and recipes in the classical shadow protocol.
+
+        Please refer to :func:`~.pennylane.classical_shadow` for detailed documentation.
+
+        .. seealso:: :func:`~pennylane.classical_shadow`
+
+        Args:
+            obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow measurement process
+            circuit (~.tapes.QuantumTape): The quantum tape that is being executed
+
+        Raises:
+            QuantumFunctionError: Classical shadow is currently unsupported on :class:`~.QutritDevice`
+        """
+        # TODO: Add support for ClassicalShadowMP
+        raise qml.QuantumFunctionError(
+            "Qutrit devices don't support classical shadow measurements."
+        )
+
+    def shadow_expval(self, obs, circuit):
+        r"""Compute expectation values using classical shadows in a differentiable manner.
+
+        Please refer to :func:`~.pennylane.shadow_expval` for detailed documentation.
+
+        .. seealso:: :func:`~pennylane.shadow_expval`
+
+        Args:
+            obs (~.pennylane.measurements.ShadowExpvalMP): The classical shadow expectation
+                value measurement process
+            circuit (~.tapes.QuantumTape): The quantum tape that is being executed
+
+        Raises:
+            QuantumFunctionError: Shadow Expectation values are currently unsupported on :class:`~.QutritDevice`
+        """
+        # TODO: Add support for ShadowExpvalMP
+        raise qml.QuantumFunctionError(
+            "Qutrit devices don't support shadow expectation value measurements."
+        )
+
     def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         """Return the estimated probability of each computational basis state
         using the generated samples.


Test patch:
diff --git a/tests/returntypes/core/test_new_return_types_qnode.py b/tests/returntypes/core/test_new_return_types_qnode.py
index 578c415b2..0f0896185 100644
--- a/tests/returntypes/core/test_new_return_types_qnode.py
+++ b/tests/returntypes/core/test_new_return_types_qnode.py
@@ -20,7 +20,18 @@ import pytest
 import pennylane as qml
 
 wires = [2, 3, 4]
-devices = ["default.qubit", "lightning.qubit", "default.mixed"]
+devices = ["default.qubit", "lightning.qubit", "default.mixed", "default.qutrit"]
+
+
+def qubit_ansatz(x):
+    qml.Hadamard(wires=[0])
+    qml.CRX(x, wires=[0, 1])
+
+
+def qutrit_ansatz(x):
+    qml.THadamard(wires=[0])
+    mat = np.exp(1j * x) * np.eye(9)
+    qml.QutritUnitary(mat, wires=[0, 1])
 
 
 class TestIntegrationSingleReturn:
@@ -32,8 +43,7 @@ class TestIntegrationSingleReturn:
         dev = qml.device("default.qubit", wires=wires)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.state()
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -48,8 +58,7 @@ class TestIntegrationSingleReturn:
         dev = qml.device("default.mixed", wires=wires)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.state()
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -63,27 +72,30 @@ class TestIntegrationSingleReturn:
     def test_density_matrix(self, d_wires, device):
         """Return density matrix with default.qubit."""
         dev = qml.device(device, wires=4)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            func(x)
             return qml.density_matrix(wires=range(0, d_wires))
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        assert res.shape == (2**d_wires, 2**d_wires)
+        dim = 3 if device == "default.qutrit" else 2
+        assert res.shape == (dim**d_wires, dim**d_wires)
         assert isinstance(res, np.ndarray)
 
     @pytest.mark.parametrize("device", devices)
     def test_expval(self, device):
         """Return a single expval."""
         dev = qml.device(device, wires=2)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.expval(qml.PauliZ(wires=1))
+            func(x)
+            return qml.expval(
+                qml.PauliZ(wires=1) if device != "default.qutrit" else qml.GellMann(1, 3)
+            )
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -95,11 +107,13 @@ class TestIntegrationSingleReturn:
     def test_var(self, device):
         """Return a single var."""
         dev = qml.device(device, wires=2)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.var(qml.PauliZ(wires=1))
+            func(x)
+            return qml.var(
+                qml.PauliZ(wires=1) if device != "default.qutrit" else qml.GellMann(1, 3)
+            )
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -110,11 +124,13 @@ class TestIntegrationSingleReturn:
     @pytest.mark.parametrize("device", devices)
     def test_vn_entropy(self, device):
         """Return a single vn entropy."""
+        if device == "default.qutrit":
+            pytest.skip("DefaultQutrit does not support VnEntropy.")
+
         dev = qml.device(device, wires=2)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.vn_entropy(wires=0)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -131,8 +147,7 @@ class TestIntegrationSingleReturn:
 
         @qml.qnode(device=dev)
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.mutual_info(wires0=[0], wires1=[1])
 
         with pytest.raises(
@@ -143,11 +158,13 @@ class TestIntegrationSingleReturn:
     @pytest.mark.parametrize("device", devices)
     def test_mutual_info(self, device):
         """Return a single mutual information."""
+        if device == "default.qutrit":
+            pytest.skip("DefaultQutrit does not support MutualInfo.")
+
         dev = qml.device(device, wires=2)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.mutual_info(wires0=[0], wires1=[1])
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -164,8 +181,7 @@ class TestIntegrationSingleReturn:
 
         @qml.qnode(device=dev)
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.mutual_info(wires0=[0], wires1=[1])
 
         with pytest.raises(
@@ -185,13 +201,14 @@ class TestIntegrationSingleReturn:
     @pytest.mark.parametrize("op,wires", probs_data)
     def test_probs(self, op, wires, device):
         """Return a single prob."""
-        if device == "lightning.qubit":
-            pytest.skip("Lightning does not support probs with unordered wires.")
+        if device == "lightning.qubit" or device == "default.qutrit":
+            pytest.skip(
+                "Skip Lightning (wire reordering unsupported) and Qutrit (unsuported observables)."
+            )
         dev = qml.device(device, wires=3)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.probs(op=op, wires=wires)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -203,17 +220,54 @@ class TestIntegrationSingleReturn:
         assert res.shape == (2 ** len(wires),)
         assert isinstance(res, np.ndarray)
 
+    probs_data_qutrit = [
+        (qml.GellMann(0, 3), None),
+        (qml.THermitian(np.eye(9), wires=[1, 0]), None),
+        (None, [0]),
+        (None, [0, 1]),
+    ]
+
+    @pytest.mark.parametrize("op,wires", probs_data_qutrit)
+    def test_probs_qutrit(self, op, wires):
+        """Return a single prob."""
+        dev = qml.device("default.qutrit", wires=3)
+
+        def circuit(x):
+            qutrit_ansatz(x)
+            return qml.probs(op=op, wires=wires)
+
+        qnode = qml.QNode(circuit, dev, diff_method=None)
+        res = qnode(0.5)
+
+        if wires is None:
+            wires = op.wires
+
+        assert res.shape == (3 ** len(wires),)
+        assert isinstance(res, np.ndarray)
+
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize(
-        "measurement", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0]), qml.sample(wires=[0, 1])]
+        "measurement",
+        [
+            qml.sample(qml.PauliZ(0)),
+            qml.sample(wires=[0]),
+            qml.sample(wires=[0, 1]),
+            qml.sample(qml.GellMann(0, 3)),
+        ],
     )
     def test_sample(self, measurement, device, shots=100):
         """Test the sample measurement."""
+        if device == "default.qutrit":
+            if isinstance(measurement.obs, qml.PauliZ):
+                pytest.skip("DefaultQutrit doesn't support qubit observables.")
+        elif isinstance(measurement.obs, qml.GellMann):
+            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+
         dev = qml.device(device, wires=2, shots=shots)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            func(x)
             return qml.apply(measurement)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -224,19 +278,31 @@ class TestIntegrationSingleReturn:
         if measurement.wires.tolist() != [0, 1]:
             assert res.shape == (shots,)
         else:
-            assert res.shape == (shots, 2)
+            assert res.shape == (shots, 2) if device != "default.qutrit" else (shots, 3)
 
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize(
-        "measurement", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0]), qml.counts(wires=[0, 1])]
+        "measurement",
+        [
+            qml.counts(qml.PauliZ(0)),
+            qml.counts(wires=[0]),
+            qml.counts(wires=[0, 1]),
+            qml.counts(qml.GellMann(0, 3)),
+        ],
     )
     def test_counts(self, measurement, device, shots=100):
         """Test the counts measurement."""
+        if device == "default.qutrit":
+            if isinstance(measurement.obs, qml.PauliZ):
+                pytest.skip("DefaultQutrit doesn't support qubit observables.")
+        elif isinstance(measurement.obs, qml.GellMann):
+            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+
         dev = qml.device(device, wires=2, shots=shots)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            func(x)
             return qml.apply(measurement)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -923,7 +989,7 @@ class TestIntegrationSingleReturnJax:
 
 wires = [([0], [1]), ([1], [0]), ([0], [0]), ([1], [1])]
 
-devices = ["default.qubit", "lightning.qubit", "default.mixed"]
+devices = ["default.qubit", "lightning.qubit", "default.mixed", "default.qutrit"]
 
 
 class TestIntegrationMultipleReturns:
@@ -936,10 +1002,17 @@ class TestIntegrationMultipleReturns:
         """Return multiple expvals."""
         dev = qml.device(device, wires=2)
 
+        obs1 = (
+            qml.Projector([0], wires=0)
+            if device != "default.qutrit"
+            else qml.THermitian(np.eye(3), wires=0)
+        )
+        obs2 = qml.PauliZ(wires=1) if device != "default.qutrit" else qml.GellMann(1, 3)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
+
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.expval(qml.Projector([0], wires=0)), qml.expval(qml.PauliZ(wires=1))
+            func(x)
+            return qml.expval(obs1), qml.expval(obs2)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -958,10 +1031,17 @@ class TestIntegrationMultipleReturns:
         """Return multiple vars."""
         dev = qml.device(device, wires=2)
 
+        obs1 = (
+            qml.Projector([0], wires=0)
+            if device != "default.qutrit"
+            else qml.THermitian(np.eye(3), wires=0)
+        )
+        obs2 = qml.PauliZ(wires=1) if device != "default.qutrit" else qml.GellMann(1, 3)
+        func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
+
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.var(qml.PauliZ(wires=0)), qml.var(qml.Hermitian([[1, 0], [0, 1]], wires=1))
+            func(x)
+            return qml.var(obs1), qml.var(obs2)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -991,11 +1071,13 @@ class TestIntegrationMultipleReturns:
     @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data)
     def test_multiple_prob(self, op1, op2, wires1, wires2, device):
         """Return multiple probs."""
+        if device == "default.qutrit":
+            pytest.skip("Separate test for DefaultQutrit.")
+
         dev = qml.device(device, wires=2)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return qml.probs(op=op1, wires=wires1), qml.probs(op=op2, wires=wires2)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
@@ -1016,16 +1098,56 @@ class TestIntegrationMultipleReturns:
         assert isinstance(res[1], np.ndarray)
         assert res[1].shape == (2 ** len(wires2),)
 
+    multi_probs_data_qutrit = [
+        (qml.GellMann(0, 3), None, qml.GellMann(1, 3), None),
+        (None, [0], qml.GellMann(1, 3), None),
+        (qml.GellMann(0, 3), None, None, [1]),
+        (qml.GellMann(1, 3), None, qml.GellMann(0, 3), None),
+        (None, [0], None, [0]),
+        (None, [0], None, [0, 1]),
+        (None, [0, 1], None, [0]),
+        (None, [0, 1], None, [0, 1]),
+    ]
+
+    @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data_qutrit)
+    def test_multiple_prob_qutrit(self, op1, op2, wires1, wires2):
+        """Return multiple probs."""
+        dev = qml.device("default.qutrit", wires=2)
+
+        def circuit(x):
+            qutrit_ansatz(x)
+            return qml.probs(op=op1, wires=wires1), qml.probs(op=op2, wires=wires2)
+
+        qnode = qml.QNode(circuit, dev, diff_method=None)
+        res = qnode(0.5)
+
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+
+        if wires1 is None:
+            wires1 = op1.wires
+
+        if wires2 is None:
+            wires2 = op2.wires
+
+        assert isinstance(res[0], np.ndarray)
+        assert res[0].shape == (3 ** len(wires1),)
+
+        assert isinstance(res[1], np.ndarray)
+        assert res[1].shape == (3 ** len(wires2),)
+
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data)
     @pytest.mark.parametrize("wires3, wires4", wires)
     def test_mix_meas(self, op1, wires1, op2, wires2, wires3, wires4, device):
         """Return multiple different measurements."""
+        if device == "default.qutrit":
+            pytest.skip("Different test for DefaultQutrit.")
+
         dev = qml.device(device, wires=2)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
+            qubit_ansatz(x)
             return (
                 qml.probs(op=op1, wires=wires1),
                 qml.vn_entropy(wires=wires3),
@@ -1057,16 +1179,67 @@ class TestIntegrationMultipleReturns:
         assert isinstance(res[3], np.ndarray)
         assert res[3].shape == ()
 
+    @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data_qutrit)
+    @pytest.mark.parametrize("wires3, wires4", wires)
+    def test_mix_meas_qutrit(self, op1, wires1, op2, wires2, wires3, wires4):
+        """Return multiple different measurements."""
+        pytest.skip("Non-commuting observables don't work correctly for qutrits yet.")
+
+        dev = qml.device("default.qutrit", wires=2)
+
+        def circuit(x):
+            qutrit_ansatz(x)
+            return (
+                qml.probs(op=op1, wires=wires1),
+                qml.var(qml.GellMann(wires3, 3)),
+                qml.probs(op=op2, wires=wires2),
+                qml.expval(qml.GellMann(wires4, 3)),
+            )
+
+        qnode = qml.QNode(circuit, dev, diff_method=None)
+        res = qnode(0.5)
+
+        if wires1 is None:
+            wires1 = op1.wires
+
+        if wires2 is None:
+            wires2 = op2.wires
+
+        assert isinstance(res, tuple)
+        assert len(res) == 4
+
+        assert isinstance(res[0], np.ndarray)
+        assert res[0].shape == (3 ** len(wires1),)
+
+        assert isinstance(res[1], np.ndarray)
+        assert res[1].shape == ()
+
+        assert isinstance(res[2], np.ndarray)
+        assert res[2].shape == (3 ** len(wires2),)
+
+        assert isinstance(res[3], np.ndarray)
+        assert res[3].shape == ()
+
     @pytest.mark.parametrize("device", devices)
-    @pytest.mark.parametrize("measurement", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0])])
+    @pytest.mark.parametrize(
+        "measurement",
+        [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0]), qml.sample(qml.GellMann(0, 3))],
+    )
     def test_expval_sample(self, measurement, device, shots=100):
         """Test the expval and sample measurements together."""
+        if device == "default.qutrit":
+            if isinstance(measurement.obs, qml.PauliZ):
+                pytest.skip("DefaultQutrit doesn't support qubit observables.")
+        elif isinstance(measurement.obs, qml.GellMann):
+            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+
         dev = qml.device(device, wires=2, shots=shots)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ(1) if device != "default.qutrit" else qml.GellMann(1, 3)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.expval(qml.PauliX(1)), qml.apply(measurement)
+            func(x)
+            return qml.expval(obs), qml.apply(measurement)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -1080,15 +1253,25 @@ class TestIntegrationMultipleReturns:
         assert res[1].shape == (shots,)
 
     @pytest.mark.parametrize("device", devices)
-    @pytest.mark.parametrize("measurement", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    @pytest.mark.parametrize(
+        "measurement",
+        [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0]), qml.counts(qml.GellMann(0, 3))],
+    )
     def test_expval_counts(self, measurement, device, shots=100):
         """Test the expval and counts measurements together."""
+        if device == "default.qutrit":
+            if isinstance(measurement.obs, qml.PauliZ):
+                pytest.skip("DefaultQutrit doesn't support qubit observables.")
+        elif isinstance(measurement.obs, qml.GellMann):
+            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+
         dev = qml.device(device, wires=2, shots=shots)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ(1) if device != "default.qutrit" else qml.GellMann(1, 3)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.expval(qml.PauliX(1)), qml.apply(measurement)
+            func(x)
+            return qml.expval(obs), qml.apply(measurement)
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -1108,11 +1291,12 @@ class TestIntegrationMultipleReturns:
     def test_list_one_expval(self, wires, device):
         """Return a comprehension list of one expvals."""
         dev = qml.device(device, wires=wires)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ(0) if device != "default.qutrit" else qml.GellMann(0, 3)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return [qml.expval(qml.PauliZ(wires=0))]
+            func(x)
+            return [qml.expval(obs)]
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -1130,11 +1314,15 @@ class TestIntegrationMultipleReturns:
     def test_list_multiple_expval(self, wires, device, shot_vector):
         """Return a comprehension list of multiple expvals."""
         dev = qml.device(device, wires=wires, shots=shot_vector)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ if device != "default.qutrit" else qml.GellMann
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return [qml.expval(qml.PauliZ(wires=i)) for i in range(0, wires)]
+            func(x)
+            return [
+                qml.expval(obs(wires=i) if device != "default.qutrit" else obs(wires=i, index=3))
+                for i in range(0, wires)
+            ]
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
@@ -1158,33 +1346,40 @@ class TestIntegrationMultipleReturns:
     @pytest.mark.parametrize("device", devices)
     def test_array_multiple(self, device):
         """Return PennyLane array of multiple measurements"""
+        if device == "default.qutrit":
+            pytest.skip("Non-commuting observables don't work correctly for qutrits yet.")
+
         dev = qml.device(device, wires=2)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ(1) if device != "default.qutrit" else qml.GellMann(1, 3)
 
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.numpy.array([qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1])])
+            func(x)
+            return qml.numpy.array([qml.expval(obs), qml.probs(wires=[0, 1])])
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
         assert isinstance(res, qml.numpy.ndarray)
         assert res[0].shape == ()
-        assert res[1].shape == (4,)
+        assert res[1].shape == (4,) if device != "default.qutrit" else (9,)
 
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize("comp_basis_sampling", [qml.sample(), qml.counts()])
     def test_sample_counts_no_obs(self, device, comp_basis_sampling):
         """Measuring qml.sample()/qml.counts() works with other measurements even with the same wire being measured."""
+        if device == "default.qutrit":
+            pytest.skip("Non-commuting observables don't work correctly for qutrits yet.")
+
         shot_num = 1000
         num_wires = 2
         dev = qml.device(device, wires=num_wires, shots=shot_num)
+        func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
+        obs = qml.PauliZ(1) if device != "default.qutrit" else qml.GellMann(1, 3)
 
-        @qml.qnode(dev, diff_method=None)
         def circuit(x):
-            qml.Hadamard(wires=[0])
-            qml.CRX(x, wires=[0, 1])
-            return qml.apply(comp_basis_sampling), qml.expval(qml.PauliZ(0)), qml.probs(wires=[0])
+            func(x)
+            return qml.apply(comp_basis_sampling), qml.expval(obs), qml.probs(wires=[0])
 
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
diff --git a/tests/test_qutrit_device.py b/tests/test_qutrit_device.py
index 30cfdd331..d8c549c0f 100644
--- a/tests/test_qutrit_device.py
+++ b/tests/test_qutrit_device.py
@@ -14,6 +14,7 @@
 """
 Unit tests for the :mod:`pennylane` :class:`QutritDevice` class.
 """
+# pylint: disable=protected-access, redefined-outer-name, too-many-arguments, too-few-public-methods, unused-variable, unused-argument
 from random import random
 
 import numpy as np
@@ -150,7 +151,7 @@ class TestOperations:
         with pytest.raises(
             ValueError, match="Cannot access the operation queue outside of the execution context!"
         ):
-            dev.op_queue
+            _ = dev.op_queue
 
     def test_op_queue_is_filled_during_execution(self, mock_qutrit_device, monkeypatch):
         """Tests that the op_queue is correctly filled when apply is called and that accessing
@@ -159,7 +160,6 @@ class TestOperations:
 
         with qml.queuing.AnnotatedQueue() as q:
             queue = [qml.QutritUnitary(U, wires=0), qml.QutritUnitary(U, wires=0)]
-            observables = [qml.expval(qml.Identity(0))]
 
         tape = QuantumScript.from_queue(q)
         call_history = []
@@ -188,15 +188,14 @@ class TestOperations:
         """Tests that the operations are properly applied and queued"""
         U = unitary_group.rvs(3, random_state=10)
 
-        with qml.queuing.AnnotatedQueue() as q:
-            queue = [
-                qml.QutritUnitary(U, wires=0),
-                qml.Hadamard(wires=1),
-                qml.QutritUnitary(U, wires=2),
-            ]
-            observables = [qml.expval(qml.Identity(0)), qml.var(qml.Identity(1))]
+        queue = [
+            qml.QutritUnitary(U, wires=0),
+            qml.Hadamard(wires=1),
+            qml.QutritUnitary(U, wires=2),
+        ]
+        observables = [qml.expval(qml.Identity(0)), qml.var(qml.Identity(1))]
 
-        tape = QuantumScript.from_queue(q)
+        tape = QuantumScript(queue, observables)
         dev = mock_qutrit_device()
         with pytest.raises(DeviceError, match="Gate Hadamard not supported on device"):
             dev.execute(tape)
@@ -233,8 +232,7 @@ class TestOperations:
             dev = mock_qutrit_device()
             dev.execute(tape, hash=tape.graph.hash)
 
-        len(call_history.items()) == 1
-        call_history["hash"] = tape.graph.hash
+        assert call_history["hash"] == tape.graph.hash
 
 
 class TestObservables:
@@ -250,17 +248,16 @@ class TestObservables:
             ValueError,
             match="Cannot access the observable value queue outside of the execution context!",
         ):
-            dev.obs_queue
+            _ = dev.obs_queue
 
     def test_unsupported_observables_raise_error(self, mock_qutrit_device):
         """Tests that the operations are properly applied and queued"""
         U = unitary_group.rvs(3, random_state=10)
 
-        with qml.queuing.AnnotatedQueue() as q:
-            queue = [qml.QutritUnitary(U, wires=0)]
-            observables = [qml.expval(qml.Hadamard(0))]
+        queue = [qml.QutritUnitary(U, wires=0)]
+        observables = [qml.expval(qml.Hadamard(0))]
 
-        tape = QuantumScript.from_queue(q)
+        tape = QuantumScript(queue, observables)
         dev = mock_qutrit_device()
         with pytest.raises(DeviceError, match="Observable Hadamard not supported on device"):
             dev.execute(tape)
@@ -300,7 +297,7 @@ class TestParameters:
             ValueError,
             match="Cannot access the free parameter mapping outside of the execution context!",
         ):
-            dev.parameters
+            _ = dev.parameters
 
 
 class TestExtractStatistics:
@@ -445,7 +442,7 @@ class TestSample:
 
         obs = SomeObservable(wires=0)
         with pytest.raises(qml.operation.EigvalsUndefinedError, match="Cannot compute samples"):
-            dev.sample(SomeObservable(wires=0))
+            dev.sample(obs)
 
     def test_samples_with_bins(self, mock_qutrit_device_with_original_statistics, monkeypatch):
         """Tests that sample works correctly when instantiating device with shot list"""
@@ -459,7 +456,7 @@ class TestSample:
         bin_size = 3
 
         out = dev.sample(obs, shot_range=shot_range, bin_size=bin_size)
-        expected_samples = samples.reshape(-1, 3, 2)
+        expected_samples = samples.reshape((-1, 3, 2))
 
         assert np.array_equal(out, expected_samples)
 
@@ -972,7 +969,10 @@ class TestBatchExecution:
 
     with qml.queuing.AnnotatedQueue() as q1:
         qml.QutritUnitary(np.eye(3), wires=0)
-        qml.expval(qml.Identity(0)), qml.expval(qml.Identity(1))
+        _ = (
+            qml.expval(qml.Identity(0)),
+            qml.expval(qml.Identity(1)),
+        )
 
     tape1 = QuantumScript.from_queue(q1)
     with qml.queuing.AnnotatedQueue() as q2:
@@ -1173,7 +1173,7 @@ class TestShotList:
 
         res = circuit(pnp.eye(9))
 
-        assert res.shape == expected_shape
+        assert res.shape == expected_shape  # pylint: disable=no-member
         assert circuit.device._shot_vector == shot_vector
         assert circuit.device.shots == total_shots
 
@@ -1226,3 +1226,21 @@ class TestUnimplemented:
 
         with pytest.raises(qml.QuantumFunctionError, match="Unsupported return type"):
             dev.mutual_info(0, 1, log_base=3)
+
+    def test_classical_shadow(self, mock_qutrit_device):
+        """Test that classical_shadow is unimplemented"""
+        dev = mock_qutrit_device()
+        qs = qml.tape.QuantumScript()
+        obs = qml.GellMann(1, 1)
+
+        with pytest.raises(qml.QuantumFunctionError, match="Qutrit devices don't support"):
+            dev.classical_shadow(obs, qs)
+
+    def test_shadow_expval(self, mock_qutrit_device):
+        """Test that shadow_expval is unimplemented"""
+        dev = mock_qutrit_device()
+        qs = qml.tape.QuantumScript()
+        obs = qml.GellMann(1, 1)
+
+        with pytest.raises(qml.QuantumFunctionError, match="Qutrit devices don't support"):
+            dev.shadow_expval(obs, qs)
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index c36103b15..de55fb311 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -53,3 +53,4 @@ tests/devices/test_default_mixed_torch.py
 tests/devices/test_default_mixed_tf.py
 tests/devices/test_null_qubit.py
 tests/fourier/test_coefficients.py
+tests/test_qutrit_device.py