Repository: PennyLaneAI/pennylane
Pull request #5049 (2024-01-12 09:55:47)

Qutrit mixed measure
Context:
Currently the qutrit_mixed device is being developed this is a necessary addition to the overall project as this measures 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 measuring a qutrit mixed-state. The new measure function can be used to measure a qutrit mixed-state.

Benefits:
Allows for measurements of 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 refactor or copied code code-smell. Methods are not the fastest possible, optimizations can still be made

Related GitHub Issues:
N/A



Gabriel-Bottrill: This is still fairly early, I meant to have it closer to finishing by now but I got behind with apply_operations. I just wanted to put it up so you all would know it's being worked on and will be finished soon! @albi3ro @glassnotes @mudit2812 @trbromley .
Gabriel-Bottrill: measurementprocess.diagonalizing_gates() is throwing errors due to the observables being for qutrits so I added methods to compute some of the main measurable values. I can leave it like this, but it slightly defeats the purpose of the new device structure. Any suggestions on getting around this would be very helpful, thanks!
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`c0ce3ce`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c0ce3ced1644cc552a3454fec5a09790ee5347fc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`96d75c2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.
> Report is 8 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5049      +/-   ##
==========================================
- Coverage   99.68%   99.68%   -0.01%     
==========================================
  Files         394      395       +1     
  Lines       35720    35750      +30     
==========================================
+ Hits        35609    35637      +28     
- Misses        111      113       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?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: @albi3ro, @glassnotes, @mudit2812, @trbromley I have the measure functionality here, I will have to remerge it once applystate is merged, but it is ready to be reviewed. 
Gabriel-Bottrill: @albi3ro, @glassnotes, @mudit2812, @trbromley now that apply_operation is done this PR is ready to be reviewed!
Gabriel-Bottrill: Also should I squash the commits or is it better to have the commit history?
Gabriel-Bottrill: changing `math.where(probs < 0, -probs, probs)` to `math.clip` broke tests for NaN values. I could remove the NaN value tests, but then I would probably need to add some testing for interfaces elsewhere. I would appreciate some input on what would be best.
Gabriel-Bottrill: Updates since last reviews (Thursday):

- Fixed suggested issues:
   - doc-string consistency and capitalization
   - Names of methods
   - Added analytical tests for expval
   - Moved `_map_indices_apply_operation` out of `apply_observable_einsum`
   - ... (a few small changes I may have forgot)
- Removed NaN tests
- Made broadcasting tests test ml_interfaces


Concerns/questions:
- [ ] Should `get_measurement_function` be defined separately?
- [ ] Should I delete `measure_func=measure` in `calculate_expval_sum_of_terms`


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`85dd14f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/85dd14f3b37bba4fbf0153e0f533ac18ba46163f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.69% compared to head [(`1fdc9ce`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.


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


```diff
@@            Coverage Diff             @@
##           master    #5049      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         394      395       +1     
  Lines       36118    35947     -171     
==========================================
- Hits        36007    35835     -172     
- Misses        111      112       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5049?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5ef82b3ae..1d4a4b5e1 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -121,9 +121,6 @@
          [1, 0, 0, 1, 1]])
   ```
 
-* 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>Improvements 🛠</h3>
 
 <h4>Faster gradients with VJPs and other performance improvements</h4>
@@ -174,6 +171,13 @@
 * The module `pennylane/math/quantum.py` has now support for the min-entropy.
   [(#3959)](https://github.com/PennyLaneAI/pennylane/pull/3959/)
 
+* 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)
+
+* A function called `measure` has been added to the new `qutrit_mixed` module found in `qml.devices` that measures device-compatible states for a collection of measurement processes.
+  [(#5049)](https://github.com/PennyLaneAI/pennylane/pull/5049)
+
+
 <h4>Other improvements</h4>
 
 * `qml.pauli.group_observables` now supports grouping `Prod` and `SProd` operators.
diff --git a/pennylane/devices/qutrit_mixed/__init__.py b/pennylane/devices/qutrit_mixed/__init__.py
index 7ab250690..ed62cd929 100644
--- a/pennylane/devices/qutrit_mixed/__init__.py
+++ b/pennylane/devices/qutrit_mixed/__init__.py
@@ -24,7 +24,9 @@ at your own discretion.
 
     create_initial_state
     apply_operation
+    measure
 """
 
 from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
+from .measure import measure
diff --git a/pennylane/devices/qutrit_mixed/measure.py b/pennylane/devices/qutrit_mixed/measure.py
new file mode 100644
index 000000000..d05bd85da
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/measure.py
@@ -0,0 +1,293 @@
+# 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.
+"""
+Code relevant for performing measurements on a qutrit mixed state.
+"""
+
+from typing import Callable
+from string import ascii_letters as alphabet
+from pennylane import math
+from pennylane.ops import Sum, Hamiltonian
+from pennylane.measurements import (
+    StateMeasurement,
+    MeasurementProcess,
+    ExpectationMP,
+    StateMP,
+    ProbabilityMP,
+    VarianceMP,
+)
+from pennylane.operation import Observable
+from pennylane.typing import TensorLike
+
+from .utils import (
+    get_einsum_mapping,
+    reshape_state_as_matrix,
+    get_num_wires,
+    get_new_state_einsum_indices,
+    QUDIT_DIM,
+)
+from .apply_operation import apply_operation
+
+
+def _map_indices_apply_operation(**kwargs):
+    """Map indices to wires.
+
+    Args:
+        **kwargs (dict): Stores indices calculated in `get_einsum_mapping`:
+            state_indices (str): Indices that are summed.
+            row_indices (str): Indices that must be replaced with sums.
+            new_row_indices (str): Tensor indices of the state.
+
+    Returns:
+        String of einsum indices to complete einsum calculations.
+    """
+    op_1_indices = f"{kwargs['new_row_indices']}{kwargs['row_indices']}"
+
+    new_state_indices = get_new_state_einsum_indices(
+        old_indices=kwargs["row_indices"],
+        new_indices=kwargs["new_row_indices"],
+        state_indices=kwargs["state_indices"],
+    )
+
+    return f"{op_1_indices},...{kwargs['state_indices']}->...{new_state_indices}"
+
+
+def apply_observable_einsum(obs: Observable, state, is_state_batched: bool = False):
+    r"""Applies an observable to a density matrix rho, giving obs@state.
+
+    Args:
+        obs (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:
+        TensorLike: the result of obs@state.
+    """
+
+    num_ch_wires = len(obs.wires)
+    einsum_indices = get_einsum_mapping(obs, state, _map_indices_apply_operation, is_state_batched)
+    obs_mat = obs.matrix()
+    obs_shape = [QUDIT_DIM] * num_ch_wires * 2
+    obs_mat = math.cast(math.reshape(obs_mat, obs_shape), complex)
+    return math.einsum(einsum_indices, obs_mat, state)
+
+
+def calculate_expval(
+    measurementprocess: ExpectationMP, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Measure the expectation value of an observable by finding the trace of obs@rho.
+
+    Args:
+        measurementprocess (ExpectationMP): measurement process to apply to the state.
+        state (TensorLike): the state to measure.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        TensorLike: expectation value of observable wrt the state.
+    """
+    obs = measurementprocess.obs
+    rho_mult_obs = apply_observable_einsum(obs, state, is_state_batched)
+
+    # using einsum since trace function axis selection parameter names
+    # are not consistent across interfaces, they don't exist for torch
+
+    num_wires = get_num_wires(state, is_state_batched)
+    rho_mult_obs_reshaped = reshape_state_as_matrix(rho_mult_obs, num_wires)
+    if is_state_batched:
+        return math.real(math.stack([math.sum(math.diagonal(dm)) for dm in rho_mult_obs_reshaped]))
+
+    return math.real(math.sum(math.diagonal(rho_mult_obs_reshaped)))
+
+
+def calculate_reduced_density_matrix(  # TODO: ask if I should have state diagonalization gates?
+    measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Get the state or reduced density matrix.
+
+    Args:
+        measurementprocess (StateMeasurement): measurement to apply to the state.
+        state (TensorLike): state to apply the measurement to.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        TensorLike: state or reduced density matrix.
+    """
+    wires = measurementprocess.wires
+    if not wires:
+        return reshape_state_as_matrix(state, get_num_wires(state, is_state_batched))
+
+    num_obs_wires = len(wires)
+    num_state_wires = get_num_wires(state, is_state_batched)
+    state_wire_indices_list = list(alphabet[:num_state_wires] * 2)
+    final_state_wire_indices_list = [""] * (2 * num_obs_wires)
+
+    for i, wire in enumerate(wires):
+        col_index = wire + num_state_wires
+        state_wire_indices_list[col_index] = alphabet[col_index]
+        final_state_wire_indices_list[i] = alphabet[wire]
+        final_state_wire_indices_list[i + num_obs_wires] = alphabet[col_index]
+
+    state_wire_indices = "".join(state_wire_indices_list)
+    final_state_wire_indices = "".join(final_state_wire_indices_list)
+
+    state = math.einsum(f"...{state_wire_indices}->...{final_state_wire_indices}", state)
+
+    return reshape_state_as_matrix(state, len(wires))
+
+
+def calculate_probability(
+    measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Find the probability of measuring states.
+
+    Args:
+        measurementprocess (StateMeasurement): measurement to apply to the state.
+        state (TensorLike): state to apply the measurement to.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        TensorLike: the probability of the state being in each measurable state.
+    """
+    for op in measurementprocess.diagonalizing_gates():
+        state = apply_operation(op, state, is_state_batched=is_state_batched)
+
+    num_state_wires = get_num_wires(state, is_state_batched)
+
+    # probs are diagonal elements
+    # stacking list since diagonal function axis selection parameter names
+    # are not consistent across interfaces
+    reshaped_state = reshape_state_as_matrix(state, num_state_wires)
+    if is_state_batched:
+        probs = math.real(math.stack([math.diagonal(dm) for dm in reshaped_state]))
+    else:
+        probs = math.real(math.diagonal(reshaped_state))
+
+    # if a probability is very small it may round to negative, undesirable.
+    # math.clip with None bounds breaks with tensorflow, using this instead:
+    probs = math.where(probs < 0, 0, probs)
+
+    if mp_wires := measurementprocess.wires:
+        expanded_shape = [QUDIT_DIM] * num_state_wires
+        new_shape = [QUDIT_DIM ** len(mp_wires)]
+        if is_state_batched:
+            batch_size = probs.shape[0]
+            expanded_shape.insert(0, batch_size)
+            new_shape.insert(0, batch_size)
+        wires_to_trace = tuple(
+            x + is_state_batched for x in range(num_state_wires) if x not in mp_wires
+        )
+
+        expanded_probs = math.reshape(probs, expanded_shape)
+        summed_probs = math.sum(expanded_probs, axis=wires_to_trace)
+        return math.reshape(summed_probs, new_shape)
+
+    return probs
+
+
+def calculate_variance(
+    measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Find variance of observable.
+
+    Args:
+        measurementprocess (StateMeasurement): measurement to apply to the state.
+        state (TensorLike): state to apply the measurement to.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        TensorLike: the variance of the observable wrt the state.
+    """
+    probs = calculate_probability(measurementprocess, state, is_state_batched)
+    eigvals = math.asarray(measurementprocess.eigvals(), dtype="float64")
+    # In case of broadcasting, `probs` has two axes and these are a matrix-vector products
+    return math.dot(probs, (eigvals**2)) - math.dot(probs, eigvals) ** 2
+
+
+def calculate_expval_sum_of_terms(
+    measurementprocess: ExpectationMP,
+    state: TensorLike,
+    is_state_batched: bool = False,
+) -> TensorLike:
+    """Measure the expectation value of the state when the measured observable is a ``Hamiltonian`` or ``Sum``
+    and it must be backpropagation compatible.
+
+    Args:
+        measurementprocess (ExpectationMP): measurement process to apply to the state.
+        state (TensorLike): the state to measure.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        TensorLike: the expectation value of the sum of Hamiltonian observable wrt the state.
+    """
+    if isinstance(measurementprocess.obs, Sum):
+        # Recursively call measure on each term, so that the best measurement method can
+        # be used for each term
+        return sum(
+            measure(ExpectationMP(term), state, is_state_batched=is_state_batched)
+            for term in measurementprocess.obs
+        )
+    # else hamiltonian
+    return sum(
+        c * measure(ExpectationMP(t), state, is_state_batched=is_state_batched)
+        for c, t in zip(*measurementprocess.obs.terms())
+    )
+
+
+# pylint: disable=too-many-return-statements
+def get_measurement_function(
+    measurementprocess: MeasurementProcess,
+) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
+    """Get the appropriate method for performing a measurement.
+
+    Args:
+        measurementprocess (MeasurementProcess): measurement process to apply to the state.
+        state (TensorLike): the state to measure.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        Callable: function that returns the measurement result.
+    """
+    if isinstance(measurementprocess, StateMeasurement):
+        if isinstance(measurementprocess, ExpectationMP):
+            # TODO add faster methods
+            # TODO add support for sparce Hamiltonians
+            if isinstance(measurementprocess.obs, (Hamiltonian, Sum)):
+                return calculate_expval_sum_of_terms
+            if measurementprocess.obs.has_matrix:
+                return calculate_expval
+        if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
+            if isinstance(measurementprocess, StateMP):
+                return calculate_reduced_density_matrix
+            if isinstance(measurementprocess, ProbabilityMP):
+                return calculate_probability
+            if isinstance(measurementprocess, VarianceMP):
+                return calculate_variance
+
+    raise NotImplementedError
+
+
+def measure(
+    measurementprocess: MeasurementProcess, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Apply a measurement process to a state.
+
+    Args:
+        measurementprocess (MeasurementProcess): measurement process to apply to the state.
+        state (TensorLike): the state to measure.
+        is_state_batched (bool): whether the state is batched or not.
+
+    Returns:
+        Tensorlike: the result of the measurement process being applied to the state.
+    """
+    return get_measurement_function(measurementprocess)(measurementprocess, state, is_state_batched)
diff --git a/pennylane/devices/qutrit_mixed/utils.py b/pennylane/devices/qutrit_mixed/utils.py
index c80ef8dc1..3bbf47b30 100644
--- a/pennylane/devices/qutrit_mixed/utils.py
+++ b/pennylane/devices/qutrit_mixed/utils.py
@@ -15,6 +15,7 @@
 import functools
 from string import ascii_letters as alphabet
 import pennylane as qml
+from pennylane import math
 from pennylane import numpy as np
 
 alphabet_array = np.array(list(alphabet))
@@ -33,7 +34,7 @@ def get_einsum_mapping(
         is_state_batched (bool): Boolean representing whether the state is batched or not
 
     Returns:
-        str: indices mapping that defines the einsum
+        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)
@@ -68,6 +69,36 @@ def get_einsum_mapping(
     )
 
 
+def reshape_state_as_matrix(state, num_wires):
+    """Given a non-flat, potentially batched state, flatten it to square matrix or matrices if batched.
+
+    Args:
+        state (TensorLike): A state that needs to be reshaped to a square matrix or matrices if batched
+        num_wires (int): The number of wires the state represents
+
+    Returns:
+        Tensorlike: A reshaped, square state, with an extra batch dimension if necessary
+    """
+    dim = QUDIT_DIM**num_wires
+    batch_size = math.get_batch_size(state, ((QUDIT_DIM,) * (num_wires * 2)), dim**2)
+    shape = (batch_size, dim, dim) if batch_size is not None else (dim, dim)
+    return math.reshape(state, shape)
+
+
+def get_num_wires(state, is_state_batched: bool = False):
+    """Finds the number of wires associated with a state
+
+    Args:
+        state (TensorLike): A device compatible state that may or may not be batched
+        is_state_batched (int): Boolean representing whether the state is batched or not
+
+    Returns:
+        int: Number of wires associated with state
+    """
+    len_row_plus_col = len(math.shape(state)) - is_state_batched
+    return int(len_row_plus_col / 2)
+
+
 def get_new_state_einsum_indices(old_indices, new_indices, state_indices):
     """Retrieves the einsum indices string for the new state
 
@@ -77,7 +108,7 @@ def get_new_state_einsum_indices(old_indices, new_indices, state_indices):
         state_indices (str): indices of the original state
 
     Returns:
-        str: the einsum indices of the new state
+        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]),


Test patch:
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py
new file mode 100644
index 000000000..d2801ffd2
--- /dev/null
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py
@@ -0,0 +1,474 @@
+# 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 measuring states in devices/qutrit_mixed."""
+
+from functools import reduce
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane import math
+
+from pennylane.devices.qutrit_mixed import measure
+from pennylane.devices.qutrit_mixed.measure import (
+    apply_observable_einsum,
+    get_measurement_function,
+    calculate_expval,
+    calculate_expval_sum_of_terms,
+    calculate_reduced_density_matrix,
+    calculate_probability,
+    calculate_variance,
+)
+
+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),
+]
+BATCH_SIZE = 2
+
+
+def get_expanded_op_mult_state(op, state):
+    """Finds the expanded matrix to multiply state by and multiplies it by a flattened state"""
+    num_qutrits = int(len(state.shape) / 2)
+    pre_wires_identity = np.eye(3 ** min(op.wires))
+    post_wires_identity = np.eye(3 ** ((num_qutrits - 1) - op.wires[-1]))
+
+    expanded_op = reduce(np.kron, (pre_wires_identity, op.matrix(), post_wires_identity))
+    flattened_state = state.reshape((3**num_qutrits,) * 2)
+    return expanded_op @ flattened_state
+
+
+def get_expval(op, state):
+    """Finds op@state and traces to find the expectation value of observable on the state"""
+    op_mult_state = get_expanded_op_mult_state(op, state)
+    return np.trace(op_mult_state)
+
+
+def test_probs_with_negative_on_diagonal():
+    """Test that if there is a negative on diagonal it is clipped to 0"""
+    state = np.array([[1 - 1e-4 + 0j, 0, 0], [0, -1e-4, 0], [0, 0, 1e-4]])
+    probs = measure(qml.probs(), state)
+    expected_probs = np.array([1 - 1e-4, 0, 1e-4])
+    assert np.allclose(probs, expected_probs)
+
+
+@pytest.mark.parametrize(
+    "mp", [qml.sample(), qml.counts(), qml.sample(wires=0), qml.counts(wires=0)]
+)
+class TestCurrentlyUnsupportedCases:
+    # pylint: disable=too-few-public-methods
+    def test_sample_based_observable(self, mp, two_qutrit_state):
+        """Test sample-only measurements raise a NotImplementedError."""
+        with pytest.raises(NotImplementedError):
+            _ = measure(mp, two_qutrit_state)
+
+
+class TestMeasurementDispatch:
+    """Test that get_measurement_function dispatchs to the correct place."""
+
+    def test_state_no_obs(self):
+        """Test that the correct internal function is used for a measurement process with no observables."""
+        # Test a case where state_measurement_process is used
+        mp1 = qml.state()
+        assert get_measurement_function(mp1) is calculate_reduced_density_matrix
+
+    def test_prod_calculate_expval_method(self):
+        """Test that the expectation value of a product uses the calculate expval method."""
+        prod = qml.prod(*(qml.GellMann(i, 1) for i in range(8)))
+        assert get_measurement_function(qml.expval(prod)) is calculate_expval
+
+    def test_hermitian_calculate_expval_method(self):
+        """Test that the expectation value of a hermitian uses the calculate expval method."""
+        mp = qml.expval(qml.THermitian(np.eye(3), wires=0))
+        assert get_measurement_function(mp) is calculate_expval
+
+    def test_hamiltonian_sum_of_terms(self):
+        """Check that the sum of terms method is used when Hamiltonian."""
+        H = qml.Hamiltonian([2], [qml.GellMann(0, 1)])
+        assert get_measurement_function(qml.expval(H)) is calculate_expval_sum_of_terms
+
+    def test_sum_sum_of_terms(self):
+        """Check that the sum of terms method is used when sum of terms"""
+        S = qml.prod(*(qml.GellMann(i, 1) for i in range(8))) + qml.prod(
+            *(qml.GellMann(i, 2) for i in range(8))
+        )
+        assert get_measurement_function(qml.expval(S)) is calculate_expval_sum_of_terms
+
+    def test_probs_compute_probabilities(self):
+        """Check that compute probabilities method is used when probs"""
+        assert get_measurement_function(qml.probs()) is calculate_probability
+
+    def test_var_compute_variance(self):
+        """Check that the compute variance method is used when variance"""
+        obs = qml.GellMann(0, 1)
+        assert get_measurement_function(qml.var(obs)) is calculate_variance
+
+
+@pytest.mark.parametrize(
+    "obs",
+    [
+        qml.GellMann(2, 2),
+        qml.GellMann(1, 8),
+        qml.GellMann(0, 5),
+        (qml.GellMann(0, 5) @ qml.GellMann(1, 8)),
+    ],
+)
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+class TestApplyObservableEinsum:
+    """Tests that observables are applied correctly for calculate expval method."""
+
+    num_qutrits = 3
+    dims = (3**num_qutrits, 3**num_qutrits)
+
+    def test_apply_observable_einsum(self, obs, three_qutrit_state, ml_framework):
+        """Tests that unbatched observables are applied correctly to a unbatched state."""
+        res = apply_observable_einsum(
+            obs, qml.math.asarray(three_qutrit_state, like=ml_framework), is_state_batched=False
+        )
+        expected = get_expanded_op_mult_state(obs, three_qutrit_state)
+        expected = expected.reshape([3] * (3 * 2))
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    def test_apply_observable_einsum_batched(self, obs, three_qutrit_batched_state, ml_framework):
+        """Tests that unbatched observables are applied correctly to a batched state."""
+        res = apply_observable_einsum(
+            obs,
+            qml.math.asarray(three_qutrit_batched_state, like=ml_framework),
+            is_state_batched=True,
+        )
+        expected = [
+            get_expanded_op_mult_state(obs, state).reshape([3] * (3 * 2))
+            for state in three_qutrit_batched_state
+        ]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+
+class TestMeasurements:
+    """Test that measurements on unbatched states work as expected."""
+
+    @pytest.mark.parametrize(
+        "measurement, get_expected",
+        [
+            (qml.state(), lambda x: math.reshape(x, newshape=(9, 9))),
+            (qml.density_matrix(wires=0), lambda x: math.trace(x, axis1=1, axis2=3)),
+            (
+                qml.probs(wires=[0]),
+                lambda x: math.real(math.diag(math.trace(x, axis1=1, axis2=3))),
+            ),
+            (
+                qml.probs(),
+                lambda x: math.real(math.diag(x.reshape(9, 9))),
+            ),
+        ],
+    )
+    def test_state_measurement_no_obs(self, measurement, get_expected, two_qutrit_state):
+        """Test that state measurements with no observable work as expected."""
+        res = measure(measurement, two_qutrit_state)
+        expected = get_expected(two_qutrit_state)
+
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "coeffs, observables",
+        [
+            ([-0.5, 2], [qml.GellMann(0, 7), qml.GellMann(1, 8)]),
+            ([-0.3, 1], [qml.GellMann(0, 2), qml.GellMann(0, 4)]),
+            ([-0.45, 2.6], [qml.GellMann(1, 6), qml.GellMann(1, 3)]),
+        ],
+    )
+    def test_hamiltonian_expval(self, coeffs, observables, two_qutrit_state):
+        """Test that measurements of hamiltonian work correctly."""
+
+        obs = qml.Hamiltonian(coeffs, observables)
+        res = measure(qml.expval(obs), two_qutrit_state)
+
+        expected = 0
+        for i, coeff in enumerate(coeffs):
+            expected += coeff * get_expval(observables[i], two_qutrit_state)
+
+        assert np.isclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "observable",
+        [
+            qml.THermitian(
+                -0.5 * qml.GellMann(0, 7).matrix() + 2 * qml.GellMann(0, 8).matrix(), wires=0
+            ),
+            qml.THermitian(
+                -0.55 * qml.GellMann(1, 4).matrix() + 2.4 * qml.GellMann(1, 5).matrix(), wires=1
+            ),
+        ],
+    )
+    def test_hermitian_expval(self, observable, two_qutrit_state):
+        """Test that measurements of qutrit hermitian work correctly."""
+        res = measure(qml.expval(observable), two_qutrit_state)
+        expected = get_expval(observable, two_qutrit_state)
+
+        assert np.isclose(res, expected)
+
+    def test_sum_expval_tensor_contraction(self):
+        """Test that `Sum` expectation values are correct when tensor contraction
+        is used for computation."""
+        summands = (qml.prod(qml.GellMann(i, 2), qml.GellMann(i + 1, 3)) for i in range(4))
+        obs = qml.sum(*summands)
+
+        @qml.qnode(qml.device("default.qutrit", wires=5))
+        def find_state(x):
+            for i in range(5):
+                qml.TRX(x, wires=i)
+            return qml.state()
+
+        rots = [0.123, 0.321]
+        schmidts = [0.7, 0.3]
+        state = np.zeros([3] * (2 * 5), dtype=complex)
+
+        for i in range(2):
+            vec = find_state(rots[i])
+            state += schmidts[i] * np.outer(vec, np.conj(vec)).reshape([3] * (2 * 5))
+
+        res = measure(qml.expval(obs), state)
+        expected = 0
+        for schmidt, theta in zip(schmidts, rots):
+            expected += schmidt * (4 * (-np.sin(theta) * np.cos(theta)))
+
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "observable",
+        [
+            qml.GellMann(0, 1),
+            qml.GellMann(1, 6),
+            (qml.GellMann(0, 6) @ qml.GellMann(1, 2)),
+        ],
+    )
+    def test_variance_measurement(self, observable, two_qutrit_state):
+        """Test that variance measurements work as expected."""
+        res = measure(qml.var(observable), two_qutrit_state)
+
+        expval_obs = get_expval(observable, two_qutrit_state)
+
+        obs_squared = qml.prod(observable, observable)
+        expval_of_squared_obs = get_expval(obs_squared, two_qutrit_state)
+
+        expected = expval_of_squared_obs - expval_obs**2
+        assert np.allclose(res, expected)
+
+
+@pytest.mark.parametrize(
+    "obs",
+    [
+        qml.Hamiltonian([-0.5, 2], [qml.GellMann(0, 5), qml.GellMann(0, 3)]),
+        qml.THermitian(
+            -0.5 * qml.GellMann(0, 5).matrix() + 2 * qml.GellMann(0, 3).matrix(), wires=0
+        ),
+    ],
+)
+class TestExpValAnalytical:
+    def test_expval_pure_state(self, obs):
+        """Test that measurements work on pure states as expected from analytical calculation."""
+        # Create TRX[0,2](0.246)TRX[0,1](0.246)|0> state
+        state_vector = np.array([np.cos(0.123) ** 2, -1j * np.sin(0.123), -1j * np.sin(0.246) / 2])
+        state = np.outer(state_vector, np.conj(state_vector))
+        res = measure(qml.expval(obs), state)
+
+        expected = 0.5 * (np.sin(0.246) * np.cos(0.123) ** 2) + 2 * (
+            np.cos(0.123) ** 4 - np.sin(0.123) ** 2
+        )
+        assert np.allclose(res, expected)
+
+    def test_expval_mixed_state(self, obs):
+        """Test that measurements work on mixed states as expected from analytical calculation."""
+        # Create TRX[0,1](0.246)|0> state mixed with TRX[0,2](0.246)|0>
+        state_vector_one = np.array([np.cos(0.123), -1j * np.sin(0.123), 0])
+        state_one = np.outer(state_vector_one, np.conj(state_vector_one))
+
+        state_vector_two = np.array([np.cos(0.123), 0, -1j * np.sin(0.123)])
+        state_two = np.outer(state_vector_two, np.conj(state_vector_two))
+
+        state = (0.33 * state_one) + (0.67 * state_two)
+
+        res = measure(qml.expval(obs), state)
+        expected_pure_state_one = 2 * np.cos(0.246)
+        expected_pure_state_two = 0.5 * np.sin(0.246) + 2 * np.cos(0.123) ** 2
+        expected = (0.33 * expected_pure_state_one) + (0.67 * expected_pure_state_two)
+        assert np.allclose(res, expected)
+
+
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+class TestBroadcasting:
+    """Test that measurements work when the state has a batch dim"""
+
+    num_qutrits = 2
+
+    @pytest.mark.parametrize(
+        "measurement, get_expected",
+        [
+            (qml.state(), lambda x: math.reshape(x, newshape=(BATCH_SIZE, 9, 9))),
+            (
+                qml.density_matrix(wires=[0, 1]),
+                lambda x: math.reshape(x, newshape=(BATCH_SIZE, 9, 9)),
+            ),
+            (qml.density_matrix(wires=[1]), lambda x: math.trace(x, axis1=1, axis2=3)),
+        ],
+    )
+    def test_state_measurement(
+        self, measurement, get_expected, ml_framework, two_qutrit_batched_state
+    ):
+        """Test that state measurements work on broadcasted state"""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        res = measure(measurement, initial_state, is_state_batched=True)
+        expected = get_expected(two_qutrit_batched_state)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "measurement, matrix_transform",
+        [
+            (qml.probs(wires=[0, 1]), lambda x: math.reshape(x, (2, 9, 9))),
+            (qml.probs(wires=[0]), lambda x: math.trace(x, axis1=2, axis2=4)),
+        ],
+    )
+    def test_probs_measurement(
+        self, measurement, matrix_transform, ml_framework, two_qutrit_batched_state
+    ):
+        """Test that probability measurements work on broadcasted state"""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        res = measure(measurement, initial_state, is_state_batched=True)
+
+        transformed_state = matrix_transform(two_qutrit_batched_state)
+
+        expected = []
+        for i in range(BATCH_SIZE):
+            expected.append(math.diag(transformed_state[i]))
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert np.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "observable",
+        [
+            qml.GellMann(1, 3),
+            qml.GellMann(0, 6),
+            qml.prod(qml.GellMann(0, 2), qml.GellMann(1, 3)),
+            qml.THermitian(
+                np.array(
+                    [
+                        [1.37770247 + 0.0j, 0.60335894 - 0.10889947j, 0.98223403 - 0.94429544j],
+                        [0.60335894 + 0.10889947j, 0.90178212 + 0.0j, 0.45529663 - 0.03054001j],
+                        [0.98223403 + 0.94429544j, 0.45529663 + 0.03054001j, 0.37721683 + 0.0j],
+                    ]
+                ),
+                wires=1,
+            ),
+        ],
+    )
+    def test_expval_measurement(self, observable, ml_framework, two_qutrit_batched_state):
+        """Test that expval measurements work on broadcasted state"""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        res = measure(qml.expval(observable), initial_state, is_state_batched=True)
+
+        expected = [get_expval(observable, two_qutrit_batched_state[i]) for i in range(BATCH_SIZE)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "observable",
+        [
+            qml.sum((2 * qml.GellMann(1, 1)), (0.4 * qml.GellMann(0, 6))),
+            qml.sum((2.4 * qml.GellMann(0, 3)), (0.2 * qml.GellMann(0, 7))),
+            qml.sum((0.9 * qml.GellMann(1, 5)), (1.2 * qml.GellMann(1, 8))),
+        ],
+    )
+    def test_expval_sum_measurement(self, observable, ml_framework, two_qutrit_batched_state):
+        """Test that expval Sum measurements work on broadcasted state"""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        res = measure(qml.expval(observable), initial_state, is_state_batched=True)
+
+        expanded_mat = np.zeros((9, 9), dtype=complex)
+        for summand in observable:
+            mat = summand.matrix()
+            expanded_mat += (
+                np.kron(np.eye(3), mat) if summand.wires[0] == 1 else np.kron(mat, np.eye(3))
+            )
+
+        expected = []
+        for i in range(BATCH_SIZE):
+            expval_sum = 0.0
+            for summand in observable:
+                expval_sum += get_expval(summand, two_qutrit_batched_state[i])
+            expected.append(expval_sum)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    def test_expval_hamiltonian_measurement(self, ml_framework, two_qutrit_batched_state):
+        """Test that expval Hamiltonian measurements work on broadcasted state"""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        observables = [qml.GellMann(1, 1), qml.GellMann(0, 6)]
+        coeffs = [2, 0.4]
+        observable = qml.Hamiltonian(coeffs, observables)
+        res = measure(qml.expval(observable), initial_state, is_state_batched=True)
+
+        expanded_mat = np.zeros((9, 9), dtype=complex)
+        for coeff, summand in zip(coeffs, observables):
+            mat = summand.matrix()
+            expanded_mat += coeff * (
+                np.kron(np.eye(3), mat) if summand.wires[0] == 1 else np.kron(mat, np.eye(3))
+            )
+
+        expected = []
+        for i in range(BATCH_SIZE):
+            expval_sum = 0.0
+            for coeff, summand in zip(coeffs, observables):
+                expval_sum += coeff * get_expval(summand, two_qutrit_batched_state[i])
+            expected.append(expval_sum)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize(
+        "observable",
+        [
+            qml.GellMann(0, 1),
+            qml.GellMann(1, 6),
+            (qml.GellMann(0, 6) @ qml.GellMann(1, 2)),
+        ],
+    )
+    def test_variance_measurement(self, observable, ml_framework, two_qutrit_batched_state):
+        """Test that variance measurements work on broadcasted state."""
+        initial_state = math.asarray(two_qutrit_batched_state, like=ml_framework)
+        res = measure(qml.var(observable), initial_state, is_state_batched=True)
+
+        obs_squared = qml.prod(observable, observable)
+
+        expected = []
+        for state in two_qutrit_batched_state:
+            expval_obs = get_expval(observable, state)
+            expval_of_squared_obs = get_expval(obs_squared, state)
+            expected.append(expval_of_squared_obs - expval_obs**2)
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert np.allclose(res, expected)
+
+
+# TODO TestSumOfTermsDifferentiability in future PR (with other differentiabilty tests)