Repository: PennyLaneAI/pennylane
Pull request #3584 (2022-12-27 12:09:05)

Add `tensordot`-based method to `default.mixed` device
**Context:**
The `_apply_channel` method of `default.mixed` devices uses `qml.math.einsum`.
For a device with `n` wires, `einsum` does not support applying operations that act on all `n` wires if `n>7`, see #3582.

**Description of the Change:**
This PR adds an `_apply_channel_tensordot` method to the device, which does not have the above problem.
In addition, preliminary benchmarks showed that the `tensordot`-based method is equally fast or faster for operations that act on more than a single qubit. Therefore, this PR not only switches the device to use the new `tensordot`-based method if an operation acts on more than 7 qubits, but already for all multi-qubit operations.

**Benefits:**
Support of many-qubit operations.
Performance upgrade for multi-qubit operations.

**Possible Drawbacks:**
A bit of code duplication/complexity. If this is a major drawback, it probably is best to even remove the old `einsum`-based method, as its performance gain compared to the new method is small even for single-qubit operations.

**Related GitHub Issues:**
#3582 


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3584?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3584](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3584?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (40fe0a3) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/880517228462ba01af35d94204e55fc9c47a7cac?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8805172) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3584   +/-   ##
=======================================
  Coverage   99.87%   99.87%           
=======================================
  Files         314      314           
  Lines       28019    28037   +18     
=======================================
+ Hits        27983    28001   +18     
  Misses         36       36           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3584?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/default\_mixed.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3584/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

ankit27kh: Hey @dwierichs, I tested this with my original use case, and the results match. Also, this is much faster. For a 10 qubit circuit, `einsum` implementation took ~120 seconds while `tensordot` implementation took ~1 second.

Note that the results do not match exactly and differ in 15th or higher digit after the decimal for probability value.
dwierichs: Hey @ankit27kh, that is great to hear!
Deviations on the 15th decimal place are very common, I'd say, as this is touching the representation precision of `float64`. As the order and number of floating point operations likely is changed by the patch, I'd be surprised if they would match to even higher precision, where the decimal places basically can be discarded anyways.
dwierichs: I performed some more comparisons, and the performance of `einsum` vs `tensordot` depends heavily on the used interface. As the mixed simulator works for all interfaces, it's difficult to make a decision here. For JAX, for example, the `einsum`-based approach is fastest for almost all operations and device qubit numbers I considered (up to 6).
dwierichs: I ran a few additional timings:
![apply_channel_variants_times](https://user-images.githubusercontent.com/20772922/215053698-d6461e0d-dc54-4495-ad38-7a8c4d6fb613.png)

As we can see, there are massive speedups, but only for `interface in [None, "autograd"]` and unitary operations that have a single Kraus matrix, and significant slowdowns for other interfaces.
A boundary to decide for `tensordot` could be `num_wires>3 and interface in [None, "autograd"] and len(kraus)>1`, but that's already rather specific.
rmoyard: Some issues with pylint in the default mixed tests, after that it will be good to be merged :+1: 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e29226639..8f758b77f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -407,6 +407,10 @@
 
 <h3>Improvements</h3>
 
+* The `default.mixed` device received a performance improvement for multi-qubit operations.
+  This also allows to apply channels that act on more than seven qubits, which was not possible before.
+  [(#3584)](https://github.com/PennyLaneAI/pennylane/pull/3584)
+
 * `qml.dot` now groups coefficients together.
   [(#3691)](https://github.com/PennyLaneAI/pennylane/pull/3691)
 
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index aba03ed3f..3a075dde5 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -334,6 +334,53 @@ class DefaultMixed(QubitDevice):
 
         self._state = qnp.einsum(einsum_indices, kraus, self._state, kraus_dagger)
 
+    def _apply_channel_tensordot(self, kraus, wires):
+        r"""Apply a quantum channel specified by a list of Kraus operators to subsystems of the
+        quantum state. For a unitary gate, there is a single Kraus operator.
+
+        Args:
+            kraus (list[array]): Kraus operators
+            wires (Wires): target wires
+        """
+        channel_wires = self.map_wires(wires)
+        num_ch_wires = len(channel_wires)
+
+        # Shape kraus operators and cast them to complex data type
+        kraus_shape = [2] * (num_ch_wires * 2)
+        kraus = [qnp.cast(qnp.reshape(k, kraus_shape), dtype=self.C_DTYPE) for k in kraus]
+
+        # row indices of the quantum state affected by this operation
+        row_wires_list = channel_wires.tolist()
+        # column indices are shifted by the number of wires
+        col_wires_list = [w + self.num_wires for w in row_wires_list]
+
+        channel_col_ids = list(range(num_ch_wires, 2 * num_ch_wires))
+        axes_left = [channel_col_ids, row_wires_list]
+        # Use column indices instead or rows to incorporate transposition of K^\dagger
+        axes_right = [col_wires_list, channel_col_ids]
+
+        # Apply the Kraus operators, and sum over all Kraus operators afterwards
+        def _conjugate_state_with(k):
+            """Perform the double tensor product k @ self._state @ k.conj().
+            The `axes_left` and `axes_right` arguments are taken from the ambient variable space
+            and `axes_right` is assumed to incorporate the tensor product and the transposition
+            of k.conj() simultaneously."""
+            return qnp.tensordot(qnp.tensordot(k, self._state, axes_left), qnp.conj(k), axes_right)
+
+        if len(kraus) == 1:
+            _state = _conjugate_state_with(kraus[0])
+        else:
+            _state = qnp.sum(qnp.stack([_conjugate_state_with(k) for k in kraus]), axis=0)
+
+        # Permute the affected axes to their destination places.
+        # The row indices of the kraus operators are moved from the beginning to the original
+        # target row locations, the column indices from the end to the target column locations
+        source_left = list(range(num_ch_wires))
+        dest_left = row_wires_list
+        source_right = list(range(-num_ch_wires, 0))
+        dest_right = col_wires_list
+        self._state = qnp.moveaxis(_state, source_left + source_right, dest_left + dest_right)
+
     def _apply_diagonal_unitary(self, eigvals, wires):
         r"""Apply a diagonal unitary gate specified by a list of eigenvalues. This method uses
         the fact that the unitary is diagonal for a more efficient implementation.
@@ -422,7 +469,7 @@ class DefaultMixed(QubitDevice):
         else:
             # generate basis states on subset of qubits via the cartesian product
             basis_states = qnp.asarray(
-                list(itertools.product([0, 1], repeat=len(device_wires))), dtype=self.C_DTYPE
+                list(itertools.product([0, 1], repeat=len(device_wires))), dtype=int
             )
 
             # get basis states to alter on full set of qubits
@@ -543,7 +590,14 @@ class DefaultMixed(QubitDevice):
         if operation in diagonal_in_z_basis:
             self._apply_diagonal_unitary(matrices, wires)
         else:
-            self._apply_channel(matrices, wires)
+            num_op_wires = len(wires)
+            interface = qml.math.get_interface(self._state, *matrices)
+            # Use tensordot for Autograd and Numpy if there are more than 2 wires
+            # Use tensordot in any case for more than 7 wires, as einsum does not support this case
+            if (num_op_wires > 2 and interface in {"autograd", "numpy"}) or num_op_wires > 7:
+                self._apply_channel_tensordot(matrices, wires)
+            else:
+                self._apply_channel(matrices, wires)
 
     # pylint: disable=arguments-differ
 
diff --git a/pennylane/ops/channel.py b/pennylane/ops/channel.py
index 0e8adf33c..c81c4f5f8 100644
--- a/pennylane/ops/channel.py
+++ b/pennylane/ops/channel.py
@@ -495,7 +495,7 @@ class ResetError(Channel):
         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]")
 
-        interface = np.get_interface([p_0, p_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(


Test patch:
diff --git a/tests/devices/test_default_mixed.py b/tests/devices/test_default_mixed.py
index d64f1b247..f972381a0 100644
--- a/tests/devices/test_default_mixed.py
+++ b/tests/devices/test_default_mixed.py
@@ -14,8 +14,12 @@
 """
 Unit tests for the :mod:`pennylane.devices.DefaultMixed` device.
 """
+# pylint: disable=protected-access
 
+import copy
 import pytest
+import numpy as np
+
 import pennylane as qml
 from pennylane import QubitStateVector, BasisState, DeviceError
 from pennylane.devices import DefaultMixed
@@ -26,32 +30,34 @@ from pennylane.ops import (
     PauliX,
     Hadamard,
     CNOT,
+    SWAP,
     ISWAP,
     AmplitudeDamping,
     DepolarizingChannel,
     ResetError,
     PauliError,
+    MultiControlledX,
 )
 from pennylane.wires import Wires
 
-import numpy as np
-
 INV_SQRT2 = 1 / np.sqrt(2)
 
 
-# functions for creating different states used in testing
 def basis_state(index, nr_wires):
+    """Generate the density matrix of the computational basis state
+    indicated by ``index``."""
     rho = np.zeros((2**nr_wires, 2**nr_wires), dtype=np.complex128)
     rho[index, index] = 1
     return rho
 
 
 def hadamard_state(nr_wires):
-    """Equal superposition state (Hadamard on all qubits)"""
+    """Generate the equal superposition state (Hadamard on all qubits)"""
     return np.ones((2**nr_wires, 2**nr_wires), dtype=np.complex128) / (2**nr_wires)
 
 
 def max_mixed_state(nr_wires):
+    """Generate the maximally mixed state."""
     return np.eye(2**nr_wires, dtype=np.complex128) / (2**nr_wires)
 
 
@@ -62,6 +68,15 @@ def root_state(nr_wires):
     return np.outer(ket, np.conj(ket))
 
 
+def random_state(num_wires):
+    """Generate a random density matrix."""
+    shape = (2**num_wires, 2**num_wires)
+    state = np.random.random(shape) + 1j * np.random.random(shape)
+    state = state @ state.T.conj()
+    state /= np.trace(state)
+    return state
+
+
 @pytest.mark.parametrize("nr_wires", [1, 2, 3])
 class TestCreateBasisState:
     """Unit tests for the method `_create_basis_state()`"""
@@ -144,6 +159,7 @@ class TestReset:
     """Unit tests for the method `reset()`"""
 
     def test_reset_basis(self, nr_wires, tol):
+        """Test the reset after creating a basis state."""
         dev = qml.device("default.mixed", wires=nr_wires)
         dev._state = dev._create_basis_state(1)
         dev.reset()
@@ -168,6 +184,9 @@ class TestReset:
             ResetError(0.1, 0.5, wires=[0]),
             PauliError("X", 0.5, wires=0),
             PauliError("ZY", 0.3, wires=[1, 0]),
+            PauliX(0),
+            PauliZ(0),
+            Hadamard(0),
         ],
     )
     def test_reset_after_channel(self, nr_wires, op, tol):
@@ -179,16 +198,6 @@ class TestReset:
 
         assert np.allclose(dev._state, dev._create_basis_state(0), atol=tol, rtol=0)
 
-    @pytest.mark.parametrize("op", [PauliX, PauliZ, Hadamard])
-    def test_reset_after_channel(self, nr_wires, op, tol):
-        """Tests that state is correctly reset after applying gates on the first
-        wire"""
-        dev = qml.device("default.mixed", wires=nr_wires)
-        dev.apply([op(wires=0)])
-        dev.reset()
-
-        assert np.allclose(dev._state, dev._create_basis_state(0), atol=tol, rtol=0)
-
 
 @pytest.mark.parametrize("nr_wires", [1, 2, 3])
 class TestAnalyticProb:
@@ -330,6 +339,7 @@ class TestKrausOps:
         assert np.allclose(dev._get_kraus(ops[0]), ops[1], atol=tol, rtol=0)
 
 
+@pytest.mark.parametrize("apply_method", ["_apply_channel", "_apply_channel_tensordot"])
 class TestApplyChannel:
     """Unit tests for the method `_apply_channel()`"""
 
@@ -354,21 +364,14 @@ class TestApplyChannel:
     ]
 
     @pytest.mark.parametrize("x", x_apply_channel_init)
-    def test_channel_init(self, x, tol):
+    def test_channel_init(self, x, tol, apply_method):
         """Tests that channels are correctly applied to the default initial state"""
         nr_wires = x[0]
         op = x[1]
         target_state = np.reshape(x[2], [2] * 2 * nr_wires)
         dev = qml.device("default.mixed", wires=nr_wires)
         kraus = dev._get_kraus(op)
-        if (
-            op.name == "CNOT"
-            or op.name == "ISWAP"
-            or (isinstance(op, PauliError) and nr_wires == 2)
-        ):
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
-        else:
-            dev._apply_channel(kraus, wires=Wires(0))
+        getattr(dev, apply_method)(kraus, wires=op.wires)
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
@@ -397,7 +400,7 @@ class TestApplyChannel:
     ]
 
     @pytest.mark.parametrize("x", x_apply_channel_mixed)
-    def test_channel_mixed(self, x, tol):
+    def test_channel_mixed(self, x, tol, apply_method):
         """Tests that channels are correctly applied to the maximally mixed state"""
         nr_wires = x[0]
         op = x[1]
@@ -406,14 +409,7 @@ class TestApplyChannel:
         max_mixed = np.reshape(max_mixed_state(nr_wires), [2] * 2 * nr_wires)
         dev._state = max_mixed
         kraus = dev._get_kraus(op)
-        if (
-            op.name == "CNOT"
-            or op.name == "ISWAP"
-            or (isinstance(op, PauliError) and nr_wires == 2)
-        ):
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
-        else:
-            dev._apply_channel(kraus, wires=Wires(0))
+        getattr(dev, apply_method)(kraus, wires=op.wires)
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
@@ -479,7 +475,7 @@ class TestApplyChannel:
     ]
 
     @pytest.mark.parametrize("x", x_apply_channel_root)
-    def test_channel_root(self, x, tol):
+    def test_channel_root(self, x, tol, apply_method):
         """Tests that channels are correctly applied to root state"""
         nr_wires = x[0]
         op = x[1]
@@ -488,14 +484,44 @@ class TestApplyChannel:
         root = np.reshape(root_state(nr_wires), [2] * 2 * nr_wires)
         dev._state = root
         kraus = dev._get_kraus(op)
-        if (
-            op.name == "CNOT"
-            or op.name == "ISWAP"
-            or (isinstance(op, PauliError) and nr_wires == 2)
-        ):
-            dev._apply_channel(kraus, wires=Wires([0, 1]))
-        else:
-            dev._apply_channel(kraus, wires=Wires(0))
+        getattr(dev, apply_method)(kraus, wires=op.wires)
+        assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
+
+    ops = [
+        PauliX(wires=0),
+        PauliX(wires=2),
+        Hadamard(wires=0),
+        CNOT(wires=[0, 1]),
+        ISWAP(wires=[0, 1]),
+        SWAP(wires=[2, 0]),
+        MultiControlledX(wires=[0, 1, 2]),
+        MultiControlledX(wires=[2, 0, 1]),
+        AmplitudeDamping(0.5, wires=0),
+        DepolarizingChannel(0.5, wires=0),
+        ResetError(0.1, 0.5, wires=0),
+        PauliError("Z", 0.3, wires=0),
+        PauliError("XY", 0.5, wires=[0, 1]),
+        PauliError("XZY", 0.1, wires=[0, 2, 1]),
+    ]
+
+    @pytest.mark.parametrize("op", ops)
+    @pytest.mark.parametrize("num_dev_wires", [1, 2, 3])
+    def test_channel_against_matmul(self, num_dev_wires, op, apply_method, tol):
+        """Test the application of a channel againt matrix multiplication."""
+        if num_dev_wires < max(op.wires) + 1:
+            pytest.skip("Need at least as many wires in the device as in the operation.")
+        np.random.seed(52)
+        dev = qml.device("default.mixed", wires=num_dev_wires)
+        init_state = random_state(num_dev_wires)
+        dev._state = qml.math.reshape(init_state, [2] * (2 * num_dev_wires))
+
+        kraus = dev._get_kraus(op)
+        full_kraus = [qml.math.expand_matrix(k, op.wires, wire_order=dev.wires) for k in kraus]
+
+        target_state = qml.math.sum([k @ init_state @ k.conj().T for k in full_kraus], axis=0)
+        target_state = qml.math.reshape(target_state, [2] * (2 * num_dev_wires))
+
+        getattr(dev, apply_method)(kraus, wires=op.wires)
 
         assert np.allclose(dev._state, target_state, atol=tol, rtol=0)
 
@@ -811,31 +837,49 @@ class TestApplyDensityMatrix:
 
 
 class TestApplyOperation:
-    """Unit tests for the method `_apply_operation()`. Since this just calls `_apply_channel()`
-    and `_apply_diagonal_unitary()`, we just check that the correct method is called"""
+    """Unit tests for the method `_apply_operation()`. Since this just calls `_apply_channel()`,
+    `_apply_diagonal_unitary()` or `_apply_channel_tensordot`, we just check
+    that the correct method is called"""
 
     def test_diag_apply_op(self, mocker):
         """Tests that when applying a diagonal gate, only `_apply_diagonal_unitary` is called,
         exactly once"""
         spy_channel = mocker.spy(DefaultMixed, "_apply_channel")
+        spy_channel_tensordot = mocker.spy(DefaultMixed, "_apply_channel_tensordot")
         spy_diag = mocker.spy(DefaultMixed, "_apply_diagonal_unitary")
         dev = qml.device("default.mixed", wires=1)
         dev._apply_operation(PauliZ(0))
 
         spy_channel.assert_not_called()
+        spy_channel_tensordot.assert_not_called()
         spy_diag.assert_called_once()
 
     def test_channel_apply_op(self, mocker):
         """Tests that when applying a non-diagonal gate, only `_apply_channel` is called,
         exactly once"""
         spy_channel = mocker.spy(DefaultMixed, "_apply_channel")
+        spy_channel_tensordot = mocker.spy(DefaultMixed, "_apply_channel_tensordot")
         spy_diag = mocker.spy(DefaultMixed, "_apply_diagonal_unitary")
         dev = qml.device("default.mixed", wires=1)
         dev._apply_operation(PauliX(0))
 
         spy_diag.assert_not_called()
+        spy_channel_tensordot.assert_not_called()
         spy_channel.assert_called_once()
 
+    def test_channel_apply_tensordot_op(self, mocker):
+        """Tests that when applying a non-diagonal gate on more than two qubits,
+        only `_apply_channel_tensordot` is called, exactly once"""
+        spy_channel = mocker.spy(DefaultMixed, "_apply_channel")
+        spy_channel_tensordot = mocker.spy(DefaultMixed, "_apply_channel_tensordot")
+        spy_diag = mocker.spy(DefaultMixed, "_apply_diagonal_unitary")
+        dev = qml.device("default.mixed", wires=3)
+        dev._apply_operation(MultiControlledX(wires=[0, 1, 2]))
+
+        spy_diag.assert_not_called()
+        spy_channel.assert_not_called()
+        spy_channel_tensordot.assert_called_once()
+
     def test_identity_skipped(self, mocker):
         """Test that applying the identity does not perform any additional computations."""
 
@@ -845,7 +889,7 @@ class TestApplyOperation:
         spy_diagonal_unitary = mocker.spy(dev, "_apply_diagonal_unitary")
         spy_apply_channel = mocker.spy(dev, "_apply_channel")
 
-        initialstate = dev.state.__copy__()
+        initialstate = copy.copy(dev.state)
 
         dev._apply_operation(op)
 
diff --git a/tests/devices/test_default_mixed_autograd.py b/tests/devices/test_default_mixed_autograd.py
index b5047f285..781a9c9d2 100644
--- a/tests/devices/test_default_mixed_autograd.py
+++ b/tests/devices/test_default_mixed_autograd.py
@@ -14,7 +14,7 @@
 """
 Tests for the ``default.mixed`` device for the Autograd interface
 """
-import re
+# pylint: disable=protected-access
 import pytest
 
 import pennylane as qml
@@ -22,8 +22,9 @@ from pennylane import numpy as np
 from pennylane.devices.default_mixed import DefaultMixed
 from pennylane import DeviceError
 
+pytestmark = pytest.mark.autograd
+
 
-@pytest.mark.autograd
 def test_analytic_deprecation():
     """Tests if the kwarg `analytic` is used and displays error message."""
     msg = "The analytic argument has been replaced by shots=None. "
@@ -36,7 +37,6 @@ def test_analytic_deprecation():
         qml.device("default.mixed", wires=1, shots=1, analytic=True)
 
 
-@pytest.mark.autograd
 class TestQNodeIntegration:
     """Integration tests for default.mixed.autograd. This test ensures it integrates
     properly with the PennyLane UI, in particular the QNode."""
@@ -67,7 +67,7 @@ class TestQNodeIntegration:
         """Test that the plugin device loads correctly"""
         dev = qml.device("default.mixed", wires=2)
         assert dev.num_wires == 2
-        assert dev.shots == None
+        assert dev.shots is None
         assert dev.short_name == "default.mixed"
         assert dev.capabilities()["passthru_devices"]["autograd"] == "default.mixed"
 
@@ -116,7 +116,6 @@ class TestQNodeIntegration:
         assert np.allclose(state, expected, atol=tol, rtol=0)
 
 
-@pytest.mark.autograd
 class TestDtypePreserved:
     """Test that the user-defined dtype of the device is preserved for QNode
     evaluation"""
@@ -131,7 +130,7 @@ class TestDtypePreserved:
             qml.probs(wires=[2, 0]),
         ],
     )
-    def test_real_dtype(self, r_dtype, measurement, tol):
+    def test_real_dtype(self, r_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with real-valued outputs"""
         p = 0.543
@@ -152,7 +151,7 @@ class TestDtypePreserved:
         "measurement",
         [qml.state(), qml.density_matrix(wires=[1]), qml.density_matrix(wires=[2, 0])],
     )
-    def test_complex_dtype(self, c_dtype, measurement, tol):
+    def test_complex_dtype(self, c_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with complex-valued outputs"""
         p = 0.543
@@ -168,7 +167,6 @@ class TestDtypePreserved:
         assert res.dtype == c_dtype
 
 
-@pytest.mark.autograd
 class TestOps:
     """Unit tests for operations supported by the default.mixed.autograd device"""
 
@@ -217,7 +215,37 @@ class TestOps:
         spy.assert_called()
 
 
-@pytest.mark.autograd
+@pytest.mark.parametrize(
+    "op, exp_method, dev_wires",
+    [
+        (qml.RX(np.array(0.2), 0), "_apply_channel", 1),
+        (qml.RX(np.array(0.2), 0), "_apply_channel", 8),
+        (qml.CNOT([0, 1]), "_apply_channel", 3),
+        (qml.CNOT([0, 1]), "_apply_channel", 8),
+        (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+        (qml.MultiControlledX(wires=list(range(3))), "_apply_channel_tensordot", 3),
+        (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+        (qml.PauliError("X", np.array(0.5), 0), "_apply_channel", 2),
+        (qml.PauliError("XXX", np.array(0.5), [0, 1, 2]), "_apply_channel_tensordot", 4),
+        (qml.PauliError("X" * 8, np.array(0.5), list(range(8))), "_apply_channel_tensordot", 8),
+    ],
+)
+def test_method_choice(mocker, op, exp_method, dev_wires):
+    """Test that the right method between _apply_channel and _apply_channel_tensordot
+    is chosen."""
+
+    methods = ["_apply_channel", "_apply_channel_tensordot"]
+    del methods[methods.index(exp_method)]
+    unexp_method = methods[0]
+    spy_exp = mocker.spy(DefaultMixed, exp_method)
+    spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+    dev = qml.device("default.mixed", wires=dev_wires)
+    dev._apply_operation(op)
+
+    spy_unexp.assert_not_called()
+    spy_exp.assert_called_once()
+
+
 class TestPassthruIntegration:
     """Tests for integration with the PassthruQNode"""
 
@@ -315,17 +343,29 @@ class TestPassthruIntegration:
         res = grad_fn(p)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize(
+        "op, wire_ids, exp_fn",
+        [
+            (qml.RY, [0], lambda a: -np.sin(a)),
+            (qml.AmplitudeDamping, [0], lambda a: -2),
+            (qml.DepolarizingChannel, [-1], lambda a: -4 / 3),
+            (lambda a, wires: qml.ResetError(p0=a, p1=0.1, wires=wires), [0], lambda a: -2),
+            (lambda a, wires: qml.ResetError(p0=0.1, p1=a, wires=wires), [0], lambda a: 0),
+        ],
+    )
     @pytest.mark.parametrize("wires", [[0], ["abc"]])
-    def test_state_differentiability(self, wires, tol):
+    def test_state_differentiability(self, wires, op, wire_ids, exp_fn, tol):
         """Test that the device state can be differentiated"""
+        # pylint: disable=too-many-arguments
         dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="autograd")
         def circuit(a):
-            qml.RY(a, wires=wires[0])
+            qml.PauliX(wires[wire_ids[0]])
+            op(a, wires=[wires[idx] for idx in wire_ids])
             return qml.state()
 
-        a = np.array(0.54, requires_grad=True)
+        a = np.array(0.23, requires_grad=True)
 
         def cost(a):
             """A function of the device quantum state, as a function
@@ -335,7 +375,7 @@ class TestPassthruIntegration:
             return res[1][1] - res[0][0]
 
         grad = qml.grad(cost)(a)
-        expected = np.sin(a)
+        expected = exp_fn(a)
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("wires", [range(2), [-12.32, "abc"]])
@@ -421,6 +461,7 @@ class TestPassthruIntegration:
 
     def test_sample_backprop_error(self):
         """Test that sampling in backpropagation mode raises an error"""
+        # pylint: disable=unused-variable
         dev = qml.device("default.mixed", wires=1, shots=100)
 
         msg = "Backpropagation is only supported when shots=None"
@@ -595,7 +636,6 @@ class TestPassthruIntegration:
         assert np.allclose(np.diag(res_b), expected_b, atol=tol, rtol=0)
 
 
-@pytest.mark.autograd
 class TestHighLevelIntegration:
     """Tests for integration with higher level components of PennyLane."""
 
@@ -619,12 +659,14 @@ class TestHighLevelIntegration:
         dev = qml.device("default.mixed", wires=2)
 
         def ansatz(weights, **kwargs):
+            # pylint: disable=unused-argument
             qml.RX(weights[0], wires=0)
             qml.RY(weights[1], wires=1)
             qml.CNOT(wires=[0, 1])
 
         obs_list = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)]
-        qnodes = qml.map(ansatz, obs_list, dev, interface="autograd")
+        with pytest.warns(UserWarning, match="The map function is deprecated"):
+            qnodes = qml.map(ansatz, obs_list, dev, interface="autograd")
 
         assert qnodes.interface == "autograd"
 
diff --git a/tests/devices/test_default_mixed_jax.py b/tests/devices/test_default_mixed_jax.py
index ba1cc6f91..de4562748 100644
--- a/tests/devices/test_default_mixed_jax.py
+++ b/tests/devices/test_default_mixed_jax.py
@@ -14,13 +14,13 @@
 """
 Tests for the ``default.mixed`` device for the JAX interface
 """
-import re
+# pylint: disable=protected-access
 import pytest
 
+import numpy as np
 import pennylane as qml
-from pennylane import numpy as np
+from pennylane import numpy as pnp
 from pennylane.devices.default_mixed import DefaultMixed
-from pennylane import DeviceError
 
 pytestmark = pytest.mark.jax
 
@@ -40,7 +40,7 @@ class TestQNodeIntegration:
         """Test that the plugin device loads correctly"""
         dev = qml.device("default.mixed", wires=2)
         assert dev.num_wires == 2
-        assert dev.shots == None
+        assert dev.shots is None
         assert dev.short_name == "default.mixed"
         assert dev.capabilities()["passthru_devices"]["jax"] == "default.mixed"
 
@@ -102,7 +102,7 @@ class TestDtypePreserved:
             qml.probs(wires=[2, 0]),
         ],
     )
-    def test_real_dtype(self, enable_x64, r_dtype, measurement, tol):
+    def test_real_dtype(self, enable_x64, r_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with real-valued outputs"""
         config.config.update("jax_enable_x64", enable_x64)
@@ -123,7 +123,7 @@ class TestDtypePreserved:
         "measurement",
         [qml.state(), qml.density_matrix(wires=[1]), qml.density_matrix(wires=[2, 0])],
     )
-    def test_complex_dtype(self, enable_x64, c_dtype, measurement, tol):
+    def test_complex_dtype(self, enable_x64, c_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with complex-valued outputs"""
         config.config.update("jax_enable_x64", enable_x64)
@@ -190,6 +190,83 @@ class TestOps:
         spy.assert_called()
 
 
+class TestApplyChannelMethodChoice:
+    """Test that the right method between _apply_channel and _apply_channel_tensordot
+    is chosen."""
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(jnp.array(0.2), 0), "_apply_channel", 1),
+            (qml.RX(jnp.array(0.2), 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel_tensordot", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", jnp.array(0.5), 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", jnp.array(0.5), [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, jnp.array(0.5), list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_numpy_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a numpy array as device state."""
+
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = np.array(state).reshape([2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(jnp.array(0.2), 0), "_apply_channel", 1),
+            (qml.RX(jnp.array(0.2), 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", jnp.array(0.5), 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", jnp.array(0.5), [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, jnp.array(0.5), list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_jax_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a JAX array as device state."""
+
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = jnp.array(state).reshape([2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+
 class TestPassthruIntegration:
     """Tests for integration with the PassthruQNode"""
 
@@ -264,7 +341,7 @@ class TestPassthruIntegration:
     def test_backprop_jacobian_agrees_parameter_shift(self, jacobian_fn, decorator, tol):
         """Test that jacobian of a QNode with an attached default.mixed device with JAX
         gives the correct result with respect to the parameter-shift method"""
-        p = np.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
+        p = pnp.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
         p_jax = jnp.array(p)
 
         def circuit(x):
@@ -290,14 +367,28 @@ class TestPassthruIntegration:
         res = decorator(jacobian_fn(circuit1, 0))(p_jax)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize(
+        "op, wire_ids, exp_fn",
+        [
+            (qml.RY, [0], lambda a: -jnp.sin(a)),
+            (qml.AmplitudeDamping, [0], lambda a: -2),
+            (qml.DepolarizingChannel, [-1], lambda a: -4 / 3),
+            (lambda a, wires: qml.ResetError(p0=a, p1=0.1, wires=wires), [0], lambda a: -2),
+            (lambda a, wires: qml.ResetError(p0=0.1, p1=a, wires=wires), [0], lambda a: 0),
+        ],
+    )
     @pytest.mark.parametrize("decorator", decorators)
-    def test_state_differentiability(self, decorator, tol):
+    def test_state_differentiability(self, decorator, op, wire_ids, exp_fn, tol):
         """Test that the device state can be differentiated"""
+        # pylint: disable=too-many-arguments
+        config.config.update("jax_enable_x64", True)
+
         dev = qml.device("default.mixed", wires=1)
 
         @qml.qnode(dev, interface="jax", diff_method="backprop")
         def circuit(a):
-            qml.RY(a, wires=0)
+            qml.PauliX(dev.wires[wire_ids[0]])
+            op(a, wires=[dev.wires[idx] for idx in wire_ids])
             return qml.state()
 
         a = jnp.array(0.54)
@@ -307,7 +398,7 @@ class TestPassthruIntegration:
             return res[1][1] - res[0][0]
 
         grad = decorator(jax.grad(cost))(a)
-        expected = np.sin(a)
+        expected = exp_fn(a)
 
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
@@ -414,6 +505,7 @@ class TestPassthruIntegration:
 
     def test_sample_backprop_error(self):
         """Test that sampling in backpropagation mode raises an error"""
+        # pylint: disable=unused-variable
         dev = qml.device("default.mixed", wires=1, shots=100)
 
         msg = "Backpropagation is only supported when shots=None"
@@ -622,13 +714,14 @@ class TestHighLevelIntegration:
         dev = qml.device("default.mixed", wires=2)
 
         obs_list = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)]
-        qnodes = qml.map(
-            qml.templates.StronglyEntanglingLayers,
-            obs_list,
-            dev,
-            interface="jax",
-            diff_method="backprop",
-        )
+        with pytest.warns(UserWarning, match="The map function is deprecated"):
+            qnodes = qml.map(
+                qml.templates.StronglyEntanglingLayers,
+                obs_list,
+                dev,
+                interface="jax",
+                diff_method="backprop",
+            )
 
         assert qnodes.interface == "jax"
 
diff --git a/tests/devices/test_default_mixed_tf.py b/tests/devices/test_default_mixed_tf.py
index ad15198b1..395b03dd4 100644
--- a/tests/devices/test_default_mixed_tf.py
+++ b/tests/devices/test_default_mixed_tf.py
@@ -14,13 +14,13 @@
 """
 Tests for the ``default.mixed`` device for the TensorFlow interface
 """
-import re
+# pylint: disable=protected-access
 import pytest
 
+import numpy as np
 import pennylane as qml
-from pennylane import numpy as np
+from pennylane import numpy as pnp
 from pennylane.devices.default_mixed import DefaultMixed
-from pennylane import DeviceError
 
 pytestmark = pytest.mark.tf
 
@@ -41,7 +41,7 @@ class TestQNodeIntegration:
         """Test that the plugin device loads correctly"""
         dev = qml.device("default.mixed", wires=2)
         assert dev.num_wires == 2
-        assert dev.shots == None
+        assert dev.shots is None
         assert dev.short_name == "default.mixed"
         assert dev.capabilities()["passthru_devices"]["tf"] == "default.mixed"
 
@@ -104,7 +104,7 @@ class TestDtypePreserved:
             qml.probs(wires=[2, 0]),
         ],
     )
-    def test_real_dtype(self, r_dtype, measurement, tol):
+    def test_real_dtype(self, r_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with real-valued outputs"""
         p = tf.constant(0.543)
@@ -124,7 +124,7 @@ class TestDtypePreserved:
         "measurement",
         [qml.state(), qml.density_matrix(wires=[1]), qml.density_matrix(wires=[2, 0])],
     )
-    def test_complex_dtype(self, c_dtype, measurement, tol):
+    def test_complex_dtype(self, c_dtype, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with complex-valued outputs"""
         p = tf.constant(0.543)
@@ -193,6 +193,83 @@ class TestOps:
         spy.assert_called()
 
 
+class TestApplyChannelMethodChoice:
+    """Test that the right method between _apply_channel and _apply_channel_tensordot
+    is chosen."""
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(tf.constant(0.2), 0), "_apply_channel", 1),
+            (qml.RX(tf.constant(0.2), 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel_tensordot", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", tf.constant(0.5), 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", tf.constant(0.5), [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, tf.constant(0.5), list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_numpy_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a numpy array as device state."""
+
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = np.array(state).reshape([2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(tf.constant(0.2), 0), "_apply_channel", 1),
+            (qml.RX(tf.constant(0.2), 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", tf.constant(0.5), 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", tf.constant(0.5), [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, tf.constant(0.5), list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_tf_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a Tensorflow array as device state."""
+
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = tf.reshape(tf.constant(state), [2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+
 class TestPassthruIntegration:
     """Tests for integration with the PassthruQNode"""
 
@@ -270,7 +347,7 @@ class TestPassthruIntegration:
     def test_backprop_jacobian_agrees_parameter_shift(self, decorator, interface, tol):
         """Test that jacobian of a QNode with an attached default.mixed.tf device
         gives the correct result with respect to the parameter-shift method"""
-        p = np.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
+        p = pnp.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
         p_tf = tf.Variable(p, trainable=True)
 
         def circuit(x):
@@ -298,16 +375,28 @@ class TestPassthruIntegration:
         res = tape.jacobian(res, p_tf)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize(
+        "op, wire_ids, exp_fn",
+        [
+            (qml.RY, [0], lambda a: -np.sin(a)),
+            (qml.AmplitudeDamping, [0], lambda a: -2),
+            (qml.DepolarizingChannel, [-1], lambda a: -4 / 3),
+            (lambda a, wires: qml.ResetError(p0=a, p1=0.1, wires=wires), [0], lambda a: -2),
+            (lambda a, wires: qml.ResetError(p0=0.1, p1=a, wires=wires), [0], lambda a: 0),
+        ],
+    )
     @pytest.mark.parametrize("wires", [[0], ["abc"]])
     @pytest.mark.parametrize("decorator, interface", decorators_interfaces)
-    def test_state_differentiability(self, decorator, interface, wires, tol):
+    def test_state_differentiability(self, decorator, interface, wires, op, wire_ids, exp_fn, tol):
         """Test that the device state can be differentiated"""
+        # pylint: disable=too-many-arguments
         dev = qml.device("default.mixed", wires=wires)
 
         @decorator
         @qml.qnode(dev, interface=interface, diff_method="backprop")
         def circuit(a):
-            qml.RY(a, wires=wires[0])
+            qml.PauliX(wires[wire_ids[0]])
+            op(a, wires=[wires[idx] for idx in wire_ids])
             return qml.state()
 
         a = tf.Variable(0.54, trainable=True)
@@ -318,7 +407,7 @@ class TestPassthruIntegration:
             res = res[1][1] - res[0][0]
 
         grad = tape.gradient(res, a)
-        expected = np.sin(a)
+        expected = exp_fn(a)
 
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
@@ -432,6 +521,7 @@ class TestPassthruIntegration:
 
     def test_sample_backprop_error(self):
         """Test that sampling in backpropagation mode raises an error"""
+        # pylint: disable=unused-variable
         dev = qml.device("default.mixed", wires=1, shots=100)
 
         msg = "Backpropagation is only supported when shots=None"
@@ -568,6 +658,7 @@ class TestPassthruIntegration:
     def test_ragged_differentiation(self, decorator, interface, dev_name, diff_method, mode, tol):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
+        # pylint: disable=too-many-arguments
 
         dev = qml.device(dev_name, wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
@@ -661,7 +752,8 @@ class TestHighLevelIntegration:
         dev = qml.device("default.mixed", wires=2)
 
         obs_list = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)]
-        qnodes = qml.map(qml.templates.StronglyEntanglingLayers, obs_list, dev, interface="tf")
+        with pytest.warns(UserWarning, match="The map function is deprecated"):
+            qnodes = qml.map(qml.templates.StronglyEntanglingLayers, obs_list, dev, interface="tf")
 
         assert qnodes.interface == "tf"
 
@@ -674,7 +766,7 @@ class TestHighLevelIntegration:
             return tf.reduce_sum(qnodes(weights))
 
         with tf.GradientTape() as tape:
-            res = qnodes(weights)
+            res = cost(weights)
 
         grad = tape.gradient(res, weights)
 
diff --git a/tests/devices/test_default_mixed_torch.py b/tests/devices/test_default_mixed_torch.py
index 52268ae70..6e3049605 100644
--- a/tests/devices/test_default_mixed_torch.py
+++ b/tests/devices/test_default_mixed_torch.py
@@ -14,13 +14,13 @@
 """
 Tests for the ``default.mixed`` device for the Torch interface.
 """
-import re
+# pylint: disable=protected-access, import-outside-toplevel
 import pytest
+import numpy as np
 
 import pennylane as qml
-from pennylane import numpy as np
+from pennylane import numpy as pnp
 from pennylane.devices.default_mixed import DefaultMixed
-from pennylane import DeviceError
 
 pytestmark = pytest.mark.torch
 
@@ -35,7 +35,7 @@ class TestQNodeIntegration:
         """Test that the plugin device loads correctly"""
         dev = qml.device("default.mixed", wires=2)
         assert dev.num_wires == 2
-        assert dev.shots == None
+        assert dev.shots is None
         assert dev.short_name == "default.mixed"
         assert dev.capabilities()["passthru_devices"]["torch"] == "default.mixed"
 
@@ -100,7 +100,7 @@ class TestDtypePreserved:
             qml.probs(wires=[2, 0]),
         ],
     )
-    def test_real_dtype(self, r_dtype, r_dtype_torch, measurement, tol):
+    def test_real_dtype(self, r_dtype, r_dtype_torch, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with real-valued outputs"""
         p = torch.tensor(0.543)
@@ -128,7 +128,7 @@ class TestDtypePreserved:
         "measurement",
         [qml.state(), qml.density_matrix(wires=[1]), qml.density_matrix(wires=[2, 0])],
     )
-    def test_complex_dtype(self, c_dtype, c_dtype_torch, measurement, tol):
+    def test_complex_dtype(self, c_dtype, c_dtype_torch, measurement):
         """Test that the user-defined dtype of the device is preserved
         for QNodes with complex-valued outputs"""
         if c_dtype_torch == "torchc64":
@@ -198,6 +198,89 @@ class TestOps:
         spy.assert_called()
 
 
+class TestApplyChannelMethodChoice:
+    """Test that the right method between _apply_channel and _apply_channel_tensordot
+    is chosen."""
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(0.2, 0), "_apply_channel", 1),
+            (qml.RX(0.2, 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel_tensordot", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", 0.5, 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", 0.5, [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, 0.5, list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_numpy_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a numpy array as device state."""
+
+        # Manually set the data of the operation to be torch data
+        # This is due to an import problem if these tests are skipped.
+        op.data = [d if isinstance(d, str) else torch.tensor(d) for d in op.data]
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = np.array(state).reshape([2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "op, exp_method, dev_wires",
+        [
+            (qml.RX(0.2, 0), "_apply_channel", 1),
+            (qml.RX(0.2, 0), "_apply_channel", 8),
+            (qml.CNOT([0, 1]), "_apply_channel", 3),
+            (qml.CNOT([0, 1]), "_apply_channel", 8),
+            (qml.MultiControlledX(wires=list(range(2))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(3))), "_apply_channel", 3),
+            (qml.MultiControlledX(wires=list(range(8))), "_apply_channel_tensordot", 8),
+            (qml.PauliError("X", 0.5, 0), "_apply_channel", 2),
+            (qml.PauliError("XXX", 0.5, [0, 1, 2]), "_apply_channel", 4),
+            (
+                qml.PauliError("X" * 8, 0.5, list(range(8))),
+                "_apply_channel_tensordot",
+                8,
+            ),
+        ],
+    )
+    def test_with_torch_state(self, mocker, op, exp_method, dev_wires):
+        """Test with a Torch array as device state."""
+
+        # Manually set the data of the operation to be torch data
+        # This is due to an import problem if these tests are skipped.
+        op.data = [d if isinstance(d, str) else torch.tensor(d) for d in op.data]
+        methods = ["_apply_channel", "_apply_channel_tensordot"]
+        del methods[methods.index(exp_method)]
+        unexp_method = methods[0]
+        spy_exp = mocker.spy(DefaultMixed, exp_method)
+        spy_unexp = mocker.spy(DefaultMixed, unexp_method)
+        dev = qml.device("default.mixed", wires=dev_wires)
+        state = np.zeros((2**dev_wires, 2**dev_wires))
+        state[0, 0] = 1.0
+        dev._state = torch.tensor(state).reshape([2] * (2 * dev_wires))
+        dev._apply_operation(op)
+
+        spy_unexp.assert_not_called()
+        spy_exp.assert_called_once()
+
+
 class TestPassthruIntegration:
     """Tests for integration with the PassthruQNode"""
 
@@ -270,7 +353,7 @@ class TestPassthruIntegration:
     def test_backprop_jacobian_agrees_parameter_shift(self, tol):
         """Test that jacobian of a QNode with an attached default.mixed.torch device
         gives the correct result with respect to the parameter-shift method"""
-        p = np.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
+        p = pnp.array([0.43316321, 0.2162158, 0.75110998, 0.94714242])
         p_torch = torch.tensor(p, dtype=torch.float64, requires_grad=True)
 
         def circuit(x):
@@ -296,14 +379,26 @@ class TestPassthruIntegration:
         grad = torch.autograd.functional.jacobian(circuit1, p_torch)
         assert qml.math.allclose(grad, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize(
+        "op, wire_ids, exp_fn",
+        [
+            (qml.RY, [0], lambda a: -torch.sin(a)),
+            (qml.AmplitudeDamping, [0], lambda a: -2.0),
+            (qml.DepolarizingChannel, [-1], lambda a: -4 / 3),
+            (lambda a, wires: qml.ResetError(p0=a, p1=0.1, wires=wires), [0], lambda a: -2.0),
+            (lambda a, wires: qml.ResetError(p0=0.1, p1=a, wires=wires), [0], lambda a: 0.0),
+        ],
+    )
     @pytest.mark.parametrize("wires", [[0], ["abc"]])
-    def test_state_differentiability(self, wires, tol):
+    def test_state_differentiability(self, wires, op, wire_ids, exp_fn, tol):
         """Test that the device state can be differentiated"""
+        # pylint: disable=too-many-arguments
         dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="torch")
         def circuit(a):
-            qml.RY(a, wires=wires[0])
+            qml.PauliX(wires[wire_ids[0]])
+            op(a, wires=[wires[idx] for idx in wire_ids])
             return qml.state()
 
         a = torch.tensor(0.54, dtype=torch.float64, requires_grad=True)
@@ -313,7 +408,7 @@ class TestPassthruIntegration:
         res = res[1][1] - res[0][0]
         res.backward()
 
-        expected = torch.sin(a)
+        expected = torch.tensor(exp_fn(a), dtype=torch.float64)
         assert torch.allclose(a.grad, expected, atol=tol, rtol=0)
 
     def test_state_vector_differentiability(self, tol):
@@ -414,6 +509,7 @@ class TestPassthruIntegration:
 
     def test_sample_backprop_error(self):
         """Test that sampling in backpropagation mode raises an error"""
+        # pylint: disable=unused-variable
         dev = qml.device("default.mixed", wires=1, shots=100)
 
         msg = "Backpropagation is only supported when shots=None"
@@ -622,7 +718,10 @@ class TestHighLevelIntegration:
         dev = qml.device("default.mixed", wires=2)
 
         obs_list = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)]
-        qnodes = qml.map(qml.templates.StronglyEntanglingLayers, obs_list, dev, interface="torch")
+        with pytest.warns(UserWarning, match="The map function is deprecated"):
+            qnodes = qml.map(
+                qml.templates.StronglyEntanglingLayers, obs_list, dev, interface="torch"
+            )
 
         assert qnodes.interface == "torch"
 
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index e317ed80f..77fc4b085 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -34,3 +34,8 @@ tests/gradients/test_spsa_gradient.py
 tests/returntypes/gradients/test_spsa_gradient_new.py
 tests/returntypes/gradients/test_spsa_gradient_shot_vec_new.py
 tests/kernels/**
+tests/devices/test_default_mixed.py
+tests/devices/test_default_mixed_autograd.py
+tests/devices/test_default_mixed_jax.py
+tests/devices/test_default_mixed_torch.py
+tests/devices/test_default_mixed_tf.py