Repository: PennyLaneAI/pennylane
Pull request #5186 (2024-02-09 07:51:56)

Qutrit mixed test differentiability
**Context:**
Measuring and applying operation functionality have been added to the `qutrit_mixed` module, these tests ensure that these operations are differentiable. 
Since support for qml.prod for qutrits has been implemented from [PR](https://github.com/PennyLaneAI/pennylane/pull/5400). It will be possible to change the method of state measurement without issue.
The way expectation value measurement was done may have been overly complicated so this reducing that complication, while still allowing for product type observables.

**Description of the Change:**
Tests have been added that ensure that the `apply_operation` and `measure` functions retain differentiability for the new qutrit mixed-state device.
Expectation value measurements have been changed to use diagonalizing gates and eigenvalues instead of trace.

**Benefits:**
Shows that the `apply_operation` and `measure` functions are differentiable.

**Possible Drawbacks:**
This PR does not remove the abstractions made in apply_operations to make the trace based measure expectation value work so there is some code left that may be overly abstract.

**Related GitHub Issues:**
N/A

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5186?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:
> Project coverage is 99.63%. Comparing base [(`3bb04be`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3bb04be42eb4788813bcc8a69b626b5156fd99e0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`7769525`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5186?src=pr&el=desc&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    #5186      +/-   ##
==========================================
- Coverage   99.63%   99.63%   -0.01%     
==========================================
  Files         400      400              
  Lines       37132    36846     -286     
==========================================
- Hits        36998    36711     -287     
- Misses        134      135       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5186?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: Hi @albi3ro, @glassnotes, @mudit2812, and @trbromley. 
These are the tests that ensure that apply operation and measurement have the correct gradients.
One small change to the code was necessary for an autograd test to work, all other tests required no changes.
Gabriel-Bottrill: There was an issue with a certain case for autograd. I fixed it by applying `math.asarray` to the output of expval. The other option is to change how expval is found, using diagonalizing gate and eigenvalues with `math.prod`. The issue is that this doesn't work with `op_math.Prod` and `op_math.SProd` simply. 

I had to rewrite 2 functions, but I implemented a working version in this [secondary PR](https://github.com/PennyLaneAI/pennylane/pull/5350). There are some pylint issues in that option, but they can be fixed. Unfortunately it requires a lot of repeated code however it has some advantages such as making  `op_math.Prod` work for variance. It also fixes issues with `op_math.Prod` in sampling.

I'm not sure what is the better choice to do here. 
Another option could be to leave it as is and add another PR to make that change later.
glassnotes: Update: just read your comment above, it sounds like the second version might be a better option, if it makes things work that we'd expect to work. Which of them is most consistent with how things work for the qubit devices? (Would be good to get the opinion of the Xanadu crew as well, though)
Gabriel-Bottrill: > Update: just read your comment above, it sounds like the second version might be a better option, if it makes things work that we'd expect to work. Which of them is most consistent with how things work for the qubit devices? (Would be good to get the opinion of the Xanadu crew as well, though)

I think the second version would be more consistent it just uses some copied code so it will need to be retested. I was thinking it may be better for it to be it's own PR.
Gabriel-Bottrill: @glassnotes I made the changes we discussed
Gabriel-Bottrill: Since this [PR](https://github.com/PennyLaneAI/pennylane/pull/5400) was merged, I updated the way measurement works for expectation values. This makes it consistent with how `default_qubit` and `default_mixed` measures expectation values and simplifies it significantly.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2713c21fc..b117f4b3f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -230,6 +230,9 @@
 * `qml.transforms.split_non_commuting` will now work with single-term operator arithmetic.
   [(#5314)](https://github.com/PennyLaneAI/pennylane/pull/5314)
 
+* Fixed differentiability for Hamiltonian measurements in new `qutrit_mixed` module. 
+  [(#5186)](https://github.com/PennyLaneAI/pennylane/pull/5186)
+  
 * Added `simulate` function to the new `qutrit_mixed` module in `qml.devices`. This allows for simulation of a 
   noisy qutrit circuit with measurement and sampling.
   [(#5213)](https://github.com/PennyLaneAI/pennylane/pull/5213)
diff --git a/pennylane/devices/qutrit_mixed/measure.py b/pennylane/devices/qutrit_mixed/measure.py
index d05bd85da..236071bfc 100644
--- a/pennylane/devices/qutrit_mixed/measure.py
+++ b/pennylane/devices/qutrit_mixed/measure.py
@@ -27,66 +27,16 @@ from pennylane.measurements import (
     ProbabilityMP,
     VarianceMP,
 )
-from pennylane.operation import Observable
 from pennylane.typing import TensorLike
 
-from .utils import (
-    get_einsum_mapping,
-    reshape_state_as_matrix,
-    get_num_wires,
-    get_new_state_einsum_indices,
-    QUDIT_DIM,
-)
+from .utils import reshape_state_as_matrix, get_num_wires, QUDIT_DIM
 from .apply_operation import apply_operation
 
 
-def _map_indices_apply_operation(**kwargs):
-    """Map indices to wires.
-
-    Args:
-        **kwargs (dict): Stores indices calculated in `get_einsum_mapping`:
-            state_indices (str): Indices that are summed.
-            row_indices (str): Indices that must be replaced with sums.
-            new_row_indices (str): Tensor indices of the state.
-
-    Returns:
-        String of einsum indices to complete einsum calculations.
-    """
-    op_1_indices = f"{kwargs['new_row_indices']}{kwargs['row_indices']}"
-
-    new_state_indices = get_new_state_einsum_indices(
-        old_indices=kwargs["row_indices"],
-        new_indices=kwargs["new_row_indices"],
-        state_indices=kwargs["state_indices"],
-    )
-
-    return f"{op_1_indices},...{kwargs['state_indices']}->...{new_state_indices}"
-
-
-def apply_observable_einsum(obs: Observable, state, is_state_batched: bool = False):
-    r"""Applies an observable to a density matrix rho, giving obs@state.
-
-    Args:
-        obs (Operator): Operator to apply to the quantum state.
-        state (array[complex]): Input quantum state.
-        is_state_batched (bool): Boolean representing whether the state is batched or not.
-
-    Returns:
-        TensorLike: the result of obs@state.
-    """
-
-    num_ch_wires = len(obs.wires)
-    einsum_indices = get_einsum_mapping(obs, state, _map_indices_apply_operation, is_state_batched)
-    obs_mat = obs.matrix()
-    obs_shape = [QUDIT_DIM] * num_ch_wires * 2
-    obs_mat = math.cast(math.reshape(obs_mat, obs_shape), complex)
-    return math.einsum(einsum_indices, obs_mat, state)
-
-
 def calculate_expval(
     measurementprocess: ExpectationMP, state: TensorLike, is_state_batched: bool = False
 ) -> TensorLike:
-    """Measure the expectation value of an observable by finding the trace of obs@rho.
+    """Measure the expectation value of an observable.
 
     Args:
         measurementprocess (ExpectationMP): measurement process to apply to the state.
@@ -96,21 +46,13 @@ def calculate_expval(
     Returns:
         TensorLike: expectation value of observable wrt the state.
     """
-    obs = measurementprocess.obs
-    rho_mult_obs = apply_observable_einsum(obs, state, is_state_batched)
-
-    # using einsum since trace function axis selection parameter names
-    # are not consistent across interfaces, they don't exist for torch
-
-    num_wires = get_num_wires(state, is_state_batched)
-    rho_mult_obs_reshaped = reshape_state_as_matrix(rho_mult_obs, num_wires)
-    if is_state_batched:
-        return math.real(math.stack([math.sum(math.diagonal(dm)) for dm in rho_mult_obs_reshaped]))
-
-    return math.real(math.sum(math.diagonal(rho_mult_obs_reshaped)))
+    probs = calculate_probability(measurementprocess, state, is_state_batched)
+    eigvals = math.asarray(measurementprocess.eigvals(), dtype="float64")
+    # In case of broadcasting, `probs` has two axes and these are a matrix-vector products
+    return math.dot(probs, eigvals)
 
 
-def calculate_reduced_density_matrix(  # TODO: ask if I should have state diagonalization gates?
+def calculate_reduced_density_matrix(
     measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False
 ) -> TensorLike:
     """Get the state or reduced density matrix.
@@ -260,8 +202,6 @@ def get_measurement_function(
     """
     if isinstance(measurementprocess, StateMeasurement):
         if isinstance(measurementprocess, ExpectationMP):
-            # TODO add faster methods
-            # TODO add support for sparce Hamiltonians
             if isinstance(measurementprocess.obs, (Hamiltonian, Sum)):
                 return calculate_expval_sum_of_terms
             if measurementprocess.obs.has_matrix:


Test patch:
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py
index c5bb0255a..30e89df56 100644
--- a/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py
@@ -20,7 +20,7 @@ from scipy.stats import unitary_group
 import pennylane as qml
 from pennylane import math
 from pennylane.operation import Channel
-from pennylane.devices.qutrit_mixed import apply_operation
+from pennylane.devices.qutrit_mixed import apply_operation, measure
 
 ml_frameworks_list = [
     "numpy",
@@ -31,6 +31,24 @@ ml_frameworks_list = [
 ]
 
 subspaces = [(0, 1), (0, 2), (1, 2)]
+kraus_matrix = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]], dtype=complex)
+
+
+class CustomChannel(Channel):  # pylint: disable=too-few-public-methods
+    num_params = 1
+    num_wires = 1
+
+    def __init__(self, p, wires, id=None):
+        super().__init__(p, wires=wires, id=id)
+
+    @staticmethod
+    def compute_kraus_matrices(p):
+        if math.get_interface(p) == "tensorflow":
+            p = math.cast_like(p, 1j)
+
+        K0 = math.sqrt(1 - p + math.eps) * math.convert_like(math.eye(3, dtype=complex), p)
+        K1 = math.sqrt(p + math.eps) * math.convert_like(kraus_matrix, p)
+        return [K0, K1]
 
 
 def test_custom_operator_with_matrix(one_qutrit_state):
@@ -320,3 +338,223 @@ class TestChannels:  # pylint: disable=too-few-public-methods
 
         assert qml.math.get_interface(res) == ml_framework
         assert qml.math.allclose(res, expected)
+
+
+@pytest.mark.parametrize("subspace", [(0, 1), (0, 2), (1, 2)])
+class TestTRXCalcGrad:
+    """Tests the application and differentiation of a TRX gate in the different interfaces."""
+
+    phi = 0.325
+
+    @staticmethod
+    def compare_expected_result(phi, state, probs, subspace, jacobian):
+        """Compare the expected result for this circuit and gradient with observed values"""
+        trx = qml.TRX.compute_matrix(phi, subspace)
+        trx_adj = qml.TRX.compute_matrix(-phi, subspace)
+        state = math.reshape(state, (9, 9))
+
+        expected_probs = math.diagonal(
+            np.kron(trx, np.eye(3)) @ state @ np.kron(trx_adj, np.eye(3))
+        )
+        assert qml.math.allclose(probs, expected_probs)
+
+        if subspace[0] == 0:
+            gell_mann_index = 1 if subspace[1] == 1 else 4
+        else:
+            gell_mann_index = 6
+        gell_mann_matrix = qml.GellMann.compute_matrix(gell_mann_index)
+        trx_derivative = -0.5j * gell_mann_matrix @ trx
+        trx_adj_derivative = 0.5j * gell_mann_matrix @ trx_adj
+
+        expected_derivative_state = (
+            np.kron(trx_derivative, np.eye(3)) @ state @ np.kron(trx_adj, np.eye(3))
+        ) + (np.kron(trx, np.eye(3)) @ state @ np.kron(trx_adj_derivative, np.eye(3)))
+        expected_derivative = np.diagonal(expected_derivative_state)
+        assert qml.math.allclose(jacobian, expected_derivative)
+
+    @pytest.mark.autograd
+    def test_trx_grad_autograd(self, two_qutrit_state, subspace):
+        """Test that the application of a trx gate is differentiable with autograd."""
+
+        state = qml.numpy.array(two_qutrit_state)
+
+        def f(phi):
+            op = qml.TRX(phi, wires=0, subspace=subspace)
+            new_state = apply_operation(op, state)
+            return measure(qml.probs(), new_state)
+
+        phi = qml.numpy.array(self.phi, requires_grad=True)
+
+        probs = f(phi)
+        jacobian = qml.jacobian(lambda x: qml.math.real(f(x)))(phi)
+        self.compare_expected_result(phi, state, probs, subspace, jacobian)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", (True, False))
+    def test_trx_grad_jax(self, use_jit, two_qutrit_state, subspace):
+        """Test that the application of a trx gate is differentiable with jax."""
+
+        import jax
+
+        state = jax.numpy.array(two_qutrit_state)
+
+        def f(phi):
+            op = qml.TRX(phi, wires=0, subspace=subspace)
+            new_state = apply_operation(op, state)
+            return measure(qml.probs(), new_state)
+
+        if use_jit:
+            f = jax.jit(f)
+
+        probs = f(self.phi)
+        jacobian = jax.jacobian(f)(self.phi)
+        self.compare_expected_result(self.phi, state, probs, subspace, jacobian)
+
+    @pytest.mark.torch
+    def test_trx_grad_torch(self, two_qutrit_state, subspace):
+        """Tests the application and differentiation of a trx gate with torch."""
+
+        import torch
+
+        state = torch.tensor(two_qutrit_state)
+
+        def f(phi):
+            op = qml.TRX(phi, wires=0, subspace=subspace)
+            new_state = apply_operation(op, state)
+            return measure(qml.probs(), new_state)
+
+        phi = torch.tensor(self.phi, requires_grad=True)
+
+        probs = f(phi)
+        jacobian = torch.autograd.functional.jacobian(f, phi)
+
+        self.compare_expected_result(
+            phi.detach().numpy(),
+            state.detach().numpy(),
+            probs.detach().numpy(),
+            subspace,
+            jacobian.detach().numpy(),
+        )
+
+    @pytest.mark.tf
+    def test_trx_grad_tf(self, two_qutrit_state, subspace):
+        """Tests the application and differentiation of a trx gate with tensorflow"""
+        import tensorflow as tf
+
+        state = tf.Variable(two_qutrit_state)
+        phi = tf.Variable(0.8589, trainable=True)
+
+        with tf.GradientTape() as grad_tape:
+            op = qml.TRX(phi, wires=0, subspace=subspace)
+            new_state = apply_operation(op, state)
+            probs = measure(qml.probs(), new_state)
+
+        jacobians = grad_tape.jacobian(probs, [phi])
+        phi_jacobian = jacobians[0]
+
+        self.compare_expected_result(phi, state, probs, subspace, phi_jacobian)
+
+
+class TestChannelCalcGrad:
+    """Tests the application and differentiation of a Channel in the different interfaces."""
+
+    p = 0.325
+
+    @staticmethod
+    def compare_expected_result(p, state, new_state, jacobian):
+        """Compare the expected result for this channel and gradient with observed values"""
+        kraus_matrix_two_qutrits = np.kron(np.eye(3), kraus_matrix)
+        kraus_matrix_two_qutrits_adj = kraus_matrix_two_qutrits.transpose()
+        state = math.reshape(state, (9, 9))
+
+        state_kraus_applied = kraus_matrix_two_qutrits @ state @ kraus_matrix_two_qutrits_adj
+
+        expected_state = (1 - p) * state + (p * state_kraus_applied)
+        expected_probs = np.diagonal(expected_state)
+        assert qml.math.allclose(new_state, expected_probs)
+
+        expected_derivative_state = state_kraus_applied - state
+        expected_derivative = np.diagonal(expected_derivative_state)
+        assert qml.math.allclose(jacobian, expected_derivative)
+
+    @pytest.mark.autograd
+    def test_channel_grad_autograd(self, two_qutrit_state):
+        """Test that the application of a channel is differentiable with autograd."""
+
+        state = qml.numpy.array(two_qutrit_state)
+
+        def f(p):
+            channel = CustomChannel(p, wires=1)
+            new_state = apply_operation(channel, state)
+            return measure(qml.probs(), new_state)
+
+        p = qml.numpy.array(self.p, requires_grad=True)
+
+        probs = f(p)
+        jacobian = qml.jacobian(lambda x: qml.math.real(f(x)))(p)
+        self.compare_expected_result(p, state, probs, jacobian)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", (True, False))
+    def test_channel_grad_jax(self, use_jit, two_qutrit_state):
+        """Test that the application of a channel is differentiable with jax."""
+
+        import jax
+
+        state = jax.numpy.array(two_qutrit_state)
+
+        def f(p):
+            op = CustomChannel(p, wires=1)
+            new_state = apply_operation(op, state)
+            return measure(qml.probs(), new_state)
+
+        if use_jit:
+            f = jax.jit(f)
+
+        probs = f(self.p)
+        jacobian = jax.jacobian(f)(self.p)
+        self.compare_expected_result(self.p, state, probs, jacobian)
+
+    @pytest.mark.torch
+    def test_channel_grad_torch(self, two_qutrit_state):
+        """Tests the application and differentiation of a channel with torch."""
+
+        import torch
+
+        state = torch.tensor(two_qutrit_state)
+
+        def f(p):
+            op = CustomChannel(p, wires=1)
+            new_state = apply_operation(op, state)
+            return measure(qml.probs(), new_state)
+
+        p = torch.tensor(self.p, requires_grad=True)
+
+        probs = f(p)
+        jacobian = torch.autograd.functional.jacobian(f, p)
+
+        self.compare_expected_result(
+            p.detach().numpy(),
+            state.detach().numpy(),
+            probs.detach().numpy(),
+            jacobian.detach().numpy(),
+        )
+
+    @pytest.mark.tf
+    def test_channel_grad_tf(self, two_qutrit_state):
+        """Tests the application and differentiation of a channel with tensorflow"""
+        import tensorflow as tf
+
+        state = tf.Variable(two_qutrit_state)
+        p = tf.Variable(0.8589 + 0j, trainable=True)
+
+        with tf.GradientTape() as grad_tape:
+            op = CustomChannel(p, wires=1)
+            new_state = apply_operation(op, state)
+            probs = measure(qml.probs(), new_state)
+
+        jacobians = grad_tape.jacobian(probs, [p])
+        # tf takes gradient with respect to conj(z), so we need to conj the gradient
+        phi_jacobian = tf.math.conj(jacobians[0])
+
+        self.compare_expected_result(p, state, probs, phi_jacobian)
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py
index d2801ffd2..58c13661a 100644
--- a/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py
@@ -20,9 +20,8 @@ import pennylane as qml
 from pennylane import numpy as np
 from pennylane import math
 
-from pennylane.devices.qutrit_mixed import measure
+from pennylane.devices.qutrit_mixed import measure, create_initial_state, apply_operation
 from pennylane.devices.qutrit_mixed.measure import (
-    apply_observable_einsum,
     get_measurement_function,
     calculate_expval,
     calculate_expval_sum_of_terms,
@@ -118,49 +117,6 @@ class TestMeasurementDispatch:
         assert get_measurement_function(qml.var(obs)) is calculate_variance
 
 
-@pytest.mark.parametrize(
-    "obs",
-    [
-        qml.GellMann(2, 2),
-        qml.GellMann(1, 8),
-        qml.GellMann(0, 5),
-        (qml.GellMann(0, 5) @ qml.GellMann(1, 8)),
-    ],
-)
-@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
-class TestApplyObservableEinsum:
-    """Tests that observables are applied correctly for calculate expval method."""
-
-    num_qutrits = 3
-    dims = (3**num_qutrits, 3**num_qutrits)
-
-    def test_apply_observable_einsum(self, obs, three_qutrit_state, ml_framework):
-        """Tests that unbatched observables are applied correctly to a unbatched state."""
-        res = apply_observable_einsum(
-            obs, qml.math.asarray(three_qutrit_state, like=ml_framework), is_state_batched=False
-        )
-        expected = get_expanded_op_mult_state(obs, three_qutrit_state)
-        expected = expected.reshape([3] * (3 * 2))
-
-        assert qml.math.get_interface(res) == ml_framework
-        assert qml.math.allclose(res, expected)
-
-    def test_apply_observable_einsum_batched(self, obs, three_qutrit_batched_state, ml_framework):
-        """Tests that unbatched observables are applied correctly to a batched state."""
-        res = apply_observable_einsum(
-            obs,
-            qml.math.asarray(three_qutrit_batched_state, like=ml_framework),
-            is_state_batched=True,
-        )
-        expected = [
-            get_expanded_op_mult_state(obs, state).reshape([3] * (3 * 2))
-            for state in three_qutrit_batched_state
-        ]
-
-        assert qml.math.get_interface(res) == ml_framework
-        assert qml.math.allclose(res, expected)
-
-
 class TestMeasurements:
     """Test that measurements on unbatched states work as expected."""
 
@@ -471,4 +427,177 @@ class TestBroadcasting:
         assert np.allclose(res, expected)
 
 
-# TODO TestSumOfTermsDifferentiability in future PR (with other differentiabilty tests)
+@pytest.mark.usefixtures("use_legacy_and_new_opmath")
+class TestSumOfTermsDifferentiability:
+    x = 0.52
+
+    @staticmethod
+    def f(scale, coeffs, n_wires=5, offset=0.1):
+        """Function to differentiate that implements a circuit with a SumOfTerms operator"""
+        ops = [qml.TRX(offset + scale * i, wires=i, subspace=(0, 2)) for i in range(n_wires)]
+        H = qml.Hamiltonian(
+            coeffs,
+            [
+                reduce(lambda x, y: x @ y, (qml.GellMann(i, 3) for i in range(n_wires))),
+                reduce(lambda x, y: x @ y, (qml.GellMann(i, 5) for i in range(n_wires))),
+            ],
+        )
+        state = create_initial_state(range(n_wires), like=math.get_interface(scale))
+        for op in ops:
+            state = apply_operation(op, state)
+        return measure(qml.expval(H), state)
+
+    @staticmethod
+    def expected(scale, coeffs, n_wires=5, offset=0.1, like="numpy"):
+        """Get the expected expval of the class' circuit."""
+        phase = offset + scale * qml.math.asarray(range(n_wires), like=like)
+        cosines = math.cos(phase / 2) ** 2
+        sines = -math.sin(phase)
+        return coeffs[0] * qml.math.prod(cosines) + coeffs[1] * qml.math.prod(sines)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize(
+        "coeffs",
+        [
+            (qml.numpy.array(2.5), qml.numpy.array(6.2)),
+            (qml.numpy.array(2.5, requires_grad=False), qml.numpy.array(6.2, requires_grad=False)),
+        ],
+    )
+    def test_autograd_backprop(self, coeffs):
+        """Test that backpropagation derivatives work in autograd with
+        Hamiltonians using new and old math."""
+
+        x = qml.numpy.array(self.x)
+        out = self.f(x, coeffs)
+        expected_out = self.expected(x, coeffs)
+        assert qml.math.allclose(out, expected_out)
+
+        gradient = qml.grad(self.f)(x, coeffs)
+        expected_gradient = qml.grad(self.expected)(x, coeffs)
+        assert qml.math.allclose(expected_gradient, gradient)
+
+    @pytest.mark.autograd
+    def test_autograd_backprop_coeffs(self):
+        """Test that backpropagation derivatives work in autograd with
+        the coefficients of Hamiltonians using new and old math."""
+
+        coeffs = qml.numpy.array((2.5, 6.2), requires_grad=True)
+        gradient = qml.grad(self.f, argnum=1)(self.x, coeffs)
+        expected_gradient = qml.grad(self.expected)(self.x, coeffs)
+
+        assert len(gradient) == 2
+        assert qml.math.allclose(expected_gradient, gradient)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", (True, False))
+    def test_jax_backprop(self, use_jit):
+        """Test that backpropagation derivatives work with jax with
+        Hamiltonians using new and old math."""
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+
+        x = jax.numpy.array(self.x, dtype=jax.numpy.float64)
+        coeffs = (5.2, 6.7)
+        f = jax.jit(self.f, static_argnums=(1, 2, 3, 4)) if use_jit else self.f
+
+        out = f(x, coeffs)
+        expected_out = self.expected(x, coeffs)
+        assert qml.math.allclose(out, expected_out)
+
+        gradient = jax.grad(f)(x, coeffs)
+        expected_gradient = jax.grad(self.expected)(x, coeffs)
+        assert qml.math.allclose(expected_gradient, gradient)
+
+    @pytest.mark.jax
+    def test_jax_backprop_coeffs(self):
+        """Test that backpropagation derivatives work with jax with
+        the coefficients of Hamiltonians using new and old math."""
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+        coeffs = jax.numpy.array((5.2, 6.7), dtype=jax.numpy.float64)
+
+        gradient = jax.grad(self.f, argnums=1)(self.x, coeffs)
+        expected_gradient = jax.grad(self.expected, argnums=1)(self.x, coeffs)
+        assert len(gradient) == 2
+        assert qml.math.allclose(expected_gradient, gradient)
+
+    @pytest.mark.torch
+    def test_torch_backprop(self):
+        """Test that backpropagation derivatives work with torch with
+        Hamiltonians using new and old math."""
+        import torch
+
+        coeffs = [
+            torch.tensor(9.2, requires_grad=False, dtype=torch.float64),
+            torch.tensor(6.2, requires_grad=False, dtype=torch.float64),
+        ]
+
+        x = torch.tensor(-0.289, requires_grad=True, dtype=torch.float64)
+        x2 = torch.tensor(-0.289, requires_grad=True, dtype=torch.float64)
+        out = self.f(x, coeffs)
+        expected_out = self.expected(x2, coeffs, like="torch")
+        assert qml.math.allclose(out, expected_out)
+
+        out.backward()
+        expected_out.backward()
+        assert qml.math.allclose(x.grad, x2.grad)
+
+    @pytest.mark.torch
+    def test_torch_backprop_coeffs(self):
+        """Test that backpropagation derivatives work with torch with
+        the coefficients of Hamiltonians using new and old math."""
+        import torch
+
+        coeffs = torch.tensor((9.2, 6.2), requires_grad=True, dtype=torch.float64)
+        coeffs_expected = torch.tensor((9.2, 6.2), requires_grad=True, dtype=torch.float64)
+
+        x = torch.tensor(-0.289, requires_grad=False, dtype=torch.float64)
+        out = self.f(x, coeffs)
+        expected_out = self.expected(x, coeffs_expected, like="torch")
+        assert qml.math.allclose(out, expected_out)
+
+        out.backward()
+        expected_out.backward()
+        assert len(coeffs.grad) == 2
+        assert qml.math.allclose(coeffs.grad, coeffs_expected.grad)
+
+    @pytest.mark.tf
+    def test_tf_backprop(self):
+        """Test that backpropagation derivatives work with tensorflow with
+        Hamiltonians using new and old math."""
+        import tensorflow as tf
+
+        x = tf.Variable(self.x)
+        coeffs = [8.3, 5.7]
+
+        with tf.GradientTape() as tape1:
+            out = self.f(x, coeffs)
+
+        with tf.GradientTape() as tape2:
+            expected_out = self.expected(x, coeffs)
+
+        assert qml.math.allclose(out, expected_out)
+        gradient = tape1.gradient(out, x)
+        expected_gradient = tape2.gradient(expected_out, x)
+        assert qml.math.allclose(expected_gradient, gradient)
+
+    @pytest.mark.tf
+    def test_tf_backprop_coeffs(self):
+        """Test that backpropagation derivatives work with tensorflow with
+        the coefficients of Hamiltonians using new and old math."""
+        import tensorflow as tf
+
+        coeffs = tf.Variable([8.3, 5.7])
+
+        with tf.GradientTape() as tape1:
+            out = self.f(self.x, coeffs)
+
+        with tf.GradientTape() as tape2:
+            expected_out = self.expected(self.x, coeffs)
+
+        gradient = tape1.gradient(out, coeffs)
+        expected_gradient = tape2.gradient(expected_out, coeffs)
+        assert len(gradient) == 2
+        assert qml.math.allclose(expected_gradient, gradient)