Repository: PennyLaneAI/pennylane
Pull request #5694 (2024-05-16 13:21:47)

Support Angstrom units with Molecule class and molecular_hamiltonian function
**Context:**
Support Angstrom units along with Bohr

**Description of the Change:**
Molecule class and molecular_hamiltonian function now works with both Bohr and Angstrom units.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


ddhawan11: [sc-62838]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5694?dropdown=coverage&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:
> :exclamation: No coverage uploaded for pull request base (`molecule_central_object@615cba6`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).


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


```diff
@@                    Coverage Diff                     @@
##             molecule_central_object    #5694   +/-   ##
==========================================================
  Coverage                           ?   99.69%           
==========================================================
  Files                              ?      412           
  Lines                              ?    38316           
  Branches                           ?        0           
==========================================================
  Hits                               ?    38198           
  Misses                             ?      118           
  Partials                           ?        0           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5694?dropdown=coverage&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/.github/workflows/install_deps/action.yml b/.github/workflows/install_deps/action.yml
index 61a0f763d..0e6fc6d27 100644
--- a/.github/workflows/install_deps/action.yml
+++ b/.github/workflows/install_deps/action.yml
@@ -82,7 +82,7 @@ runs:
       if: inputs.install_jax == 'true'
       env:
         JAX_VERSION: ${{ inputs.jax_version != '' && format('=={0}', inputs.jax_version) || '' }}
-      run: pip install "jax${{ env.JAX_VERSION}}" "jaxlib${{ env.JAX_VERSION }}"
+      run: pip install "jax${{ env.JAX_VERSION}}" "jaxlib${{ env.JAX_VERSION }}" scipy~=1.12.0
 
     - name: Install additional PIP packages
       shell: bash
diff --git a/doc/introduction/circuits.rst b/doc/introduction/circuits.rst
index bb637fc1f..74f8ff8bf 100644
--- a/doc/introduction/circuits.rst
+++ b/doc/introduction/circuits.rst
@@ -88,7 +88,7 @@ instantiated using the :func:`device <pennylane.device>` loader.
     dev = qml.device('default.qubit', wires=2, shots=1000)
 
 PennyLane offers some basic devices such as the ``'default.qubit'``, ``'default.mixed'``, ``lightning.qubit``,
-``'default.gaussian'``, and ``'default.clifford'`` simulators; additional devices can be installed as plugins
+``'default.gaussian'``, ``'default.clifford'``, and ``'default.tensor'`` simulators; additional devices can be installed as plugins
 (see `available plugins <https://pennylane.ai/plugins.html>`_ for more details). Note that the
 choice of a device significantly determines the speed of your computation, as well as
 the available options that can be passed to the device loader.
diff --git a/doc/introduction/operations.rst b/doc/introduction/operations.rst
index 970e7b4ec..b4b9b0d3c 100644
--- a/doc/introduction/operations.rst
+++ b/doc/introduction/operations.rst
@@ -507,6 +507,7 @@ Qutrit noisy channels
     :nosignatures:
 
     ~pennylane.QutritDepolarizingChannel
+    ~pennylane.QutritAmplitudeDamping
 
 :html:`</div>`
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b5dfab96a..a0418b58e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,8 +4,17 @@
 
 <h3>New features since last release</h3>
 
+* The `default.tensor` device is introduced to perform tensor network simulation of a quantum circuit.
+  [(#5699)](https://github.com/PennyLaneAI/pennylane/pull/5699)
+
 <h3>Improvements 🛠</h3>
 
+* `ctrl` now works with tuple-valued `control_values` when applied to any already controlled operation.
+  [(#5725)](https://github.com/PennyLaneAI/pennylane/pull/5725)
+
+* Add support for 3 new pytest markers: `unit`, `integration` and `system`.
+  [(#5517)](https://github.com/PennyLaneAI/pennylane/pull/5517)
+
 * The sorting order of parameter-shift terms is now guaranteed to resolve ties in the absolute value with the sign of the shifts.
   [(#5582)](https://github.com/PennyLaneAI/pennylane/pull/5582)
 
@@ -71,6 +80,10 @@
   `qml.devices.Device`, which follows the new device API.
   [(#5581)](https://github.com/PennyLaneAI/pennylane/pull/5581)
 
+* `qml.qchem.Molecule` and qchem functions that take `Molecule` object as an argument now work with coordinates
+  provided in units Angstrom along with Bohr.
+  [(#5694)](https://github.com/PennyLaneAI/pennylane/pull/5694)
+
 * The `dtype` for `eigvals` of `X`, `Y`, `Z` and `Hadamard` is changed from `int` to `float`, making them 
   consistent with the other observables. The `dtype` of the returned values when sampling these observables 
   (e.g. `qml.sample(X(0))`) is also changed to `float`. 
@@ -116,6 +129,9 @@
   returning a list of `QuantumTape`s and a post-processing function instead of simply the transformed circuit.
   [(#5693)](https://github.com/PennyLaneAI/pennylane/pull/5693)
 
+* `qml.QutritAmplitudeDamping` channel has been added, allowing for noise processes modelled by amplitude damping to be simulated on the `default.qutrit.mixed` device.
+  [(#5503)](https://github.com/PennyLaneAI/pennylane/pull/5503)
+
 <h3>Deprecations 👋</h3>
 
 * The `simplify` argument in `qml.Hamiltonian` and `qml.ops.LinearCombination` is deprecated. 
@@ -132,6 +148,12 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* The decomposition of `StronglyEntanglingLayers` is now compatible with broadcasting.
+  [(#5716)](https://github.com/PennyLaneAI/pennylane/pull/5716)
+
+* `qml.cond` can now be applied to `ControlledOp` operations when deferring measurements.
+  [(#5725)](https://github.com/PennyLaneAI/pennylane/pull/5725)
+
 * The legacy `Tensor` class can now handle a `Projector` with abstract tracer input.
   [(#5720)](https://github.com/PennyLaneAI/pennylane/pull/5720)
 
@@ -172,5 +194,6 @@ Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,
 Vincent Michaud-Rioux,
+Lee James O'Riordan,
 Kenya Sakka,
 David Wierichs.
diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py
index ecbb9ef5f..f82976388 100644
--- a/pennylane/devices/__init__.py
+++ b/pennylane/devices/__init__.py
@@ -158,4 +158,5 @@ from .default_mixed import DefaultMixed
 from .default_clifford import DefaultClifford
 from .null_qubit import NullQubit
 from .default_qutrit_mixed import DefaultQutritMixed
+from .default_tensor import DefaultTensor
 from .._device import Device as LegacyDevice
diff --git a/pennylane/devices/default_tensor.py b/pennylane/devices/default_tensor.py
new file mode 100644
index 000000000..5c850329f
--- /dev/null
+++ b/pennylane/devices/default_tensor.py
@@ -0,0 +1,595 @@
+# 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.
+"""
+This module contains the default.tensor device to perform tensor network simulation of a quantum circuit using ``quimb``.
+"""
+import copy
+from dataclasses import replace
+from numbers import Number
+from typing import Callable, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+import pennylane as qml
+from pennylane.devices import DefaultExecutionConfig, Device, ExecutionConfig
+from pennylane.devices.modifiers import simulator_tracking, single_tape_support
+from pennylane.devices.preprocess import (
+    decompose,
+    validate_device_wires,
+    validate_measurements,
+    validate_observables,
+)
+from pennylane.measurements import ExpectationMP, MeasurementProcess, StateMeasurement, VarianceMP
+from pennylane.tape import QuantumScript, QuantumTape
+from pennylane.transforms.core import TransformProgram
+from pennylane.typing import Result, ResultBatch, TensorLike
+
+Result_or_ResultBatch = Union[Result, ResultBatch]
+QuantumTapeBatch = Sequence[QuantumTape]
+QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
+PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
+
+has_quimb = True
+try:
+    import quimb.tensor as qtn
+except (ModuleNotFoundError, ImportError) as import_error:  # pragma: no cover
+    has_quimb = False
+
+_operations = frozenset(
+    {
+        "Identity",
+        "QubitUnitary",
+        "ControlledQubitUnitary",
+        "MultiControlledX",
+        "DiagonalQubitUnitary",
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "MultiRZ",
+        "GlobalPhase",
+        "Hadamard",
+        "S",
+        "T",
+        "SX",
+        "CNOT",
+        "SWAP",
+        "ISWAP",
+        "PSWAP",
+        "SISWAP",
+        "SQISW",
+        "CSWAP",
+        "Toffoli",
+        "CY",
+        "CZ",
+        "PhaseShift",
+        "ControlledPhaseShift",
+        "CPhase",
+        "RX",
+        "RY",
+        "RZ",
+        "Rot",
+        "CRX",
+        "CRY",
+        "CRZ",
+        "CRot",
+        "IsingXX",
+        "IsingYY",
+        "IsingZZ",
+        "IsingXY",
+        "SingleExcitation",
+        "SingleExcitationPlus",
+        "SingleExcitationMinus",
+        "DoubleExcitation",
+        "QubitCarry",
+        "QubitSum",
+        "OrbitalRotation",
+        "QFT",
+        "ECR",
+        "BlockEncode",
+    }
+)
+# The set of supported operations.
+
+
+_observables = frozenset(
+    {
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "Hadamard",
+        "Hermitian",
+        "Identity",
+        "Projector",
+        "SparseHamiltonian",
+        "Hamiltonian",
+        "LinearCombination",
+        "Sum",
+        "SProd",
+        "Prod",
+        "Exp",
+    }
+)
+# The set of supported observables.
+
+_methods = frozenset({"mps"})
+# The set of supported methods.
+
+
+def accepted_methods(method: str) -> bool:
+    """A function that determines whether or not a method is supported by ``default.tensor``."""
+    return method in _methods
+
+
+def stopping_condition(op: qml.operation.Operator) -> bool:
+    """A function that determines if an operation is supported by ``default.tensor``."""
+    return op.name in _operations
+
+
+def accepted_observables(obs: qml.operation.Operator) -> bool:
+    """A function that determines if an observable is supported by ``default.tensor``."""
+    return obs.name in _observables
+
+
+@simulator_tracking
+@single_tape_support
+class DefaultTensor(Device):
+    """A PennyLane device to perform tensor network operations on a quantum circuit using
+    `quimb <https://github.com/jcmgray/quimb/>`_.
+
+    Args:
+        wires (int, Iterable[Number, str]): Number of wires present on the device, or iterable that
+            contains unique labels for the wires as numbers (i.e., ``[-1, 0, 2]``) or strings
+            (``['aux_wire', 'q1', 'q2']``).
+        method (str): Supported method. Currently, only ``"mps"`` is supported.
+        dtype (type): Datatype for the tensor representation. Must be one of ``np.complex64`` or ``np.complex128``.
+            Default is ``np.complex128``.
+        **kwargs: keyword arguments. The following options are currently supported:
+
+            ``max_bond_dim`` (int): Maximum bond dimension for the MPS simulator.
+                It corresponds to the number of Schmidt coefficients retained at the end of the SVD algorithm when applying gates. Default is ``None``.
+            ``cutoff`` (float): Truncation threshold for the Schmidt coefficients in a MPS simulator. Default is ``np.finfo(dtype).eps``.
+            ``contract`` (str): The contraction method for applying gates. It can be either ``auto-mps`` or ``nonlocal``.
+                ``nonlocal`` turns each gate into a MPO and applies it directly to the MPS, while ``auto-mps`` swaps nonlocal qubits in 2-qubit gates to be next
+                    to each other before applying the gate, then swaps them back. Default is ``auto-mps``.
+    """
+
+    # pylint: disable=too-many-instance-attributes
+
+    # So far we just consider the options for MPS simulator
+    _device_options = (
+        "contract",
+        "cutoff",
+        "dtype",
+        "method",
+        "max_bond_dim",
+    )
+
+    def __init__(
+        self,
+        wires,
+        *,
+        method="mps",
+        dtype=np.complex128,
+        **kwargs,
+    ) -> None:
+
+        if wires is None:
+            raise TypeError("Wires must be provided for the default.tensor device.")
+
+        if not has_quimb:
+            raise ImportError(
+                "This feature requires quimb, a library for tensor network manipulations. "
+                "It can be installed with:\n\npip install quimb"
+            )  # pragma: no cover
+
+        if not accepted_methods(method):
+            raise ValueError(
+                f"Unsupported method: {method}. The only currently supported method is mps."
+            )
+
+        if dtype not in [np.complex64, np.complex128]:
+            raise TypeError(
+                f"Unsupported type: {dtype}. Supported types are np.complex64 and np.complex128."
+            )
+
+        super().__init__(wires=wires, shots=None)
+
+        self._method = method
+        self._dtype = dtype
+
+        # options for MPS
+        self._max_bond_dim = kwargs.get("max_bond_dim", None)
+        self._cutoff = kwargs.get("cutoff", np.finfo(self._dtype).eps)
+        self._contract = kwargs.get("contract", "auto-mps")
+
+        device_options = self._setup_execution_config().device_options
+
+        self._init_state_opts = {
+            "binary": "0" * (len(self._wires) if self._wires else 1),
+            "dtype": self._dtype.__name__,
+            "tags": [str(l) for l in self._wires.labels] if self._wires else None,
+        }
+
+        self._gate_opts = {
+            "parametrize": None,
+            "contract": device_options["contract"],
+            "cutoff": device_options["cutoff"],
+            "max_bond": device_options["max_bond_dim"],
+        }
+
+        self._expval_opts = {
+            "dtype": self._dtype.__name__,
+            "simplify_sequence": "ADCRS",
+            "simplify_atol": 0.0,
+        }
+
+        self._circuitMPS = qtn.CircuitMPS(psi0=self._initial_mps())
+
+        for arg in kwargs:
+            if arg not in self._device_options:
+                raise TypeError(
+                    f"Unexpected argument: {arg} during initialization of the default.tensor device."
+                )
+
+    @property
+    def name(self):
+        """The name of the device."""
+        return "default.tensor"
+
+    @property
+    def method(self):
+        """Supported method."""
+        return self._method
+
+    @property
+    def dtype(self):
+        """Tensor complex data type."""
+        return self._dtype
+
+    def _reset_state(self) -> None:
+        """Reset the MPS."""
+        self._circuitMPS = qtn.CircuitMPS(psi0=self._initial_mps())
+
+    def _initial_mps(self) -> "qtn.MatrixProductState":
+        r"""
+        Return an initial state to :math:`\ket{0}`.
+
+        Internally, it uses `quimb`'s `MPS_computational_state` method.
+
+        Returns:
+            MatrixProductState: The initial MPS of a circuit.
+        """
+        return qtn.MPS_computational_state(**self._init_state_opts)
+
+    def _setup_execution_config(
+        self, config: Optional[ExecutionConfig] = DefaultExecutionConfig
+    ) -> ExecutionConfig:
+        """
+        Update the execution config with choices for how the device should be used and the device options.
+        """
+        # TODO: add options for gradients next quarter
+
+        updated_values = {}
+
+        new_device_options = dict(config.device_options)
+        for option in self._device_options:
+            if option not in new_device_options:
+                new_device_options[option] = getattr(self, f"_{option}", None)
+
+        return replace(config, **updated_values, device_options=new_device_options)
+
+    def preprocess(
+        self,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """This function defines the device transform program to be applied and an updated device configuration.
+
+        Args:
+            execution_config (Union[ExecutionConfig, Sequence[ExecutionConfig]]): A data structure describing the
+                parameters needed to fully describe the execution.
+
+        Returns:
+            TransformProgram, ExecutionConfig: A transform program that when called returns :class:`~.QuantumTape`'s that the
+            device can natively execute as well as a postprocessing function to be called after execution, and a configuration
+            with unset specifications filled in.
+
+        This device currently:
+
+        * Does not support finite shots.
+        * Does not support derivatives.
+        * Does not support vector-Jacobian products.
+        """
+
+        config = self._setup_execution_config(execution_config)
+
+        program = TransformProgram()
+
+        program.add_transform(validate_measurements, name=self.name)
+        program.add_transform(validate_observables, accepted_observables, name=self.name)
+        program.add_transform(validate_device_wires, self._wires, name=self.name)
+        program.add_transform(
+            decompose,
+            stopping_condition=stopping_condition,
+            skip_initial_state_prep=True,
+            name=self.name,
+        )
+        program.add_transform(qml.transforms.broadcast_expand)
+
+        return program, config
+
+    def execute(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ) -> Result_or_ResultBatch:
+        """Execute a circuit or a batch of circuits and turn it into results.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the quantum circuits to be executed.
+            execution_config (ExecutionConfig): a data structure with additional information required for execution.
+
+        Returns:
+            TensorLike, tuple[TensorLike], tuple[tuple[TensorLike]]: A numeric result of the computation.
+        """
+
+        results = []
+        for circuit in circuits:
+            # we need to check if the wires of the circuit are compatible with the wires of the device
+            # since the initial tensor state is created with the wires of the device
+            if not self.wires.contains_wires(circuit.wires):
+                raise AttributeError(
+                    f"Circuit has wires {circuit.wires.tolist()}. "
+                    f"Tensor on device has wires {self.wires.tolist()}"
+                )
+            circuit = circuit.map_to_standard_wires()
+            results.append(self.simulate(circuit))
+
+        return tuple(results)
+
+    def simulate(self, circuit: QuantumScript) -> Result:
+        """Simulate a single quantum script. This function assumes that all operations provide matrices.
+
+        Args:
+            circuit (QuantumScript): The single circuit to simulate.
+
+        Returns:
+            Tuple[TensorLike]: The results of the simulation.
+        """
+
+        self._reset_state()
+
+        for op in circuit.operations:
+            self._apply_operation(op)
+
+        if not circuit.shots:
+            if len(circuit.measurements) == 1:
+                return self.measurement(circuit.measurements[0])
+            return tuple(self.measurement(mp) for mp in circuit.measurements)
+
+        raise NotImplementedError  # pragma: no cover
+
+    def _apply_operation(self, op: qml.operation.Operator) -> None:
+        """Apply a single operator to the circuit, keeping the state always in a MPS form.
+
+        Internally it uses `quimb`'s `apply_gate` method.
+
+        Args:
+            op (Operator): The operation to apply.
+        """
+
+        self._circuitMPS.apply_gate(op.matrix().astype(self._dtype), *op.wires, **self._gate_opts)
+
+    def measurement(self, measurementprocess: MeasurementProcess) -> TensorLike:
+        """Measure the measurement required by the circuit over the MPS.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement to apply to the state.
+
+        Returns:
+            TensorLike: the result of the measurement.
+        """
+
+        return self._get_measurement_function(measurementprocess)(measurementprocess)
+
+    def _get_measurement_function(
+        self, measurementprocess: MeasurementProcess
+    ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
+        """Get the appropriate method for performing a measurement.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement process to apply to the state
+
+        Returns:
+            Callable: function that returns the measurement result
+        """
+        if isinstance(measurementprocess, StateMeasurement):
+            if isinstance(measurementprocess, ExpectationMP):
+                return self.expval
+
+            if isinstance(measurementprocess, VarianceMP):
+                return self.var
+
+        raise NotImplementedError
+
+    def expval(self, measurementprocess: MeasurementProcess) -> float:
+        """Expectation value of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the MPS.
+
+        Returns:
+            Expectation value of the observable.
+        """
+
+        obs = measurementprocess.obs
+
+        result = self._local_expectation(obs.matrix(), tuple(obs.wires))
+
+        return result
+
+    def var(self, measurementprocess: MeasurementProcess) -> float:
+        """Variance of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the MPS.
+
+        Returns:
+            Variance of the observable.
+        """
+
+        obs = measurementprocess.obs
+
+        obs_mat = obs.matrix()
+        expect_op = self.expval(measurementprocess)
+        expect_squar_op = self._local_expectation(obs_mat @ obs_mat.conj().T, tuple(obs.wires))
+
+        return expect_squar_op - np.square(expect_op)
+
+    def _local_expectation(self, matrix, wires) -> float:
+        """Compute the local expectation value of a matrix on the MPS.
+
+        Internally, it uses `quimb`'s `local_expectation` method.
+
+        Args:
+            matrix (array): the matrix to compute the expectation value of.
+            wires (tuple[int]): the wires the matrix acts on.
+
+        Returns:
+            Local expectation value of the matrix on the MPS.
+        """
+
+        # We need to copy the MPS to avoid modifying the original state
+        qc = copy.deepcopy(self._circuitMPS)
+
+        exp_val = qc.local_expectation(
+            matrix,
+            wires,
+            **self._expval_opts,
+        )
+
+        return float(np.real(exp_val))
+
+    # pylint: disable=unused-argument
+    def supports_derivatives(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[qml.tape.QuantumTape] = None,
+    ) -> bool:
+        """Check whether or not derivatives are available for a given configuration and circuit.
+
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation.
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+
+        Returns:
+            Bool: Whether or not a derivative can be calculated provided the given information.
+
+        """
+        return False
+
+    def compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Calculate the Jacobian of either a single or a batch of circuits on the device.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuits to calculate derivatives for.
+            execution_config (ExecutionConfig): a data structure with all additional information required for execution.
+
+        Returns:
+            Tuple: The Jacobian for each trainable parameter.
+        """
+        raise NotImplementedError(
+            "The computation of derivatives has yet to be implemented for the default.tensor device."
+        )
+
+    def execute_and_compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Compute the results and Jacobians of circuits at the same time.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuits or batch of circuits.
+            execution_config (ExecutionConfig): a data structure with all additional information required for execution.
+
+        Returns:
+            tuple: A numeric result of the computation and the gradient.
+        """
+        raise NotImplementedError(
+            "The computation of derivatives has yet to be implemented for the default.tensor device."
+        )
+
+    # pylint: disable=unused-argument
+    def supports_vjp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom vector-Jacobian product.
+
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation.
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+
+        Returns:
+            Bool: Whether or not a derivative can be calculated provided the given information.
+        """
+        return False
+
+    def compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        r"""The vector-Jacobian product used in reverse-mode differentiation.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits.
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit. If the circuit has a single output, ``cotangents`` may be a single number, not an iterable
+                of numbers.
+            execution_config (ExecutionConfig): a data structure with all additional information required for execution.
+
+        Returns:
+            tensor-like: A numeric result of computing the vector-Jacobian product.
+        """
+        raise NotImplementedError(
+            "The computation of vector-Jacobian product has yet to be implemented for the default.tensor device."
+        )
+
+    def execute_and_compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Calculate both the results and the vector-Jacobian product used in reverse-mode differentiation.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits to be executed.
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit.
+            execution_config (ExecutionConfig): a data structure with all additional information required for execution.
+
+        Returns:
+            Tuple, Tuple: the result of executing the scripts and the numeric result of computing the vector-Jacobian product
+        """
+        raise NotImplementedError(
+            "The computation of vector-Jacobian product has yet to be implemented for the default.tensor device."
+        )
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index bf01f4fc8..ca47e76aa 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -295,20 +295,22 @@ def _no_trainable_grad(tape):
     return [], lambda _: tuple(qml.math.zeros([0]) for _ in range(len(tape.measurements)))
 
 
-def _swap_first_two_axes(grads, first_axis_size, second_axis_size):
+def _swap_first_two_axes(grads, first_axis_size, second_axis_size, squeeze=True):
     """Transpose the first two axes of an iterable of iterables, returning
-    a tuple of tuples."""
-    if first_axis_size == 1:
+    a tuple of tuples. Tuple version of ``np.moveaxis(grads, 0, 1)``"""
+    if first_axis_size == 1 and squeeze:
         return tuple(grads[0][i] for i in range(second_axis_size))
     return tuple(
         tuple(grads[j][i] for j in range(first_axis_size)) for i in range(second_axis_size)
     )
 
 
-def _move_first_axis_to_third_pos(grads, first_axis_size, second_axis_size, third_axis_size):
+def _move_first_axis_to_third_pos(
+    grads, first_axis_size, second_axis_size, third_axis_size, squeeze=True
+):
     """Transpose the first two axes of an iterable of iterables, returning
-    a tuple of tuples."""
-    if first_axis_size == 1:
+    a tuple of tuples. Tuple version of ``np.moveaxis(grads, 0, 2)``"""
+    if first_axis_size == 1 and squeeze:
         return tuple(
             tuple(grads[0][i][j] for j in range(third_axis_size)) for i in range(second_axis_size)
         )
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 9036378f7..b9f558c2e 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -24,7 +24,6 @@ import numpy as np
 
 import pennylane as qml
 from pennylane import transform
-from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
 from pennylane.measurements import VarianceMP
 
 from .finite_difference import finite_diff
@@ -36,7 +35,10 @@ from .general_shift_rules import (
 )
 from .gradient_transform import (
     _all_zero_grad,
+    _contract_qjac_with_cjac,
+    _move_first_axis_to_third_pos,
     _no_trainable_grad,
+    _swap_first_two_axes,
     assert_multimeasure_not_broadcasted,
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
@@ -152,7 +154,7 @@ def _single_meas_grad(result, coeffs, unshifted_coeff, r0):
 
     If an unshifted term exists, its contribution is added to the gradient.
     """
-    if isinstance(result, list) and result == []:
+    if isinstance(result, tuple) and result == ():
         if unshifted_coeff is None:
             raise ValueError(
                 "This gradient component neither has a shifted nor an unshifted component. "
@@ -173,19 +175,14 @@ def _single_meas_grad(result, coeffs, unshifted_coeff, r0):
 def _multi_meas_grad(res, coeffs, r0, unshifted_coeff, num_measurements):
     """Compute the gradient for multiple measurements by taking the linear combination of
     the coefficients and each measurement result."""
-    g = []
     if r0 is None:
         r0 = [None] * num_measurements
-    for meas_idx in range(num_measurements):
-        # Gather the measurement results
-        meas_result = [param_result[meas_idx] for param_result in res]
-        g_component = _single_meas_grad(meas_result, coeffs, unshifted_coeff, r0[meas_idx])
-        g.append(g_component)
+    if res == ():
+        res = tuple(() for _ in range(num_measurements))
+    return tuple(_single_meas_grad(r, coeffs, unshifted_coeff, r0_) for r, r0_ in zip(res, r0))
 
-    return tuple(g)
 
-
-def _evaluate_gradient(tape, res, data, r0):
+def _evaluate_gradient(tape_specs, res, data, r0, batch_size):
     """Use shifted tape evaluations and parameter-shift rule coefficients to evaluate
     a gradient result. If res is an empty list, ``r0`` and ``data[3]``, which is the
     coefficient for the unshifted term, must be given and not None.
@@ -197,42 +194,46 @@ def _evaluate_gradient(tape, res, data, r0):
     if fn is not None:
         res = fn(res)
 
-    num_measurements = len(tape.measurements)
+    *_, num_measurements, shots = tape_specs
+    scalar_shots, len_shot_vec = not shots.has_partitioned_shots, shots.num_copies
+
+    if r0 is None and not scalar_shots:
+        r0 = [None] * int(len_shot_vec)
 
     if num_measurements == 1:
-        if not tape.shots.has_partitioned_shots:
+        if scalar_shots:
+            # Res has axes (parameters,)
             return _single_meas_grad(res, coeffs, unshifted_coeff, r0)
-        g = []
-        len_shot_vec = tape.shots.num_copies
-        # Res has order of axes:
-        # 1. Number of parameters
-        # 2. Shot vector
-        if r0 is None:
-            r0 = [None] * int(len_shot_vec)
-        for i in range(len_shot_vec):
-            shot_comp_res = [r[i] for r in res]
-            shot_comp_res = _single_meas_grad(shot_comp_res, coeffs, unshifted_coeff, r0[i])
-            g.append(shot_comp_res)
-        return tuple(g)
-
-    g = []
-    if not tape.shots.has_partitioned_shots:
+        # Res has axes (parameters, shots) or with broadcasting (shots, parameters)
+        if batch_size is None:
+            # Move shots to first position
+            res = _swap_first_two_axes(res, len(res), len_shot_vec, squeeze=False)
+        # _single_meas_grad expects axis (parameters,), iterate over shot vector
+        return tuple(_single_meas_grad(r, coeffs, unshifted_coeff, r0_) for r, r0_ in zip(res, r0))
+
+    if scalar_shots:
+        # Res has axes (parameters, measurements) or with broadcasting (measurements, parameters)
+        if batch_size is None and len(res) > 0:
+            # Move measurements to first position
+            res = _swap_first_two_axes(res, len(res), num_measurements, squeeze=False)
+        # _multi_meas_grad expects axes (measurements, parameters)
         return _multi_meas_grad(res, coeffs, r0, unshifted_coeff, num_measurements)
 
-    # Res has order of axes:
-    # 1. Number of parameters
-    # 2. Shot vector
-    # 3. Number of measurements
-    for idx_shot_comp in range(tape.shots.num_copies):
-        single_shot_component_result = [
-            result_for_each_param[idx_shot_comp] for result_for_each_param in res
-        ]
-        multi_meas_grad = _multi_meas_grad(
-            single_shot_component_result, coeffs, r0, unshifted_coeff, num_measurements
-        )
-        g.append(multi_meas_grad)
-
-    return tuple(g)
+    # Res has axes (parameters, shots, measurements)
+    # or with broadcasting (shots, measurements, parameters)
+    if batch_size is None:
+        if len(res) > 0:
+            # Move first axis (parameters) to last position
+            res = _move_first_axis_to_third_pos(
+                res, len(res), len_shot_vec, num_measurements, squeeze=False
+            )
+        else:
+            res = (() for _ in range(len_shot_vec))
+    # _multi_meas_grad expects (measurements, parameters), so we iterate over shot vector
+    return tuple(
+        _multi_meas_grad(r, coeffs, r0_, unshifted_coeff, num_measurements)
+        for r, r0_ in zip(res, r0)
+    )
 
 
 def _get_operation_recipe(tape, t_idx, shifts, order=1):
@@ -424,14 +425,14 @@ def expval_param_shift(
                     grads.append(None)
                     continue
                 # The gradient for this parameter is computed from r0 alone.
-                g = _evaluate_gradient(tape, [], data, r0)
+                g = _evaluate_gradient(tape_specs, (), data, r0, batch_size)
                 grads.append(g)
                 continue
 
             res = results[start : start + num_tapes] if batch_size is None else results[start]
             start = start + num_tapes
 
-            g = _evaluate_gradient(tape, res, data, r0)
+            g = _evaluate_gradient(tape_specs, res, data, r0, batch_size)
             grads.append(g)
 
         # g will have been defined at least once (because otherwise all gradients would have
diff --git a/pennylane/ops/op_math/controlled.py b/pennylane/ops/op_math/controlled.py
index 262897628..33d23ca4e 100644
--- a/pennylane/ops/op_math/controlled.py
+++ b/pennylane/ops/op_math/controlled.py
@@ -140,6 +140,8 @@ def create_controlled_op(op, control, control_values=None, work_wires=None):
         control_values = [control_values]
     elif control_values is None:
         control_values = [True] * len(control)
+    elif isinstance(control_values, tuple):
+        control_values = list(control_values)
 
     ctrl_op = _try_wrap_in_custom_ctrl_op(
         op, control, control_values=control_values, work_wires=work_wires
diff --git a/pennylane/ops/qutrit/__init__.py b/pennylane/ops/qutrit/__init__.py
index 670f950ba..da3ee27ca 100644
--- a/pennylane/ops/qutrit/__init__.py
+++ b/pennylane/ops/qutrit/__init__.py
@@ -26,7 +26,7 @@ from .non_parametric_ops import *
 from .observables import *
 from .parametric_ops import *
 from .state_preparation import *
-from .channel import QutritDepolarizingChannel
+from .channel import *
 
 # TODO: Change `qml.Identity` for qutrit support or add `qml.TIdentity` for qutrits
 __ops__ = {
@@ -50,6 +50,7 @@ __obs__ = {
 }
 __channels__ = {
     "QutritDepolarizingChannel",
+    "QutritAmplitudeDamping",
 }
 
 __all__ = list(__ops__ | __obs__ | __channels__)
diff --git a/pennylane/ops/qutrit/channel.py b/pennylane/ops/qutrit/channel.py
index 8443fd2de..c50572cf2 100644
--- a/pennylane/ops/qutrit/channel.py
+++ b/pennylane/ops/qutrit/channel.py
@@ -222,3 +222,94 @@ class QutritDepolarizingChannel(Channel):
         )
 
         return [identity] + Ks
+
+
+class QutritAmplitudeDamping(Channel):
+    r"""
+    Single-qutrit amplitude damping error channel.
+
+    Interaction with the environment can lead to changes in the state populations of a qutrit.
+    This can be modelled by the qutrit amplitude damping channel with the following Kraus matrices:
+
+    .. math::
+        K_0 = \begin{bmatrix}
+                1 & 0 & 0\\
+                0 & \sqrt{1-\gamma_1} & 0 \\
+                0 & 0 & \sqrt{1-\gamma_2}
+                \end{bmatrix}, \quad
+        K_1 = \begin{bmatrix}
+                0 & \sqrt{\gamma_1} & 0 \\
+                0 & 0 & 0 \\
+                0 & 0 & 0
+                \end{bmatrix}, \quad
+        K_2 = \begin{bmatrix}
+                0 & 0 & \sqrt{\gamma_2} \\
+                0 & 0 & 0 \\
+                0 & 0 & 0
+                \end{bmatrix}
+
+    where :math:`\gamma_1 \in [0, 1]` and :math:`\gamma_2 \in [0, 1]` are the amplitude damping
+    probabilities for subspaces (0,1) and (0,2) respectively.
+
+    .. note::
+
+        The Kraus operators :math:`\{K_0, K_1, K_2\}` are adapted from [`1 <https://doi.org/10.48550/arXiv.1905.10481>`_] (Eq. 8).
+
+    **Details:**
+
+    * Number of wires: 1
+    * Number of parameters: 2
+
+    Args:
+        gamma_1 (float): :math:`|1 \rangle \rightarrow |0 \rangle` amplitude damping probability.
+        gamma_2 (float): :math:`|2 \rangle \rightarrow |0 \rangle` amplitude damping probability.
+        wires (Sequence[int] or int): the wire the channel acts on
+        id (str or None): String representing the operation (optional)
+    """
+
+    num_params = 2
+    num_wires = 1
+    grad_method = "F"
+
+    def __init__(self, gamma_1, gamma_2, wires, id=None):
+        # Verify gamma_1 and gamma_2
+        for gamma in (gamma_1, gamma_2):
+            if not (math.is_abstract(gamma_1) or math.is_abstract(gamma_2)):
+                if not 0.0 <= gamma <= 1.0:
+                    raise ValueError("Each probability must be in the interval [0,1]")
+        super().__init__(gamma_1, gamma_2, wires=wires, id=id)
+
+    @staticmethod
+    def compute_kraus_matrices(gamma_1, gamma_2):  # pylint:disable=arguments-differ
+        r"""Kraus matrices representing the ``QutritAmplitudeDamping`` channel.
+
+        Args:
+            gamma_1 (float): :math:`|1\rangle \rightarrow |0\rangle` amplitude damping probability.
+            gamma_2 (float): :math:`|2\rangle \rightarrow |0\rangle` amplitude damping probability.
+
+        Returns:
+            list(array): list of Kraus matrices
+
+        **Example**
+
+        >>> qml.QutritAmplitudeDamping.compute_kraus_matrices(0.5, 0.25)
+        [
+        array([ [1.        , 0.        , 0.        ],
+                [0.        , 0.70710678, 0.        ],
+                [0.        , 0.        , 0.8660254 ]]),
+        array([ [0.        , 0.70710678, 0.        ],
+                [0.        , 0.        , 0.        ],
+                [0.        , 0.        , 0.        ]]),
+        array([ [0.        , 0.        , 0.5       ],
+                [0.        , 0.        , 0.        ],
+                [0.        , 0.        , 0.        ]])
+        ]
+        """
+        K0 = math.diag([1, math.sqrt(1 - gamma_1 + math.eps), math.sqrt(1 - gamma_2 + math.eps)])
+        K1 = math.sqrt(gamma_1 + math.eps) * math.convert_like(
+            math.cast_like(math.array([[0, 1, 0], [0, 0, 0], [0, 0, 0]]), gamma_1), gamma_1
+        )
+        K2 = math.sqrt(gamma_2 + math.eps) * math.convert_like(
+            math.cast_like(math.array([[0, 0, 1], [0, 0, 0], [0, 0, 0]]), gamma_2), gamma_2
+        )
+        return [K0, K1, K2]
diff --git a/pennylane/qchem/molecule.py b/pennylane/qchem/molecule.py
index 24e67d43d..7f6877a0f 100644
--- a/pennylane/qchem/molecule.py
+++ b/pennylane/qchem/molecule.py
@@ -27,6 +27,9 @@ from .basis_data import atomic_numbers
 from .basis_set import BasisFunction, mol_basis_data
 from .integrals import contracted_norm, primitive_norm
 
+# Bohr-Angstrom correlation coefficient (https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0)
+bohr_angs = 0.529177210903
+
 
 class Molecule:
     r"""Create a molecule object that stores molecular information and default basis set parameters.
@@ -51,6 +54,7 @@ class Molecule:
         coeff (array[float]): coefficients of the contracted Gaussian functions
         r (array[float]): positions of the Gaussian functions
         normalize (bool): if True, the basis functions get normalized
+        unit (str): unit of atomic coordinates. Available options are ``unit="bohr"`` and ``unit="angstrom"``.
 
     **Example**
 
@@ -75,6 +79,7 @@ class Molecule:
         alpha=None,
         coeff=None,
         normalize=True,
+        unit="bohr",
     ):
         if (
             basis_name.lower()
@@ -97,12 +102,22 @@ class Molecule:
 
         self.symbols = symbols
         self.coordinates = coordinates
+        self.unit = unit.strip().lower()
         self.charge = charge
         self.mult = mult
         self.basis_name = basis_name.lower()
         self.name = name
         self.n_basis, self.basis_data = mol_basis_data(self.basis_name, self.symbols, load_data)
 
+        if self.unit not in ("angstrom", "bohr"):
+            raise ValueError(
+                f"The provided unit '{unit}' is not supported. "
+                f"Please set 'unit' to 'bohr' or 'angstrom'."
+            )
+
+        if self.unit == "angstrom":
+            self.coordinates = self.coordinates / bohr_angs
+
         self.nuclear_charges = [atomic_numbers[s] for s in self.symbols]
 
         self.n_electrons = sum(self.nuclear_charges) - self.charge
diff --git a/pennylane/qchem/openfermion_obs.py b/pennylane/qchem/openfermion_obs.py
index 69716d790..a9252e055 100644
--- a/pennylane/qchem/openfermion_obs.py
+++ b/pennylane/qchem/openfermion_obs.py
@@ -997,6 +997,7 @@ def _(
 def _(
     symbols,
     coordinates,
+    unit="bohr",
     name="molecule",
     charge=0,
     mult=1,
@@ -1013,6 +1014,16 @@ def _(
     load_data=False,
     convert_tol=1e12,
 ):
+
+    if (coord_unit := unit.strip().lower()) not in ("angstrom", "bohr"):
+        raise ValueError(
+            f"The provided unit '{unit}' is not supported. "
+            f"Please set 'unit' to 'bohr' or 'angstrom'."
+        )
+
+    if coord_unit == "angstrom":
+        coordinates = coordinates / bohr_angs
+
     return _molecular_hamiltonian(
         symbols,
         coordinates=coordinates,
diff --git a/pennylane/templates/layers/strongly_entangling.py b/pennylane/templates/layers/strongly_entangling.py
index fd07f05ac..554f100e8 100644
--- a/pennylane/templates/layers/strongly_entangling.py
+++ b/pennylane/templates/layers/strongly_entangling.py
@@ -200,7 +200,7 @@ class StronglyEntanglingLayers(Operation):
         CNOT(wires=['a', 'a']),
         CNOT(wires=['b', 'b'])]
         """
-        n_layers = qml.math.shape(weights)[0]
+        n_layers = qml.math.shape(weights)[-3]
         wires = qml.wires.Wires(wires)
         op_list = []
 
diff --git a/requirements-ci.txt b/requirements-ci.txt
index 6ff2c5afc..932820a26 100644
--- a/requirements-ci.txt
+++ b/requirements-ci.txt
@@ -1,5 +1,5 @@
 numpy
-scipy
+scipy<1.13.0
 cvxpy
 cvxopt
 networkx
@@ -8,7 +8,7 @@ autograd
 toml
 appdirs
 semantic_version
-autoray>=0.6.11
+autoray>=0.6.1,<0.6.10
 matplotlib
 requests
 tomli # Drop once minimum Python version is 3.11
diff --git a/setup.py b/setup.py
index 15863f6e8..892e962ec 100644
--- a/setup.py
+++ b/setup.py
@@ -60,6 +60,7 @@ info = {
             "default.qutrit = pennylane.devices.default_qutrit:DefaultQutrit",
             "default.clifford = pennylane.devices.default_clifford:DefaultClifford",
             "default.qutrit.mixed = pennylane.devices.default_qutrit_mixed:DefaultQutritMixed",
+            "default.tensor = pennylane.devices.default_tensor:DefaultTensor",
         ],
         "console_scripts": ["pl-device-test=pennylane.devices.tests:cli"],
     },


Test patch:
diff --git a/.github/workflows/interface-unit-tests.yml b/.github/workflows/interface-unit-tests.yml
index 41f2f3434..ba2fc3815 100644
--- a/.github/workflows/interface-unit-tests.yml
+++ b/.github/workflows/interface-unit-tests.yml
@@ -366,7 +366,7 @@ jobs:
       install_pennylane_lightning_master: false
       pytest_coverage_flags: ${{ inputs.pytest_coverage_flags }}
       pytest_markers: external
-      additional_pip_packages: pyzx pennylane-catalyst matplotlib stim
+      additional_pip_packages: pyzx pennylane-catalyst matplotlib stim quimb
       requirements_file: ${{ strategy.job-index == 0 && 'external.txt' || '' }}
       disable_new_opmath: ${{ inputs.disable_new_opmath }}
 
diff --git a/tests/devices/default_tensor/test_default_tensor.py b/tests/devices/default_tensor/test_default_tensor.py
new file mode 100644
index 000000000..4fdf300ee
--- /dev/null
+++ b/tests/devices/default_tensor/test_default_tensor.py
@@ -0,0 +1,352 @@
+# 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 the DefaultTensor class.
+"""
+
+
+import numpy as np
+import pytest
+from scipy.sparse import csr_matrix
+
+import pennylane as qml
+
+quimb = pytest.importorskip("quimb")
+
+pytestmark = pytest.mark.external
+
+# gates for which device support is tested
+operations_list = {
+    "Identity": qml.Identity(wires=[0]),
+    "BlockEncode": qml.BlockEncode([[0.1, 0.2], [0.3, 0.4]], wires=[0, 1]),
+    "CNOT": qml.CNOT(wires=[0, 1]),
+    "CRX": qml.CRX(0, wires=[0, 1]),
+    "CRY": qml.CRY(0, wires=[0, 1]),
+    "CRZ": qml.CRZ(0, wires=[0, 1]),
+    "CRot": qml.CRot(0, 0, 0, wires=[0, 1]),
+    "CSWAP": qml.CSWAP(wires=[0, 1, 2]),
+    "CZ": qml.CZ(wires=[0, 1]),
+    "CCZ": qml.CCZ(wires=[0, 1, 2]),
+    "CY": qml.CY(wires=[0, 1]),
+    "CH": qml.CH(wires=[0, 1]),
+    "DiagonalQubitUnitary": qml.DiagonalQubitUnitary(np.array([1, 1]), wires=[0]),
+    "Hadamard": qml.Hadamard(wires=[0]),
+    "MultiRZ": qml.MultiRZ(0, wires=[0]),
+    "PauliX": qml.X(0),
+    "PauliY": qml.Y(0),
+    "PauliZ": qml.Z(0),
+    "X": qml.X([0]),
+    "Y": qml.Y([0]),
+    "Z": qml.Z([0]),
+    "PhaseShift": qml.PhaseShift(0, wires=[0]),
+    "PCPhase": qml.PCPhase(0, 1, wires=[0, 1]),
+    "ControlledPhaseShift": qml.ControlledPhaseShift(0, wires=[0, 1]),
+    "CPhaseShift00": qml.CPhaseShift00(0, wires=[0, 1]),
+    "CPhaseShift01": qml.CPhaseShift01(0, wires=[0, 1]),
+    "CPhaseShift10": qml.CPhaseShift10(0, wires=[0, 1]),
+    "QubitUnitary": qml.QubitUnitary(np.eye(2), wires=[0]),
+    "SpecialUnitary": qml.SpecialUnitary(np.array([0.2, -0.1, 2.3]), wires=1),
+    "ControlledQubitUnitary": qml.ControlledQubitUnitary(np.eye(2), control_wires=[1], wires=[0]),
+    "MultiControlledX": qml.MultiControlledX(wires=[1, 2, 0]),
+    "IntegerComparator": qml.IntegerComparator(1, geq=True, wires=[0, 1, 2]),
+    "RX": qml.RX(0, wires=[0]),
+    "RY": qml.RY(0, wires=[0]),
+    "RZ": qml.RZ(0, wires=[0]),
+    "Rot": qml.Rot(0, 0, 0, wires=[0]),
+    "S": qml.S(wires=[0]),
+    "Adjoint(S)": qml.adjoint(qml.S(wires=[0])),
+    "SWAP": qml.SWAP(wires=[0, 1]),
+    "ISWAP": qml.ISWAP(wires=[0, 1]),
+    "PSWAP": qml.PSWAP(0, wires=[0, 1]),
+    "ECR": qml.ECR(wires=[0, 1]),
+    "Adjoint(ISWAP)": qml.adjoint(qml.ISWAP(wires=[0, 1])),
+    "T": qml.T(wires=[0]),
+    "Adjoint(T)": qml.adjoint(qml.T(wires=[0])),
+    "SX": qml.SX(wires=[0]),
+    "Adjoint(SX)": qml.adjoint(qml.SX(wires=[0])),
+    "Toffoli": qml.Toffoli(wires=[0, 1, 2]),
+    "QFT": qml.templates.QFT(wires=[0, 1, 2]),
+    "IsingXX": qml.IsingXX(0, wires=[0, 1]),
+    "IsingYY": qml.IsingYY(0, wires=[0, 1]),
+    "IsingZZ": qml.IsingZZ(0, wires=[0, 1]),
+    "IsingXY": qml.IsingXY(0, wires=[0, 1]),
+    "SingleExcitation": qml.SingleExcitation(0, wires=[0, 1]),
+    "SingleExcitationPlus": qml.SingleExcitationPlus(0, wires=[0, 1]),
+    "SingleExcitationMinus": qml.SingleExcitationMinus(0, wires=[0, 1]),
+    "DoubleExcitation": qml.DoubleExcitation(0, wires=[0, 1, 2, 3]),
+    "QubitCarry": qml.QubitCarry(wires=[0, 1, 2, 3]),
+    "QubitSum": qml.QubitSum(wires=[0, 1, 2]),
+    "PauliRot": qml.PauliRot(0, "XXYY", wires=[0, 1, 2, 3]),
+    "U1": qml.U1(0, wires=0),
+    "U2": qml.U2(0, 0, wires=0),
+    "U3": qml.U3(0, 0, 0, wires=0),
+    "SISWAP": qml.SISWAP(wires=[0, 1]),
+    "Adjoint(SISWAP)": qml.adjoint(qml.SISWAP(wires=[0, 1])),
+    "OrbitalRotation": qml.OrbitalRotation(0, wires=[0, 1, 2, 3]),
+    "FermionicSWAP": qml.FermionicSWAP(0, wires=[0, 1]),
+    "GlobalPhase": qml.GlobalPhase(0.123, wires=[0, 1]),
+}
+
+all_ops = operations_list.keys()
+
+# observables for which device support is tested
+observables_list = {
+    "Identity": qml.Identity(wires=[0]),
+    "Hadamard": qml.Hadamard(wires=[0]),
+    "Hermitian": qml.Hermitian(np.eye(2), wires=[0]),
+    "PauliX": qml.PauliX(0),
+    "PauliY": qml.PauliY(0),
+    "PauliZ": qml.PauliZ(0),
+    "X": qml.X(0),
+    "Y": qml.Y(0),
+    "Z": qml.Z(0),
+    "Projector": [
+        qml.Projector(np.array([1]), wires=[0]),
+        qml.Projector(np.array([0, 1]), wires=[0]),
+    ],
+    "SparseHamiltonian": qml.SparseHamiltonian(csr_matrix(np.eye(8)), wires=[0, 1, 2]),
+    "Hamiltonian": qml.Hamiltonian([1, 1], [qml.Z(0), qml.X(0)]),
+    "LinearCombination": qml.ops.LinearCombination([1, 1], [qml.Z(0), qml.X(0)]),
+}
+
+all_obs = observables_list.keys()
+
+
+def test_name():
+    """Test the name of DefaultTensor."""
+    assert qml.device("default.tensor", wires=0).name == "default.tensor"
+
+
+def test_wires():
+    """Test that a device can be created with wires."""
+    assert qml.device("default.tensor", wires=0).wires is not None
+    assert qml.device("default.tensor", wires=2).wires == qml.wires.Wires([0, 1])
+    assert qml.device("default.tensor", wires=[0, 2]).wires == qml.wires.Wires([0, 2])
+
+    with pytest.raises(AttributeError):
+        qml.device("default.tensor", wires=0).wires = [0, 1]
+
+
+def test_wires_error():
+    """Test that an error is raised if the wires are not provided."""
+    with pytest.raises(TypeError):
+        qml.device("default.tensor")
+
+    with pytest.raises(TypeError):
+        qml.device("default.tensor", wires=None)
+
+
+def test_wires_execution_error():
+    """Test that this device cannot execute a tape if its wires do not match the wires on the device."""
+    dev = qml.device("default.tensor", wires=3)
+    ops = [
+        qml.Identity(0),
+        qml.Identity((0, 1)),
+        qml.RX(2, 0),
+        qml.RY(1, 5),
+        qml.RX(2, 1),
+    ]
+    measurements = [qml.expval(qml.PauliZ(15))]
+    tape = qml.tape.QuantumScript(ops, measurements)
+
+    with pytest.raises(AttributeError):
+        dev.execute(tape)
+
+
+@pytest.mark.parametrize("max_bond_dim", [None, 10])
+@pytest.mark.parametrize("cutoff", [1e-16, 1e-12])
+@pytest.mark.parametrize("contract", ["auto-mps", "nonlocal"])
+def test_kwargs(max_bond_dim, cutoff, contract):
+    """Test the class initialization with different arguments and returned properties."""
+
+    kwargs = {"max_bond_dim": max_bond_dim, "cutoff": cutoff, "contract": contract}
+
+    dev = qml.device("default.tensor", wires=0, **kwargs)
+
+    _, config = dev.preprocess()
+    assert config.device_options["method"] == "mps"
+    assert config.device_options["max_bond_dim"] == max_bond_dim
+    assert config.device_options["cutoff"] == cutoff
+    assert config.device_options["contract"] == contract
+
+
+def test_invalid_kwarg():
+    """Test an invalid keyword argument."""
+    with pytest.raises(
+        TypeError,
+        match="Unexpected argument: fake_arg during initialization of the default.tensor device.",
+    ):
+        qml.device("default.tensor", wires=0, fake_arg=None)
+
+
+def test_method():
+    """Test the device method."""
+    assert qml.device("default.tensor", wires=0).method == "mps"
+
+
+def test_invalid_method():
+    """Test an invalid method."""
+    method = "invalid_method"
+    with pytest.raises(ValueError, match=f"Unsupported method: {method}"):
+        qml.device("default.tensor", wires=0, method=method)
+
+
+@pytest.mark.parametrize("dtype", [np.complex64, np.complex128])
+def test_data_type(dtype):
+    """Test the data type."""
+    assert qml.device("default.tensor", wires=0, dtype=dtype).dtype == dtype
+
+
+def test_ivalid_data_type():
+    """Test that data type can only be np.complex64 or np.complex128."""
+    with pytest.raises(TypeError):
+        qml.device("default.tensor", wires=0, dtype=float)
+
+
+class TestSupportedGatesAndObservables:
+    """Test that the DefaultTensor device supports all gates and observables that it claims to support."""
+
+    @pytest.mark.parametrize("operation", all_ops)
+    def test_supported_gates_can_be_implemented(self, operation):
+        """Test that the device can implement all its supported gates."""
+
+        dev = qml.device("default.tensor", wires=4, method="mps")
+
+        tape = qml.tape.QuantumScript(
+            [operations_list[operation]],
+            [qml.expval(qml.Identity(wires=0))],
+        )
+
+        result = dev.execute(circuits=tape)
+        assert np.allclose(result, 1.0)
+
+    @pytest.mark.parametrize("observable", all_obs)
+    def test_supported_observables_can_be_implemented(self, observable):
+        """Test that the device can implement all its supported observables."""
+
+        dev = qml.device("default.tensor", wires=3, method="mps")
+
+        if observable == "Projector":
+            for o in observables_list[observable]:
+                tape = qml.tape.QuantumScript(
+                    [qml.PauliX(0)],
+                    [qml.expval(o)],
+                )
+                result = dev.execute(circuits=tape)
+                assert isinstance(result, (float, np.ndarray))
+
+        else:
+            tape = qml.tape.QuantumScript(
+                [qml.PauliX(0)],
+                [qml.expval(observables_list[observable])],
+            )
+            result = dev.execute(circuits=tape)
+            assert isinstance(result, (float, np.ndarray))
+
+    def test_not_implemented_meas(self):
+        """Tests that support only exists for `qml.expval` and `qml.var` so far."""
+
+        op = [qml.Identity(0)]
+        measurements = [qml.probs(qml.PauliZ(0))]
+        tape = qml.tape.QuantumScript(op, measurements)
+
+        dev = qml.device("default.tensor", wires=tape.wires)
+
+        with pytest.raises(NotImplementedError):
+            dev.execute(tape)
+
+
+class TestSupportsDerivatives:
+    """Test that DefaultTensor states what kind of derivatives it supports."""
+
+    def test_support_derivatives(self):
+        """Test that the device does not support derivatives yet."""
+        dev = qml.device("default.tensor", wires=0)
+        assert not dev.supports_derivatives()
+
+    def test_compute_derivatives(self):
+        """Test that an error is raised if the `compute_derivatives` method is called."""
+        dev = qml.device("default.tensor", wires=0)
+        with pytest.raises(
+            NotImplementedError,
+            match="The computation of derivatives has yet to be implemented for the default.tensor device.",
+        ):
+            dev.compute_derivatives(circuits=None)
+
+    def test_execute_and_compute_derivatives(self):
+        """Test that an error is raised if `execute_and_compute_derivative` method is called."""
+        dev = qml.device("default.tensor", wires=0)
+        with pytest.raises(
+            NotImplementedError,
+            match="The computation of derivatives has yet to be implemented for the default.tensor device.",
+        ):
+            dev.execute_and_compute_derivatives(circuits=None)
+
+    def test_supports_vjp(self):
+        """Test that the device does not support VJP yet."""
+        dev = qml.device("default.tensor", wires=0)
+        assert not dev.supports_vjp()
+
+    def test_compute_vjp(self):
+        """Test that an error is raised if `compute_vjp` method is called."""
+        dev = qml.device("default.tensor", wires=0)
+        with pytest.raises(
+            NotImplementedError,
+            match="The computation of vector-Jacobian product has yet to be implemented for the default.tensor device.",
+        ):
+            dev.compute_vjp(circuits=None, cotangents=None)
+
+    def test_execute_and_compute_vjp(self):
+        """Test that an error is raised if `execute_and_compute_vjp` method is called."""
+        dev = qml.device("default.tensor", wires=0)
+        with pytest.raises(
+            NotImplementedError,
+            match="The computation of vector-Jacobian product has yet to be implemented for the default.tensor device.",
+        ):
+            dev.execute_and_compute_vjp(circuits=None, cotangents=None)
+
+    @pytest.mark.jax
+    def test_jax(self):
+        """Test the device with JAX."""
+
+        jax = pytest.importorskip("jax")
+        dev = qml.device("default.tensor", wires=1)
+        ref_dev = qml.device("default.qubit.jax", wires=1)
+
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.Rot(x[0], x[1], x[2], wires=0)
+            return qml.expval(qml.Z(0))
+
+        weights = jax.numpy.array([0.2, 0.5, 0.1])
+        print(isinstance(dev, qml.Device))
+        qnode = qml.QNode(circuit, dev, interface="jax")
+        ref_qnode = qml.QNode(circuit, ref_dev, interface="jax")
+
+        assert np.allclose(qnode(weights), ref_qnode(weights))
+
+    @pytest.mark.jax
+    def test_jax_jit(self):
+        """Test the device with JAX's JIT compiler."""
+
+        jax = pytest.importorskip("jax")
+        dev = qml.device("default.tensor", wires=1)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit():
+            qml.Hadamard(0)
+            return qml.expval(qml.Z(0))
+
+        assert np.allclose(circuit(), 0.0)
diff --git a/tests/devices/default_tensor/test_tensor_expval.py b/tests/devices/default_tensor/test_tensor_expval.py
new file mode 100644
index 000000000..663dba1ee
--- /dev/null
+++ b/tests/devices/default_tensor/test_tensor_expval.py
@@ -0,0 +1,399 @@
+# 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.
+"""
+Tests for the expectation value calculations on the DefaultTensor device.
+"""
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.devices import DefaultQubit
+
+THETA = np.linspace(0.11, 1, 3)
+PHI = np.linspace(0.32, 1, 3)
+VARPHI = np.linspace(0.02, 1, 3)
+
+quimb = pytest.importorskip("quimb")
+
+pytestmark = pytest.mark.external
+
+# pylint: disable=too-many-arguments, redefined-outer-name
+
+
+@pytest.fixture(params=[np.complex64, np.complex128])
+def dev(request):
+    """Device fixture."""
+    return qml.device("default.tensor", wires=3, dtype=request.param)
+
+
+def calculate_reference(tape):
+    """Calculate the reference value of the tape using DefaultQubit."""
+    dev = DefaultQubit(max_workers=1)
+    program, _ = dev.preprocess()
+    tapes, transf_fn = program([tape])
+    results = dev.execute(tapes)
+    return transf_fn(results)
+
+
+def execute(dev, tape):
+    """Execute the tape on the device and return the result."""
+    results = dev.execute(tape)
+    return results
+
+
+@pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
+class TestExpval:
+    """Test expectation value calculations"""
+
+    def test_Identity(self, theta, phi, dev, tol):
+        """Tests applying identities."""
+
+        ops = [
+            qml.Identity(0),
+            qml.Identity((0, 1)),
+            qml.RX(theta, 0),
+            qml.Identity((1, 2)),
+            qml.RX(phi, 1),
+        ]
+        measurements = [qml.expval(qml.PauliZ(0))]
+        tape = qml.tape.QuantumScript(ops, measurements)
+
+        result = dev.execute(tape)
+        expected = np.cos(theta)
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(result, expected, tol)
+
+    def test_identity_expectation(self, theta, phi, dev, tol):
+        """Tests identity expectations."""
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(theta, wires=[0]), qml.RX(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.expval(qml.Identity(wires=[0])), qml.expval(qml.Identity(wires=[1]))],
+        )
+        result = dev.execute(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(1.0, result, tol)
+
+    def test_multi_wire_identity_expectation(self, theta, phi, dev, tol):
+        """Tests multi-wire identity."""
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(theta, wires=[0]), qml.RX(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.expval(qml.Identity(wires=[0, 1]))],
+        )
+        result = dev.execute(tape)
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(1.0, result, tol)
+
+    @pytest.mark.parametrize(
+        "wires",
+        [([0, 1]), (["a", 1]), (["b", "a"]), ([-1, 2.5])],
+    )
+    def test_custom_wires(self, theta, phi, tol, wires):
+        """Tests custom wires."""
+        dev = qml.device("default.tensor", wires=wires, dtype=np.complex128)
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(theta, wires=wires[0]),
+                qml.RX(phi, wires=wires[1]),
+                qml.CNOT(wires=wires),
+            ],
+            [
+                qml.expval(qml.PauliZ(wires=wires[0])),
+                qml.expval(qml.PauliZ(wires=wires[1])),
+            ],
+        )
+
+        calculated_val = execute(dev, tape)
+        reference_val = np.array([np.cos(theta), np.cos(theta) * np.cos(phi)])
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize(
+        "Obs, Op, expected_fn",
+        [
+            (
+                [qml.PauliX(wires=[0]), qml.PauliX(wires=[1])],
+                qml.RY,
+                lambda theta, phi: np.array([np.sin(theta) * np.sin(phi), np.sin(phi)]),
+            ),
+            (
+                [qml.PauliY(wires=[0]), qml.PauliY(wires=[1])],
+                qml.RX,
+                lambda theta, phi: np.array([0, -np.cos(theta) * np.sin(phi)]),
+            ),
+            (
+                [qml.PauliZ(wires=[0]), qml.PauliZ(wires=[1])],
+                qml.RX,
+                lambda theta, phi: np.array([np.cos(theta), np.cos(theta) * np.cos(phi)]),
+            ),
+            (
+                [qml.Hadamard(wires=[0]), qml.Hadamard(wires=[1])],
+                qml.RY,
+                lambda theta, phi: np.array(
+                    [
+                        np.sin(theta) * np.sin(phi) + np.cos(theta),
+                        np.cos(theta) * np.cos(phi) + np.sin(phi),
+                    ]
+                )
+                / np.sqrt(2),
+            ),
+        ],
+    )
+    def test_single_wire_observables_expectation(self, Obs, Op, expected_fn, theta, phi, tol, dev):
+        """Test that expectation values for single wire observables are correct"""
+
+        tape = qml.tape.QuantumScript(
+            [Op(theta, wires=[0]), Op(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.expval(Obs[0]), qml.expval(Obs[1])],
+        )
+        result = execute(dev, tape)
+        expected = expected_fn(theta, phi)
+
+        assert np.allclose(result, expected, tol)
+
+    def test_hermitian_expectation(self, theta, phi, tol, dev):
+        """Tests an Hermitian operator."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+            qml.RX(theta + phi, wires=2)
+
+            for idx in range(3):
+                qml.expval(qml.Hermitian([[1, 0], [0, -1]], wires=[idx]))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_hamiltonian_expectation(self, theta, phi, tol, dev):
+        """Tests a Hamiltonian."""
+
+        ham = qml.Hamiltonian(
+            [1.0, 0.3, 0.3, 0.4],
+            [
+                qml.PauliX(0) @ qml.PauliX(1),
+                qml.PauliZ(0),
+                qml.PauliZ(1),
+                qml.PauliX(0) @ qml.PauliY(1),
+            ],
+        )
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+            qml.RX(theta + phi, wires=2)
+
+            qml.expval(ham)
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_sparse_hamiltonian_expectation(self, theta, phi, tol, dev):
+        """Tests a Hamiltonian."""
+
+        ham = qml.SparseHamiltonian(
+            qml.Hamiltonian(
+                [1.0, 0.3, 0.3, 0.4],
+                [
+                    qml.PauliX(0) @ qml.PauliX(1),
+                    qml.PauliZ(0),
+                    qml.PauliZ(1),
+                    qml.PauliX(0) @ qml.PauliY(1),
+                ],
+            ).sparse_matrix(),
+            wires=[0, 1],
+        )
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+
+            qml.expval(ham)
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+
+@pytest.mark.parametrize("phi", PHI)
+class TestOperatorArithmetic:
+    """Test integration with SProd, Prod, and Sum."""
+
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            qml.s_prod(0.5, qml.PauliZ(0)),
+            qml.prod(qml.PauliZ(0), qml.PauliX(1)),
+            qml.sum(qml.PauliZ(0), qml.PauliX(1)),
+        ],
+    )
+    def test_op_math(self, phi, dev, obs, tol):
+        """Tests the `SProd`, `Prod`, and `Sum` classes."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(phi, wires=[0]),
+                qml.Hadamard(wires=[1]),
+                qml.PauliZ(wires=[1]),
+                qml.RX(-1.1 * phi, wires=[1]),
+            ],
+            [qml.expval(obs)],
+        )
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_integration(self, phi, dev, tol):
+        """Test a Combination of `Sum`, `SProd`, and `Prod`."""
+
+        obs = qml.sum(
+            qml.s_prod(2.3, qml.PauliZ(0)),
+            -0.5 * qml.prod(qml.PauliY(0), qml.PauliZ(1)),
+        )
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(phi, wires=[0]), qml.RX(-1.1 * phi, wires=[0])],
+            [qml.expval(obs)],
+        )
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+
+@pytest.mark.parametrize("theta, phi, varphi", list(zip(THETA, PHI, VARPHI)))
+class TestTensorExpval:
+    """Test tensor expectation values"""
+
+    def test_PauliX_PauliY(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliX and PauliY."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.expval(qml.PauliX(0) @ qml.PauliY(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_PauliZ_identity(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliZ and Identity."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.Identity(wires=[0])
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.expval(qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_PauliZ_hadamard_PauliY(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliY, PauliZ and Hadamard."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.expval(qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+
+@pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
+def test_multi_qubit_gates(theta, phi, dev):
+    """Tests a simple circuit with multi-qubit gates."""
+
+    ops = [
+        qml.PauliX(wires=[0]),
+        qml.RX(theta, wires=[0]),
+        qml.CSWAP(wires=[7, 0, 5]),
+        qml.RX(phi, wires=[1]),
+        qml.CNOT(wires=[3, 4]),
+        qml.DoubleExcitation(phi, wires=[1, 2, 3, 4]),
+        qml.CZ(wires=[4, 5]),
+        qml.Hadamard(wires=[4]),
+        qml.CCZ(wires=[0, 1, 2]),
+        qml.CSWAP(wires=[2, 3, 4]),
+        qml.QFT(wires=[0, 1, 2]),
+        qml.CNOT(wires=[2, 4]),
+        qml.Toffoli(wires=[0, 1, 2]),
+        qml.DoubleExcitation(phi, wires=[0, 1, 3, 4]),
+    ]
+
+    meas = [
+        qml.expval(qml.PauliY(wires=[2])),
+        qml.expval(qml.Hamiltonian([1, 5, 6], [qml.Z(6), qml.X(0), qml.Hadamard(4)])),
+        qml.expval(
+            qml.Hamiltonian(
+                [4, 5, 7],
+                [
+                    qml.Z(6) @ qml.Y(4),
+                    qml.X(7),
+                    qml.Hadamard(4),
+                ],
+            )
+        ),
+    ]
+
+    tape = qml.tape.QuantumScript(ops=ops, measurements=meas)
+
+    reference_val = calculate_reference(tape)
+    dev = qml.device("default.tensor", wires=tape.wires, dtype=np.complex128)
+    calculated_val = dev.execute(tape)
+
+    assert np.allclose(calculated_val, reference_val)
diff --git a/tests/devices/default_tensor/test_tensor_var.py b/tests/devices/default_tensor/test_tensor_var.py
new file mode 100644
index 000000000..a4aa35e40
--- /dev/null
+++ b/tests/devices/default_tensor/test_tensor_var.py
@@ -0,0 +1,406 @@
+# 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.
+"""
+Tests for the variance calculation on the DefaultTensor device.
+"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.tape import QuantumScript
+
+THETA = np.linspace(0.11, 1, 3)
+PHI = np.linspace(0.32, 1, 3)
+VARPHI = np.linspace(0.02, 1, 3)
+
+quimb = pytest.importorskip("quimb")
+
+pytestmark = pytest.mark.external
+
+# pylint: disable=too-many-arguments, redefined-outer-name
+
+
+@pytest.fixture(params=[np.complex64, np.complex128])
+def dev(request):
+    """Device fixture."""
+    return qml.device("default.tensor", wires=3, dtype=request.param)
+
+
+def calculate_reference(tape):
+    """Calculate the reference value of the tape using DefaultQubit."""
+    dev = qml.device("default.qubit", max_workers=1)
+    program, _ = dev.preprocess()
+    tapes, transf_fn = program([tape])
+    results = dev.execute(tapes)
+    return transf_fn(results)
+
+
+def execute(dev, tape):
+    """Execute the tape on the device and return the result."""
+    results = dev.execute(tape)
+    return results
+
+
+@pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
+class TestVar:
+    """Tests for the variance"""
+
+    def test_Identity(self, theta, phi, dev):
+        """Tests applying identities."""
+
+        ops = [
+            qml.Identity(0),
+            qml.Identity((0, 1)),
+            qml.RX(theta, 0),
+            qml.Identity((1, 2)),
+            qml.RX(phi, 1),
+        ]
+        measurements = [qml.var(qml.PauliZ(0))]
+        tape = qml.tape.QuantumScript(ops, measurements)
+
+        result = dev.execute(tape)
+        expected = 1 - np.cos(theta) ** 2
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(result, expected, atol=tol, rtol=0)
+
+    def test_identity_variance(self, theta, phi, dev):
+        """Tests identity variances."""
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(theta, wires=[0]), qml.RX(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.var(qml.Identity(wires=[0])), qml.var(qml.Identity(wires=[1]))],
+        )
+        result = dev.execute(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(0.0, result, atol=tol, rtol=0)
+
+    def test_multi_wire_identity_variance(self, theta, phi, dev):
+        """Tests multi-wire identity."""
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(theta, wires=[0]), qml.RX(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.var(qml.Identity(wires=[0, 1]))],
+        )
+        result = dev.execute(tape)
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(0.0, result, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize(
+        "wires",
+        [([0, 1]), (["a", 1]), (["b", "a"]), ([-1, 2.5])],
+    )
+    def test_custom_wires(self, theta, phi, wires):
+        """Tests custom wires."""
+        device = qml.device("default.tensor", wires=wires, dtype=np.complex128)
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(theta, wires=wires[0]),
+                qml.RX(phi, wires=wires[1]),
+                qml.CNOT(wires=wires),
+            ],
+            [qml.var(qml.PauliZ(wires=wires[0])), qml.var(qml.PauliZ(wires=wires[1]))],
+        )
+
+        calculated_val = execute(device, tape)
+        reference_val = np.array(
+            [1 - np.cos(theta) ** 2, 1 - np.cos(theta) ** 2 * np.cos(phi) ** 2]
+        )
+
+        tol = 1e-5 if device.dtype == np.complex64 else 1e-7
+        assert np.allclose(calculated_val, reference_val, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize(
+        "Obs, Op, expected_fn",
+        [
+            (
+                [qml.PauliX(wires=[0]), qml.PauliX(wires=[1])],
+                qml.RY,
+                lambda theta, phi: np.array(
+                    [1 - np.sin(theta) ** 2 * np.sin(phi) ** 2, 1 - np.sin(phi) ** 2]
+                ),
+            ),
+            (
+                [qml.PauliY(wires=[0]), qml.PauliY(wires=[1])],
+                qml.RX,
+                lambda theta, phi: np.array([1, 1 - np.cos(theta) ** 2 * np.sin(phi) ** 2]),
+            ),
+            (
+                [qml.PauliZ(wires=[0]), qml.PauliZ(wires=[1])],
+                qml.RX,
+                lambda theta, phi: np.array(
+                    [1 - np.cos(theta) ** 2, 1 - np.cos(theta) ** 2 * np.cos(phi) ** 2]
+                ),
+            ),
+            (
+                [qml.Hadamard(wires=[0]), qml.Hadamard(wires=[1])],
+                qml.RY,
+                lambda theta, phi: np.array(
+                    [
+                        1 - (np.sin(theta) * np.sin(phi) + np.cos(theta)) ** 2 / 2,
+                        1 - (np.cos(theta) * np.cos(phi) + np.sin(phi)) ** 2 / 2,
+                    ]
+                ),
+            ),
+        ],
+    )
+    def test_single_wire_observables_variance(self, Obs, Op, expected_fn, theta, phi, dev):
+        """Test that variance values for single wire observables are correct"""
+
+        tape = qml.tape.QuantumScript(
+            [Op(theta, wires=[0]), Op(phi, wires=[1]), qml.CNOT(wires=[0, 1])],
+            [qml.var(Obs[0]), qml.var(Obs[1])],
+        )
+        result = execute(dev, tape)
+        expected = expected_fn(theta, phi)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(result, expected, atol=tol, rtol=0)
+
+    def test_hermitian_variance(self, theta, phi, dev):
+        """Tests an Hermitian operator."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+            qml.RX(theta + phi, wires=2)
+
+            for idx in range(3):
+                qml.var(qml.Hermitian([[1, 0], [0, -1]], wires=[idx]))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(calculated_val, reference_val, atol=tol, rtol=0)
+
+    def test_hamiltonian_variance(self, theta, phi, dev):
+        """Tests a Hamiltonian."""
+
+        ham = qml.Hamiltonian(
+            [1.0, 0.3, 0.3, 0.4],
+            [
+                qml.PauliX(0) @ qml.PauliX(1),
+                qml.PauliZ(0),
+                qml.PauliZ(1),
+                qml.PauliX(0) @ qml.PauliY(1),
+            ],
+        )
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+            qml.RX(theta + phi, wires=2)
+
+            qml.var(ham)
+
+        tape2 = QuantumScript(tape1.operations, [qml.var(qml.dot(ham.coeffs, ham.ops))])
+
+        calculated_val = execute(dev, tape1)
+        reference_val = calculate_reference(tape2)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(calculated_val, reference_val, atol=tol, rtol=0)
+
+    def test_sparse_hamiltonian_variance(self, theta, phi, dev):
+        """Tests a Hamiltonian."""
+
+        ham = qml.SparseHamiltonian(
+            qml.Hamiltonian(
+                [1.0, 0.3, 0.3, 0.4],
+                [
+                    qml.PauliX(0) @ qml.PauliX(1),
+                    qml.PauliZ(0),
+                    qml.PauliZ(1),
+                    qml.PauliX(0) @ qml.PauliY(1),
+                ],
+            ).sparse_matrix(),
+            wires=[0, 1],
+        )
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(theta, wires=0)
+            qml.RX(phi, wires=1)
+
+            qml.var(ham)
+
+        tape2 = QuantumScript(
+            tape1.operations, [qml.var(qml.Hermitian(ham.matrix(), wires=[0, 1]))]
+        )
+
+        calculated_val = execute(dev, tape1)
+        reference_val = calculate_reference(tape2)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+        assert np.allclose(calculated_val, reference_val, atol=tol, rtol=0)
+
+
+@pytest.mark.parametrize("phi", PHI)
+class TestOperatorArithmetic:
+    """Test integration with SProd, Prod, and Sum."""
+
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            qml.s_prod(0.5, qml.PauliZ(0)),
+            qml.prod(qml.PauliZ(0), qml.PauliX(1)),
+            qml.sum(qml.PauliZ(0), qml.PauliX(1)),
+        ],
+    )
+    def test_op_math(self, phi, dev, obs, tol):
+        """Tests the `SProd`, `Prod`, and `Sum` classes."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(phi, wires=[0]),
+                qml.Hadamard(wires=[1]),
+                qml.PauliZ(wires=[1]),
+                qml.RX(-1.1 * phi, wires=[1]),
+            ],
+            [qml.var(obs)],
+        )
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_integration(self, phi, dev, tol):
+        """Test a Combination of `Sum`, `SProd`, and `Prod`."""
+
+        obs = qml.sum(
+            qml.s_prod(2.3, qml.PauliZ(0)),
+            -0.5 * qml.prod(qml.PauliY(0), qml.PauliZ(1)),
+        )
+
+        tape = qml.tape.QuantumScript(
+            [qml.RX(phi, wires=[0]), qml.RX(-1.1 * phi, wires=[0])],
+            [qml.var(obs)],
+        )
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+
+@pytest.mark.parametrize("theta, phi, varphi", list(zip(THETA, PHI, VARPHI)))
+class TestTensorVar:
+    """Test tensor variances"""
+
+    def test_PauliX_PauliY(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliX and PauliY."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.var(qml.PauliX(0) @ qml.PauliY(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_PauliZ_identity(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliZ and Identity."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.Identity(wires=[0])
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.var(qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+    def test_PauliZ_hadamard_PauliY(self, theta, phi, varphi, dev, tol):
+        """Tests a tensor product involving PauliY, PauliZ and Hadamard."""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(theta, wires=[0])
+            qml.RX(phi, wires=[1])
+            qml.RX(varphi, wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.var(qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2))
+
+        calculated_val = execute(dev, tape)
+        reference_val = calculate_reference(tape)
+
+        tol = 1e-5 if dev.dtype == np.complex64 else 1e-7
+
+        assert np.allclose(calculated_val, reference_val, tol)
+
+
+@pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
+def test_multi_qubit_gates(theta, phi, dev):
+    """Tests a simple circuit with multi-qubit gates."""
+
+    ops = [
+        qml.PauliX(wires=[0]),
+        qml.RX(theta, wires=[0]),
+        qml.CSWAP(wires=[7, 0, 5]),
+        qml.RX(phi, wires=[1]),
+        qml.CNOT(wires=[3, 4]),
+        qml.DoubleExcitation(phi, wires=[1, 2, 3, 4]),
+        qml.CZ(wires=[4, 5]),
+        qml.Hadamard(wires=[4]),
+        qml.CCZ(wires=[0, 1, 2]),
+        qml.CSWAP(wires=[2, 3, 4]),
+        qml.QFT(wires=[0, 1, 2]),
+        qml.CNOT(wires=[2, 4]),
+        qml.Toffoli(wires=[0, 1, 2]),
+        qml.DoubleExcitation(phi, wires=[0, 1, 3, 4]),
+    ]
+
+    meas = [
+        qml.var(qml.PauliY(wires=[2])),
+        qml.var(qml.Hamiltonian([1, 5, 6], [qml.Z(6), qml.X(0), qml.Hadamard(4)])),
+        qml.var(
+            qml.Hamiltonian(
+                [4, 5, 7],
+                [
+                    qml.Z(6) @ qml.Y(4),
+                    qml.X(0),
+                    qml.Hadamard(7),
+                ],
+            )
+        ),
+    ]
+
+    tape = qml.tape.QuantumScript(ops=ops, measurements=meas)
+
+    reference_val = calculate_reference(tape)
+    dev = qml.device("default.tensor", wires=tape.wires, dtype=np.complex128)
+    calculated_val = dev.execute(tape)
+
+    assert np.allclose(calculated_val, reference_val)
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index 4a33f631b..6823a0a3a 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -20,12 +20,221 @@ from pennylane import numpy as np
 from pennylane.devices import DefaultQubitLegacy
 from pennylane.gradients import param_shift
 from pennylane.gradients.parameter_shift import (
+    _evaluate_gradient,
     _get_operation_recipe,
     _make_zero_rep,
     _put_zeros_in_pdA2_involutory,
 )
+from pennylane.measurements.shots import Shots
 from pennylane.operation import AnyWires, Observable
 
+# Constants for TestEvaluateGradient
+# Coefficients and expectation values
+X = np.arange(1, 5)
+# Expected "shift rule" result
+Z = np.sum(-np.arange(1, 5) ** 2)
+# Single coefficient/expectation value that leads to the same result as X
+w = np.sqrt(30)
+# Prefactors to emulate a shot vector
+shv = np.array([0.1, 0.4, 0.7])
+# Fake probability vector (just a 1d array)
+p = np.array([0.01, 0.06, -0.2, 0.5, -0.1, 0.7, -0.09])
+# Second fake probability vector (just a 1d array)
+p2 = p[1:5]
+# shifted probability evaluations
+P = np.outer(X, p)
+# shifted probability evaluations for p2
+P2 = np.outer(X, p2)
+# Single unshifted result that lead to the same result as P
+v = w * p
+# Single unshifted result that lead to the same result as P2
+v2 = w * p2
+# Prefactors to emulate different shot values and multi measurement
+shv_m = np.outer([0.1, 0.4, 0.7], [1, 2])
+
+
+class TestEvaluateGradient:
+    """Test _evaluate_gradient."""
+
+    # pylint: disable=too-many-arguments
+
+    # We could theoretically compute the required res, r0 and expected from the parametrization of coeffs,
+    # unshifted_coeff and batch_size, but that turned out to take lots of effort and edge case logic
+
+    test_cases_single_shots_single_meas = [
+        # Expectation value
+        (X, None, None, tuple(-X), None, Z),
+        (X, None, 4, -X, None, Z),
+        (X[:-1], X[-1], None, tuple(-X[:-1]), -X[-1], Z),
+        (X[:-1], X[-1], 4, -X[:-1], -X[-1], Z),
+        (np.ones(0), w, None, (), -w, Z),
+        (np.ones(0), w, 4, (), -w, Z),
+        # Probability
+        (X, None, None, tuple(-P), None, p * Z),
+        (X, None, 4, -P, None, p * Z),
+        (X[:-1], X[-1], None, tuple(-P[:-1]), -P[-1], p * Z),
+        (X[:-1], X[-1], 4, -P[:-1], -P[-1], p * Z),
+        (np.ones(0), w, None, (), -v, p * Z),
+        (np.ones(0), w, 4, (), -v, p * Z),
+    ]
+
+    @pytest.mark.parametrize(
+        "coeffs, unshifted_coeff, batch_size, res, r0, expected",
+        test_cases_single_shots_single_meas,
+    )
+    def test_single_shots_single_meas(self, coeffs, unshifted_coeff, batch_size, res, r0, expected):
+        """Test that a single shots, single measurement gradient is evaluated correctly."""
+
+        shots = Shots(100)
+        tape_specs = (None, None, 1, shots)
+        data = [None, coeffs, None, unshifted_coeff, None]
+        grad = _evaluate_gradient(tape_specs, res, data, r0, batch_size)
+
+        assert isinstance(grad, np.ndarray)
+        assert grad.shape == expected.shape
+        assert np.allclose(grad, expected)
+
+    exp_probs = (p2 * Z, 2 * p * Z)
+    test_cases_single_shots_multi_meas = [
+        # Expectation values
+        (X, None, None, tuple(zip(-X, -2 * X)), None, (Z, 2 * Z)),
+        (X, None, 4, (-X, -2 * X), None, (Z, 2 * Z)),
+        (X[:-1], X[-1], None, tuple(zip(-X[:-1], -2 * X[:-1])), (-X[-1], -2 * X[-1]), (Z, 2 * Z)),
+        (X[:-1], X[-1], 4, (-X[:-1], -2 * X[:-1]), (-X[-1], -2 * X[-1]), (Z, 2 * Z)),
+        (np.ones(0), w, None, (), (-w, -2 * w), (Z, 2 * Z)),
+        (np.ones(0), w, 4, (), (-w, -2 * w), (Z, 2 * Z)),
+        # Expval and Probability
+        (X, None, None, tuple(zip(-X, -2 * P)), None, (Z, 2 * p * Z)),
+        (X, None, 4, (-X, -2 * P), None, (Z, 2 * p * Z)),
+        (X[:-1], X[-1], None, tuple(zip(-X, -2 * P))[:-1], (-X[-1], -2 * P[-1]), (Z, 2 * p * Z)),
+        (X[:-1], X[-1], 4, (-X[:-1], -2 * P[:-1]), (-X[-1], -2 * P[-1]), (Z, 2 * p * Z)),
+        (np.ones(0), w, None, (), (-w, -2 * v), (Z, 2 * p * Z)),
+        (np.ones(0), w, 4, (), (-w, -2 * v), (Z, 2 * p * Z)),
+        # Probabilities
+        (X, None, None, tuple(zip(-P2, -2 * P)), None, exp_probs),
+        (X, None, 4, (-P2, -2 * P), None, exp_probs),
+        (X[:-1], X[-1], None, tuple(zip(-P2, -2 * P))[:-1], (-P2[-1], -2 * P[-1]), exp_probs),
+        (X[:-1], X[-1], 4, (-P2[:-1], -2 * P[:-1]), (-P2[-1], -2 * P[-1]), exp_probs),
+        (np.ones(0), w, None, (), (-v2, -2 * v), exp_probs),
+        (np.ones(0), w, 4, (), (-v2, -2 * v), exp_probs),
+    ]
+
+    @pytest.mark.parametrize(
+        "coeffs, unshifted_coeff, batch_size, res, r0, expected",
+        test_cases_single_shots_multi_meas,
+    )
+    def test_single_shots_multi_meas(self, coeffs, unshifted_coeff, batch_size, res, r0, expected):
+        """Test that a single shots, multiple measurements gradient is evaluated correctly."""
+
+        shots = Shots(100)
+        tape_specs = (None, None, 2, shots)
+        data = [None, coeffs, None, unshifted_coeff, None]
+        grad = _evaluate_gradient(tape_specs, res, data, r0, batch_size)
+
+        assert isinstance(grad, tuple) and len(grad) == 2
+        for g, e in zip(grad, expected):
+            assert isinstance(g, np.ndarray) and g.shape == e.shape
+            assert np.allclose(g, e)
+
+    shot_vec_X = tuple(zip(*(-c * X for c in shv)))
+    shot_vec_P = tuple(zip(*(-c * P for c in shv)))
+    shot_vec_P_partial = tuple(-c * P[:-1] for c in shv)
+
+    exp_shot_vec_prob = np.outer(shv, p) * Z
+    test_cases_multi_shots_single_meas = [
+        # Expectation value
+        (X, None, None, shot_vec_X, None, shv * Z),
+        (X, None, 4, tuple(-c * X for c in shv), None, shv * Z),
+        (X[:-1], X[-1], None, shot_vec_X[:-1], shot_vec_X[-1], shv * Z),
+        (X[:-1], X[-1], 4, tuple(-c * X[:-1] for c in shv), tuple(-shv * X[-1]), shv * Z),
+        (np.ones(0), w, None, (), tuple(-c * w for c in shv), shv * Z),
+        (np.ones(0), w, 4, ((), (), ()), tuple(-c * w for c in shv), shv * Z),
+        # Probability
+        (X, None, None, shot_vec_P, None, exp_shot_vec_prob),
+        (X, None, 4, tuple(-c * P for c in shv), None, exp_shot_vec_prob),
+        (X[:-1], X[-1], None, shot_vec_P[:-1], shot_vec_P[-1], exp_shot_vec_prob),
+        (X[:-1], X[-1], 4, shot_vec_P_partial, tuple(np.outer(-shv, P[-1])), exp_shot_vec_prob),
+        (np.ones(0), w, None, (), tuple(-c * v for c in shv), exp_shot_vec_prob),
+        (np.ones(0), w, 4, ((), (), ()), tuple(-c * v for c in shv), exp_shot_vec_prob),
+    ]
+
+    @pytest.mark.parametrize(
+        "coeffs, unshifted_coeff, batch_size, res, r0, expected",
+        test_cases_multi_shots_single_meas,
+    )
+    def test_multi_shots_single_meas(self, coeffs, unshifted_coeff, batch_size, res, r0, expected):
+        """Test that a shot vector, single measurements gradient is evaluated correctly."""
+
+        shots = Shots((100, 101, 102))
+        tape_specs = (None, None, 1, shots)
+        data = [None, coeffs, None, unshifted_coeff, None]
+        grad = _evaluate_gradient(tape_specs, res, data, r0, batch_size)
+
+        assert isinstance(grad, tuple) and len(grad) == 3
+        for g, e in zip(grad, expected):
+            assert isinstance(g, np.ndarray) and g.shape == e.shape
+            assert np.allclose(g, e)
+
+    multi_X = tuple(tuple((-c * x, -2 * c * x) for c in shv) for x in X)
+    batched_multi_X = tuple((-c * X, -2 * c * X) for c in shv)
+    partial_multi_X = tuple((-c * X[:-1], -2 * c * X[:-1]) for c in shv)
+    expvals_r0 = tuple((-c * w, -2 * c * w) for c in shv)
+
+    multi_X_P = tuple(tuple((-c * _p, -2 * c * x) for c in shv) for x, _p in zip(X, P))
+    batched_multi_X_P = tuple((-c * P, -2 * c * X) for c in shv)
+    partial_multi_X_P = tuple((-c * P[:-1], -2 * c * X[:-1]) for c in shv)
+    prob_expval_r0 = tuple((-c * v, -2 * c * w) for c in shv)
+
+    multi_P_P = tuple(tuple((-c * _p, -2 * c * _q) for c in shv) for _q, _p in zip(P2, P))
+    batched_multi_P_P = tuple((-c * P, -2 * c * P2) for c in shv)
+    partial_multi_P_P = tuple((-c * P[:-1], -2 * c * P2[:-1]) for c in shv)
+    probs_r0 = tuple((-c * v, -2 * c * v2) for c in shv)
+
+    exp_shot_vec_prob_expval = tuple((c * p * Z, 2 * c * Z) for c in shv)
+    exp_shot_vec_probs = tuple((c * p * Z, 2 * c * p2 * Z) for c in shv)
+    test_cases_multi_shots_multi_meas = [
+        # Expectation values
+        (X, None, None, multi_X, None, shv_m * Z),
+        (X, None, 4, batched_multi_X, None, shv_m * Z),
+        (X[:-1], X[-1], None, multi_X[:-1], multi_X[-1], shv_m * Z),
+        (X[:-1], X[-1], 4, partial_multi_X, multi_X[-1], shv_m * Z),
+        (np.ones(0), w, None, (), expvals_r0, shv_m * Z),
+        (np.ones(0), w, 4, ((), (), ()), expvals_r0, shv_m * Z),
+        # Probability and expectation
+        (X, None, None, multi_X_P, None, exp_shot_vec_prob_expval),
+        (X, None, 4, batched_multi_X_P, None, exp_shot_vec_prob_expval),
+        (X[:-1], X[-1], None, multi_X_P[:-1], multi_X_P[-1], exp_shot_vec_prob_expval),
+        (X[:-1], X[-1], 4, partial_multi_X_P, multi_X_P[-1], exp_shot_vec_prob_expval),
+        (np.ones(0), w, None, (), prob_expval_r0, exp_shot_vec_prob_expval),
+        (np.ones(0), w, 4, ((), (), ()), prob_expval_r0, exp_shot_vec_prob_expval),
+        # Probabilities
+        (X, None, None, multi_P_P, None, exp_shot_vec_probs),
+        (X, None, 4, batched_multi_P_P, None, exp_shot_vec_probs),
+        (X[:-1], X[-1], None, multi_P_P[:-1], multi_P_P[-1], exp_shot_vec_probs),
+        (X[:-1], X[-1], 4, partial_multi_P_P, multi_P_P[-1], exp_shot_vec_probs),
+        (np.ones(0), w, None, (), probs_r0, exp_shot_vec_probs),
+        (np.ones(0), w, 4, ((), (), ()), probs_r0, exp_shot_vec_probs),
+    ]
+
+    @pytest.mark.parametrize(
+        "coeffs, unshifted_coeff, batch_size, res, r0, expected",
+        test_cases_multi_shots_multi_meas,
+    )
+    def test_multi_shots_multi_meas(self, coeffs, unshifted_coeff, batch_size, res, r0, expected):
+        """Test that a shot vector, multiple measurements gradient is evaluated correctly."""
+
+        shots = Shots((100, 101, 102))
+        tape_specs = (None, None, 2, shots)
+        data = [None, coeffs, None, unshifted_coeff, None]
+        grad = _evaluate_gradient(tape_specs, res, data, r0, batch_size)
+
+        assert isinstance(grad, tuple) and len(grad) == 3
+        for g, e in zip(grad, expected):
+            assert isinstance(g, tuple) and len(g) == 2
+            for _g, _e in zip(g, e):
+                assert isinstance(_g, np.ndarray) and _g.shape == _e.shape
+                assert np.allclose(_g, _e)
+
 
 # pylint: disable=too-few-public-methods
 class RY_with_F(qml.RY):
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_cv.py b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
index 064251126..8064bb518 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_cv.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
@@ -14,7 +14,7 @@
 """Tests for the gradients.parameter_shift_cv module."""
 # pylint: disable=protected-access, no-self-use, not-callable, no-value-for-parameter
 
-import unittest.mock as mock
+from unittest import mock
 
 import pytest
 
diff --git a/tests/numpy/test_numpy_random.py b/tests/numpy/test_numpy_random.py
index d8db4ff0e..6cfd9e483 100644
--- a/tests/numpy/test_numpy_random.py
+++ b/tests/numpy/test_numpy_random.py
@@ -43,6 +43,7 @@ bit_generator_classes = [random.MT19937, random.PCG64, random.Philox, random.SFC
 general_gen = random.default_rng()
 
 
+@pytest.mark.unit
 class TestGeneratorDistributions:
     @pytest.mark.parametrize("distribution", distributions_no_extra_input)
     def test_generator_distributions(self, distribution):
@@ -61,6 +62,7 @@ class TestGeneratorDistributions:
         assert output.requires_grad is False
 
 
+@pytest.mark.unit
 class Test_default_rng:
     def test_no_input(self):
         """Tests that np.random.default_rng() returns a generator object when
diff --git a/tests/numpy/test_numpy_wrapper.py b/tests/numpy/test_numpy_wrapper.py
index 0c28828e8..f89bc715b 100644
--- a/tests/numpy/test_numpy_wrapper.py
+++ b/tests/numpy/test_numpy_wrapper.py
@@ -25,6 +25,7 @@ from pennylane import numpy as np
 from pennylane.numpy.tensor import tensor_to_arraybox
 
 
+@pytest.mark.unit
 class TestExtractTensors:
     """Tests for the extract_tensors function"""
 
@@ -56,6 +57,7 @@ class TestExtractTensors:
         assert res[1] is arr2
 
 
+@pytest.mark.unit
 class TestTensor:
     """Tests for the Tensor(ndarray) subclass"""
 
@@ -141,6 +143,7 @@ ARRAY_SHAPE_FNS = [
 ]
 
 
+@pytest.mark.unit
 class TestNumpyIntegration:
     """Test that the wrapped NumPy functionality integrates well
     with standard NumPy functions."""
@@ -428,6 +431,7 @@ class TestNumpyIntegration:
         assert len(res) == 2
 
 
+@pytest.mark.integration
 class TestAutogradIntegration:
     """Test autograd works with the new tensor subclass"""
 
@@ -458,6 +462,7 @@ class TestAutogradIntegration:
             grad_fn(arr1)
 
 
+@pytest.mark.unit
 class TestScalarHashing:
     """Test for the hashing capability of scalar arrays."""
 
@@ -505,6 +510,7 @@ class TestScalarHashing:
 class TestNumpyConversion:
     """Tests for the tensor.unwrap() and tensor.numpy() methods"""
 
+    @pytest.mark.unit
     def test_convert_scalar_array(self):
         """Test that a scalar array converts to a python literal"""
         data = np.array(1.543)
@@ -512,6 +518,7 @@ class TestNumpyConversion:
         assert res == data.item()
         assert isinstance(res, float)
 
+    @pytest.mark.unit
     def test_convert_array(self):
         """Test that a numpy array successfully converts"""
         data = np.array([1, 2, 3])
@@ -522,6 +529,7 @@ class TestNumpyConversion:
         assert isinstance(res, np.ndarray)
         assert not isinstance(res, np.tensor)
 
+    @pytest.mark.system
     def test_single_gate_parameter(self):
         """Test that when supplied a PennyLane tensor, a QNode passes an
         unwrapped tensor as the argument to a gate taking a single parameter"""
@@ -545,6 +553,7 @@ class TestNumpyConversion:
             assert op.name == "RX"
             assert op.parameters == [p]
 
+    @pytest.mark.system
     def test_multiple_gate_parameter(self):
         """Test that when supplied a PennyLane tensor, a QNode passes arguments
         as unwrapped tensors to a gate taking multiple parameters"""
diff --git a/tests/ops/op_math/test_controlled.py b/tests/ops/op_math/test_controlled.py
index b9630d206..2d4234a77 100644
--- a/tests/ops/op_math/test_controlled.py
+++ b/tests/ops/op_math/test_controlled.py
@@ -163,6 +163,17 @@ class TestControlledInit:
         op = Controlled(self.temp_op, (0, 1), control_values=[0, 1])
         assert op.control_values == [False, True]
 
+    @pytest.mark.parametrize("control_values", [True, False, 0, 1])
+    def test_scalar_control_values(self, control_values):
+        """Test assignment of provided control_values."""
+        op = Controlled(self.temp_op, 0, control_values=control_values)
+        assert op.control_values == [control_values]
+
+    def test_tuple_control_values(self):
+        """Test assignment of provided control_values."""
+        op = Controlled(self.temp_op, (0, 1), control_values=(0, 1))
+        assert op.control_values == [False, True]
+
     def test_non_boolean_control_values(self):
         """Test control values are converted to booleans."""
         op = Controlled(self.temp_op, (0, 1, 2), control_values=["", None, 5])
diff --git a/tests/ops/op_math/test_controlled_ops.py b/tests/ops/op_math/test_controlled_ops.py
index a0bbc879d..415502e96 100644
--- a/tests/ops/op_math/test_controlled_ops.py
+++ b/tests/ops/op_math/test_controlled_ops.py
@@ -737,3 +737,15 @@ def test_controlled_method(base, cbase):
     """Tests the _controlled method for parametric ops."""
     # pylint: disable=protected-access
     assert qml.equal(base._controlled("a"), cbase)
+
+
+@pytest.mark.parametrize(
+    "control, control_values",
+    [([0, 1], [True, False]), ([10, "a"], (0, 0)), ([2], 1), (2, (True,))],
+)
+@pytest.mark.parametrize(
+    "base_op", [qml.CRX(0.2, [21, 22]), qml.CNOT([21, 22]), qml.CPhase(0.6, [21, 22])]
+)
+def test_controlling_a_controlled_operation(control, control_values, base_op):
+    """Test that a controlled op can be controlled again."""
+    qml.ctrl(base_op, control=control, control_values=control_values)
diff --git a/tests/ops/qutrit/test_qutrit_channel_ops.py b/tests/ops/qutrit/test_qutrit_channel_ops.py
index be1025d0c..585f623f4 100644
--- a/tests/ops/qutrit/test_qutrit_channel_ops.py
+++ b/tests/ops/qutrit/test_qutrit_channel_ops.py
@@ -19,6 +19,7 @@ import pytest
 from numpy.linalg import matrix_power
 
 import pennylane as qml
+from pennylane import math
 from pennylane import numpy as pnp
 from pennylane.ops.qutrit import channel
 
@@ -74,7 +75,7 @@ class TestQutritDepolarizingChannel:
         dev = qml.device("default.qutrit.mixed")
         prob = pnp.array(0.5, requires_grad=True)
 
-        @qml.qnode(dev)
+        @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(p):
             qml.TRX(angle, wires=0, subspace=(0, 1))
             qml.TRX(angle, wires=0, subspace=(1, 2))
@@ -113,28 +114,24 @@ class TestQutritDepolarizingChannel:
     @staticmethod
     def kraus_fn(p):
         """Gets a matrix of the Kraus matrices to be tested."""
-        return qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
+        return math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
 
     @staticmethod
     def kraus_fn_real(p):
         """Gets a matrix of the real part of the Kraus matrices to be tested."""
-        return qml.math.real(
-            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
-        )
+        return math.real(math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices()))
 
     @staticmethod
     def kraus_fn_imag(p):
         """Gets a matrix of the imaginary part of the Kraus matrices to be tested."""
-        return qml.math.imag(
-            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
-        )
+        return math.imag(math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices()))
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
         """Tests Jacobian of Kraus matrices using autograd."""
         p = pnp.array(0.43, requires_grad=True)
         jac = qml.jacobian(self.kraus_fn_real)(p) + 1j * qml.jacobian(self.kraus_fn_imag)(p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+        assert math.allclose(jac, self.expected_jac_fn(p))
 
     @pytest.mark.torch
     def test_kraus_jac_torch(self):
@@ -144,7 +141,7 @@ class TestQutritDepolarizingChannel:
         p = torch.tensor(0.43, requires_grad=True)
         jacobian = torch.autograd.functional.jacobian
         jac = jacobian(self.kraus_fn_real, p) + 1j * jacobian(self.kraus_fn_imag, p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p.detach().numpy()))
+        assert math.allclose(jac, self.expected_jac_fn(p.detach().numpy()))
 
     @pytest.mark.tf
     def test_kraus_jac_tf(self):
@@ -157,10 +154,10 @@ class TestQutritDepolarizingChannel:
         with tf.GradientTape() as imag_tape:
             imag_out = self.kraus_fn_imag(p)
 
-        real_jac = qml.math.cast(real_tape.jacobian(real_out, p), complex)
-        imag_jac = qml.math.cast(imag_tape.jacobian(imag_out, p), complex)
+        real_jac = math.cast(real_tape.jacobian(real_out, p), complex)
+        imag_jac = math.cast(imag_tape.jacobian(imag_out, p), complex)
         jac = real_jac + 1j * imag_jac
-        assert qml.math.allclose(jac, self.expected_jac_fn(0.43))
+        assert math.allclose(jac, self.expected_jac_fn(0.43))
 
     @pytest.mark.jax
     def test_kraus_jac_jax(self):
@@ -171,4 +168,98 @@ class TestQutritDepolarizingChannel:
 
         p = jax.numpy.array(0.43, dtype=jax.numpy.complex128)
         jac = jax.jacobian(self.kraus_fn, holomorphic=True)(p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+        assert math.allclose(jac, self.expected_jac_fn(p))
+
+
+class TestQutritAmplitudeDamping:
+    """Tests for the qutrit quantum channel QutritAmplitudeDamping"""
+
+    def test_gamma_zero(self, tol):
+        """Test gamma_1=gamma_2=0 gives correct Kraus matrices"""
+        kraus_mats = qml.QutritAmplitudeDamping(0, 0, wires=0).kraus_matrices()
+        assert np.allclose(kraus_mats[0], np.eye(3), atol=tol, rtol=0)
+        assert np.allclose(kraus_mats[1], np.zeros((3, 3)), atol=tol, rtol=0)
+        assert np.allclose(kraus_mats[2], np.zeros((3, 3)), atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("gamma1,gamma2", ((0.1, 0.2), (0.75, 0.75)))
+    def test_gamma_arbitrary(self, gamma1, gamma2, tol):
+        """Test the correct Kraus matrices are returned, also ensures that the sum of gammas can be over 1."""
+        K_0 = np.diag((1, np.sqrt(1 - gamma1), np.sqrt(1 - gamma2)))
+
+        K_1 = np.zeros((3, 3))
+        K_1[0, 1] = np.sqrt(gamma1)
+
+        K_2 = np.zeros((3, 3))
+        K_2[0, 2] = np.sqrt(gamma2)
+
+        expected = [K_0, K_1, K_2]
+        damping_channel = qml.QutritAmplitudeDamping(gamma1, gamma2, wires=0)
+        assert np.allclose(damping_channel.kraus_matrices(), expected, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("gamma1,gamma2", ((1.5, 0.0), (0.0, 1.0 + math.eps)))
+    def test_gamma_invalid_parameter(self, gamma1, gamma2):
+        """Ensures that error is thrown when gamma_1 or gamma_2 are outside [0,1]"""
+        with pytest.raises(ValueError, match="Each probability must be in the interval"):
+            channel.QutritAmplitudeDamping(gamma1, gamma2, wires=0).kraus_matrices()
+
+    @staticmethod
+    def expected_jac_fn(gamma_1, gamma_2):
+        """Gets the expected Jacobian of Kraus matrices"""
+        partial_1 = [math.zeros((3, 3)) for _ in range(3)]
+        partial_1[0][1, 1] = -1 / (2 * math.sqrt(1 - gamma_1))
+        partial_1[1][0, 1] = 1 / (2 * math.sqrt(gamma_1))
+
+        partial_2 = [math.zeros((3, 3)) for _ in range(3)]
+        partial_2[0][2, 2] = -1 / (2 * math.sqrt(1 - gamma_2))
+        partial_2[2][0, 2] = 1 / (2 * math.sqrt(gamma_2))
+
+        return [partial_1, partial_2]
+
+    @staticmethod
+    def kraus_fn(gamma_1, gamma_2):
+        """Gets the Kraus matrices of QutritAmplitudeDamping channel, used for differentiation."""
+        damping_channel = qml.QutritAmplitudeDamping(gamma_1, gamma_2, wires=0)
+        return math.stack(damping_channel.kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        """Tests Jacobian of Kraus matrices using autograd."""
+        gamma_1 = pnp.array(0.43, requires_grad=True)
+        gamma_2 = pnp.array(0.12, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(gamma_1, gamma_2)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        """Tests Jacobian of Kraus matrices using PyTorch."""
+        import torch
+
+        gamma_1 = torch.tensor(0.43, requires_grad=True)
+        gamma_2 = torch.tensor(0.12, requires_grad=True)
+
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, (gamma_1, gamma_2))
+        expected = self.expected_jac_fn(gamma_1.detach().numpy(), gamma_2.detach().numpy())
+        assert math.allclose(jac[0].detach().numpy(), expected[0])
+        assert math.allclose(jac[1].detach().numpy(), expected[1])
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        """Tests Jacobian of Kraus matrices using TensorFlow."""
+        import tensorflow as tf
+
+        gamma_1 = tf.Variable(0.43)
+        gamma_2 = tf.Variable(0.12)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(gamma_1, gamma_2)
+        jac = tape.jacobian(out, (gamma_1, gamma_2))
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        """Tests Jacobian of Kraus matrices using JAX."""
+        import jax
+
+        gamma_1 = jax.numpy.array(0.43)
+        gamma_2 = jax.numpy.array(0.12)
+        jac = jax.jacobian(self.kraus_fn, argnums=[0, 1])(gamma_1, gamma_2)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))
diff --git a/tests/pytest.ini b/tests/pytest.ini
index 85bd26efb..4c2812d4a 100644
--- a/tests/pytest.ini
+++ b/tests/pytest.ini
@@ -1,5 +1,8 @@
 [pytest]
 markers =
+    unit: marks tests as unit-level tests (select with '-m "unit"')
+    integration: marks tests as integration-level tests (select with '-m "integration"')
+    system: marks tests as system-level tests (select with '-m "system"')
     core: marks tests for core testing (select with '-m "core"')
     autograd: marks tests for autograd testing (select with '-m "autograd"')
     torch: marks tests for torch testing (select with '-m "core"')
diff --git a/tests/qchem/of_tests/test_molecular_hamiltonian.py b/tests/qchem/of_tests/test_molecular_hamiltonian.py
index 8d8d68508..355a6b19b 100644
--- a/tests/qchem/of_tests/test_molecular_hamiltonian.py
+++ b/tests/qchem/of_tests/test_molecular_hamiltonian.py
@@ -782,3 +782,77 @@ def test_error_raised_for_missing_molecule_information():
         match="The provided arguments do not contain information about symbols in the molecule.",
     ):
         qchem.molecular_hamiltonian(charge=0, mult=1, method="dhf")
+
+
+@pytest.mark.parametrize(
+    ("method"),
+    [
+        "pyscf",
+        "dhf",
+        "openfermion",
+    ],
+)
+def test_coordinate_units_for_molecular_hamiltonian(method, tmpdir):
+    r"""Test that molecular_hamiltonian generates the Hamiltonian for both Bohr and Angstrom units."""
+
+    symbols = ["H", "H"]
+    geometry_bohr = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+    geometry_ang = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.529177210903]])
+
+    hamiltonian_bohr, _ = qchem.molecular_hamiltonian(
+        symbols,
+        geometry_bohr,
+        unit="bohr",
+        method=method,
+        outpath=tmpdir.strpath,
+    )
+
+    hamiltonian_ang, _ = qchem.molecular_hamiltonian(
+        symbols,
+        geometry_ang,
+        unit="angstrom",
+        method=method,
+        outpath=tmpdir.strpath,
+    )
+    assert qml.equal(hamiltonian_ang, hamiltonian_bohr)
+
+
+@pytest.mark.parametrize(
+    ("method"),
+    [
+        "pyscf",
+        "dhf",
+        "openfermion",
+    ],
+)
+def test_coordinate_units_for_molecular_hamiltonian_molecule_class(method, tmpdir):
+    r"""Test that molecular_hamiltonian generates the Hamiltonian for both Bohr and Angstrom units."""
+
+    symbols = ["H", "H"]
+    geometry_bohr = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+    geometry_ang = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.529177210903]])
+
+    molecule_bohr = qchem.Molecule(symbols, geometry_bohr, unit="bohr")
+    hamiltonian_bohr, _ = qchem.molecular_hamiltonian(
+        molecule_bohr,
+        method=method,
+        outpath=tmpdir.strpath,
+    )
+
+    molecule_ang = qchem.Molecule(symbols, geometry_ang, unit="angstrom")
+    hamiltonian_ang, _ = qchem.molecular_hamiltonian(
+        molecule_ang,
+        method=method,
+        outpath=tmpdir.strpath,
+    )
+    assert qml.equal(hamiltonian_ang, hamiltonian_bohr)
+
+
+def test_unit_error_molecular_hamiltonian():
+    r"""Test that an error is raised if a wrong/not-supported unit for coordinates is entered."""
+
+    symbols = ["H", "H"]
+    geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+
+    with pytest.raises(ValueError, match="The provided unit 'degrees' is not supported."):
+        qchem.molecular_hamiltonian(symbols, geometry, unit="degrees")
diff --git a/tests/qchem/test_molecule.py b/tests/qchem/test_molecule.py
index 5191f3445..bd01511cb 100644
--- a/tests/qchem/test_molecule.py
+++ b/tests/qchem/test_molecule.py
@@ -353,3 +353,12 @@ class TestMolecule:
         assert l == l_ref
         assert alpha == alpha_ref
         assert coeff == coeff_ref
+
+    def test_unit_error(self):
+        r"""Test that an error is raised if a wrong/not-supported unit for coordinates is entered."""
+
+        symbols = ["H", "H"]
+        geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+
+        with pytest.raises(ValueError, match="The provided unit 'degrees' is not supported."):
+            qchem.Molecule(symbols, geometry, unit="degrees")
diff --git a/tests/templates/test_layers/test_strongly_entangling.py b/tests/templates/test_layers/test_strongly_entangling.py
index d02bbe67e..ed9839b84 100644
--- a/tests/templates/test_layers/test_strongly_entangling.py
+++ b/tests/templates/test_layers/test_strongly_entangling.py
@@ -36,6 +36,7 @@ def test_standard_validity():
     qml.ops.functions.assert_valid(op)
 
 
+@pytest.mark.parametrize("batch_dim", [None, 1, 3])
 class TestDecomposition:
     """Tests that the template defines the correct decomposition."""
 
@@ -57,23 +58,39 @@ class TestDecomposition:
     ]
 
     @pytest.mark.parametrize("n_wires, weight_shape, expected_names, expected_wires", QUEUES)
-    def test_expansion(self, n_wires, weight_shape, expected_names, expected_wires):
+    def test_expansion(self, n_wires, weight_shape, expected_names, expected_wires, batch_dim):
         """Checks the queue for the default settings."""
+        # pylint: disable=too-many-arguments
 
+        if batch_dim is not None:
+            weight_shape = (batch_dim,) + weight_shape
         weights = np.random.random(size=weight_shape)
 
         op = qml.StronglyEntanglingLayers(weights, wires=range(n_wires))
         tape = op.expand()
 
+        if batch_dim is None:
+            param_sets = iter(weights.reshape((-1, 3)))
+        else:
+            param_sets = iter(weights.reshape((batch_dim, -1, 3)).transpose([1, 2, 0]))
+
         for i, gate in enumerate(tape.operations):
             assert gate.name == expected_names[i]
+            if gate.name == "Rot":
+                assert gate.batch_size == batch_dim
+                assert qml.math.allclose(gate.data, next(param_sets))
+            else:
+                assert gate.batch_size is None
             assert gate.wires.labels == tuple(expected_wires[i])
 
     @pytest.mark.parametrize("n_layers, n_wires", [(2, 2), (1, 3), (2, 4)])
-    def test_uses_correct_imprimitive(self, n_layers, n_wires):
+    def test_uses_correct_imprimitive(self, n_layers, n_wires, batch_dim):
         """Test that correct number of entanglers are used in the circuit."""
 
-        weights = np.random.randn(n_layers, n_wires, 3)
+        shape = (n_layers, n_wires, 3)
+        if batch_dim is not None:
+            shape = (batch_dim,) + shape
+        weights = np.random.randn(*shape)
 
         op = qml.StronglyEntanglingLayers(weights=weights, wires=range(n_wires), imprimitive=qml.CZ)
         ops = op.expand().operations
@@ -81,9 +98,10 @@ class TestDecomposition:
         gate_names = [gate.name for gate in ops]
         assert gate_names.count("CZ") == n_wires * n_layers
 
-    def test_custom_wire_labels(self, tol):
+    def test_custom_wire_labels(self, tol, batch_dim):
         """Test that template can deal with non-numeric, nonconsecutive wire labels."""
-        weights = np.random.random(size=(1, 3, 3))
+        shape = (1, 3, 3) if batch_dim is None else (batch_dim, 1, 3, 3)
+        weights = np.random.random(size=shape)
 
         dev = qml.device("default.qubit", wires=3)
         dev2 = qml.device("default.qubit", wires=["z", "a", "k"])
@@ -107,10 +125,13 @@ class TestDecomposition:
     @pytest.mark.parametrize(
         "n_layers, n_wires, ranges", [(2, 2, [1, 1]), (1, 3, [2]), (4, 4, [2, 3, 1, 3])]
     )
-    def test_custom_range_sequence(self, n_layers, n_wires, ranges):
+    def test_custom_range_sequence(self, n_layers, n_wires, ranges, batch_dim):
         """Test that correct sequence of custom ranges are used in the circuit."""
 
-        weights = np.random.randn(n_layers, n_wires, 3)
+        shape = (n_layers, n_wires, 3)
+        if batch_dim is not None:
+            shape = (batch_dim,) + shape
+        weights = np.random.randn(*shape)
 
         op = qml.StronglyEntanglingLayers(weights=weights, wires=range(n_wires), ranges=ranges)
         ops = op.expand().operations
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index cdc1fd5f1..4d4faea65 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -1063,7 +1063,7 @@ class TestConditionalOperations:
     )
     def test_cond_qfunc(self, device):
         """Test that a qfunc can also used with qml.cond."""
-        dev = qml.device(device, wires=3)
+        dev = qml.device(device, wires=4)
 
         r = 2.324
 
@@ -1074,12 +1074,14 @@ class TestConditionalOperations:
             qml.CNOT(wires=[0, 1])
             qml.CRY(rads, wires=[0, 1])
             qml.CZ(wires=[0, 1])
-            return qml.probs(wires=1)
+            qml.ctrl(qml.CRX, control=0, control_values=[1])(0.5, [1, 2])
+            return qml.probs(wires=[1, 2])
 
         def f(x):
             qml.PauliX(1)
             qml.RY(x, wires=1)
             qml.PauliZ(1)
+            qml.CRX(0.5, [1, 2])
 
         @qml.defer_measurements
         @qml.qnode(dev)
@@ -1087,7 +1089,7 @@ class TestConditionalOperations:
             qml.Hadamard(0)
             m_0 = qml.measure(0)
             qml.cond(m_0, f)(r)
-            return qml.probs(wires=1)
+            return qml.probs(wires=[1, 2])
 
         exp = normal_circuit(r)
         cond_probs = quantum_control_circuit(r)