From 46afaa5c748c8bc0f0ad174dd71af755f0fec6b2 Mon Sep 17 00:00:00 2001
From: Gabriel Bottrill <78718539+Gabriel-Bottrill@users.noreply.github.com>
Date: Fri, 12 Apr 2024 06:54:54 -0700
Subject: [PATCH] Qutrit mixed test differentiability (#5186)

**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

---------

Co-authored-by: Gabe PC <bottrill@student.ubc.ca>
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>
Co-authored-by: Christina Lee <chrissie.c.l@gmail.com>
Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>
---
 doc/releases/changelog-dev.md                 |   3 +
 pennylane/devices/qutrit_mixed/measure.py     |  74 +-----
 .../test_qutrit_mixed_apply_operation.py      | 240 +++++++++++++++++-
 .../qutrit_mixed/test_qutrit_mixed_measure.py | 221 ++++++++++++----
 4 files changed, 424 insertions(+), 114 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2713c21fc81..b117f4b3f82 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 d05bd85da40..236071bfce0 100644
--- a/pennylane/devices/qutrit_mixed/measure.py
+++ b/pennylane/devices/qutrit_mixed/measure.py
@@ -27,66 +27,16 @@
     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:
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 c5bb0255ab2..30e89df56cd 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 @@
 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 @@
 ]
 
 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 @@ def test_broadcasted_state(self, ml_framework, wire, two_qutrit_batched_state):
 
         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 d2801ffd213..58c13661aa4 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 @@
 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 @@ def test_var_compute_variance(self):
         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 @@ def test_variance_measurement(self, observable, ml_framework, two_qutrit_batched
         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)