Repository: PennyLaneAI/pennylane
Pull request #5082 (2024-01-19 05:26:04)

qutrit mixed state sampling
**Context:**
Necessary next step in `qutrit_mixed` device development, allows for sampling a quantum circuit on the mixed state. This will be necessary to simulate sampling a qutrit device.

**Description of the Change:**
Added functionality for sampling a qutrit mixed-state via the new `sample_state` and `measure_with_samples` functions.

**Benefits:**
Allows for sampling a qutrit mixed-state, allowing for noisy qutrit simulation sampling

**Possible Drawbacks:**
Abstracting for qubits and more generally qutrits may have added challenges and will require a refactor or copied code code-smell. Methods are not the fastest possible, optimizations can still be made.

**Related GitHub Issues:**
N/A


Gabriel-Bottrill: Current state of sampling PR, I could remove it and make it closer to what is in `default.qutrit`. I plan to cut it into two to also include expval, var, ... . This however should be a good start to get what is most necessesary while also leaving abstractions to make it easy to include other functionality.
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5082?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `21 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`0428922`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/0428922017c61c6bc0405582653eae2a9ecbd8cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`641bdc2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5082?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.62%.

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

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


```diff
@@            Coverage Diff             @@
##           master    #5082      +/-   ##
==========================================
- Coverage   99.68%   99.62%   -0.07%     
==========================================
  Files         394      396       +2     
  Lines       36298    36247      -51     
==========================================
- Hits        36183    36110      -73     
- Misses        115      137      +22     
```



</details>

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

Gabriel-Bottrill: @albi3ro @glassnotes @mudit2812 @trbromley  I cleaned up this PR since measure is close to being merged. Hope you like it!
Gabriel-Bottrill: Measure has been merged, so this is ready for review.
Gabriel-Bottrill: Fixed broken tests. Should be good to go!
trbromley: Thanks! @mudit2812 do you have the ability to merge in `master` to the PR branch so that we can do a final round of tests?
mudit2812: > Thanks! @mudit2812 do you have the ability to merge in `master` to the PR branch so that we can do a final round of tests?

I will check on the next PR, but I think I should as I'm part of the org that owns the fork.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2e201a45f..7138095e0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -98,6 +98,12 @@
 
 * Upgraded `null.qubit` to the new device API. Also, added support for all measurements and various modes of differentiation.
   [(#5211)](https://github.com/PennyLaneAI/pennylane/pull/5211)
+  
+<h4>Community contributions 🥳</h4>
+
+* Functions `measure_with_samples` and `sample_state` have been added to the new `qutrit_mixed` module found in
+ `qml.devices`. These functions are used to sample device-compatible states, returning either the final measured state or value of an observable.
+  [(#5082)](https://github.com/PennyLaneAI/pennylane/pull/5082)
 
 * The `QNode` now defers `diff_method` validation to the device under the new device api `qml.devices.Device`.
   [(#5176)](https://github.com/PennyLaneAI/pennylane/pull/5176)
@@ -137,6 +143,7 @@ This release contains contributions from (in alphabetical order):
 
 Korbinian Kottmann,
 Guillermo Alonso,
+Gabriel Bottrill,
 Astral Cai,
 Amintor Dusko,
 Pietropaolo Frisoni,
diff --git a/pennylane/devices/qutrit_mixed/__init__.py b/pennylane/devices/qutrit_mixed/__init__.py
index ed62cd929..68ec69e72 100644
--- a/pennylane/devices/qutrit_mixed/__init__.py
+++ b/pennylane/devices/qutrit_mixed/__init__.py
@@ -25,8 +25,11 @@ at your own discretion.
     create_initial_state
     apply_operation
     measure
+    measure_with_samples
+    sample_state
 """
 
 from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
 from .measure import measure
+from .sampling import sample_state, measure_with_samples
diff --git a/pennylane/devices/qutrit_mixed/sampling.py b/pennylane/devices/qutrit_mixed/sampling.py
new file mode 100644
index 000000000..89aec1378
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/sampling.py
@@ -0,0 +1,366 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Code relevant for sampling a qutrit mixed state.
+"""
+import functools
+import numpy as np
+import pennylane as qml
+from pennylane import math
+from pennylane.ops import Sum, Hamiltonian
+from pennylane.measurements import (
+    Shots,
+    SampleMeasurement,
+    SampleMP,
+    CountsMP,
+    ExpectationMP,
+    VarianceMP,
+)
+from pennylane.typing import TensorLike
+from .utils import QUDIT_DIM, get_num_wires
+from .measure import measure
+from .apply_operation import apply_operation
+
+
+def _apply_diagonalizing_gates(
+    mp: SampleMeasurement, state: np.ndarray, is_state_batched: bool = False
+):
+    """Applies diagonalizing gates when necessary"""
+    if mp.obs:
+        for op in mp.diagonalizing_gates():
+            state = apply_operation(op, state, is_state_batched=is_state_batched)
+
+    return state
+
+
+def _process_samples(
+    mp,
+    samples,
+    wire_order,
+):
+    """Processes samples like SampleMP.process_samples, but fixed for qutrits"""
+    wire_map = dict(zip(wire_order, range(len(wire_order))))
+    mapped_wires = [wire_map[w] for w in mp.wires]
+
+    if mapped_wires:
+        # if wires are provided, then we only return samples from those wires
+        samples = samples[..., mapped_wires]
+
+    num_wires = samples.shape[-1]  # wires is the last dimension
+
+    if mp.obs is None:
+        # if no observable was provided then return the raw samples
+        return samples
+
+    # Replace the basis state in the computational basis with the correct eigenvalue.
+    # Extract only the columns of the basis samples required based on ``wires``.
+    powers_of_three = QUDIT_DIM ** qml.math.arange(num_wires)[::-1]
+    indices = qml.math.array(samples @ powers_of_three)
+    return mp.eigvals()[indices]
+
+
+def _process_counts_samples(processed_sample, mp_has_obs):
+    """Processes a set of samples and counts the results."""
+    observables, counts = math.unique(processed_sample, return_counts=True, axis=0)
+    if not mp_has_obs:
+        observables = ["".join(observable.astype("str")) for observable in observables]
+    return dict(zip(observables, counts))
+
+
+def _process_expval_samples(processed_sample):
+    """Processes a set of samples and returns the expectation value of an observable."""
+    eigvals, counts = math.unique(processed_sample, return_counts=True)
+    probs = counts / math.sum(counts)
+    return math.dot(probs, eigvals)
+
+
+def _process_variance_samples(processed_sample):
+    """Processes a set of samples and returns the variance of an observable."""
+    eigvals, counts = math.unique(processed_sample, return_counts=True)
+    probs = counts / math.sum(counts)
+    return math.dot(probs, (eigvals**2)) - math.dot(probs, eigvals) ** 2
+
+
+# pylint:disable = too-many-arguments
+def _measure_with_samples_diagonalizing_gates(
+    mp: SampleMeasurement,
+    state: np.ndarray,
+    shots: Shots,
+    is_state_batched: bool = False,
+    rng=None,
+    prng_key=None,
+) -> TensorLike:
+    """Returns the samples of the measurement process performed on the given state,
+    by rotating the state into the measurement basis using the diagonalizing gates
+    given by the measurement process.
+
+    Args:
+        mp (~.measurements.SampleMeasurement): The sample measurement to perform
+        state (np.ndarray[complex]): The state vector to sample from
+        shots (~.measurements.Shots): The number of samples to take
+        is_state_batched (bool): whether the state is batched or not
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. Only for simulation using JAX.
+
+    Returns:
+        TensorLike[Any]: Sample measurement results
+    """
+    # apply diagonalizing gates
+    state = _apply_diagonalizing_gates(mp, state, is_state_batched)
+
+    total_indices = get_num_wires(state, is_state_batched)
+    wires = qml.wires.Wires(range(total_indices))
+
+    def _process_single_shot(samples):
+        samples_processed = _process_samples(mp, samples, wires)
+        if isinstance(mp, SampleMP):
+            return math.squeeze(samples_processed)
+        if isinstance(mp, CountsMP):
+            process_func = functools.partial(_process_counts_samples, mp_has_obs=mp.obs is not None)
+        elif isinstance(mp, ExpectationMP):
+            process_func = _process_expval_samples
+        elif isinstance(mp, VarianceMP):
+            process_func = _process_variance_samples
+        else:
+            raise NotImplementedError
+
+        if is_state_batched:
+            ret = []
+            for processed_sample in samples_processed:
+                ret.append(process_func(processed_sample))
+            return math.squeeze(ret)
+        return process_func(samples_processed)
+
+    # if there is a shot vector, build a list containing results for each shot entry
+    if shots.has_partitioned_shots:
+        processed_samples = []
+        for s in shots:
+            # Like default.qubit currently calling sample_state for each shot entry,
+            # but it may be better to call sample_state just once with total_shots,
+            # then use the shot_range keyword argument
+            samples = sample_state(
+                state,
+                shots=s,
+                is_state_batched=is_state_batched,
+                wires=wires,
+                rng=rng,
+                prng_key=prng_key,
+            )
+            processed_samples.append(_process_single_shot(samples))
+
+        return tuple(processed_samples)
+
+    samples = sample_state(
+        state,
+        shots=shots.total_shots,
+        is_state_batched=is_state_batched,
+        wires=wires,
+        rng=rng,
+        prng_key=prng_key,
+    )
+
+    return _process_single_shot(samples)
+
+
+def _measure_sum_with_samples(
+    mp: SampleMeasurement,
+    state: np.ndarray,
+    shots: Shots,
+    is_state_batched: bool = False,
+    rng=None,
+    prng_key=None,
+):
+    """Compute expectation values of Sum or Hamiltonian Observables"""
+    # mp.obs returns is the list of observables for Sum,
+    # mp.obs.terms()[1] returns is the list of observables for Hamiltonian
+    obs_terms = mp.obs if isinstance(mp.obs, Sum) else mp.obs.terms()[1]
+
+    def _sum_for_single_shot(s):
+        results = []
+        for term in obs_terms:
+            results.append(
+                measure_with_samples(
+                    ExpectationMP(term),
+                    state,
+                    s,
+                    is_state_batched=is_state_batched,
+                    rng=rng,
+                    prng_key=prng_key,
+                )
+            )
+
+        if isinstance(mp.obs, Hamiltonian):
+            # If Hamiltonian apply coefficients
+            return sum((c * res for c, res in zip(mp.obs.terms()[0], results)))
+        return sum(results)
+
+    if shots.has_partitioned_shots:
+        return tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
+    return _sum_for_single_shot(shots)
+
+
+def _sample_state_jax(
+    state,
+    shots: int,
+    prng_key,
+    is_state_batched: bool = False,
+    wires=None,
+) -> np.ndarray:
+    """Returns a series of samples of a state for the JAX interface based on the PRNG.
+
+    Args:
+        state (array[complex]): A state vector to be sampled
+        shots (int): The number of samples to take
+        prng_key (jax.random.PRNGKey): A``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator.
+        is_state_batched (bool): whether the state is batched or not
+        wires (Sequence[int]): The wires to sample
+
+    Returns:
+        ndarray[int]: Sample values of the shape (shots, num_wires)
+    """
+    # pylint: disable=import-outside-toplevel
+    import jax
+    import jax.numpy as jnp
+
+    key = prng_key
+
+    total_indices = get_num_wires(state, is_state_batched)
+    state_wires = qml.wires.Wires(range(total_indices))
+
+    wires_to_sample = wires or state_wires
+    num_wires = len(wires_to_sample)
+    basis_states = np.arange(QUDIT_DIM**num_wires)
+
+    with qml.queuing.QueuingManager.stop_recording():
+        probs = measure(qml.probs(wires=wires_to_sample), state, is_state_batched)
+
+    if is_state_batched:
+        # Produce separate keys for each of the probabilities along the broadcasted axis
+        keys = []
+        for _ in state:
+            key, subkey = jax.random.split(key)
+            keys.append(subkey)
+        samples = jnp.array(
+            [
+                jax.random.choice(_key, basis_states, shape=(shots,), p=prob)
+                for _key, prob in zip(keys, probs)
+            ]
+        )
+    else:
+        samples = jax.random.choice(key, basis_states, shape=(shots,), p=probs)
+
+    res = np.zeros(samples.shape + (num_wires,), dtype=np.int64)
+    for i in range(num_wires):
+        res[..., -(i + 1)] = (samples // (QUDIT_DIM**i)) % QUDIT_DIM
+    return res
+
+
+def sample_state(
+    state,
+    shots: int,
+    is_state_batched: bool = False,
+    wires=None,
+    rng=None,
+    prng_key=None,
+) -> np.ndarray:
+    """Returns a series of computational basis samples of a state.
+
+    Args:
+        state (array[complex]): A state vector to be sampled
+        shots (int): The number of samples to take
+        is_state_batched (bool): whether the state is batched or not
+        wires (Sequence[int]): The wires to sample
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]):
+            A seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. Only for simulation using JAX.
+
+    Returns:
+        ndarray[int]: Sample values of the shape (shots, num_wires)
+    """
+    if prng_key is not None:
+        return _sample_state_jax(
+            state, shots, prng_key, is_state_batched=is_state_batched, wires=wires
+        )
+
+    rng = np.random.default_rng(rng)
+
+    total_indices = get_num_wires(state, is_state_batched)
+    state_wires = qml.wires.Wires(range(total_indices))
+
+    wires_to_sample = wires or state_wires
+    num_wires = len(wires_to_sample)
+    basis_states = np.arange(QUDIT_DIM**num_wires)
+
+    with qml.queuing.QueuingManager.stop_recording():
+        probs = measure(qml.probs(wires=wires_to_sample), state, is_state_batched)
+
+    if is_state_batched:
+        # rng.choice doesn't support broadcasting
+        samples = np.stack([rng.choice(basis_states, shots, p=p) for p in probs])
+    else:
+        samples = rng.choice(basis_states, shots, p=probs)
+
+    res = np.zeros(samples.shape + (num_wires,), dtype=np.int64)
+    for i in range(num_wires):
+        res[..., -(i + 1)] = (samples // (QUDIT_DIM**i)) % QUDIT_DIM
+    return res
+
+
+def measure_with_samples(
+    mp: SampleMeasurement,
+    state: np.ndarray,
+    shots: Shots,
+    is_state_batched: bool = False,
+    rng=None,
+    prng_key=None,
+) -> TensorLike:
+    """Returns the samples of the measurement process performed on the given state.
+    This function assumes that the user-defined wire labels in the measurement process
+    have already been mapped to integer wires used in the device.
+
+    Args:
+        mp (SampleMeasurement): The sample measurement to perform
+        state (np.ndarray[complex]): The state vector to sample from
+        shots (Shots): The number of samples to take
+        is_state_batched (bool): whether the state is batched or not
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. Only for simulation using JAX.
+
+    Returns:
+        TensorLike[Any]: Sample measurement results
+    """
+
+    if isinstance(mp, ExpectationMP) and isinstance(mp.obs, (Hamiltonian, Sum)):
+        measure_fn = _measure_sum_with_samples
+    else:
+        # measure with the usual method (rotate into the measurement basis)
+        measure_fn = _measure_with_samples_diagonalizing_gates
+
+    return measure_fn(
+        mp,
+        state,
+        shots,
+        is_state_batched=is_state_batched,
+        rng=rng,
+        prng_key=prng_key,
+    )


Test patch:
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py
new file mode 100644
index 000000000..4f15954fa
--- /dev/null
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py
@@ -0,0 +1,688 @@
+# 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 sampling states in devices/qutrit_mixed."""
+import pytest
+import numpy as np
+from flaky import flaky
+
+import pennylane as qml
+from pennylane import math
+from pennylane.devices.qutrit_mixed import (
+    sample_state,
+    measure_with_samples,
+    create_initial_state,
+    apply_operation,
+)
+from pennylane.devices.qutrit_mixed.sampling import _sample_state_jax
+from pennylane.measurements import Shots
+
+
+APPROX_ATOL = 0.01
+QUDIT_DIM = 3
+ONE_QUTRIT = 1
+TWO_QUTRITS = 2
+THREE_QUTRITS = 3
+
+ml_frameworks_list = [
+    "numpy",
+    pytest.param("autograd", marks=pytest.mark.autograd),
+    pytest.param("jax", marks=pytest.mark.jax),
+    pytest.param("torch", marks=pytest.mark.torch),
+    pytest.param("tensorflow", marks=pytest.mark.tf),
+]
+
+shots_and_copies = (
+    [
+        [(100,), 1],
+        [((100, 1),), 1],
+        [((100, 2),), 2],
+        [(100, 100), 2],
+        [(100, 200), 2],
+        [(100, 100, 200), 3],
+        [(200, (100, 2)), 3],
+    ],
+)
+
+
+def get_dm_of_state(state_vector, num_qudits, normalization=1):
+    state = np.outer(state_vector, np.conj(state_vector)) / normalization
+    return state.reshape((QUDIT_DIM,) * num_qudits * 2)
+
+
+@pytest.fixture(name="two_qutrit_pure_state")
+def fixture_two_qutrit_pure_state():
+    state_vector = np.array([0, 0, 1.0j, -1.0, 0, 0, 0, 1.0, 0], dtype=np.complex128)
+    return get_dm_of_state(state_vector, TWO_QUTRITS, 3)
+
+
+@pytest.fixture(name="batched_two_qutrit_pure_state")
+def fixture_batched_two_qutrit_pure_state():
+    state_vectors = [
+        np.array([0, 0, 0, 0, 1, 0, 0, 0, 0]),
+        np.array([1, 0, 0, 0, 0, 1, 0, 0, 0]) / np.sqrt(2),
+        np.array([0, 1, 0, 1, 0, 0, 0, 1, 1]) / 2,
+    ]
+    states = [get_dm_of_state(state_vector, TWO_QUTRITS) for state_vector in state_vectors]
+    return math.stack(states)
+
+
+def samples_to_probs(samples, num_wires):
+    """Converts samples to probs"""
+    samples_decimal = [np.ravel_multi_index(sample, [QUDIT_DIM] * num_wires) for sample in samples]
+    counts = [0] * (QUDIT_DIM**num_wires)
+
+    for sample in samples_decimal:
+        counts[sample] += 1
+
+    return np.array(counts, dtype=np.float64) / len(samples)
+
+
+def assert_correct_sampled_two_qutrit_pure_state(samples):
+    """Asserts that the returned samples of the two qutrit pure state only contains expected states"""
+    for sample in samples:
+        assert (
+            math.allequal(sample, [0, 2])
+            or math.allequal(sample, [1, 0])
+            or math.allequal(sample, [2, 1])
+        )
+
+
+def assert_correct_sampled_batched_two_qutrit_pure_state(samples):
+    """Asserts that the returned samples of the two qutrit batched state only contains expected states"""
+    # first batch of samples is always |11>
+    assert np.all(samples[0] == 1)
+
+    # second batch of samples is either |00> or |12>
+    assert np.all(np.logical_or(samples[1] == [0, 0], samples[1] == [1, 2]))
+
+    # third batch of samples can be any of |01>, |10>, |21>, or |22>
+    third_batch_bool_1 = np.logical_or(samples[2] == [0, 1], samples[2] == [1, 0])
+    third_batch_bool_2 = np.logical_or(samples[2] == [2, 1], samples[2] == [2, 2])
+    assert np.all(np.logical_or(third_batch_bool_1, third_batch_bool_2))
+
+
+class TestSampleState:
+    """Test that the sample_state function works as expected"""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_sample_state_basic(self, interface, two_qutrit_pure_state):
+        """Tests that the returned samples are as expected."""
+        state = math.array(two_qutrit_pure_state, like=interface)
+        samples = sample_state(state, 10)
+        assert samples.shape == (10, 2)
+        assert samples.dtype == np.int64
+        assert_correct_sampled_two_qutrit_pure_state(samples)
+
+    @pytest.mark.jax
+    def test_prng_key_as_seed_uses_sample_state_jax(self, mocker, two_qutrit_state):
+        """Tests that sample_state calls _sample_state_jax if the seed is a JAX PRNG key"""
+        import jax
+
+        spy = mocker.spy(qml.devices.qutrit_mixed.sampling, "_sample_state_jax")
+        state = qml.math.array(two_qutrit_state, like="jax")
+
+        # prng_key specified, should call _sample_state_jax
+        _ = sample_state(state, 10, prng_key=jax.random.PRNGKey(15))
+        # prng_key defaults to None, should NOT call _sample_state_jax
+        _ = sample_state(state, 10, rng=15)
+
+        spy.assert_called_once()
+
+    @pytest.mark.jax
+    def test_sample_state_jax(self, two_qutrit_pure_state):
+        """Tests that the returned samples are as expected when explicitly calling _sample_state_jax."""
+        import jax
+
+        state = qml.math.array(two_qutrit_pure_state, like="jax")
+
+        samples = _sample_state_jax(state, 10, prng_key=jax.random.PRNGKey(84))
+
+        assert samples.shape == (10, 2)
+        assert samples.dtype == np.int64
+        assert_correct_sampled_two_qutrit_pure_state(samples)
+
+    @pytest.mark.jax
+    def test_prng_key_determines_sample_state_jax_results(self, two_qutrit_pure_state):
+        """Test that setting the seed as a JAX PRNG key determines the results for _sample_state_jax"""
+        import jax
+
+        state = qml.math.array(two_qutrit_pure_state, like="jax")
+
+        samples = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(12))
+        samples2 = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(12))
+        samples3 = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(13))
+
+        assert np.all(samples == samples2)
+        assert not np.allclose(samples, samples3)
+
+    def test_sample_state_custom_rng(self, two_qutrit_state):
+        """Tests that a custom RNG can be used with sample_state."""
+        custom_rng = np.random.default_rng(12345)
+        samples = sample_state(two_qutrit_state, 4, rng=custom_rng)
+        expected = [[0, 2], [1, 0], [2, 1], [1, 2]]
+        assert qml.math.allequal(samples, expected)
+
+    @flaky
+    def test_entangled_qutrit_samples_always_match(self):
+        """Tests that entangled qutrits are always in the same state."""
+        num_samples = 10000
+
+        bell_state_vector = np.array([[1, 0, 0, 0, 1, 0, 0, 0, 1]])
+        bell_state = get_dm_of_state(bell_state_vector, 2, 3)
+
+        samples = sample_state(bell_state, num_samples)
+        assert samples.shape == (num_samples, 2)
+        assert not any(samples[:, 0] ^ samples[:, 1])  # all samples are entangled
+
+        # all samples are approximately equivalently sampled
+        assert np.isclose(
+            math.count_nonzero(samples[:, 0] == 0) / num_samples,
+            1 / 3,
+            atol=APPROX_ATOL,
+        )
+        assert np.isclose(
+            math.count_nonzero(samples[:, 0] == 1) / num_samples,
+            1 / 3,
+            atol=APPROX_ATOL,
+        )
+        assert np.isclose(
+            math.count_nonzero(samples[:, 0] == 2) / num_samples,
+            1 / 3,
+            atol=APPROX_ATOL,
+        )
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize(
+        "num_wires,shuffled_wires",
+        (
+            (8, [7, 5, 3, 1, 6, 0, 4, 2]),
+            (9, [4, 0, 2, 1, 5, 7, 6, 8, 3]),
+        ),
+    )
+    def test_sample_state_many_wires(self, num_wires, shuffled_wires):
+        """Tests that sample_state works as expected with many wires, and with re-ordering."""
+        shots = 10000
+        probs_shape = (QUDIT_DIM,) * num_wires
+        state_shape = probs_shape * 2
+        diagonal = np.arange(1, QUDIT_DIM**num_wires + 1, dtype=np.float64)
+        diagonal = diagonal / np.sum(diagonal)
+        state = np.diag(diagonal).reshape(state_shape)
+
+        ordered_samples = sample_state(state, shots)
+        ordered_probs = samples_to_probs(ordered_samples, num_wires)
+        assert np.allclose(ordered_probs, diagonal, atol=APPROX_ATOL)
+
+        random_samples = sample_state(state, shots, wires=shuffled_wires)
+        random_probs = samples_to_probs(random_samples, num_wires)
+
+        reordered_probs = ordered_probs.reshape(probs_shape).transpose(shuffled_wires).flatten()
+        assert np.allclose(reordered_probs, random_probs, atol=APPROX_ATOL)
+
+
+class TestMeasureWithSamples:
+    """Test that the measure_with_samples function works as expected"""
+
+    def test_sample_measure(self, two_qutrit_pure_state):
+        """Test that a sample measurement works as expected"""
+        shots = qml.measurements.Shots(100)
+        mp = qml.sample(wires=range(2))
+
+        result = measure_with_samples(mp, two_qutrit_pure_state, shots=shots)
+
+        assert result.shape == (shots.total_shots, 2)
+        assert result.dtype == np.int64
+        assert_correct_sampled_two_qutrit_pure_state(result)
+
+    def test_sample_measure_single_wire(self):
+        """Test that a sample measurement on a single wire works as expected"""
+        state_vector = np.array([1, -1j, 1, 0, 0, 0, 0, 0, 0])
+        state = get_dm_of_state(state_vector, 2, 3)
+        shots = qml.measurements.Shots(100)
+
+        mp0 = qml.sample(wires=0)
+        mp1 = qml.sample(wires=1)
+
+        result0 = measure_with_samples(mp0, state, shots=shots)
+        result1 = measure_with_samples(mp1, state, shots=shots)
+
+        assert result0.shape == (shots.total_shots,)
+        assert result0.dtype == np.int64
+        assert np.all(result0 == 0)
+
+        assert result1.shape == (shots.total_shots,)
+        assert result1.dtype == np.int64
+        assert len(np.unique(result1)) == 3
+
+    def test_approximate_sample_measure(self, two_qutrit_pure_state):
+        """Test that a sample measurement returns approximately the correct distribution"""
+        shots = qml.measurements.Shots(10000)
+        mp = qml.sample(wires=range(2))
+
+        result = measure_with_samples(mp, two_qutrit_pure_state, shots=shots, rng=1234)
+
+        one_or_two_prob = np.count_nonzero(result[:, 0]) / result.shape[0]
+        one_prob = np.count_nonzero(result[:, 0] == 1) / result.shape[0]
+        assert np.allclose(one_or_two_prob, 2 / 3, atol=APPROX_ATOL)
+        assert np.allclose(one_prob, 1 / 3, atol=APPROX_ATOL)
+
+    def test_approximate_expval_measure(self, two_qutrit_state):
+        """Test that an expval measurement works as expected"""
+        state = qml.math.array(two_qutrit_state)
+        shots = qml.measurements.Shots(10000)
+        mp = qml.expval(qml.GellMann(0, 1) @ qml.GellMann(1, 1))
+
+        result = measure_with_samples(mp, state, shots=shots, rng=1234)
+
+        gellmann_1_matrix = qml.GellMann.compute_matrix(1)
+        observable_matrix = np.kron(gellmann_1_matrix, gellmann_1_matrix)
+        expected = np.trace(observable_matrix @ state.reshape((9, 9)))
+
+        assert isinstance(result, np.float64)
+        assert np.allclose(result, expected, atol=APPROX_ATOL)
+
+    def test_approximate_var_measure(self, two_qutrit_state):
+        """Test that a variance measurement works as expected"""
+        state = qml.math.array(two_qutrit_state)
+        shots = qml.measurements.Shots(10000)
+        mp = qml.var(qml.GellMann(0, 1) @ qml.GellMann(1, 1))
+
+        result = measure_with_samples(mp, state, shots=shots, rng=123)
+
+        gellmann_1_matrix = qml.GellMann.compute_matrix(1)
+        obs_mat = np.kron(gellmann_1_matrix, gellmann_1_matrix)
+        reshaped_state = state.reshape((9, 9))
+        obs_squared = np.linalg.matrix_power(obs_mat, 2)
+        expected = np.trace(obs_squared @ reshaped_state) - np.trace(obs_mat @ reshaped_state) ** 2
+
+        assert isinstance(result, np.float64)
+        assert np.allclose(result, expected, atol=APPROX_ATOL)
+
+    @flaky
+    def test_counts_measure(self, two_qutrit_pure_state):
+        """Test that a counts measurement works as expected"""
+        num_shots = 10000
+        shots = qml.measurements.Shots(num_shots)
+        mp = qml.counts()
+
+        result = measure_with_samples(mp, two_qutrit_pure_state, shots=shots)
+
+        assert isinstance(result, dict)
+        assert sorted(result.keys()) == ["02", "10", "21"]
+        assert np.isclose(result["02"] / num_shots, 1 / 3, atol=APPROX_ATOL)
+        assert np.isclose(result["10"] / num_shots, 1 / 3, atol=APPROX_ATOL)
+        assert np.isclose(result["21"] / num_shots, 1 / 3, atol=APPROX_ATOL)
+
+    @flaky
+    def test_counts_measure_single_wire(self):
+        """Test that a counts measurement on a single wire works as expected"""
+        state_vector = np.sqrt(np.array([5, -2j, 1, 0, 0, 0, 0, 0, 0]) / 8)
+        state = get_dm_of_state(state_vector, 2)
+
+        num_shots = 10000
+        shots = qml.measurements.Shots(num_shots)
+
+        mp0 = qml.counts(wires=0)
+        mp1 = qml.counts(wires=1)
+
+        result0 = measure_with_samples(mp0, state, shots=shots)
+        result1 = measure_with_samples(mp1, state, shots=shots)
+
+        assert isinstance(result1, dict)
+        assert result0 == {"0": num_shots}
+
+        assert isinstance(result1, dict)
+        assert sorted(result1.keys()) == ["0", "1", "2"]
+        assert np.isclose(result1["0"] / num_shots, 5 / 8, atol=APPROX_ATOL)
+        assert np.isclose(result1["1"] / num_shots, 1 / 4, atol=APPROX_ATOL)
+        assert np.isclose(result1["2"] / num_shots, 1 / 8, atol=APPROX_ATOL)
+
+    def test_sample_observables(self):
+        """Test that counts measurements properly counts samples of an observable"""
+        state_vector = np.sqrt(np.array([0, 0, 0, 0, 2, 0, 1, 0, 1]) / 4)
+        state = get_dm_of_state(state_vector, 2)
+        num_shots = 100
+        shots = qml.measurements.Shots(num_shots)
+
+        results_gel_3 = measure_with_samples(qml.sample(qml.GellMann(0, 3)), state, shots=shots)
+        assert results_gel_3.shape == (shots.total_shots,)
+        assert results_gel_3.dtype == np.int64
+        assert sorted(np.unique(results_gel_3)) == [-1, 0]
+
+        results_gel_1s = measure_with_samples(
+            qml.sample(qml.GellMann(0, 1) @ qml.GellMann(1, 1)), state, shots=shots
+        )
+        assert results_gel_1s.shape == (shots.total_shots,)
+        assert results_gel_1s.dtype == np.int64
+        assert sorted(np.unique(results_gel_1s)) == [-1, 0, 1]
+
+    @flaky
+    def test_counts_observables(self):
+        """Test that a set of sample and counts measurements works as expected"""
+        state_vector = np.sqrt(np.array([0, 0, 0, 0, 3, 0, 1, 0, 1]) / 5)
+        state = get_dm_of_state(state_vector, 2)
+        num_shots = 10000
+        shots = qml.measurements.Shots(num_shots)
+
+        results_gel_3 = measure_with_samples(qml.counts(qml.GellMann(0, 3)), state, shots=shots)
+
+        assert isinstance(results_gel_3, dict)
+        assert sorted(results_gel_3.keys()) == [-1, 0]
+        assert np.isclose(results_gel_3[-1] / num_shots, 3 / 5, atol=APPROX_ATOL)
+        assert np.isclose(results_gel_3[0] / num_shots, 2 / 5, atol=APPROX_ATOL)
+
+        results_gel_1s = measure_with_samples(
+            qml.counts(qml.GellMann(0, 1) @ qml.GellMann(1, 1)), state, shots=shots
+        )
+        assert isinstance(results_gel_1s, dict)
+        assert sorted(results_gel_1s.keys()) == [-1, 0, 1]
+        assert np.isclose(results_gel_1s[-1] / num_shots, 0.3, atol=APPROX_ATOL)
+        assert np.isclose(results_gel_1s[0] / num_shots, 0.4, atol=APPROX_ATOL)
+        assert np.isclose(results_gel_1s[1] / num_shots, 0.3, atol=APPROX_ATOL)
+
+
+class TestInvalidSampling:
+    """Tests for non-expected states and inputs."""
+
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_only_catch_nan_errors(self, shots):
+        """Test that when probabilities don't add to 1 Error is thrown."""
+        state = np.zeros((3,) * QUDIT_DIM * 2).astype(np.complex128)
+        mp = qml.sample(wires=range(2))
+        _shots = Shots(shots)
+
+        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+            _ = measure_with_samples(mp, state, _shots)
+
+    @pytest.mark.parametrize("mp", [qml.probs(0), qml.probs(op=qml.GellMann(0, 1))])
+    def test_currently_unsupported_observable(self, mp, two_qutrit_state):
+        """Test sample measurements that are not sample, counts, expval,
+        or var raise a NotImplementedError."""
+        shots = qml.measurements.Shots(1)
+        with pytest.raises(NotImplementedError):
+            _ = measure_with_samples(mp, two_qutrit_state, shots)
+
+
+shots_to_test_samples = [
+    ((100, 2),),
+    (100, 100),
+    (100, 100),
+    (100, 100, 200),
+    (200, (100, 2)),
+]
+
+shots_to_test_nonsamples = [
+    ((10000, 2),),
+    (10000, 10000),
+    (10000, 20000),
+    (10000, 10000, 20000),
+    (20000, (10000, 2)),
+]
+
+
+class TestBroadcasting:
+    """Test that measurements work when the state has a batch dim"""
+
+    def test_sample_measure(self, batched_two_qutrit_pure_state):
+        """Test that broadcasting works for qml.sample and single shots"""
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(100)
+        state = batched_two_qutrit_pure_state
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(measurement, state, shots, is_state_batched=True, rng=rng)
+
+        assert res.shape == (3, shots.total_shots, 2)
+        assert res.dtype == np.int64
+
+    def test_counts_measure(self, batched_two_qutrit_pure_state):
+        """Test that broadcasting works for qml.sample and single shots"""
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(100)
+        state = batched_two_qutrit_pure_state
+
+        measurement = qml.counts(wires=[0, 1])
+        res = measure_with_samples(measurement, state, shots, is_state_batched=True, rng=rng)
+
+        assert res.shape == (3,)
+        assert list(res[0].keys()) == ["11"]
+        assert list(res[1].keys()) == ["00", "12"]
+        assert list(res[2].keys()) == ["01", "10", "21", "22"]
+
+    @pytest.mark.parametrize("shots", shots_to_test_samples)
+    def test_sample_measure_shot_vector(self, shots, batched_two_qutrit_pure_state):
+        """Test that broadcasting works for qml.sample and shot vectors"""
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(shots)
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(
+            measurement,
+            batched_two_qutrit_pure_state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+        )
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for s, r in zip(shots, res):
+            assert r.shape == (3, s, 2)
+            assert r.dtype == np.int64
+
+            assert_correct_sampled_batched_two_qutrit_pure_state(r)
+
+    # pylint:disable = too-many-arguments
+    @pytest.mark.parametrize(
+        "shots",
+        shots_to_test_nonsamples,
+    )
+    @pytest.mark.parametrize(
+        "measurement, expected",
+        [
+            (qml.expval(qml.GellMann(1, 3)), np.array([1 / 2, 1, 0])),
+            (qml.var(qml.GellMann(1, 3)), np.array([1 / 4, 0, 1])),
+        ],
+    )
+    def test_nonsample_measure_shot_vector(
+        self,
+        shots,
+        measurement,
+        expected,
+    ):
+        """Test that broadcasting works for the other sample measurements and shot vectors"""
+
+        rng = np.random.default_rng(123456)
+        shots = qml.measurements.Shots(shots)
+
+        state = [
+            get_dm_of_state(np.array([[0, 0, 0, 1, 0, 1, 0, 0, 0]]), 2, 2),
+            get_dm_of_state(np.array([1, 0, 0, 1, 0, 0, 1, 0, 0]), 2, 3),
+            get_dm_of_state(np.array([[1, 1, 0, 0, 0, 0, 1, 1, 0]]), 2, 4),
+        ]
+
+        state = np.stack(state)
+
+        res = measure_with_samples(measurement, state, shots, is_state_batched=True, rng=rng)
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for r in res:
+            assert r.shape == expected.shape
+            assert np.allclose(r, expected, atol=APPROX_ATOL)
+
+
+@pytest.mark.jax
+class TestBroadcastingPRNG:
+    """Test that measurements work and use _sample_state_jax when the state has a batch dim
+    and a PRNG key is provided"""
+
+    def test_sample_measure(self, mocker, batched_two_qutrit_pure_state):
+        """Test that broadcasting works for qml.sample and single shots"""
+        import jax
+
+        spy = mocker.spy(qml.devices.qutrit_mixed.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(100)
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(
+            measurement,
+            batched_two_qutrit_pure_state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+
+        assert res.shape == (3, shots.total_shots, 2)
+        assert res.dtype == np.int64
+
+        # convert to numpy array because prng_key -> JAX -> ArrayImpl -> angry vanilla numpy below
+        res = [np.array(r) for r in res]
+
+        assert_correct_sampled_batched_two_qutrit_pure_state(res)
+
+    @pytest.mark.parametrize("shots", shots_to_test_samples)
+    def test_sample_measure_shot_vector(self, mocker, shots, batched_two_qutrit_pure_state):
+        """Test that broadcasting works for qml.sample and shot vectors"""
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+
+        spy = mocker.spy(qml.devices.qutrit_mixed.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(shots)
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(
+            measurement,
+            batched_two_qutrit_pure_state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for s, r in zip(shots, res):
+            assert r.shape == (3, s, 2)
+            assert res[0][0].dtype == np.int64
+
+            # convert to numpy array because prng_key -> JAX -> ArrayImpl -> angry vanilla numpy below
+            r = [np.array(i) for i in r]
+
+            assert_correct_sampled_batched_two_qutrit_pure_state(r)
+
+    # pylint:disable = too-many-arguments
+    @pytest.mark.parametrize("shots", shots_to_test_nonsamples)
+    @pytest.mark.parametrize(
+        "measurement, expected",
+        [
+            (qml.expval(qml.GellMann(1, 3)), np.array([1 / 2, 1, 0])),
+            (qml.var(qml.GellMann(1, 3)), np.array([1 / 4, 0, 1])),
+        ],
+    )
+    def test_nonsample_measure_shot_vector(self, mocker, shots, measurement, expected):
+        """Test that broadcasting works for the other sample measurements and shot vectors"""
+
+        import jax
+
+        spy = mocker.spy(qml.devices.qutrit_mixed.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(shots)
+
+        state = [
+            get_dm_of_state(np.array([[0, 0, 0, 1, 0, 1, 0, 0, 0]]), 2, 2),
+            get_dm_of_state(np.array([1, 0, 0, 1, 0, 0, 1, 0, 0]), 2, 3),
+            get_dm_of_state(np.array([[1, 1, 0, 0, 0, 0, 1, 1, 0]]), 2, 4),
+        ]
+        state = np.stack(state)
+
+        res = measure_with_samples(
+            measurement,
+            state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for r in res:
+            assert r.shape == expected.shape
+            assert np.allclose(r, expected, atol=0.01)
+
+
+@pytest.mark.parametrize(
+    "obs",
+    [
+        qml.Hamiltonian([0.8, 0.5], [qml.GellMann(0, 3), qml.GellMann(0, 1)]),
+        qml.s_prod(0.8, qml.GellMann(0, 3)) + qml.s_prod(0.5, qml.GellMann(0, 1)),
+    ],
+)
+class TestHamiltonianSamples:
+    """Test that the measure_with_samples function works as expected for
+    Hamiltonian and Sum observables"""
+
+    def test_hamiltonian_expval(self, obs):
+        """Test that sampling works well for Hamiltonian and Sum observables"""
+        shots = qml.measurements.Shots(10000)
+
+        x, y = np.array(0.67), np.array(0.95)
+        ops = [qml.TRY(x, wires=0), qml.TRZ(y, wires=0)]
+        state = create_initial_state((0,))
+        for op in ops:
+            state = apply_operation(op, state)
+
+        res = measure_with_samples(qml.expval(obs), state, shots=shots, rng=300)
+
+        expected = 0.8 * np.cos(x) + 0.5 * np.cos(y) * np.sin(x)
+        assert isinstance(res, np.float64)
+        assert np.allclose(res, expected, atol=APPROX_ATOL)
+
+    def test_hamiltonian_expval_shot_vector(self, obs):
+        """Test that sampling works well for Hamiltonian and Sum observables with a shot vector"""
+        shots = qml.measurements.Shots((10000, 100000))
+
+        x, y = np.array(0.67), np.array(0.95)
+        ops = [qml.TRY(x, wires=0), qml.TRZ(y, wires=0)]
+        state = create_initial_state((0,))
+        for op in ops:
+            state = apply_operation(op, state)
+
+        res = measure_with_samples(qml.expval(obs), state, shots=shots, rng=300)
+
+        expected = 0.8 * np.cos(x) + 0.5 * np.cos(y) * np.sin(x)
+
+        assert len(res) == 2
+        assert isinstance(res, tuple)
+        assert np.allclose(res[0], expected, atol=APPROX_ATOL)
+        assert np.allclose(res[1], expected, atol=APPROX_ATOL)
diff --git a/tests/fourier/test_reconstruct.py b/tests/fourier/test_reconstruct.py
index 843c3f0cf..8f8fca439 100644
--- a/tests/fourier/test_reconstruct.py
+++ b/tests/fourier/test_reconstruct.py
@@ -424,6 +424,8 @@ class TestReconstructGen:
         functions are differentiable for JAX input variables."""
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         # Convert fun to have integer frequencies
         rec = _reconstruct_gen(fun, spectrum, interface="jax")
         grad = jax.grad(rec)
@@ -918,6 +920,8 @@ class TestReconstruct:
         """Tests the reconstruction and differentiability with JAX."""
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         params = tuple(jax.numpy.array(par) for par in params)
         qnode = qml.QNode(qnode, dev_1, interface="jax")
         with qml.Tracker(qnode.device) as tracker:
diff --git a/tests/gradients/core/test_adjoint_metric_tensor.py b/tests/gradients/core/test_adjoint_metric_tensor.py
index c7ce5b689..5a6439037 100644
--- a/tests/gradients/core/test_adjoint_metric_tensor.py
+++ b/tests/gradients/core/test_adjoint_metric_tensor.py
@@ -407,6 +407,8 @@ class TestAdjointMetricTensorQNode:
 
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         j_params = tuple(jax.numpy.array(p) for p in params)
         dev = qml.device("default.qubit", wires=self.num_wires)