Repository: PennyLaneAI/pennylane
Pull request #5502 (2024-04-11 23:35:25)

Qutrit channel depolarizing
**Context:**
`default.qutrit.mixed` device has been added, but no channels have been added so far, this adds the first channel to the device so that it may be used for it's intended purpose of simulating noise.

**Description of the Change:**
Adds new channel module to `qml.ops.qutrit` package. Adding the first qutrit channel `QutritDepolarizingChannel`.

**Benefits:**
Allows for `defualt.qutrit.mixed` to simulate depolarizing noise. Makes `default.qutrit.mixed`, much more useful.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A


Gabriel-Bottrill: Hi @albi3ro, @glassnotes, @mudit2812, @trbromley. This PR is not finished, but I wanted to put it up as a draft to let you know it's here. I realized I probably should have been @ing you all in previous drafts, but I thought I would now. Also after this PR ther will be another Channel which I have drafted in this [PR](https://github.com/PennyLaneAI/pennylane/pull/5503).
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5502?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.68%. Comparing base [(`9b13456`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9b13456185dbca8f0e5b578b7134f5daabe8ecd5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e098c4e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5502?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5502      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         412      413       +1     
  Lines       38663    38404     -259     
==========================================
- Hits        38544    38284     -260     
- Misses        119      120       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5502?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

Gabriel-Bottrill: Hi @albi3ro, @glassnotes, @mudit2812, @trbromley I have finished this PR, except for the failing test, I am a bit confused why that error is occurring, so I thought I would ask you all. I will add some comments to specific parts of the code now too.

Patch:
diff --git a/doc/introduction/operations.rst b/doc/introduction/operations.rst
index 9666fc581..970e7b4ec 100644
--- a/doc/introduction/operations.rst
+++ b/doc/introduction/operations.rst
@@ -495,6 +495,21 @@ Qutrit State preparation
 
 :html:`</div>`
 
+.. _intro_ref_ops_qutrit_channels:
+
+Qutrit noisy channels
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+:html:`<div class="summary-table">`
+
+.. autosummary::
+    :nosignatures:
+
+    ~pennylane.QutritDepolarizingChannel
+
+:html:`</div>`
+
 .. _intro_ref_ops_qutrit_obs:
 
 Qutrit Observables
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0b6587eac..01db3008d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -98,6 +98,11 @@
 * ``qml.load`` has been removed in favour of more specific functions, such as ``qml.from_qiskit``, etc.
   [(#5654)](https://github.com/PennyLaneAI/pennylane/pull/5654)
 
+<h4>Community contributions 🥳</h4>
+
+* ``qml.QutritDepolarizingChannel`` has been added, allowing for depolarizing noise to be simulated on the `default.qutrit.mixed` device.
+  [(#5502)](https://github.com/PennyLaneAI/pennylane/pull/5502)
+
 <h3>Deprecations 👋</h3>
 
 * ``qml.transforms.map_batch_transform`` is deprecated, since a transform can be applied directly to a batch of tapes.
@@ -131,6 +136,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Gabriel Bottrill,
 Isaac De Vlugt,
 Pietropaolo Frisoni,
 Soran Jahangiri,
diff --git a/pennylane/devices/default_qutrit_mixed.py b/pennylane/devices/default_qutrit_mixed.py
index 7f98a9497..7c264a9d7 100644
--- a/pennylane/devices/default_qutrit_mixed.py
+++ b/pennylane/devices/default_qutrit_mixed.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """The default.qutrit.mixed device is PennyLane's standard qutrit simulator for mixed-state
 computations."""
-
 import inspect
 import logging
 from dataclasses import replace
@@ -22,6 +21,7 @@ from typing import Callable, Optional, Sequence, Tuple, Union
 import numpy as np
 
 import pennylane as qml
+from pennylane.ops import _qutrit__channel__ops__ as channels
 from pennylane.tape import QuantumTape
 from pennylane.transforms.core import TransformProgram
 from pennylane.typing import Result, ResultBatch
@@ -49,7 +49,6 @@ QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
 # always a function from a resultbatch to either a result or a result batch
 PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
 
-channels = set()
 observables = {
     "THermitian",
     "GellMann",
diff --git a/pennylane/ops/__init__.py b/pennylane/ops/__init__.py
index ee1a9eefa..6beae6b80 100644
--- a/pennylane/ops/__init__.py
+++ b/pennylane/ops/__init__.py
@@ -37,6 +37,7 @@ from .qutrit import *
 from .qutrit import __all__ as _qutrit__all__
 from .qutrit import __obs__ as _qutrit__obs__
 from .qutrit import __ops__ as _qutrit__ops__
+from .qutrit import __channels__ as _qutrit__channel__ops__
 
 _qubit__ops__ = _qubit__ops__ | _controlled_qubit__ops__
 _qubit__all__ = _qubit__all__ + list(_controlled_qubit__ops__)
diff --git a/pennylane/ops/qutrit/__init__.py b/pennylane/ops/qutrit/__init__.py
index 36aa7f227..670f950ba 100644
--- a/pennylane/ops/qutrit/__init__.py
+++ b/pennylane/ops/qutrit/__init__.py
@@ -26,6 +26,7 @@ from .non_parametric_ops import *
 from .observables import *
 from .parametric_ops import *
 from .state_preparation import *
+from .channel import QutritDepolarizingChannel
 
 # TODO: Change `qml.Identity` for qutrit support or add `qml.TIdentity` for qutrits
 __ops__ = {
@@ -47,5 +48,8 @@ __obs__ = {
     "THermitian",
     "GellMann",
 }
+__channels__ = {
+    "QutritDepolarizingChannel",
+}
 
-__all__ = list(__ops__ | __obs__)
+__all__ = list(__ops__ | __obs__ | __channels__)
diff --git a/pennylane/ops/qutrit/channel.py b/pennylane/ops/qutrit/channel.py
new file mode 100644
index 000000000..8443fd2de
--- /dev/null
+++ b/pennylane/ops/qutrit/channel.py
@@ -0,0 +1,224 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# pylint: disable=too-many-arguments
+"""
+This module contains the available built-in noisy qutrit
+quantum channels supported by PennyLane, as well as their conventions.
+"""
+import numpy as np
+
+from pennylane import math
+from pennylane.operation import Channel
+
+QUDIT_DIM = 3
+
+
+class QutritDepolarizingChannel(Channel):
+    r"""
+    Single-qutrit symmetrically depolarizing error channel.
+    This channel is modelled by the Kraus matrices generated by the following relationship:
+
+    .. math::
+        K_0 = K_{0,0} = \sqrt{1-p} \begin{bmatrix}
+                1 & 0 & 0\\
+                0 & 1 & 0\\
+                0 & 0 & 1
+                \end{bmatrix}, \quad
+        K_{i,j} = \sqrt{\frac{p}{8}}X^iZ^j
+
+    Where:
+
+    .. math::
+        X = \begin{bmatrix}
+                0 & 1 & 0 \\
+                0 & 0 & 1 \\
+                1 & 0 & 0
+                \end{bmatrix}, \quad
+        Z = \begin{bmatrix}
+                1 & 0 & 0\\
+                0 & \omega & 0\\
+                0 & 0 & \omega^2
+                \end{bmatrix}
+
+
+    These relations create the following Kraus matrices:
+
+    .. math::
+        \begin{matrix}
+            K_0 = K_{0,0} = \sqrt{1-p} \begin{bmatrix}
+                1 & 0 & 0\\
+                0 & 1 & 0\\
+                0 & 0 & 1
+                \end{bmatrix}&
+            K_1 = K_{0,1} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                1 & 0 & 0\\
+                0 & \omega & 0\\
+                0 & 0 & \omega^2
+                \end{bmatrix}&
+            K_2 = K_{0,2} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                1 & 0 & 0\\
+                0 & \omega^2 & 0\\
+                0 & 0 & \omega
+                \end{bmatrix}\\
+            K_3 = K_{1,0} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & 1 & 0 \\
+                0 & 0 & 1 \\
+                1 & 0 & 0
+                \end{bmatrix}&
+            K_4 = K_{1,1} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & \omega & 0 \\
+                0 & 0 & \omega^2 \\
+                1 & 0 & 0
+                \end{bmatrix}&
+            K_5 = K_{1,2} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & \omega^2 & 0 \\
+                0 & 0 & \omega \\
+                1 & 0 & 0
+                \end{bmatrix}\\
+            K_6 = K_{2,0} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & 0 & 1 \\
+                1 & 0 & 0 \\
+                0 & 1 & 0
+                \end{bmatrix}&
+            K_7 = K_{2,1} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & 0 & \omega^2 \\
+                1 & 0 & 0 \\
+                0 & \omega & 0
+                \end{bmatrix}&
+            K_8 = K_{2,2} = \sqrt{\frac{p}{8}}\begin{bmatrix}
+                0 & 0 & \omega \\
+                1 & 0 & 0 \\
+                0 & \omega^2 & 0
+                \end{bmatrix}
+        \end{matrix}
+
+    Where :math:`\omega=\exp(\frac{2\pi}{3})`  is the third root of unity,
+    and :math:`p \in [0, 1]` is the depolarization probability, equally
+    divided in the application of all qutrit Pauli operators.
+
+    .. note::
+
+        The Kraus operators :math:`\{K_0 \ldots K_8\}` used are the representations of the single qutrit Pauli group.
+        These Pauli group operators are defined in [`1 <https://doi.org/10.48550/arXiv.quant-ph/9802007>`_] (Eq. 5).
+        The Kraus Matrices we use are adapted from [`2 <https://doi.org/10.48550/arXiv.1905.10481>`_] (Eq. 5).
+        For this definition, please make a note of the following:
+
+        * For :math:`p = 0`, the channel will be an Identity channel, i.e., a noise-free channel.
+        * For :math:`p = \frac{8}{9}`, the channel will be a fully depolarizing channel.
+        * For :math:`p = 1`, the channel will be a uniform error channel.
+
+    **Details:**
+
+    * Number of wires: 1
+    * Number of parameters: 1
+
+    Args:
+        p (float): Each qutrit Pauli operator is applied with probability :math:`\frac{p}{8}`
+        wires (Sequence[int] or int): the wire the channel acts on
+        id (str or None): String representing the operation (optional)
+    """
+
+    num_params = 1
+    num_wires = 1
+    grad_method = "A"
+    grad_recipe = ([[1, 0, 1], [-1, 0, 0]],)
+
+    def __init__(self, p, wires, id=None):
+        super().__init__(p, wires=wires, id=id)
+
+    @staticmethod
+    def compute_kraus_matrices(p):  # pylint:disable=arguments-differ
+        r"""Kraus matrices representing the qutrit depolarizing channel.
+
+         Args:
+             p (float): each qutrit Pauli gate is applied with probability :math:`\frac{p}{8}`
+
+         Returns:
+             list (array): list of Kraus matrices
+
+         **Example**
+
+         >>> np.round(qml.QutritDepolarizingChannel.compute_kraus_matrices(0.5), 3)
+         array([[[ 0.707+0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.707+0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   ,  0.707+0.j   ]],
+
+        [[ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   , -0.125+0.217j,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   , -0.125-0.217j]],
+
+        [[ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   , -0.125-0.217j,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   , -0.125+0.217j]],
+
+        [[ 0.   +0.j   ,  0.25 +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   ,  0.25 +0.j   ],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ]],
+
+        [[ 0.   +0.j   , -0.125+0.217j,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   , -0.125-0.217j],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ]],
+
+        [[ 0.   +0.j   , -0.125-0.217j,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.   +0.j   , -0.125+0.217j],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ]],
+
+        [[ 0.   +0.j   ,  0.   +0.j   ,  0.25 +0.j   ],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   ,  0.25 +0.j   ,  0.   +0.j   ]],
+
+        [[ 0.   +0.j   ,  0.   +0.j   , -0.125-0.217j],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   , -0.125+0.217j,  0.   +0.j   ]],
+
+        [[ 0.   +0.j   ,  0.   +0.j   , -0.125+0.217j],
+         [ 0.25 +0.j   ,  0.   +0.j   ,  0.   +0.j   ],
+         [ 0.   +0.j   , -0.125-0.217j,  0.   +0.j   ]]])
+        """
+        if not math.is_abstract(p) and not 0.0 <= p <= 1.0:
+            raise ValueError("p must be in the interval [0,1]")
+
+        interface = math.get_interface(p)
+
+        w = math.exp(2j * np.pi / 3)
+        one = 1
+        z = 0
+
+        if interface == "tensorflow":
+            p = math.cast_like(p, 1j)
+            w = math.cast_like(w, p)
+            one = math.cast_like(one, p)
+            z = math.cast_like(z, p)
+
+        w2 = w**2
+
+        # The matrices are explicitly written, not generated to ensure PyTorch differentiation.
+        depolarizing_mats = [
+            [[one, z, z], [z, w, z], [z, z, w2]],
+            [[one, z, z], [z, w2, z], [z, z, w]],
+            [[z, one, z], [z, z, one], [one, z, z]],
+            [[z, w, z], [z, z, w2], [one, z, z]],
+            [[z, w2, z], [z, z, w], [one, z, z]],
+            [[z, z, one], [one, z, z], [z, one, z]],
+            [[z, z, w2], [one, z, z], [z, w, z]],
+            [[z, z, w], [one, z, z], [z, w2, z]],
+        ]
+
+        normalization = math.sqrt(p / 8 + math.eps)
+        Ks = [normalization * math.array(m, like=interface) for m in depolarizing_mats]
+        identity = math.sqrt(1 - p + math.eps) * math.array(
+            math.eye(QUDIT_DIM, dtype=complex), like=interface
+        )
+
+        return [identity] + Ks


Test patch:
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_preprocessing.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_preprocessing.py
index 27bf1331d..9574fdbec 100644
--- a/tests/devices/qutrit_mixed/test_qutrit_mixed_preprocessing.py
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_preprocessing.py
@@ -125,9 +125,10 @@ class TestPreprocessing:
             (qml.GellMann(0, 1), False),
             (qml.Snapshot(), True),
             (qml.TRX(1.1, 0), True),
+            (qml.QutritDepolarizingChannel(0.4, 0), True),
         ],
     )
-    def test_accepted_observables(self, op, expected):
+    def test_accepted_operator(self, op, expected):
         """Test that stopping_condition works correctly"""
         res = stopping_condition(op)
         assert res == expected
@@ -136,6 +137,7 @@ class TestPreprocessing:
         "obs, expected",
         [
             (qml.TShift(0), False),
+            (qml.QutritDepolarizingChannel(0.4, 0), False),
             (qml.GellMann(0, 1), True),
             (qml.Snapshot(), False),
             (qml.operation.Tensor(qml.GellMann(0, 1), qml.GellMann(3, 3)), True),
@@ -144,7 +146,7 @@ class TestPreprocessing:
             (qml.ops.op_math.Prod(qml.GellMann(0, 1), qml.GellMann(3, 3)), True),
         ],
     )
-    def test_accepted_operator(self, obs, expected):
+    def test_accepted_observable(self, obs, expected):
         """Test that observable_stopping_condition works correctly"""
         res = observable_stopping_condition(obs)
         assert res == expected
diff --git a/tests/ops/qutrit/test_qutrit_channel_ops.py b/tests/ops/qutrit/test_qutrit_channel_ops.py
new file mode 100644
index 000000000..be1025d0c
--- /dev/null
+++ b/tests/ops/qutrit/test_qutrit_channel_ops.py
@@ -0,0 +1,174 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the available built-in qutrit quantum channels.
+"""
+import numpy as np
+import pytest
+from numpy.linalg import matrix_power
+
+import pennylane as qml
+from pennylane import numpy as pnp
+from pennylane.ops.qutrit import channel
+
+QUDIT_DIM = 3
+
+
+class TestQutritDepolarizingChannel:
+    """Tests for the qutrit quantum channel QutritDepolarizingChannel"""
+
+    @staticmethod
+    def get_expected_kraus_matrices(p):
+        """Gets the expected Kraus matrices given probability p."""
+        w = np.exp(2j * np.pi / 3)
+
+        X = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]])
+        Z = np.diag([1, w, w**2])
+
+        Ks = [np.sqrt(1 - p) * np.eye(QUDIT_DIM)]
+        for i in range(3):
+            for j in range(3):
+                if i == 0 and j == 0:
+                    continue
+                Ks.append(np.sqrt(p / 8) * matrix_power(X, i) @ matrix_power(Z, j))
+        return Ks
+
+    def test_p_zero(self, tol):
+        """Test p=0 gives correct Kraus matrices"""
+        op = qml.QutritDepolarizingChannel
+        kraus_matrices = op(0, wires=0).kraus_matrices()
+
+        assert len(kraus_matrices) == 9
+        assert np.allclose(kraus_matrices[0], np.eye(QUDIT_DIM), atol=tol, rtol=0)
+        assert np.allclose(np.array(kraus_matrices[1:]), 0, atol=tol, rtol=0)
+
+    def test_p_arbitrary(self, tol):
+        """Test p=0.1 gives correct Kraus matrices"""
+        p = 0.1
+        kraus_matrices = qml.QutritDepolarizingChannel(p, wires=0).kraus_matrices()
+        expected_matrices = self.get_expected_kraus_matrices(p)
+        for kraus_matrix, expected_matrix in zip(kraus_matrices, expected_matrices):
+            assert np.allclose(kraus_matrix, expected_matrix, atol=tol, rtol=0)
+
+    def test_p_invalid_parameter(self):
+        """Test that error is raised given an inappropriate p value."""
+        with pytest.raises(ValueError, match="p must be in the interval"):
+            qml.QutritDepolarizingChannel(1.5, wires=0).kraus_matrices()
+
+    @pytest.mark.parametrize("angle", np.linspace(0, 2 * np.pi, 7))
+    def test_grad_depolarizing(self, angle):
+        """Test that analytical gradient is computed correctly for different states. Channel
+        grad recipes are independent of channel parameter"""
+
+        dev = qml.device("default.qutrit.mixed")
+        prob = pnp.array(0.5, requires_grad=True)
+
+        @qml.qnode(dev)
+        def circuit(p):
+            qml.TRX(angle, wires=0, subspace=(0, 1))
+            qml.TRX(angle, wires=0, subspace=(1, 2))
+            qml.QutritDepolarizingChannel(p, wires=0)
+            return qml.expval(qml.GellMann(0, 3) + qml.GellMann(0, 8))
+
+        expected_errorless = (
+            (np.sqrt(3) - 3) * (1 - np.cos(2 * angle)) / 24
+            - 2 / np.sqrt(3) * np.sin(angle / 2) ** 4
+            + (np.sqrt(1 / 3) + 1) * np.cos(angle / 2) ** 2
+        )
+
+        assert np.allclose(circuit(prob), ((prob - (1 / 9)) / (8 / 9)) * expected_errorless)
+
+        gradient = np.squeeze(qml.grad(circuit)(prob))
+        assert np.allclose(gradient, circuit(1) - circuit(0))
+        assert np.allclose(gradient, -(9 / 8) * expected_errorless)
+
+    @staticmethod
+    def expected_jac_fn(p):
+        """Gets the expected Jacobian of Kraus matrices given probability p."""
+        w = np.exp(2j * np.pi / 3)
+
+        X = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]])
+        Z = np.diag([1, w, w**2])
+
+        jacs = [-1 / (2 * np.sqrt(1 - p)) * np.eye(QUDIT_DIM)]
+        for i in range(3):
+            for j in range(3):
+                if i == 0 and j == 0:
+                    continue
+                jacs.append(np.sqrt(1 / (32 * p)) * matrix_power(X, i) @ matrix_power(Z, j))
+
+        return jacs
+
+    @staticmethod
+    def kraus_fn(p):
+        """Gets a matrix of the Kraus matrices to be tested."""
+        return qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
+
+    @staticmethod
+    def kraus_fn_real(p):
+        """Gets a matrix of the real part of the Kraus matrices to be tested."""
+        return qml.math.real(
+            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
+        )
+
+    @staticmethod
+    def kraus_fn_imag(p):
+        """Gets a matrix of the imaginary part of the Kraus matrices to be tested."""
+        return qml.math.imag(
+            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
+        )
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        """Tests Jacobian of Kraus matrices using autograd."""
+        p = pnp.array(0.43, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn_real)(p) + 1j * qml.jacobian(self.kraus_fn_imag)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        """Tests Jacobian of Kraus matrices using torch."""
+        import torch
+
+        p = torch.tensor(0.43, requires_grad=True)
+        jacobian = torch.autograd.functional.jacobian
+        jac = jacobian(self.kraus_fn_real, p) + 1j * jacobian(self.kraus_fn_imag, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p.detach().numpy()))
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        """Tests Jacobian of Kraus matrices using tensorflow."""
+        import tensorflow as tf
+
+        p = tf.Variable(0.43)
+        with tf.GradientTape() as real_tape:
+            real_out = self.kraus_fn_real(p)
+        with tf.GradientTape() as imag_tape:
+            imag_out = self.kraus_fn_imag(p)
+
+        real_jac = qml.math.cast(real_tape.jacobian(real_out, p), complex)
+        imag_jac = qml.math.cast(imag_tape.jacobian(imag_out, p), complex)
+        jac = real_jac + 1j * imag_jac
+        assert qml.math.allclose(jac, self.expected_jac_fn(0.43))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        """Tests Jacobian of Kraus matrices using jax."""
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+
+        p = jax.numpy.array(0.43, dtype=jax.numpy.complex128)
+        jac = jax.jacobian(self.kraus_fn, holomorphic=True)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))