Repository: PennyLaneAI/pennylane

Issue:
[BUG] `DepolarizingChannel` is not compatible with all interfaces
### Expected behavior

The channel `DepolarizingChannel` works correctly with all interfaces.

### Actual behavior

1. With `torch`, the channel does not work when computing a functional jacobian, because in the creation of the Kraus matrices, `torch.tensor` objects are multiplied with simple `numpy` arrays, raising the error 
```python
RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
```
2. With `Tensorflow`, the channel leads to wrong results: A prefactor of `0.5` is introduced somehow.

**edit**
- `ResetError` and `PauliError` seem to have the same problem with `torch`.
- The bug with Tensorflow stems from multiplying a tensorflow `float` with a complex-valued `numpy` array, resulting in the imaginary part of `K2` being ignored. For the presented test case, the contributions of `K1` and `K2` are equal, hence the observed prefactor of `0.5`.

### Additional information

_No response_

### Source code

```shell
def jacobian(fn, interface=None):
    if interface=="torch":
        return partial(torch.autograd.functional.jacobian, fn)
    elif interface=="jax":
        return jax.jacobian(fn, argnums=list(range(len(inspect.signature(fn).parameters))))
    elif interface in ["tensorflow", "tf"]:
        def wrap(*args, **kwargs):
            with tf.GradientTape() as t:
                out = fn(*args, **kwargs)
            return t.jacobian(out, *args)
        return wrap
    elif interface in ["autograd", None]:
        return qml.jacobian(fn)
    
def array(x, interface):
    if interface=="torch":
        return torch.tensor(x, requires_grad=True)
    if interface=="jax":
        return jnp.array(x)
    if interface in ["tensorflow", "tf"]:
        return tf.Variable(x)
    if interface in ["autograd", None]:
        return pnp.array(x, requires_grad=True)

outs = []
grads = []
for interface in [None, "tensorflow", "autograd",  "torch", "jax" ]:
    print(interface)
    dev = qml.device("default.mixed", wires=2)
    @qml.qnode(dev, interface=interface)
    def circuit(x):
        qml.PauliX(0)
        qml.DepolarizingChannel(x, wires=[0])
        return qml.probs(wires=[0])
    
    x = array(0.2, interface)
    outs.append(circuit(x))
    grads.append(jacobian(circuit, interface)(x))
    
print(all(qml.math.allclose(o, outs[0]) for o in outs[1:]))
print(outs)
print(all(qml.math.allclose(o, grads[0]) for o in grads[1:]))
print(grads)
```


### Tracebacks

```shell
**For tensorflow:**

False
[tensor([0.13333333, 0.86666667], requires_grad=True), <tf.Tensor: shape=(2,), dtype=float64, numpy=array([0.06666667, 0.86666665])>, tensor([0.13333333, 0.86666667], requires_grad=True), Array([0.13333333, 0.8666666 ], dtype=float32)]
False
[array([ 0.66666667, -0.66666667]), <tf.Tensor: shape=(2,), dtype=float32, numpy=array([ 0.33333334, -0.6666666 ], dtype=float32)>, array([ 0.66666667, -0.66666667]), (Array([ 0.6666667, -0.6666666], dtype=float32, weak_type=True),)]

**For torch:**

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[20], line 13
     10         return qml.probs(wires=[0])
     12     x = array(0.2, interface)
---> 13     outs.append(circuit(x))
     14     grads.append(jacobian(circuit, interface)(x))
     16 print(all(qml.math.allclose(o, outs[0]) for o in outs[1:]))

File ~/repos/pennylane/pennylane/qnode.py:822, in QNode.__call__(self, *args, **kwargs)
    819 self._tape_cached = using_custom_cache and self.tape.hash in cache
    821 if qml.active_return():
--> 822     res = qml.execute(
    823         [self.tape],
    824         device=self.device,
    825         gradient_fn=self.gradient_fn,
    826         interface=self.interface,
    827         gradient_kwargs=self.gradient_kwargs,
    828         override_shots=override_shots,
    829         **self.execute_kwargs,
    830     )
    832     res = res[0]
    834     if old_interface == "auto":

File ~/repos/pennylane/pennylane/interfaces/execution.py:593, in execute(tapes, device, gradient_fn, interface, mode, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform)
    492 """Execute a batch of tapes on a device in an autodifferentiable-compatible manner.
    493 
    494 Args:
   (...)
    590        [ 0.        ,  0.        , -0.95533649]])
    591 """
    592 if qml.active_return():
--> 593     return _execute_new(
    594         tapes,
    595         device,
    596         gradient_fn,
    597         interface=interface,
    598         mode=mode,
    599         gradient_kwargs=gradient_kwargs,
    600         cache=cache,
    601         cachesize=cachesize,
    602         max_diff=max_diff,
    603         override_shots=override_shots,
    604         expand_fn=expand_fn,
    605         max_expansion=max_expansion,
    606         device_batch_transform=device_batch_transform,
    607     )
    609 if interface == "auto":
    610     params = []

File ~/repos/pennylane/pennylane/interfaces/execution.py:388, in _execute_new(tapes, device, gradient_fn, interface, mode, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform)
    384     return batch_fn(res)
    386 if gradient_fn == "backprop" or interface is None:
    387     return batch_fn(
--> 388         qml.interfaces.cache_execute(
    389             batch_execute, cache, return_tuple=False, expand_fn=expand_fn
    390         )(tapes)
    391     )
    393 # the default execution function is batch_execute
    394 execute_fn = qml.interfaces.cache_execute(batch_execute, cache, expand_fn=expand_fn)

File ~/repos/pennylane/pennylane/interfaces/execution.py:206, in cache_execute.<locals>.wrapper(tapes, **kwargs)
    202         return (res, []) if return_tuple else res
    204 else:
    205     # execute all unique tapes that do not exist in the cache
--> 206     res = fn(execution_tapes.values(), **kwargs)
    208 final_res = []
    210 for i, tape in enumerate(tapes):

File ~/repos/pennylane/pennylane/interfaces/execution.py:131, in cache_execute.<locals>.fn(tapes, **kwargs)
    129 def fn(tapes: Sequence[QuantumTape], **kwargs):  # pylint: disable=function-redefined
    130     tapes = [expand_fn(tape) for tape in tapes]
--> 131     return original_fn(tapes, **kwargs)

File /usr/lib/python3.10/contextlib.py:79, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     76 @wraps(func)
     77 def inner(*args, **kwds):
     78     with self._recreate_cm():
---> 79         return func(*args, **kwds)

File ~/repos/pennylane/pennylane/_qubit_device.py:647, in QubitDevice.batch_execute(self, circuits)
    644 # TODO: This method and the tests can be globally implemented by Device
    645 # once it has the same signature in the execute() method
    646 if qml.active_return():
--> 647     return self._batch_execute_new(circuits=circuits)
    649 results = []
    650 for circuit in circuits:
    651     # we need to reset the device here, else it will
    652     # not start the next computation in the zero state

File ~/repos/pennylane/pennylane/_qubit_device.py:620, in QubitDevice._batch_execute_new(self, circuits)
    615 for circuit in circuits:
    616     # we need to reset the device here, else it will
    617     # not start the next computation in the zero state
    618     self.reset()
--> 620     res = self._execute_new(circuit)
    621     results.append(res)
    623 if self.tracker.active:

File ~/repos/pennylane/pennylane/_qubit_device.py:375, in QubitDevice._execute_new(self, circuit, **kwargs)
    372 self.check_validity(circuit.operations, circuit.observables)
    374 # apply all circuit operations
--> 375 self.apply(circuit.operations, rotations=circuit.diagonalizing_gates, **kwargs)
    377 # generate computational basis samples
    378 if self.shots is not None:

File ~/repos/pennylane/pennylane/devices/default_mixed.py:620, in DefaultMixed.apply(self, operations, rotations, **kwargs)
    614         raise DeviceError(
    615             f"Operation {operation.name} cannot be used after other Operations have already been applied "
    616             f"on a {self.short_name} device."
    617         )
    619 for operation in operations:
--> 620     self._apply_operation(operation)
    622 # store the pre-rotated state
    623 self._pre_rotated_state = self._state

File ~/repos/pennylane/pennylane/devices/default_mixed.py:542, in DefaultMixed._apply_operation(self, operation)
    539             self._debugger.snapshots[len(self._debugger.snapshots)] = density_matrix
    540     return
--> 542 matrices = self._get_kraus(operation)
    544 if operation in diagonal_in_z_basis:
    545     self._apply_diagonal_unitary(matrices, wires)

File ~/repos/pennylane/pennylane/devices/default_mixed.py:278, in DefaultMixed._get_kraus(self, operation)
    275     return operation.eigvals()
    277 if isinstance(operation, Channel):
--> 278     return operation.kraus_matrices()
    280 return [operation.matrix()]

File ~/repos/pennylane/pennylane/operation.py:1633, in Channel.kraus_matrices(self)
   1619 def kraus_matrices(self):
   1620     r"""Kraus matrices of an instantiated channel
   1621     in the computational basis.
   1622 
   (...)
   1631      array([[0., 0.31622777], [0., 0.]])]
   1632     """
-> 1633     return self.compute_kraus_matrices(*self.parameters, **self.hyperparameters)

File ~/repos/pennylane/pennylane/ops/channel.py:323, in DepolarizingChannel.compute_kraus_matrices(p)
    320 if not np.is_abstract(p) and not 0.0 <= p <= 1.0:
    321     raise ValueError("p must be in the interval [0,1]")
--> 323 K0 = np.sqrt(1 - p + np.eps) * np.eye(2)
    324 K1 = np.sqrt(p / 3 + np.eps) * np.array([[0, 1], [1, 0]])
    325 K2 = np.sqrt(p / 3 + np.eps) * np.array([[0, -1j], [1j, 0]])

File ~/venvs/dev/lib/python3.10/site-packages/torch/_tensor.py:956, in Tensor.__array__(self, dtype)
    954     return handle_torch_function(Tensor.__array__, (self,), self, dtype=dtype)
    955 if dtype is None:
--> 956     return self.numpy()
    957 else:
    958     return self.numpy().astype(dtype, copy=False)

RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
```


### System information

```shell
Name: PennyLane
Version: 0.29.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/david/venvs/dev/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           Linux-5.15.0-57-generic-x86_64-with-glibc2.35
Python version:          3.10.6
Numpy version:           1.23.5
Scipy version:           1.9.3
Installed devices:
- default.gaussian (PennyLane-0.29.0.dev0)
- default.mixed (PennyLane-0.29.0.dev0)
- default.qubit (PennyLane-0.29.0.dev0)
- default.qubit.autograd (PennyLane-0.29.0.dev0)
- default.qubit.jax (PennyLane-0.29.0.dev0)
- default.qubit.tf (PennyLane-0.29.0.dev0)
- default.qubit.torch (PennyLane-0.29.0.dev0)
- default.qutrit (PennyLane-0.29.0.dev0)
- null.qubit (PennyLane-0.29.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.28.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/ops/channel.py b/pennylane/ops/channel.py
--- a/pennylane/ops/channel.py
+++ b/pennylane/ops/channel.py
@@ -83,7 +83,9 @@ def compute_kraus_matrices(gamma):  # pylint:disable=arguments-differ
             raise ValueError("gamma must be in the interval [0,1].")
 
         K0 = np.diag([1, np.sqrt(1 - gamma + np.eps)])
-        K1 = np.sqrt(gamma + np.eps) * np.array([[0, 1], [0, 0]])
+        K1 = np.sqrt(gamma + np.eps) * np.convert_like(
+            np.cast_like(np.array([[0, 1], [0, 0]]), gamma), gamma
+        )
         return [K0, K1]
 
 
@@ -167,9 +169,17 @@ def compute_kraus_matrices(gamma, p):  # pylint:disable=arguments-differ
             raise ValueError("p must be in the interval [0,1].")
 
         K0 = np.sqrt(p + np.eps) * np.diag([1, np.sqrt(1 - gamma + np.eps)])
-        K1 = np.sqrt(p + np.eps) * np.sqrt(gamma) * np.array([[0, 1], [0, 0]])
+        K1 = (
+            np.sqrt(p + np.eps)
+            * np.sqrt(gamma)
+            * np.convert_like(np.cast_like(np.array([[0, 1], [0, 0]]), gamma), gamma)
+        )
         K2 = np.sqrt(1 - p + np.eps) * np.diag([np.sqrt(1 - gamma + np.eps), 1])
-        K3 = np.sqrt(1 - p + np.eps) * np.sqrt(gamma) * np.array([[0, 0], [1, 0]])
+        K3 = (
+            np.sqrt(1 - p + np.eps)
+            * np.sqrt(gamma)
+            * np.convert_like(np.cast_like(np.array([[0, 0], [1, 0]]), gamma), gamma)
+        )
         return [K0, K1, K2, K3]
 
 
@@ -320,10 +330,17 @@ def compute_kraus_matrices(p):  # pylint:disable=arguments-differ
         if not np.is_abstract(p) and not 0.0 <= p <= 1.0:
             raise ValueError("p must be in the interval [0,1]")
 
-        K0 = np.sqrt(1 - p + np.eps) * np.eye(2)
-        K1 = np.sqrt(p / 3 + np.eps) * np.array([[0, 1], [1, 0]])
-        K2 = np.sqrt(p / 3 + np.eps) * np.array([[0, -1j], [1j, 0]])
-        K3 = np.sqrt(p / 3 + np.eps) * np.array([[1, 0], [0, -1]])
+        if np.get_interface(p) == "tensorflow":
+            p = np.cast_like(p, 1j)
+
+        K0 = np.sqrt(1 - p + np.eps) * np.convert_like(np.eye(2, dtype=complex), p)
+        K1 = np.sqrt(p / 3 + np.eps) * np.convert_like(np.array([[0, 1], [1, 0]], dtype=complex), p)
+        K2 = np.sqrt(p / 3 + np.eps) * np.convert_like(
+            np.array([[0, -1j], [1j, 0]], dtype=complex), p
+        )
+        K3 = np.sqrt(p / 3 + np.eps) * np.convert_like(
+            np.array([[1, 0], [0, -1]], dtype=complex), p
+        )
         return [K0, K1, K2, K3]
 
 
@@ -386,8 +403,8 @@ def compute_kraus_matrices(p):  # pylint:disable=arguments-differ
         if not np.is_abstract(p) and not 0.0 <= p <= 1.0:
             raise ValueError("p must be in the interval [0,1]")
 
-        K0 = np.sqrt(1 - p + np.eps) * np.eye(2)
-        K1 = np.sqrt(p + np.eps) * np.array([[0, 1], [1, 0]])
+        K0 = np.sqrt(1 - p + np.eps) * np.convert_like(np.cast_like(np.eye(2), p), p)
+        K1 = np.sqrt(p + np.eps) * np.convert_like(np.cast_like(np.array([[0, 1], [1, 0]]), p), p)
         return [K0, K1]
 
 
@@ -479,11 +496,21 @@ def compute_kraus_matrices(p_0, p_1):  # pylint:disable=arguments-differ
         if not np.is_abstract(p_0 + p_1) and not 0.0 <= p_0 + p_1 <= 1.0:
             raise ValueError("p_0 + p_1 must be in the interval [0,1]")
 
-        K0 = np.sqrt(1 - p_0 - p_1 + np.eps) * np.eye(2)
-        K1 = np.sqrt(p_0 + np.eps) * np.array([[1, 0], [0, 0]])
-        K2 = np.sqrt(p_0 + np.eps) * np.array([[0, 1], [0, 0]])
-        K3 = np.sqrt(p_1 + np.eps) * np.array([[0, 0], [1, 0]])
-        K4 = np.sqrt(p_1 + np.eps) * np.array([[0, 0], [0, 1]])
+        interface = np.get_interface([p_0, p_1])
+        p_0, p_1 = np.coerce([p_0, p_1], like=interface)
+        K0 = np.sqrt(1 - p_0 - p_1 + np.eps) * np.convert_like(np.cast_like(np.eye(2), p_0), p_0)
+        K1 = np.sqrt(p_0 + np.eps) * np.convert_like(
+            np.cast_like(np.array([[1, 0], [0, 0]]), p_0), p_0
+        )
+        K2 = np.sqrt(p_0 + np.eps) * np.convert_like(
+            np.cast_like(np.array([[0, 1], [0, 0]]), p_0), p_0
+        )
+        K3 = np.sqrt(p_1 + np.eps) * np.convert_like(
+            np.cast_like(np.array([[0, 0], [1, 0]]), p_0), p_0
+        )
+        K4 = np.sqrt(p_1 + np.eps) * np.convert_like(
+            np.cast_like(np.array([[0, 0], [0, 1]]), p_0), p_0
+        )
 
         return [K0, K1, K2, K3, K4]
 
@@ -584,16 +611,23 @@ def compute_kraus_matrices(operators, p):  # pylint:disable=arguments-differ
         nq = len(operators)
 
         # K0 is sqrt(1-p) * Identity
-        K0 = np.sqrt(1 - p + np.eps) * np.eye(2**nq)
+        K0 = np.sqrt(1 - p + np.eps) * np.convert_like(np.cast_like(np.eye(2**nq), p), p)
+
+        interface = np.get_interface(p)
+        if interface == "tensorflow" or "Y" in operators:
+            if interface == "numpy":
+                p = (1 + 0j) * p
+            else:
+                p = np.cast_like(p, 1j)
 
         ops = {
-            "X": np.array([[0, 1], [1, 0]]),
-            "Y": np.array([[0, -1j], [1j, 0]]),
-            "Z": np.array([[1, 0], [0, -1]]),
+            "X": np.convert_like(np.cast_like(np.array([[0, 1], [1, 0]]), p), p),
+            "Y": np.convert_like(np.cast_like(np.array([[0, -1j], [1j, 0]]), p), p),
+            "Z": np.convert_like(np.cast_like(np.array([[1, 0], [0, -1]]), p), p),
         }
 
         # K1 is composed by Kraus matrices of operators
-        K1 = np.sqrt(p + np.eps) * np.array([[1]])
+        K1 = np.sqrt(p + np.eps) * np.convert_like(np.cast_like(np.eye(1), p), p)
         for op in operators[::-1]:
             K1 = np.multi_dispatch()(np.kron)(ops[op], K1)
 
@@ -659,8 +693,8 @@ def compute_kraus_matrices(p):  # pylint:disable=arguments-differ
         if not np.is_abstract(p) and not 0.0 <= p <= 1.0:
             raise ValueError("p must be in the interval [0,1]")
 
-        K0 = np.sqrt(1 - p + np.eps) * np.eye(2)
-        K1 = np.sqrt(p + np.eps) * np.array([[1, 0], [0, -1]])
+        K0 = np.sqrt(1 - p + np.eps) * np.convert_like(np.cast_like(np.eye(2), p), p)
+        K1 = np.sqrt(p + np.eps) * np.convert_like(np.cast_like(np.diag([1, -1]), p), p)
         return [K0, K1]
 
 


Test patch:
diff --git a/tests/ops/test_channel_ops.py b/tests/ops/test_channel_ops.py
--- a/tests/ops/test_channel_ops.py
+++ b/tests/ops/test_channel_ops.py
@@ -14,6 +14,7 @@
 """
 Unit tests for the available built-in quantum channels.
 """
+from itertools import product
 import pytest
 import numpy as np
 import pennylane as qml
@@ -25,42 +26,79 @@
 Y = np.array([[0, -1j], [1j, 0]])
 Z = np.array([[1, 0], [0, -1]])
 
-ch_list = [
+one_arg_args = [(0.0,), (0.1,), (1.0,)]
+
+channels_with_one_arg = [
     channel.AmplitudeDamping,
-    channel.GeneralizedAmplitudeDamping,
     channel.PhaseDamping,
     channel.BitFlip,
     channel.PhaseFlip,
     channel.DepolarizingChannel,
-    channel.ResetError,
-    channel.PauliError,
-    channel.ThermalRelaxationError,
+    channel.PhaseDamping,
 ]
 
+two_arg_args = [
+    (0.0, 0.0),
+    (0.0, 0.1),
+    (0.1, 0.0),
+    (1.0, 1.0),
+    (0.0, 1.0),
+    (1.0, 0.0),
+    (0.3, 0.2),
+]
+
+channels_with_two_args = [channel.GeneralizedAmplitudeDamping, channel.ResetError]
+
+channels_and_args = (
+    list(product(channels_with_one_arg, one_arg_args))
+    + list(product(channels_with_two_args, two_arg_args))
+    + [
+        (channel.PauliError, ("X", 0.0)),
+        (channel.PauliError, ("X", 0.41)),
+        (channel.PauliError, ("X", 1.0)),
+        (channel.PauliError, ("YY", 0.0)),
+        (channel.PauliError, ("YX", 0.41)),
+        (channel.PauliError, ("ZZ", 1.0)),
+        (channel.ThermalRelaxationError, (0.0, 1e-4, 1e-4, 2e-8)),
+        (channel.ThermalRelaxationError, (0.1, 1e-4, 1e-4, 2e-8)),
+        (channel.ThermalRelaxationError, (1.0, 1e-4, 1e-4, 2e-8)),
+        (channel.ThermalRelaxationError, (0.0, 1e-4, 1.2e-4, 2e-8)),
+        (channel.ThermalRelaxationError, (0.1, 1e-4, 1.2e-4, 2e-8)),
+        (channel.ThermalRelaxationError, (1.0, 1e-4, 1.2e-4, 2e-8)),
+    ]
+)
+
 
 class TestChannels:
     """Tests for the quantum channels"""
 
-    @pytest.mark.parametrize("ops", ch_list)
-    @pytest.mark.parametrize("p", [0, 0.1, 1])
-    @pytest.mark.parametrize("tr_args", [[100e-6, 100e-6, 20e-9], [100e-6, 120e-6, 20e-9]])
-    def test_kraus_matrices_sum_identity(self, ops, p, tr_args, tol):
+    @pytest.mark.parametrize(
+        "interface",
+        [
+            None,
+            pytest.param("autograd", marks=pytest.mark.autograd),
+            pytest.param("tensorflow", marks=pytest.mark.tf),
+            pytest.param("jax", marks=pytest.mark.jax),
+            pytest.param("torch", marks=pytest.mark.torch),
+        ],
+    )
+    @pytest.mark.parametrize("ch, args", channels_and_args)
+    def test_kraus_matrices_sum_identity(self, ch, args, interface, tol):
         """Test channels are trace-preserving"""
-        if ops.__name__ == "GeneralizedAmplitudeDamping":
-            op = [ops(p, p, wires=0)]
-        elif ops.__name__ == "ResetError":
-            op = [ops(p / 2, p / 3, wires=0)]
-        elif ops.__name__ == "PauliError":
-            op = [ops("X", p, wires=0), ops("XX", p, wires=[0, 1])]
-        elif ops.__name__ == "ThermalRelaxationError":
-            op = [ops(p, *tr_args, wires=0)]
+        if ch is channel.ResetError:
+            args = (args[0] / 2, args[1] / 3)
+        args = tuple(
+            arg if isinstance(arg, str) else qml.math.array(arg, like=interface) for arg in args
+        )
+        if ch is channel.PauliError and len(args[0]) > 1:
+            wires = [0, 1]
         else:
-            op = [ops(p, wires=0)]
-        for operation in op:
-            K_list = operation.kraus_matrices()
-            K_arr = np.array(K_list)
-            Kraus_sum = np.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
-            assert np.allclose(Kraus_sum, np.eye(K_list[0].shape[0]), atol=tol, rtol=0)
+            wires = [0]
+        op = ch(*args, wires=wires)
+        K_list = op.kraus_matrices()
+        K_arr = qml.math.stack(K_list)
+        Kraus_sum = qml.math.einsum("ajk,ajl->kl", qml.math.conj(K_arr), K_arr)
+        assert qml.math.allclose(Kraus_sum, np.eye(K_list[0].shape[0]), atol=tol, rtol=0)
 
 
 class TestAmplitudeDamping:
@@ -85,6 +123,44 @@ def test_gamma_invalid_parameter(self):
         with pytest.raises(ValueError, match="gamma must be in the interval"):
             channel.AmplitudeDamping(1.5, wires=0).kraus_matrices()
 
+    expected_jac_fn = lambda self, gamma: [
+        qml.math.array([[0, 0], [0, -1 / (2 * qml.math.sqrt(1 - gamma))]]),
+        qml.math.array([[0, 1 / (2 * qml.math.sqrt(gamma))], [0, 0]]),
+    ]
+    kraus_fn = lambda self, x: qml.math.stack(channel.AmplitudeDamping(x, wires=0).kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        gamma = pnp.array(0.43, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        gamma = torch.tensor(0.43, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, gamma)
+        assert qml.math.allclose(jac.detach().numpy(), self.expected_jac_fn(gamma.detach().numpy()))
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        gamma = tf.Variable(0.43)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(gamma)
+        jac = tape.jacobian(out, gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        gamma = jax.numpy.array(0.43)
+        jac = jax.jacobian(self.kraus_fn)(gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
 
 class TestGeneralizedAmplitudeDamping:
     """Tests for the quantum channel GeneralizedAmplitudeDamping"""
@@ -117,6 +193,67 @@ def test_p_invalid_parameter(self):
         with pytest.raises(ValueError, match="p must be in the interval"):
             channel.GeneralizedAmplitudeDamping(0.0, 1.5, wires=0).kraus_matrices()
 
+    expected_jac_fn = lambda self, gamma, p: (
+        [
+            qml.math.sqrt(p) * qml.math.array([[0, 0], [0, -1 / (2 * qml.math.sqrt(1 - gamma))]]),
+            qml.math.sqrt(p) * qml.math.array([[0, 1 / (2 * qml.math.sqrt(gamma))], [0, 0]]),
+            qml.math.sqrt(1 - p)
+            * qml.math.array([[-1 / (2 * qml.math.sqrt(1 - gamma)), 0], [0, 0]]),
+            qml.math.sqrt(1 - p) * qml.math.array([[0, 0], [1 / (2 * qml.math.sqrt(gamma)), 0]]),
+        ],
+        [
+            1 / (2 * qml.math.sqrt(p)) * qml.math.array([[1, 0], [0, qml.math.sqrt(1 - gamma)]]),
+            1 / (2 * qml.math.sqrt(p)) * qml.math.array([[0, qml.math.sqrt(gamma)], [0, 0]]),
+            -1
+            / (2 * qml.math.sqrt(1 - p))
+            * qml.math.array([[qml.math.sqrt(1 - gamma), 0], [0, 1]]),
+            -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.array([[0, 0], [qml.math.sqrt(gamma), 0]]),
+        ],
+    )
+
+    kraus_fn = lambda self, gamma, p: qml.math.stack(
+        channel.GeneralizedAmplitudeDamping(gamma, p, wires=0).kraus_matrices()
+    )
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        gamma = pnp.array(0.43, requires_grad=True)
+        p = pnp.array(0.3, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(gamma, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma, p))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        gamma = torch.tensor(0.43, requires_grad=True)
+        p = torch.tensor(0.3, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, (gamma, p))
+        exp_jac = self.expected_jac_fn(gamma.detach().numpy(), p.detach().numpy())
+        assert len(jac) == len(exp_jac) == 2
+        for j, exp_j in zip(jac, exp_jac):
+            assert qml.math.allclose(j.detach().numpy(), exp_j)
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        gamma = tf.Variable(0.43)
+        p = tf.Variable(0.3)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(gamma, p)
+        jac = tape.jacobian(out, (gamma, p))
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma, p))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        gamma = jax.numpy.array(0.43)
+        p = jax.numpy.array(0.3)
+        jac = jax.jacobian(self.kraus_fn, argnums=[0, 1])(gamma, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma, p))
+
 
 class TestPhaseDamping:
     """Tests for the quantum channel PhaseDamping"""
@@ -140,6 +277,45 @@ def test_gamma_invalid_parameter(self):
         with pytest.raises(ValueError, match="gamma must be in the interval"):
             channel.PhaseDamping(1.5, wires=0).kraus_matrices()
 
+    expected_jac_fn = lambda self, gamma: [
+        qml.math.array([[0, 0], [0, -1 / (2 * qml.math.sqrt(1 - gamma))]]),
+        qml.math.array([[0, 0], [0, 1 / (2 * qml.math.sqrt(gamma))]]),
+    ]
+
+    kraus_fn = lambda self, x: qml.math.stack(channel.PhaseDamping(x, wires=0).kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        gamma = pnp.array(0.43, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        gamma = torch.tensor(0.43, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, gamma)
+        assert qml.math.allclose(jac.detach().numpy(), self.expected_jac_fn(gamma.detach().numpy()))
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        gamma = tf.Variable(0.43)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(gamma)
+        jac = tape.jacobian(out, gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        gamma = jax.numpy.array(0.43)
+        jac = jax.jacobian(self.kraus_fn)(gamma)
+        assert qml.math.allclose(jac, self.expected_jac_fn(gamma))
+
 
 class TestBitFlip:
     """Tests for the quantum channel BitFlipChannel"""
@@ -177,6 +353,45 @@ def test_p_invalid_parameter(self):
         with pytest.raises(ValueError, match="p must be in the interval"):
             channel.BitFlip(1.5, wires=0).kraus_matrices()
 
+    expected_jac_fn = lambda self, p: [
+        -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(2),
+        1 / (2 * qml.math.sqrt(p)) * qml.math.array([[0, 1], [1, 0]]),
+    ]
+
+    kraus_fn = lambda self, x: qml.math.stack(channel.BitFlip(x, wires=0).kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        p = pnp.array(0.43, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        p = torch.tensor(0.43, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, p)
+        assert qml.math.allclose(jac.detach().numpy(), self.expected_jac_fn(p.detach().numpy()))
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        p = tf.Variable(0.43)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(p)
+        jac = tape.jacobian(out, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        p = jax.numpy.array(0.43)
+        jac = jax.jacobian(self.kraus_fn)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
 
 class TestPhaseFlip:
     """Test that various values of p give correct Kraus matrices"""
@@ -214,20 +429,58 @@ def test_p_invalid_parameter(self):
         with pytest.raises(ValueError, match="p must be in the interval"):
             channel.PhaseFlip(1.5, wires=0).kraus_matrices()
 
+    expected_jac_fn = lambda self, p: [
+        -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(2),
+        1 / (2 * qml.math.sqrt(p)) * qml.math.diag([1, -1]),
+    ]
+    kraus_fn = lambda self, x: qml.math.stack(channel.PhaseFlip(x, wires=0).kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        p = pnp.array(0.43, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        p = torch.tensor(0.43, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, p)
+        assert qml.math.allclose(jac.detach().numpy(), self.expected_jac_fn(p.detach().numpy()))
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        p = tf.Variable(0.43)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(p)
+        jac = tape.jacobian(out, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        p = jax.numpy.array(0.43)
+        jac = jax.jacobian(self.kraus_fn)(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
 
 class TestDepolarizingChannel:
     """Tests for the quantum channel DepolarizingChannel"""
 
     def test_p_zero(self, tol):
         """Test p=0 gives correct Kraus matrices"""
-        op = channel.DepolarizingChannel
+        op = qml.DepolarizingChannel
         assert np.allclose(op(0, wires=0).kraus_matrices()[0], np.eye(2), atol=tol, rtol=0)
         assert np.allclose(op(0, wires=0).kraus_matrices()[1], np.zeros((2, 2)), atol=tol, rtol=0)
 
     def test_p_arbitrary(self, tol):
         """Test p=0.1 gives correct Kraus matrices"""
         p = 0.1
-        op = channel.DepolarizingChannel
+        op = qml.DepolarizingChannel
         expected = np.sqrt(p / 3) * X
         assert np.allclose(op(0.1, wires=0).kraus_matrices()[1], expected, atol=tol, rtol=0)
 
@@ -251,7 +504,60 @@ def circuit(p):
 
     def test_p_invalid_parameter(self):
         with pytest.raises(ValueError, match="p must be in the interval"):
-            channel.DepolarizingChannel(1.5, wires=0).kraus_matrices()
+            qml.DepolarizingChannel(1.5, wires=0).kraus_matrices()
+
+    expected_jac_fn = lambda self, p: [
+        -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(2),
+        1 / (6 * qml.math.sqrt(p / 3)) * qml.math.array([[0, 1], [1, 0]]),
+        1 / (6 * qml.math.sqrt(p / 3)) * qml.math.array([[0, -1j], [1j, 0]]),
+        1 / (6 * qml.math.sqrt(p / 3)) * qml.math.diag([1, -1]),
+    ]
+
+    kraus_fn = lambda self, x: qml.math.stack(
+        channel.DepolarizingChannel(x, wires=0).kraus_matrices()
+    )
+
+    kraus_fn_real = lambda self, x: qml.math.real(
+        qml.math.stack(channel.DepolarizingChannel(x, wires=0).kraus_matrices())
+    )
+    kraus_fn_imag = lambda self, x: qml.math.imag(
+        qml.math.stack(channel.DepolarizingChannel(x, wires=0).kraus_matrices())
+    )
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        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):
+        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):
+        import tensorflow as tf
+
+        p = tf.Variable(0.43)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(p)
+        jac = tape.jacobian(out, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        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))
 
 
 class TestResetError:
@@ -323,6 +629,67 @@ def circuit(p_0, p_1):
             ),
         )
 
+    expected_jac_fn = lambda self, p0, p1: (
+        [
+            -1 / (2 * qml.math.sqrt(1 - p0 - p1)) * qml.math.eye(2),
+            1 / (2 * qml.math.sqrt(p0)) * qml.math.array([[1, 0], [0, 0]]),
+            1 / (2 * qml.math.sqrt(p0)) * qml.math.array([[0, 1], [0, 0]]),
+            qml.math.zeros((2, 2)),
+            qml.math.zeros((2, 2)),
+        ],
+        [
+            -1 / (2 * qml.math.sqrt(1 - p0 - p1)) * qml.math.eye(2),
+            qml.math.zeros((2, 2)),
+            qml.math.zeros((2, 2)),
+            1 / (2 * qml.math.sqrt(p1)) * qml.math.array([[0, 0], [1, 0]]),
+            1 / (2 * qml.math.sqrt(p1)) * qml.math.array([[0, 0], [0, 1]]),
+        ],
+    )
+
+    kraus_fn = lambda self, p0, p1: qml.math.stack(
+        channel.ResetError(p0, p1, wires=0).kraus_matrices()
+    )
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        p0 = pnp.array(0.43, requires_grad=True)
+        p1 = pnp.array(0.12, requires_grad=True)
+
+        jac = qml.jacobian(self.kraus_fn)(p0, p1)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p0, p1))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        import torch
+
+        p0 = torch.tensor(0.43, requires_grad=True)
+        p1 = torch.tensor(0.12, requires_grad=True)
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, (p0, p1))
+        exp_jac = self.expected_jac_fn(p0.detach().numpy(), p1.detach().numpy())
+        assert len(jac) == len(exp_jac) == 2
+        for j, exp_j in zip(jac, exp_jac):
+            assert qml.math.allclose(j.detach().numpy(), exp_j)
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        import tensorflow as tf
+
+        p0 = tf.Variable(0.43)
+        p1 = tf.Variable(0.12)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(p0, p1)
+        jac = tape.jacobian(out, (p0, p1))
+        assert qml.math.allclose(jac, self.expected_jac_fn(p0, p1))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        import jax
+
+        p0 = jax.numpy.array(0.43)
+        p1 = jax.numpy.array(0.12)
+        jac = jax.jacobian(self.kraus_fn, argnums=[0, 1])(p0, p1)
+        assert qml.math.allclose(jac, self.expected_jac_fn(p0, p1))
+
 
 class TestPauliError:
     """Tests for the quantum channel PauliError"""
@@ -419,6 +786,76 @@ def test_kraus_matrix(self, tol, operators, wires, expected_Ks):
 
         assert np.allclose(c.kraus_matrices(), expected_Ks, atol=tol, rtol=0)
 
+    expected_jac_fn = {
+        "X": lambda p: [
+            -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(2),
+            1 / (2 * qml.math.sqrt(p)) * qml.math.array([[0, 1], [1, 0]]),
+        ],
+        "XY": lambda p: [
+            -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(4),
+            1 / (2 * qml.math.sqrt(p)) * (qml.math.diag([1j, -1j, 1j, -1j])[::-1]),
+        ],
+    }
+
+    @pytest.mark.parametrize("ops", ["X", "XY"])
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self, ops):
+        p = pnp.array(0.43, requires_grad=True)
+        wires = list(range(len(ops)))
+        fn_real = lambda x: qml.math.real(
+            qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+        )
+        fn_imag = lambda x: qml.math.imag(
+            qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+        )
+        jac_fn_real = qml.jacobian(fn_real)
+        jac_fn_imag = qml.jacobian(fn_imag)
+        jac = jac_fn_real(p) + 1j * jac_fn_imag(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn[ops](p))
+
+    @pytest.mark.parametrize("ops", ["X", "XY"])
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self, ops):
+        import torch
+
+        p = torch.tensor(0.43, requires_grad=True)
+        wires = list(range(len(ops)))
+        fn_real = lambda x: qml.math.real(
+            qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+        )
+        fn_imag = lambda x: qml.math.imag(
+            qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+        )
+        jac_real = torch.autograd.functional.jacobian(fn_real, p).detach().numpy()
+        jac_imag = torch.autograd.functional.jacobian(fn_imag, p).detach().numpy()
+        assert qml.math.allclose(
+            jac_real + 1j * jac_imag, self.expected_jac_fn[ops](p.detach().numpy())
+        )
+
+    @pytest.mark.parametrize("ops", ["X", "XY"])
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self, ops):
+        import tensorflow as tf
+
+        p = tf.Variable(0.43)
+        wires = list(range(len(ops)))
+        with tf.GradientTape() as tape:
+            out = qml.math.stack(channel.PauliError(ops, p, wires=wires).kraus_matrices())
+        jac = tape.jacobian(out, p)
+        assert qml.math.allclose(jac, self.expected_jac_fn[ops](p))
+
+    @pytest.mark.parametrize("ops", ["X", "XY"])
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self, ops):
+        import jax
+
+        p = jax.numpy.array(0.43, dtype=jax.numpy.complex128)
+        wires = list(range(len(ops)))
+        fn = lambda x: qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+        jac_fn = jax.jacobian(fn, holomorphic=True)
+        jac = jac_fn(p)
+        assert qml.math.allclose(jac, self.expected_jac_fn[ops](p))
+
 
 class TestQubitChannel:
     """Tests for the quantum channel QubitChannel"""