Repository: PennyLaneAI/pennylane

Issue:
[BUG] Incompatibility qml.QubitChannel and the mixed device with jax.jit
### Expected behavior

I should be able to work with custom noise with jax.jit.

### Actual behavior

Tracing error.

### Additional information

Not sure if it's a bug, but I guess it should be working. In order to put custom noise in a circuit qml.QubitChannel is used. It works correctly without jax.jit, but it gives a tracing error with it. Below you can find a short example.

### Source code

```shell
import pennylane as qml
import jax.numpy as jnp
import jax

dev = qml.device('default.mixed', wires=1)
@jax.jit
@qml.qnode(dev, interface="jax")
def noise_channel(p):
    k0 = jnp.sqrt(1-p)*jnp.eye(2)
    k1 = jnp.sqrt(p)*jnp.eye(2)
    qml.QubitChannel([k0,k1], wires=[0])
    return qml.state()

noise_channel(0.1)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TracerArrayConversionError                Traceback (most recent call last)
File d:\Documents\1. GitHub\qurriculum_learning\test.py:15
     12     qml.QubitChannel([k0,k1], wires=[0])
     13     return qml.state()
---> 15 noise_channel(0.1)

    [... skipping hidden 12 frame]

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\workflow\qnode.py:1002, in QNode.__call__(self, *args, **kwargs)
    999         kwargs["shots"] = _get_device_shots(self._original_device)
   1001 # construct the tape
-> 1002 self.construct(args, kwargs)
   1004 cache = self.execute_kwargs.get("cache", False)
   1005 using_custom_cache = (
   1006     hasattr(cache, "__getitem__")
   1007     and hasattr(cache, "__setitem__")
   1008     and hasattr(cache, "__delitem__")
   1009 )

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\workflow\qnode.py:888, in QNode.construct(self, args, kwargs)
    885     self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
    887 with qml.queuing.AnnotatedQueue() as q:
--> 888     self._qfunc_output = self.func(*args, **kwargs)
    890 self._tape = QuantumScript.from_queue(q, shots)
    892 params = self.tape.get_parameters(trainable_only=False)

File d:\Documents\1. GitHub\qurriculum_learning\test.py:12
     10 k0 = jnp.sqrt(1-p)*jnp.eye(2)
     11 k1 = jnp.sqrt(p)*jnp.eye(2)
---> 12 qml.QubitChannel([k0,k1], wires=[0])
     13 return qml.state()

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\ops\channel.py:736, in QubitChannel.__init__(self, K_list, wires, id)
    731     raise ValueError(
    732         "Dimension of all Kraus matrices must be (2**num_wires, 2**num_wires)."
    733     )
    735 # check that the channel represents a trace-preserving map
--> 736 K_arr = np.array(K_list)
    737 Kraus_sum = np.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
    738 if not np.allclose(Kraus_sum, np.eye(K_list[0].shape[0])):

File d:\Documents\1. GitHub\venv\lib\site-packages\pennylane\math\multi_dispatch.py:39, in array(like, *args, **kwargs)
     30 def array(*args, like=None, **kwargs):
     31     """Creates an array or tensor object of the target framework.
     32 
     33     If the PyTorch interface is specified, this method preserves the Torch device used.
   (...)
     37         tensor_like: the tensor_like object of the framework
     38     """
---> 39     res = np.array(*args, like=like, **kwargs)
     40     if like is not None and get_interface(like) == "torch":
     41         res = res.to(device=like.device)

File d:\Documents\1. GitHub\venv\lib\site-packages\autoray\autoray.py:80, in do(fn, like, *args, **kwargs)
     31 """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
     32 dispatch to retrieve ``fn`` based on whichever library defines the class of
     33 the ``args[0]``, or the ``like`` keyword argument if specified.
   (...)
     77     <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
     78 """
     79 backend = choose_backend(fn, *args, like=like, **kwargs)
---> 80 return get_lib_fn(backend, fn)(*args, **kwargs)

File d:\Documents\1. GitHub\venv\lib\site-packages\jax\_src\core.py:710, in Tracer.__array__(self, *args, **kw)
    709 def __array__(self, *args, **kw):
--> 710   raise TracerArrayConversionError(self)

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[2,2].
The error occurred while tracing the function noise_channel at <ipython-input-17-22f3c77fa2c5>:7 for jit. This concrete value was not available in Python because it depends on the value of the argument p.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
```


### System information

```shell
Name: PennyLane
Version: 0.36.0.dev0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: d:\documents\1. github\venv\lib\site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.10.11
Numpy version:           1.26.2
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.36.0.dev0)
- default.gaussian (PennyLane-0.36.0.dev0)
- default.mixed (PennyLane-0.36.0.dev0)
- default.qubit (PennyLane-0.36.0.dev0)
- default.qubit.autograd (PennyLane-0.36.0.dev0)
- default.qubit.jax (PennyLane-0.36.0.dev0)
- default.qubit.legacy (PennyLane-0.36.0.dev0)
- default.qubit.tf (PennyLane-0.36.0.dev0)
- default.qubit.torch (PennyLane-0.36.0.dev0)
- default.qutrit (PennyLane-0.36.0.dev0)
- null.qubit (PennyLane-0.36.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.34.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
@@ -18,7 +18,7 @@
 """
 import warnings
 
-import pennylane.math as np
+from pennylane import math as np
 from pennylane.operation import AnyWires, Channel
 
 
@@ -717,26 +717,27 @@ def __init__(self, K_list, wires=None, id=None):
         super().__init__(*K_list, wires=wires, id=id)
 
         # check all Kraus matrices are square matrices
-        if not all(K.shape[0] == K.shape[1] for K in K_list):
+        if any(K.shape[0] != K.shape[1] for K in K_list):
             raise ValueError(
                 "Only channels with the same input and output Hilbert space dimensions can be applied."
             )
 
         # check all Kraus matrices have the same shape
-        if not all(K.shape == K_list[0].shape for K in K_list):
+        if any(K.shape != K_list[0].shape for K in K_list):
             raise ValueError("All Kraus matrices must have the same shape.")
 
         # check the dimension of all Kraus matrices are valid
-        if not all(K.ndim == 2 for K in K_list):
+        if any(K.ndim != 2 for K in K_list):
             raise ValueError(
                 "Dimension of all Kraus matrices must be (2**num_wires, 2**num_wires)."
             )
 
         # check that the channel represents a trace-preserving map
-        K_arr = np.array(K_list)
-        Kraus_sum = np.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
-        if not np.allclose(Kraus_sum, np.eye(K_list[0].shape[0])):
-            raise ValueError("Only trace preserving channels can be applied.")
+        if not any(np.is_abstract(K) for K in K_list):
+            K_arr = np.array(K_list)
+            Kraus_sum = np.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
+            if not np.allclose(Kraus_sum, np.eye(K_list[0].shape[0])):
+                raise ValueError("Only trace preserving channels can be applied.")
 
     def _flatten(self):
         return (self.data,), (self.wires, ())


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
@@ -124,11 +124,16 @@ 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())
+    @staticmethod
+    def expected_jac_fn(gamma):
+        return [
+            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]]),
+        ]
+
+    @staticmethod
+    def kraus_fn(x):
+        return qml.math.stack(channel.AmplitudeDamping(x, wires=0).kraus_matrices())
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -194,27 +199,37 @@ 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]]),
-        ],
-    )
+    @staticmethod
+    def expected_jac_fn(gamma, p):
+        return (
+            [
+                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()
-    )
+    @staticmethod
+    def kraus_fn(gamma, p):
+        return qml.math.stack(
+            channel.GeneralizedAmplitudeDamping(gamma, p, wires=0).kraus_matrices()
+        )
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -278,12 +293,16 @@ 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))]]),
-    ]
+    @staticmethod
+    def expected_jac_fn(gamma):
+        return [
+            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())
+    @staticmethod
+    def kraus_fn(x):
+        return qml.math.stack(channel.PhaseDamping(x, wires=0).kraus_matrices())
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -354,12 +373,16 @@ 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]]),
-    ]
+    @staticmethod
+    def expected_jac_fn(p):
+        return [
+            -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())
+    @staticmethod
+    def kraus_fn(x):
+        return qml.math.stack(channel.BitFlip(x, wires=0).kraus_matrices())
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -430,11 +453,16 @@ 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())
+    @staticmethod
+    def expected_jac_fn(p):
+        return [
+            -1 / (2 * qml.math.sqrt(1 - p)) * qml.math.eye(2),
+            1 / (2 * qml.math.sqrt(p)) * qml.math.diag([1, -1]),
+        ]
+
+    @staticmethod
+    def kraus_fn(x):
+        return qml.math.stack(channel.PhaseFlip(x, wires=0).kraus_matrices())
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -507,23 +535,30 @@ def test_p_invalid_parameter(self):
         with pytest.raises(ValueError, match="p must be in the interval"):
             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]),
-    ]
+    @staticmethod
+    def expected_jac_fn(p):
+        return [
+            -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()
-    )
+    @staticmethod
+    def kraus_fn(x):
+        return 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())
-    )
+    @staticmethod
+    def kraus_fn_real(x):
+        return qml.math.real(
+            qml.math.stack(channel.DepolarizingChannel(x, wires=0).kraus_matrices())
+        )
+
+    @staticmethod
+    def kraus_fn_imag(x):
+        return qml.math.imag(
+            qml.math.stack(channel.DepolarizingChannel(x, wires=0).kraus_matrices())
+        )
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -630,26 +665,28 @@ 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]]),
-        ],
-    )
+    @staticmethod
+    def expected_jac_fn(p0, p1):
+        return (
+            [
+                -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()
-    )
+    @staticmethod
+    def kraus_fn(p0, p1):
+        return qml.math.stack(channel.ResetError(p0, p1, wires=0).kraus_matrices())
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
@@ -804,12 +841,17 @@ def test_kraus_matrix(self, tol, operators, wires, expected_Ks):
     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())
-        )
+
+        def fn_real(x):
+            return qml.math.real(
+                qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+            )
+
+        def fn_imag(x):
+            return 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)
@@ -822,12 +864,17 @@ def test_kraus_jac_torch(self, ops):
 
         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())
-        )
+
+        def fn_real(x):
+            return qml.math.real(
+                qml.math.stack(channel.PauliError(ops, x, wires=wires).kraus_matrices())
+            )
+
+        def fn_imag(x):
+            return 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(
@@ -853,7 +900,10 @@ def test_kraus_jac_jax(self, ops):
 
         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())
+
+        def fn(x):
+            return 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))
@@ -910,6 +960,25 @@ def test_channel_trace_preserving(self):
         with pytest.raises(ValueError, match="Only trace preserving channels can be applied."):
             channel.QubitChannel(K_list2 * 2, wires=0)
 
+    @pytest.mark.jax
+    def test_jit_compatibility(self):
+        """Test that QubitChannel can be jitted."""
+
+        import jax
+
+        dev = qml.device("default.mixed", wires=1)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def noise_channel(p):
+            k0 = jax.numpy.sqrt(1 - p) * jax.numpy.eye(2)
+            k1 = jax.numpy.sqrt(p) * jax.numpy.eye(2)
+            qml.QubitChannel([k0, k1], wires=[0])
+            return qml.expval(qml.PauliZ(0))
+
+        # just checking it runs
+        noise_channel(0.1)
+
 
 class TestThermalRelaxationError:
     """Tests for the quantum channel ThermalRelaxationError"""