Repository: PennyLaneAI/pennylane
Pull request #5032 (2024-01-10 00:27:52)

Qutrit mixed apply operation
**Context:**
Currently the qutrit_mixed device is being developed this is a necessary addition to the overall project as this applies operations to a qutrit mixed state. This is a prerequisite for other functionality relating to the qutrit mixed device for noisy qutrit simulation.

**Description of the Change:**
Added functionality for applying operations to a qutrit mixed-state.  The new ``apply_operation`` function can be used to apply gates and Channels to a qutrit mixed-state.

**Benefits:**
Allows for Channels and operations to be applied to a mixed state, will be used to add execute functionality to qutrit mixed-state device allowing for noisy qutrit simulation

**Possible Drawbacks:**
Abstracting for qubits and more generally qutrits may have added challenges and will require a reforctor or copied code code-smell.

**Related GitHub Issues:**
N/A

Gabriel-Bottrill: I'm confused why these JAX tests are failing, I can look into it, but I thought I could ask you all now if you  all may know off the top ya'lls heads.
Gabriel-Bottrill: This is the next installment of the qutrit mixed state simulator sorry for the wait!
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5032?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `4 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`3710433`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3710433d3bbbb7f6f92fc8a478cbaadabc646616?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`f0780ee`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5032?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5032?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/qutrit\_mixed/apply\_operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5032?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXV0cml0X21peGVkL2FwcGx5X29wZXJhdGlvbi5weQ==) | 95.87% | [4 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5032?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #5032      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         392      393       +1     
  Lines       35572    35402     -170     
==========================================
- Hits        35457    35282     -175     
- Misses        115      120       +5     
```



</details>

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

Gabriel-Bottrill: @glassnotes  @mudit2812  @albi3ro  New PR!!

Gabriel-Bottrill: I noticed that the tests seem to be stuck in Expected, anyone know why that may be?
albi3ro: > Removed tensordot as it was not necessesary for functionality.
> 
> Moved some functionality to new utils module so that it can be used throughout.

@Gabriel-Bottrill it could potentially be about not being updated with master. I'd try merging in master and seeing if that fixes it.  Unfortunately, it looks like this update will need to be done manually.
Gabriel-Bottrill: Thanks @mudit2812 I don't know what is causing the issues, the failing tests should not be related to my changes. I'm wondering if it's from something else. Is this happening in other PRs?
albi3ro: @Gabriel-Bottrill  I think we are somehow leaving float64 mode accidentally.  I encountered something similar in a different PR recently.  You can see the fixes I added to that PR:

[tests/devices/qubit/test_sampling.py](https://github.com/PennyLaneAI/pennylane/pull/4914/files#diff-0b818ae73b0de8ddd1ebcb8d78fbf8e7ea65893d5d17f68531db8a70b5554163)


Updating with master will give you those fixes and might solve the test failures. If not, you can just keep adding more `jax.config.update("jax_enable_x64", True)` to the failing tests to ensure we always use float64 mode.
Gabriel-Bottrill: @albi3ro, I had to use your suggestion to input `jax.config.update("jax_enable_x64", True)` at the start of failing tests, if you are ok with that and the PR, would you please merge it or let me know what to do. Thanks!
trbromley: @Gabriel-Bottrill, we're good to go to merge this! Would you mind updating your branch as it's a blocker for merging.
glassnotes: Thanks @trbromley  - it looks like there are some test failures that I don't think are related to this PR - in the details I'm seeing errors thrown by tests related to Catalyst and `default_qubit`. Have you encountered these before?

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c0b5fd397..0cf92d86b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -128,6 +128,9 @@
 * The transform `split_non_commuting` now accepts measurements of type `probs`, `sample` and `counts` which accept both wires and observables.
   [(#4972)](https://github.com/PennyLaneAI/pennylane/pull/4972)
 
+* A function called `apply_operation` has been added to the new `qutrit_mixed` module found in `qml.devices` that applies operations to device-compatible states.
+  [(#5032)](https://github.com/PennyLaneAI/pennylane/pull/5032)
+
 <h3>Breaking changes 💔</h3>
 
 * Pin Black to `v23.12` to prevent unnecessary formatting changes.
@@ -262,6 +265,7 @@ This release contains contributions from (in alphabetical order):
 
 Abhishek Abhishek,
 Utkarsh Azad,
+Gabriel Bottrill,
 Astral Cai,
 Isaac De Vlugt,
 Korbinian Kottmann,
diff --git a/pennylane/devices/qutrit_mixed/__init__.py b/pennylane/devices/qutrit_mixed/__init__.py
index e6a696b2f..7ab250690 100644
--- a/pennylane/devices/qutrit_mixed/__init__.py
+++ b/pennylane/devices/qutrit_mixed/__init__.py
@@ -17,11 +17,14 @@ Submodule for performing mixed state simulations of qutrit-based quantum circuit
 This submodule is internal and subject to change without a deprecation cycle. Use
 at your own discretion.
 
+
 .. currentmodule:: pennylane.devices.qutrit_mixed
 .. autosummary::
     :toctree: api
 
     create_initial_state
+    apply_operation
 """
 
+from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
diff --git a/pennylane/devices/qutrit_mixed/apply_operation.py b/pennylane/devices/qutrit_mixed/apply_operation.py
new file mode 100644
index 000000000..8a9b94a6e
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/apply_operation.py
@@ -0,0 +1,184 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions to apply operations to a qutrit mixed state."""
+# pylint: disable=unused-argument
+
+from functools import singledispatch
+from string import ascii_letters as alphabet
+import pennylane as qml
+from pennylane import math
+from pennylane import numpy as np
+from pennylane.operation import Channel
+from .utils import QUDIT_DIM, get_einsum_mapping, get_new_state_einsum_indices
+
+alphabet_array = np.array(list(alphabet))
+
+
+def _map_indices_apply_channel(**kwargs):
+    """Map indices to einsum string
+    Args:
+        **kwargs (dict): Stores indices calculated in `get_einsum_mapping`
+
+    Returns:
+        String of einsum indices to complete einsum calculations
+    """
+    op_1_indices = f"{kwargs['kraus_index']}{kwargs['new_row_indices']}{kwargs['row_indices']}"
+    op_2_indices = f"{kwargs['kraus_index']}{kwargs['col_indices']}{kwargs['new_col_indices']}"
+
+    new_state_indices = get_new_state_einsum_indices(
+        old_indices=kwargs["col_indices"] + kwargs["row_indices"],
+        new_indices=kwargs["new_col_indices"] + kwargs["new_row_indices"],
+        state_indices=kwargs["state_indices"],
+    )
+    # index mapping for einsum, e.g., '...iga,...abcdef,...idh->...gbchef'
+    return (
+        f"...{op_1_indices},...{kwargs['state_indices']},...{op_2_indices}->...{new_state_indices}"
+    )
+
+
+def apply_operation_einsum(op: qml.operation.Operator, state, is_state_batched: bool = False):
+    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:
+        op (Operator): Operator to apply to the quantum state
+        state (array[complex]): Input quantum state
+        is_state_batched (bool): Boolean representing whether the state is batched or not
+
+    Returns:
+        array[complex]: output_state
+    """
+    einsum_indices = get_einsum_mapping(op, state, _map_indices_apply_channel, is_state_batched)
+
+    num_ch_wires = len(op.wires)
+
+    # This could be pulled into separate function if tensordot is added
+    if isinstance(op, Channel):
+        kraus = op.kraus_matrices()
+    else:
+        kraus = [op.matrix()]
+
+    # Shape kraus operators
+    kraus_shape = [len(kraus)] + [QUDIT_DIM] * num_ch_wires * 2
+    if not isinstance(op, Channel):
+        mat = op.matrix()
+        dim = QUDIT_DIM**num_ch_wires
+        batch_size = math.get_batch_size(mat, (dim, dim), dim**2)
+        if batch_size is not None:
+            # Add broadcasting dimension to shape
+            kraus_shape = [batch_size] + kraus_shape
+            if op.batch_size is None:
+                op._batch_size = batch_size  # pylint:disable=protected-access
+
+    kraus = math.stack(kraus)
+    kraus_transpose = math.stack(math.moveaxis(kraus, source=-1, destination=-2))
+    # Torch throws error if math.conj is used before stack
+    kraus_dagger = math.conj(kraus_transpose)
+
+    kraus = math.cast(math.reshape(kraus, kraus_shape), complex)
+    kraus_dagger = math.cast(math.reshape(kraus_dagger, kraus_shape), complex)
+
+    return math.einsum(einsum_indices, kraus, state, kraus_dagger)
+
+
+@singledispatch
+def apply_operation(
+    op: qml.operation.Operator, state, is_state_batched: bool = False, debugger=None
+):
+    """Apply an operation to a given state.
+
+    Args:
+        op (Operator): The operation to apply to ``state``
+        state (TensorLike): The starting state.
+        is_state_batched (bool): Boolean representing whether the state is batched or not
+        debugger (_Debugger): The debugger to use
+
+    Returns:
+        ndarray: output state
+
+    .. warning::
+
+        ``apply_operation`` is an internal function, and thus subject to change without a deprecation cycle.
+
+    .. warning::
+        ``apply_operation`` applies no validation to its inputs.
+
+        This function assumes that the wires of the operator correspond to indices
+        of the state. See :func:`~.map_wires` to convert operations to integer wire labels.
+
+        The shape of state should be ``[QUDIT_DIM]*(num_wires * 2)``, where ``QUDIT_DIM`` is
+        the dimension of the system.
+
+    This is a ``functools.singledispatch`` function, so additional specialized kernels
+    for specific operations can be registered like:
+
+    .. code-block:: python
+
+        @apply_operation.register
+        def _(op: type_op, state):
+            # custom op application method here
+
+    **Example:**
+
+    >>> state = np.zeros((3,3))
+    >>> state[0][0] = 1
+    >>> state
+    tensor([[1., 0., 0.],
+        [0., 0., 0.],
+        [0., 0., 0.]], requires_grad=True)
+    >>> apply_operation(qml.TShift(0), state)
+    tensor([[0., 0., 0.],
+        [0., 1., 0],
+        [0., 0., 0.],], requires_grad=True)
+
+    """
+    return _apply_operation_default(op, state, is_state_batched, debugger)
+
+
+def _apply_operation_default(op, state, is_state_batched, debugger):
+    """The default behaviour of apply_operation, accessed through the standard dispatch
+    of apply_operation, as well as conditionally in other dispatches.
+    """
+
+    return apply_operation_einsum(op, state, is_state_batched=is_state_batched)
+    # TODO add tensordot and benchmark for performance
+
+
+# TODO add diagonal for speed up.
+
+
+@apply_operation.register
+def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debugger=None):
+    """Take a snapshot of the mixed state"""
+    if debugger and debugger.active:
+        measurement = op.hyperparameters["measurement"]
+        if measurement:
+            # TODO replace with: measure once added
+            raise NotImplementedError  # TODO
+        if is_state_batched:
+            dim = int(math.sqrt(math.size(state[0])))
+            flat_shape = [math.shape(state)[0], dim, dim]
+        else:
+            dim = int(math.sqrt(math.size(state)))
+            flat_shape = [dim, dim]
+
+        snapshot = math.reshape(state, flat_shape)
+        if op.tag:
+            debugger.snapshots[op.tag] = snapshot
+        else:
+            debugger.snapshots[len(debugger.snapshots)] = snapshot
+    return state
+
+
+# TODO add special case speedups
diff --git a/pennylane/devices/qutrit_mixed/initialize_state.py b/pennylane/devices/qutrit_mixed/initialize_state.py
index 23c21b318..d247a5f95 100644
--- a/pennylane/devices/qutrit_mixed/initialize_state.py
+++ b/pennylane/devices/qutrit_mixed/initialize_state.py
@@ -16,8 +16,7 @@
 from typing import Iterable, Union
 import pennylane as qml
 from pennylane.operation import StatePrepBase
-
-qudit_dim = 3  # specifies qudit dimension
+from .utils import QUDIT_DIM
 
 
 def create_initial_state(
@@ -55,18 +54,18 @@ def _apply_state_vector(state, num_wires):  # function is easy to abstract for q
 
     Args:
         state (array[complex]): normalized input state of length
-            ``qudit_dim**num_wires``, where ``qudit_dim`` is the dimension of the system.
+            ``QUDIT_DIM**num_wires``, where ``QUDIT_DIM`` is the dimension of the system.
         num_wires (int): number of wires that get initialized in the state
 
     Returns:
-        array[complex]: complex array of shape ``[qudit_dim] * (2 * num_wires)``
-        representing the density matrix of this state, where ``qudit_dim`` is
+        array[complex]: complex array of shape ``[QUDIT_DIM] * (2 * num_wires)``
+        representing the density matrix of this state, where ``QUDIT_DIM`` is
         the dimension of the system.
     """
 
     # Initialize the entire set of wires with the state
     rho = qml.math.outer(state, qml.math.conj(state))
-    return qml.math.reshape(rho, [qudit_dim] * 2 * num_wires)
+    return qml.math.reshape(rho, [QUDIT_DIM] * 2 * num_wires)
 
 
 def _create_basis_state(num_wires, index):  # function is easy to abstract for qudit
@@ -77,10 +76,10 @@ def _create_basis_state(num_wires, index):  # function is easy to abstract for q
         index (int): integer representing the computational basis state.
 
     Returns:
-        array[complex]: complex array of shape ``[qudit_dim] * (2 * num_wires)``
-        representing the density matrix of the basis state, where ``qudit_dim`` is
+        array[complex]: complex array of shape ``[QUDIT_DIM] * (2 * num_wires)``
+        representing the density matrix of the basis state, where ``QUDIT_DIM`` is
         the dimension of the system.
     """
-    rho = qml.math.zeros((qudit_dim**num_wires, qudit_dim**num_wires))
+    rho = qml.math.zeros((QUDIT_DIM**num_wires, QUDIT_DIM**num_wires))
     rho[index, index] = 1
-    return qml.math.reshape(rho, [qudit_dim] * (2 * num_wires))
+    return qml.math.reshape(rho, [QUDIT_DIM] * (2 * num_wires))
diff --git a/pennylane/devices/qutrit_mixed/utils.py b/pennylane/devices/qutrit_mixed/utils.py
new file mode 100644
index 000000000..c80ef8dc1
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/utils.py
@@ -0,0 +1,86 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions and variables to be utilized by qutrit mixed state simulator."""
+import functools
+from string import ascii_letters as alphabet
+import pennylane as qml
+from pennylane import numpy as np
+
+alphabet_array = np.array(list(alphabet))
+QUDIT_DIM = 3  # specifies qudit dimension
+
+
+def get_einsum_mapping(
+    op: qml.operation.Operator, state, map_indices, is_state_batched: bool = False
+):
+    r"""Finds the indices for einsum to apply kraus operators to a mixed state
+
+    Args:
+        op (Operator): Operator to apply to the quantum state
+        state (array[complex]): Input quantum state
+        map_indices (function): Maps the calculated indices to an einsum indices string
+        is_state_batched (bool): Boolean representing whether the state is batched or not
+
+    Returns:
+        str: indices mapping that defines the einsum
+    """
+    num_ch_wires = len(op.wires)
+    num_wires = int((len(qml.math.shape(state)) - is_state_batched) / 2)
+    rho_dim = 2 * num_wires
+
+    # Tensor indices of the state. For each qutrit, need an index for rows *and* columns
+    state_indices = alphabet[:rho_dim]
+
+    # row indices of the quantum state affected by this operation
+    row_wires_list = op.wires.tolist()
+    row_indices = "".join(alphabet_array[row_wires_list].tolist())
+
+    # column indices are shifted by the number of wires
+    col_wires_list = [w + num_wires for w in row_wires_list]
+    col_indices = "".join(alphabet_array[col_wires_list].tolist())
+
+    # indices in einsum must be replaced with new ones
+    new_row_indices = alphabet[rho_dim : rho_dim + num_ch_wires]
+    new_col_indices = alphabet[rho_dim + num_ch_wires : rho_dim + 2 * num_ch_wires]
+
+    # index for summation over Kraus operators
+    kraus_index = alphabet[rho_dim + 2 * num_ch_wires : rho_dim + 2 * num_ch_wires + 1]
+
+    # apply mapping function
+    return map_indices(
+        state_indices=state_indices,
+        kraus_index=kraus_index,
+        row_indices=row_indices,
+        new_row_indices=new_row_indices,
+        col_indices=col_indices,
+        new_col_indices=new_col_indices,
+    )
+
+
+def get_new_state_einsum_indices(old_indices, new_indices, state_indices):
+    """Retrieves the einsum indices string for the new state
+
+    Args:
+        old_indices (str): indices that are summed
+        new_indices (str): indices that must be replaced with sums
+        state_indices (str): indices of the original state
+
+    Returns:
+        str: the einsum indices of the new state
+    """
+    return functools.reduce(
+        lambda old_string, idx_pair: old_string.replace(idx_pair[0], idx_pair[1]),
+        zip(old_indices, new_indices),
+        state_indices,
+    )


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 9135ec0a1..df2cf6d7b 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -1062,6 +1062,8 @@ class TestApplyGroverOperator:
         when applying it to a Jax state."""
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         batched = batch_dim is not None
         shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
         # Input state
diff --git a/tests/devices/qutrit_mixed/conftest.py b/tests/devices/qutrit_mixed/conftest.py
new file mode 100644
index 000000000..296317407
--- /dev/null
+++ b/tests/devices/qutrit_mixed/conftest.py
@@ -0,0 +1,63 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Pytest configuration file for PennyLane qutrit mixed state test suite."""
+import numpy as np
+from scipy.stats import unitary_group
+import pytest
+
+
+def get_random_mixed_state(num_qutrits):
+    dim = 3**num_qutrits
+
+    rng = np.random.default_rng(seed=4774)
+    basis = unitary_group(dim=dim, seed=584545).rvs()
+    schmidt_weights = rng.dirichlet(np.ones(dim), size=1).astype(complex)[0]
+    mixed_state = np.zeros((dim, dim)).astype(complex)
+    for i in range(dim):
+        mixed_state += schmidt_weights[i] * np.outer(np.conj(basis[i]), basis[i])
+
+    return mixed_state.reshape([3] * (2 * num_qutrits))
+
+
+# 1 qutrit states
+@pytest.fixture(scope="package")
+def one_qutrit_state():
+    return get_random_mixed_state(1)
+
+
+@pytest.fixture(scope="package")
+def one_qutrit_batched_state():
+    return np.array([get_random_mixed_state(1) for _ in range(2)])
+
+
+# 2 qutrit states
+@pytest.fixture(scope="package")
+def two_qutrit_state():
+    return get_random_mixed_state(2)
+
+
+@pytest.fixture(scope="package")
+def two_qutrit_batched_state():
+    return np.array([get_random_mixed_state(2) for _ in range(2)])
+
+
+# 3 qutrit states
+@pytest.fixture(scope="package")
+def three_qutrit_state():
+    return get_random_mixed_state(3)
+
+
+@pytest.fixture(scope="package")
+def three_qutrit_batched_state():
+    return np.array([get_random_mixed_state(3) for _ in range(2)])
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py
new file mode 100644
index 000000000..c5bb0255a
--- /dev/null
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py
@@ -0,0 +1,322 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for apply_operation in devices/qutrit_mixed/apply_operation."""
+
+from functools import reduce
+import pytest
+import numpy as np
+from scipy.stats import unitary_group
+import pennylane as qml
+from pennylane import math
+from pennylane.operation import Channel
+from pennylane.devices.qutrit_mixed import apply_operation
+
+ml_frameworks_list = [
+    "numpy",
+    pytest.param("autograd", marks=pytest.mark.autograd),
+    pytest.param("jax", marks=pytest.mark.jax),
+    pytest.param("torch", marks=pytest.mark.torch),
+    pytest.param("tensorflow", marks=pytest.mark.tf),
+]
+
+subspaces = [(0, 1), (0, 2), (1, 2)]
+
+
+def test_custom_operator_with_matrix(one_qutrit_state):
+    """Test that apply_operation works with any operation that defines a matrix."""
+    mat = np.array(
+        [
+            [-0.35546532 - 0.03636115j, -0.19051888 - 0.38049108j, 0.07943913 - 0.8276115j],
+            [-0.2766807 - 0.71617593j, -0.1227771 + 0.61271557j, -0.0872488 - 0.11150285j],
+            [-0.2312502 - 0.47894201j, -0.04564929 - 0.65295532j, -0.3629075 + 0.3962342j],
+        ]
+    )
+
+    # pylint: disable=too-few-public-methods
+    class CustomOp(qml.operation.Operation):
+        num_wires = 1
+
+        def matrix(self):
+            return mat
+
+    new_state = apply_operation(CustomOp(0), one_qutrit_state)
+    assert qml.math.allclose(new_state, mat @ one_qutrit_state @ np.conj(mat).T)
+
+
+# TODO add tests for special cases as they are added
+
+
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+@pytest.mark.parametrize(
+    "state,shape", [("two_qutrit_state", (9, 9)), ("two_qutrit_batched_state", (2, 9, 9))]
+)
+class TestSnapshot:
+    """Test that apply_operation works for Snapshot ops"""
+
+    class Debugger:  # pylint: disable=too-few-public-methods
+        """A dummy debugger class"""
+
+        def __init__(self):
+            self.active = True
+            self.snapshots = {}
+
+    @pytest.mark.usefixtures("two_qutrit_state")
+    def test_no_debugger(
+        self, ml_framework, state, shape, request
+    ):  # pylint: disable=unused-argument
+        """Test nothing happens when there is no debugger"""
+        state = request.getfixturevalue(state)
+        initial_state = math.asarray(state, like=ml_framework)
+
+        new_state = apply_operation(qml.Snapshot(), initial_state, is_state_batched=len(shape) != 2)
+        assert new_state.shape == initial_state.shape
+        assert math.allclose(new_state, initial_state)
+
+    def test_empty_tag(self, ml_framework, state, shape, request):
+        """Test a snapshot is recorded properly when there is no tag"""
+        state = request.getfixturevalue(state)
+        initial_state = math.asarray(state, like=ml_framework)
+
+        debugger = self.Debugger()
+        new_state = apply_operation(
+            qml.Snapshot(), initial_state, debugger=debugger, is_state_batched=len(shape) != 2
+        )
+
+        assert new_state.shape == initial_state.shape
+        assert math.allclose(new_state, initial_state)
+
+        assert list(debugger.snapshots.keys()) == [0]
+        assert debugger.snapshots[0].shape == shape
+        assert math.allclose(debugger.snapshots[0], math.reshape(initial_state, shape))
+
+    def test_provided_tag(self, ml_framework, state, shape, request):
+        """Test a snapshot is recorded properly when provided a tag"""
+        state = request.getfixturevalue(state)
+        initial_state = math.asarray(state, like=ml_framework)
+
+        debugger = self.Debugger()
+        tag = "dense"
+        new_state = apply_operation(
+            qml.Snapshot(tag), initial_state, debugger=debugger, is_state_batched=len(shape) != 2
+        )
+
+        assert new_state.shape == initial_state.shape
+        assert math.allclose(new_state, initial_state)
+
+        assert list(debugger.snapshots.keys()) == [tag]
+        assert debugger.snapshots[tag].shape == shape
+        assert math.allclose(debugger.snapshots[tag], math.reshape(initial_state, shape))
+
+    def test_snapshot_with_measurement(self, ml_framework, state, shape, request):
+        """Test a snapshot with measurement throws NotImplementedError"""
+        state = request.getfixturevalue(state)
+        initial_state = math.asarray(state, like=ml_framework)
+
+        debugger = self.Debugger()
+        with pytest.raises(NotImplementedError):
+            _ = apply_operation(
+                qml.Snapshot(measurement=qml.expval(qml.GellMann(0, 1))),
+                initial_state,
+                debugger=debugger,
+                is_state_batched=len(shape) != 2,
+            )
+
+
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+class TestOperation:  # pylint: disable=too-few-public-methods
+    """Tests that broadcasted operations (not channels) are applied correctly."""
+
+    broadcasted_ops = [
+        qml.TRX(np.array([np.pi, np.pi / 2]), wires=0, subspace=(0, 1)),
+        qml.TRY(np.array([np.pi, np.pi / 2]), wires=1, subspace=(0, 1)),
+        qml.TRZ(np.array([np.pi, np.pi / 2]), wires=2, subspace=(1, 2)),
+        qml.QutritUnitary(
+            np.array([unitary_group.rvs(27), unitary_group.rvs(27)]),
+            wires=[0, 1, 2],
+        ),
+    ]
+    unbroadcasted_ops = [
+        qml.THadamard(wires=0),
+        qml.TClock(wires=1),
+        qml.TShift(wires=2),
+        qml.TAdd(wires=[1, 2]),
+        qml.TRX(np.pi / 3, wires=0, subspace=(0, 2)),
+        qml.TRY(2 * np.pi / 3, wires=1, subspace=(1, 2)),
+        qml.TRZ(np.pi / 6, wires=2, subspace=(0, 1)),
+        qml.QutritUnitary(unitary_group.rvs(27), wires=[0, 1, 2]),
+    ]
+    num_qutrits = 3
+    num_batched = 2
+
+    @classmethod
+    def expand_matrices(cls, op, batch_size=0):
+        """Find expanded operator matrices, since qml.matrix isn't working for qutrits #4367"""
+        pre_wires_identity = np.eye(3 ** op.wires[0])
+        post_wires_identity = np.eye(3 ** ((cls.num_qutrits - 1) - op.wires[-1]))
+        mat = op.matrix()
+
+        def expand_matrix(matrix):
+            return reduce(np.kron, (pre_wires_identity, matrix, post_wires_identity))
+
+        if batch_size:
+            return [expand_matrix(mat[i]) for i in range(batch_size)]
+        return expand_matrix(mat)
+
+    @classmethod
+    def get_expected_state(cls, expanded_operator, state):
+        """Finds expected state after applying operator"""
+        flattened_state = state.reshape((3**cls.num_qutrits,) * 2)
+        adjoint_matrix = np.conj(expanded_operator).T
+        new_state = expanded_operator @ flattened_state @ adjoint_matrix
+        return new_state.reshape([3] * (cls.num_qutrits * 2))
+
+    @pytest.mark.parametrize("op", unbroadcasted_ops)
+    def test_no_broadcasting(self, op, ml_framework, three_qutrit_state):
+        """Tests that unbatched operations are applied correctly to an unbatched state."""
+        state = three_qutrit_state
+        res = apply_operation(op, qml.math.asarray(state, like=ml_framework))
+
+        expanded_operator = self.expand_matrices(op)
+        expected = self.get_expected_state(expanded_operator, state)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize("op", broadcasted_ops)
+    def test_broadcasted_op(self, op, ml_framework, three_qutrit_state):
+        """Tests that batched operations are applied correctly to an unbatched state."""
+        state = three_qutrit_state
+        res = apply_operation(op, qml.math.asarray(state, like=ml_framework))
+        expanded_operators = self.expand_matrices(op, self.num_batched)
+
+        def get_expected(m):
+            return self.get_expected_state(m, state)
+
+        expected = [get_expected(expanded_operators[i]) for i in range(self.num_batched)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize("op", unbroadcasted_ops)
+    def test_broadcasted_state(self, op, ml_framework, three_qutrit_batched_state):
+        """Tests that unbatched operations are applied correctly to a batched state."""
+        state = three_qutrit_batched_state
+        res = apply_operation(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
+        expanded_operator = self.expand_matrices(op)
+
+        def get_expected(s):
+            return self.get_expected_state(expanded_operator, s)
+
+        expected = [get_expected(state[i]) for i in range(self.num_batched)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize("op", broadcasted_ops)
+    def test_broadcasted_op_broadcasted_state(self, op, ml_framework, three_qutrit_batched_state):
+        """Tests that batched operations are applied correctly to a batched state."""
+        state = three_qutrit_batched_state
+        res = apply_operation(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
+        expanded_operators = self.expand_matrices(op, self.num_batched)
+
+        expected = [
+            self.get_expected_state(expanded_operators[i], state[i])
+            for i in range(self.num_batched)
+        ]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    def test_batch_size_set_if_missing(self, ml_framework, one_qutrit_state):
+        """Tests that the batch_size is set on an operator if it was missing before."""
+        param = qml.math.asarray([0.1, 0.2], like=ml_framework)
+        state = one_qutrit_state
+        op = qml.TRX(param, 0)
+        op._batch_size = None  # pylint:disable=protected-access
+        state = apply_operation(op, state)
+        assert state.shape == (self.num_batched, 3, 3)
+        assert op.batch_size == self.num_batched
+
+
+@pytest.mark.parametrize("wire", [0, 1])
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+class TestChannels:  # pylint: disable=too-few-public-methods
+    """Tests that Channel operations are applied correctly."""
+
+    num_qutrits = 2
+    num_batched = 2
+
+    @classmethod
+    def expand_krons(cls, op):
+        """Find expanded operator kraus matrices"""
+        pre_wires_identity = np.eye(3 ** op.wires[0])
+        post_wires_identity = np.eye(3 ** ((cls.num_qutrits - 1) - op.wires[-1]))
+        krons = op.kraus_matrices()
+        return [reduce(np.kron, (pre_wires_identity, kron, post_wires_identity)) for kron in krons]
+
+    @classmethod
+    def get_expected_state(cls, expanded_krons, state):
+        """Finds expected state after applying channel"""
+        flattened_state = state.reshape((3**cls.num_qutrits,) * 2)
+        adjoint_krons = np.conj(np.transpose(expanded_krons, (0, 2, 1)))
+        new_state = np.sum(
+            [
+                expanded_kron @ flattened_state @ adjoint_krons[i]
+                for i, expanded_kron in enumerate(expanded_krons)
+            ],
+            axis=0,
+        )
+        return new_state.reshape([3] * (cls.num_qutrits * 2))
+
+    class CustomChannel(Channel):
+        num_params = 1
+        num_wires = 1
+
+        def __init__(self, p, wires, id=None):
+            super().__init__(p, wires=wires, id=id)
+
+        @staticmethod
+        def compute_kraus_matrices(p):
+            K0 = (np.sqrt(1 - p) * math.cast_like(np.eye(3), p)).astype(complex)
+            K1 = (
+                np.sqrt(p) * math.cast_like(np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]), p)
+            ).astype(complex)
+            return [K0, K1]
+
+    def test_non_broadcasted_state(self, ml_framework, wire, two_qutrit_state):
+        """Tests that Channel operations are applied correctly to a state."""
+        state = two_qutrit_state
+        test_channel = self.CustomChannel(0.3, wires=wire)
+        res = apply_operation(test_channel, math.asarray(state, like=ml_framework))
+
+        expanded_krons = self.expand_krons(test_channel)
+        expected = self.get_expected_state(expanded_krons, state)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    def test_broadcasted_state(self, ml_framework, wire, two_qutrit_batched_state):
+        """Tests that Channel operations are applied correctly to a batched state."""
+        state = two_qutrit_batched_state
+
+        test_channel = self.CustomChannel(0.3, wires=wire)
+        res = apply_operation(test_channel, math.asarray(state, like=ml_framework))
+
+        expanded_krons = self.expand_krons(test_channel)
+        expected = [
+            self.get_expected_state(expanded_krons, state[i]) for i in range(self.num_batched)
+        ]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
diff --git a/tests/fourier/test_reconstruct.py b/tests/fourier/test_reconstruct.py
index 1a611e833..843c3f0cf 100644
--- a/tests/fourier/test_reconstruct.py
+++ b/tests/fourier/test_reconstruct.py
@@ -194,6 +194,8 @@ class TestReconstructEqu:
         functions are differentiable for JAX input variables."""
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         # Convert fun to have integer frequencies
         _fun = lambda x: fun(x / base_f)
         _rec = _reconstruct_equ(_fun, num_frequency, interface="jax")
diff --git a/tests/gradients/core/test_metric_tensor.py b/tests/gradients/core/test_metric_tensor.py
index ef07dbd02..cf01151df 100644
--- a/tests/gradients/core/test_metric_tensor.py
+++ b/tests/gradients/core/test_metric_tensor.py
@@ -1147,6 +1147,9 @@ class TestFullMetricTensor:
     @pytest.mark.parametrize("interface", ["auto", "jax"])
     def test_correct_output_jax(self, ansatz, params, interface):
         from jax import numpy as jnp
+        from jax import config
+
+        config.update("jax_enable_x64", True)
 
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         dev = qml.device("default.qubit.jax", wires=self.num_wires + 1)