Repository: PennyLaneAI/pennylane
Pull request #1608 (2021-08-30 09:57:22)

Add support for the new VQE workflows to the differentiable batch execution pipeline
**Context:** In #1483, #1596, and #1551, support was added for:

- Directly specifying `expval(H)` inside of a QNode
- Differentiating Hamiltonian coefficients
- Allowing certain devices, including `default.qubit`, to support Hamiltonians directly.

However, these changes were not ported over to the new differentiable batch-execute pipeline.

**Description of the Change:**

- Add support for differentiating (one or more) Hamiltonians to `qml.gradients.param_shift`.

- Modify `hamiltonian_grad` to support multiple Hamiltonians (`return expval(H1), expval(H2)`), as this is supported execution behaviour.

- Fixes a couple of bugs with the Autograd and TensorFlow batch interfaces

- Adds tests

**Benefits:** returning `expval(H)` from QNodes is now supported within the new batch-differentiable pipeline.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** n/a


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1608](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (002ec42) into [grad-docs](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c2689592194ed1133aec00c9792022d8b107b5ed?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c268959) will **decrease** coverage by `0.00%`.
> The diff coverage is `97.95%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@              Coverage Diff              @@
##           grad-docs    #1608      +/-   ##
=============================================
- Coverage      99.13%   99.13%   -0.01%     
=============================================
  Files            196      196              
  Lines          14148    14181      +33     
=============================================
+ Hits           14026    14058      +32     
- Misses           122      123       +1     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfY3YucHk=) | `99.13% <ø> (ø)` | |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/hamiltonian\_grad.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9oYW1pbHRvbmlhbl9ncmFkLnB5) | `96.29% <95.00%> (-3.71%)` | :arrow_down: |
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `98.09% <100.00%> (ø)` | |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/tensorflow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdGVuc29yZmxvdy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/tape/tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvdGFwZS5weQ==) | `98.91% <100.00%> (+0.01%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [c268959...002ec42](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1608?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: [ch8557]
anthayes92: Awesome addition @josh146 ! Very thorough testing

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index fcb81facb..89b41eab5 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -271,6 +271,7 @@
   [(#1508)](https://github.com/PennyLaneAI/pennylane/pull/1508)
   [(#1542)](https://github.com/PennyLaneAI/pennylane/pull/1542)
   [(#1549)](https://github.com/PennyLaneAI/pennylane/pull/1549)
+  [(#1608)](https://github.com/PennyLaneAI/pennylane/pull/1608)
 
   For example:
 
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 32e165264..f0285da06 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -527,7 +527,7 @@ class Device(abc.ABC):
             # Evaluations and gradients are paired, so that
             # devices can re-use the device state for the
             # gradient computation (if applicable).
-            res.append(circuit.execute(self))
+            res.append(self.batch_execute([circuit])[0])
             jacs.append(gradient_method(circuit, **kwargs))
 
         return res, jacs
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index cb53a6ca6..7f04c6d51 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -472,7 +472,10 @@ class DefaultQubit(QubitDevice):
         if observable.name in ("Hamiltonian", "SparseHamiltonian"):
             assert self.shots is None, f"{observable.name} must be used with shots=None"
 
-            backprop_mode = not isinstance(self.state, np.ndarray)
+            backprop_mode = (
+                not isinstance(self.state, np.ndarray)
+                or any(not isinstance(d, (float, np.ndarray)) for d in observable.data)
+            ) and observable.name == "Hamiltonian"
 
             if backprop_mode:
                 # We must compute the expectation value assuming that the Hamiltonian
@@ -513,6 +516,9 @@ class DefaultQubit(QubitDevice):
                     coo_matrix.dot(Hmat, coo_matrix(self.state.reshape(len(self.state), 1))),
                 ).toarray()[0]
 
+            if observable.name == "Hamiltonian":
+                res = qml.math.squeeze(res)
+
             return qml.math.real(res)
 
         return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index a4fcb76ae..0feec6572 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-
 Quantum gradient transforms are strategies for computing the gradient of a quantum
 circuit that work by **transforming** the quantum circuit into one or more gradient circuits.
 These gradient circuits, once executed and post-processed, return the gradient
diff --git a/pennylane/gradients/hamiltonian_grad.py b/pennylane/gradients/hamiltonian_grad.py
index 092ee65e5..fce76293f 100644
--- a/pennylane/gradients/hamiltonian_grad.py
+++ b/pennylane/gradients/hamiltonian_grad.py
@@ -25,18 +25,32 @@ def hamiltonian_grad(tape, idx, params=None):
         idx (int): index of parameter that we differentiate with respect to
         params (array): explicit parameters to set
     """
+    op, p_idx = tape.get_operation(idx)
+    new_tape = tape.copy(copy_operations=True)
 
-    t_idx = list(tape.trainable_params)[idx]
-    op = tape._par_info[t_idx]["op"]
-    p_idx = tape._par_info[t_idx]["p_idx"]
-
-    new_tape = tape.copy(copy_operations=True, tape_cls=qml.tape.QuantumTape)
     if params is not None:
         new_tape.set_parameters(params=params)
 
-    new_tape._measurements = [qml.expval(op.ops[p_idx])]
+    # get position in queue
+
+    queue_position = tape.observables.index(op)
+    new_tape._measurements[queue_position] = qml.expval(op.ops[p_idx])
 
     new_tape._par_info = {}
     new_tape._update()
 
-    return [new_tape], lambda x: qml.math.squeeze(x)
+    if len(tape.measurements) > 1:
+
+        def processing_fn(results):
+            res = results[0][queue_position]
+            zeros = qml.math.zeros_like(res)
+
+            final = []
+            for i, _ in enumerate(tape.measurements):
+                final.append(res if i == queue_position else zeros)
+
+            return qml.math.expand_dims(qml.math.stack(final), 0)
+
+        return [new_tape], processing_fn
+
+    return [new_tape], lambda x: x
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 9ccae1303..61dce8ef7 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -15,7 +15,7 @@
 This module contains functions for computing the parameter-shift gradient
 of a qubit-based quantum tape.
 """
-# pylint: disable=protected-access,too-many-arguments
+# pylint: disable=protected-access,too-many-arguments,too-many-statements
 import numpy as np
 
 import pennylane as qml
@@ -66,19 +66,8 @@ def _get_operation_recipe(tape, t_idx, shift=np.pi / 2):
     If the corresponding operation has grad_recipe=None, then
     the default two-term parameter-shift rule is assumed.
     """
-    # get the index of the parameter in the tape
-    parameter_idx = list(tape.trainable_params)[t_idx]
-
-    # get the corresponding operation
-    op = tape._par_info[parameter_idx]["op"]
-
-    # get the corresponding operation parameter index
-    # (that is, index of the parameter within the operation)
-    op_p_idx = tape._par_info[parameter_idx]["p_idx"]
-
-    # return the parameter-shift gradient for that
-    # operation parameter.
-    return op.get_parameter_shift(op_p_idx, shift=shift)
+    op, p_idx = tape.get_operation(t_idx)
+    return op.get_parameter_shift(p_idx, shift=shift)
 
 
 def _process_gradient_recipe(gradient_recipe, tol=1e-10):
@@ -150,12 +139,32 @@ def expval_param_shift(tape, argnum=None, shift=np.pi / 2, gradient_recipes=None
     shapes = []
     unshifted_coeffs = []
 
+    fns = []
+
     for idx, _ in enumerate(tape.trainable_params):
 
         if idx not in argnum:
             # parameter has zero gradient
             shapes.append(0)
             gradient_coeffs.append([])
+            fns.append(None)
+            continue
+
+        op, _ = tape.get_operation(idx)
+
+        if op.name == "Hamiltonian":
+            # operation is a Hamiltonian
+            if op.return_type is not qml.operation.Expectation:
+                raise ValueError(
+                    "Can only differentiate Hamiltonian "
+                    f"coefficients for expectations, not {op.return_type.value}"
+                )
+
+            g_tapes, h_fn = qml.gradients.hamiltonian_grad(tape, idx)
+            gradient_tapes.extend(g_tapes)
+            shapes.append(1)
+            gradient_coeffs.append(np.array([1.0]))
+            fns.append(h_fn)
             continue
 
         # get the gradient recipe for the trainable parameter
@@ -163,6 +172,7 @@ def expval_param_shift(tape, argnum=None, shift=np.pi / 2, gradient_recipes=None
         recipe = recipe or _get_operation_recipe(tape, idx, shift=shift)
         recipe = _process_gradient_recipe(recipe)
         coeffs, multipliers, shifts = recipe
+        fns.append(None)
 
         if shifts[0] == 0 and multipliers[0] == 1:
             # Gradient recipe includes a term with zero shift.
@@ -189,7 +199,7 @@ def expval_param_shift(tape, argnum=None, shift=np.pi / 2, gradient_recipes=None
         start = 1 if unshifted_coeffs and f0 is None else 0
         r0 = f0 or results[0]
 
-        for i, s in enumerate(shapes):
+        for i, (s, f) in enumerate(zip(shapes, fns)):
 
             if s == 0:
                 # parameter has zero gradient
@@ -200,6 +210,9 @@ def expval_param_shift(tape, argnum=None, shift=np.pi / 2, gradient_recipes=None
             res = results[start : start + s]
             start = start + s
 
+            if f is not None:
+                res = f(res)
+
             # compute the linear combination of results and coefficients
             res = qml.math.stack(res)
             g = qml.math.tensordot(res, qml.math.convert_like(gradient_coeffs[i], res), [[0], [0]])
diff --git a/pennylane/interfaces/batch/autograd.py b/pennylane/interfaces/batch/autograd.py
index 56bf3f276..98e05dc2f 100644
--- a/pennylane/interfaces/batch/autograd.py
+++ b/pennylane/interfaces/batch/autograd.py
@@ -190,7 +190,7 @@ def vjp(
                             gradient_kwargs=gradient_kwargs,
                         )
 
-                    vjps = processing_fn(execute_fn(vjp_tapes)[0])
+                        vjps = processing_fn(execute_fn(vjp_tapes)[0])
 
                 else:
                     vjp_tapes, processing_fn = qml.gradients.batch_vjp(
diff --git a/pennylane/interfaces/batch/tensorflow.py b/pennylane/interfaces/batch/tensorflow.py
index 107232c7e..8b1ad9943 100644
--- a/pennylane/interfaces/batch/tensorflow.py
+++ b/pennylane/interfaces/batch/tensorflow.py
@@ -122,7 +122,7 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
                                 gradient_kwargs=gradient_kwargs,
                             )
 
-                        vjps = processing_fn(execute_fn(vjp_tapes)[0])
+                            vjps = processing_fn(execute_fn(vjp_tapes)[0])
 
                     else:
                         vjp_tapes, processing_fn = qml.gradients.batch_vjp(
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index c91c3de2d..66391c7f2 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -738,6 +738,29 @@ class QuantumTape(AnnotatedQueue):
 
         self._trainable_params = param_indices
 
+    def get_operation(self, idx):
+        """Returns the trainable operation, and the corresponding operation argument
+        index, for a specified trainable parameter index.
+
+        Args:
+            idx (int): the trainable parameter index
+
+        Returns:
+            tuple[.Operation, int]: tuple containing the corresponding
+            operation, and an integer representing the argument index,
+            for the provided trainable parameter.
+        """
+        # get the index of the parameter in the tape
+        t_idx = list(self.trainable_params)[idx]
+
+        # get the corresponding operation
+        op = self._par_info[t_idx]["op"]
+
+        # get the corresponding operation parameter index
+        # (that is, index of the parameter within the operation)
+        p_idx = self._par_info[t_idx]["p_idx"]
+        return op, p_idx
+
     def get_parameters(self, trainable_only=True, **kwargs):  # pylint:disable=unused-argument
         """Return the parameters incident on the tape operations.
 
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index 612f4d3be..b19e80860 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -96,7 +96,6 @@ def classical_jacobian(qnode):
 
         def _jacobian(*args, **kwargs):
             with tf.GradientTape() as tape:
-                tape.watch(args)
                 gate_params = classical_preprocessing(*args, **kwargs)
 
             return tape.jacobian(gate_params, args)


Test patch:
diff --git a/tests/gradients/test_parameter_shift.py b/tests/gradients/test_parameter_shift.py
index 753710095..361adda29 100644
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -914,3 +914,286 @@ class TestParamShiftGradients:
             ]
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestHamiltonianExpvalGradients:
+    """Test that tapes ending with expval(H) can be
+    differentiated"""
+
+    def test_not_expval_error(self):
+        """Test that if the variance of the Hamiltonian is requested,
+        an error is raised"""
+        dev = qml.device("default.qubit", wires=2)
+
+        obs = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+        coeffs = np.array([0.1, 0.2, 0.3])
+        H = qml.Hamiltonian(coeffs, obs)
+
+        weights = np.array([0.4, 0.5])
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.var(H)
+
+        tape.trainable_params = {2, 3, 4}
+
+        with pytest.raises(ValueError, match="for expectations, not var"):
+            qml.gradients.param_shift(tape)
+
+    def test_no_trainable_coeffs(self, mocker, tol):
+        """Test no trainable Hamiltonian coefficients"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.gradients, "hamiltonian_grad")
+
+        obs = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+        coeffs = np.array([0.1, 0.2, 0.3])
+        H = qml.Hamiltonian(coeffs, obs)
+
+        weights = np.array([0.4, 0.5])
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        a, b, c = coeffs
+        x, y = weights
+        tape.trainable_params = {0, 1}
+
+        res = dev.batch_execute([tape])
+        expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        tapes, fn = qml.gradients.param_shift(tape)
+        # two shifts per rotation gate, one circuit per trainable H term
+        assert len(tapes) == 2 * 2
+        spy.assert_not_called()
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 2)
+
+        expected = [
+            -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+            b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_trainable_coeffs(self, mocker, tol):
+        """Test trainable Hamiltonian coefficients"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.gradients, "hamiltonian_grad")
+
+        obs = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+        coeffs = np.array([0.1, 0.2, 0.3])
+        H = qml.Hamiltonian(coeffs, obs)
+
+        weights = np.array([0.4, 0.5])
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        a, b, c = coeffs
+        x, y = weights
+        tape.trainable_params = {0, 1, 2, 4}
+
+        res = dev.batch_execute([tape])
+        expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        tapes, fn = qml.gradients.param_shift(tape)
+        # two shifts per rotation gate, one circuit per trainable H term
+        assert len(tapes) == 2 * 2 + 2
+        spy.assert_called()
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 4)
+
+        expected = [
+            -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+            b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+            np.cos(x),
+            -(np.sin(x) * np.sin(y)),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_multiple_hamiltonians(self, mocker, tol):
+        """Test multiple trainable Hamiltonian coefficients"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.gradients, "hamiltonian_grad")
+
+        obs = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+        coeffs = np.array([0.1, 0.2, 0.3])
+        a, b, c = coeffs
+        H1 = qml.Hamiltonian(coeffs, obs)
+
+        obs = [qml.PauliZ(0)]
+        coeffs = np.array([0.7])
+        d = coeffs[0]
+        H2 = qml.Hamiltonian(coeffs, obs)
+
+        weights = np.array([0.4, 0.5])
+        x, y = weights
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H1)
+            qml.expval(H2)
+
+        tape.trainable_params = {0, 1, 2, 4, 5}
+
+        res = dev.batch_execute([tape])
+        expected = [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        tapes, fn = qml.gradients.param_shift(tape)
+        # two shifts per rotation gate, one circuit per trainable H term
+        assert len(tapes) == 2 * 2 + 3
+        spy.assert_called()
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2, 5)
+
+        expected = [
+            [
+                -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                np.cos(x),
+                -(np.sin(x) * np.sin(y)),
+                0,
+            ],
+            [-d * np.sin(x), 0, 0, 0, np.cos(x)],
+        ]
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @staticmethod
+    def cost_fn(weights, coeffs1, coeffs2, dev=None):
+        """Cost function for gradient tests"""
+        obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+        H1 = qml.Hamiltonian(coeffs1, obs1)
+
+        obs2 = [qml.PauliZ(0)]
+        H2 = qml.Hamiltonian(coeffs2, obs2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H1)
+            qml.expval(H2)
+
+        tape.trainable_params = {0, 1, 2, 3, 4, 5}
+        tapes, fn = qml.gradients.param_shift(tape)
+        jac = fn(dev.batch_execute(tapes))
+        return jac
+
+    @staticmethod
+    def cost_fn_expected(weights, coeffs1, coeffs2):
+        """Analytic jacobian of cost_fn above"""
+        a, b, c = coeffs1
+        d = coeffs2[0]
+        x, y = weights
+        return [
+            [
+                -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                np.cos(x),
+                np.cos(x) * np.sin(y),
+                -(np.sin(x) * np.sin(y)),
+                0,
+            ],
+            [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)],
+        ]
+
+    def test_autograd(self, tol):
+        """Test gradient of multiple trainable Hamiltonian coefficients
+        using autograd"""
+        coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        coeffs2 = np.array([0.7], requires_grad=True)
+        weights = np.array([0.4, 0.5], requires_grad=True)
+        dev = qml.device("default.qubit.autograd", wires=2)
+
+        res = self.cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        # second derivative wrt to Hamiltonian coefficients should be zero
+        res = qml.jacobian(self.cost_fn)(weights, coeffs1, coeffs2, dev=dev)
+        assert np.allclose(res[1][:, 2:5], np.zeros([2, 3, 3]), atol=tol, rtol=0)
+        assert np.allclose(res[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
+
+    def test_tf(self, tol):
+        """Test gradient of multiple trainable Hamiltonian coefficients
+        using tf"""
+        tf = pytest.importorskip("tensorflow")
+
+        coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64)
+        coeffs2 = tf.Variable([0.7], dtype=tf.float64)
+        weights = tf.Variable([0.4, 0.5], dtype=tf.float64)
+
+        dev = qml.device("default.qubit.tf", wires=2)
+
+        with tf.GradientTape() as t:
+            jac = self.cost_fn(weights, coeffs1, coeffs2, dev=dev)
+
+        expected = self.cost_fn_expected(weights.numpy(), coeffs1.numpy(), coeffs2.numpy())
+        assert np.allclose(jac, expected, atol=tol, rtol=0)
+
+        # second derivative wrt to Hamiltonian coefficients should be zero
+        hess = t.jacobian(jac, [coeffs1, coeffs2])
+        assert np.allclose(hess[0][:, 2:5], np.zeros([2, 3, 3]), atol=tol, rtol=0)
+        assert np.allclose(hess[1][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
+
+    def test_torch(self, tol):
+        """Test gradient of multiple trainable Hamiltonian coefficients
+        using torch"""
+        torch = pytest.importorskip("torch")
+
+        coeffs1 = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float64, requires_grad=True)
+        coeffs2 = torch.tensor([0.7], dtype=torch.float64, requires_grad=True)
+        weights = torch.tensor([0.4, 0.5], dtype=torch.float64, requires_grad=True)
+
+        dev = qml.device("default.qubit.torch", wires=2)
+
+        res = self.cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(
+            weights.detach().numpy(), coeffs1.detach().numpy(), coeffs2.detach().numpy()
+        )
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+        # second derivative wrt to Hamiltonian coefficients should be zero
+        hess = torch.autograd.functional.jacobian(
+            lambda *args: self.cost_fn(*args, dev=dev), (weights, coeffs1, coeffs2)
+        )
+        assert np.allclose(hess[1][:, 2:5], np.zeros([2, 3, 3]), atol=tol, rtol=0)
+        assert np.allclose(hess[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
+
+    def test_jax(self, tol):
+        """Test gradient of multiple trainable Hamiltonian coefficients
+        using JAX"""
+        jax = pytest.importorskip("jax")
+        jnp = jax.numpy
+
+        coeffs1 = jnp.array([0.1, 0.2, 0.3])
+        coeffs2 = jnp.array([0.7])
+        weights = jnp.array([0.4, 0.5])
+        dev = qml.device("default.qubit.jax", wires=2)
+
+        res = self.cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        # second derivative wrt to Hamiltonian coefficients should be zero
+        res = jax.jacobian(self.cost_fn, argnums=1)(weights, coeffs1, coeffs2, dev=dev)
+        assert np.allclose(res[:, 2:5], np.zeros([2, 3, 3]), atol=tol, rtol=0)
+
+        res = jax.jacobian(self.cost_fn, argnums=1)(weights, coeffs1, coeffs2, dev=dev)
+        assert np.allclose(res[:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index e91733811..0dfa04f85 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -18,7 +18,7 @@ import pytest
 from pennylane import numpy as np
 
 import pennylane as qml
-from pennylane.gradients import param_shift
+from pennylane.gradients import finite_diff, param_shift
 from pennylane.interfaces.batch import execute
 
 
@@ -808,3 +808,93 @@ class TestHigherOrderDerivatives:
 
         expected = np.zeros([2, 2])
         assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+execute_kwargs = [
+    {"gradient_fn": param_shift},
+    {"gradient_fn": finite_diff},
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestHamiltonianWorkflows:
+    """Test that tapes ending with expectations
+    of Hamiltonians provide correct results and gradients"""
+
+    @pytest.fixture
+    def cost_fn(self, execute_kwargs):
+        """Cost function for gradient tests"""
+
+        def _cost_fn(weights, coeffs1, coeffs2, dev=None):
+            obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+            H1 = qml.Hamiltonian(coeffs1, obs1)
+
+            obs2 = [qml.PauliZ(0)]
+            H2 = qml.Hamiltonian(coeffs2, obs2)
+
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H1)
+                qml.expval(H2)
+
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        return _cost_fn
+
+    @staticmethod
+    def cost_fn_expected(weights, coeffs1, coeffs2):
+        """Analytic value of cost_fn above"""
+        a, b, c = coeffs1
+        d = coeffs2[0]
+        x, y = weights
+        return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
+
+    @staticmethod
+    def cost_fn_jacobian(weights, coeffs1, coeffs2):
+        """Analytic jacobian of cost_fn above"""
+        a, b, c = coeffs1
+        d = coeffs2[0]
+        x, y = weights
+        return np.array(
+            [
+                [
+                    -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                    b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                    np.cos(x),
+                    np.cos(x) * np.sin(y),
+                    -(np.sin(x) * np.sin(y)),
+                    0,
+                ],
+                [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)],
+            ]
+        )
+
+    def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False)
+        coeffs2 = np.array([0.7], requires_grad=False)
+        weights = np.array([0.4, 0.5], requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        coeffs2 = np.array([0.7], requires_grad=True)
+        weights = np.array([0.4, 0.5], requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev))
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/interfaces/test_batch_tensorflow.py b/tests/interfaces/test_batch_tensorflow.py
index 3e7519f30..c8744b3f4 100644
--- a/tests/interfaces/test_batch_tensorflow.py
+++ b/tests/interfaces/test_batch_tensorflow.py
@@ -18,7 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.gradients import param_shift
+from pennylane.gradients import finite_diff, param_shift
 from pennylane.interfaces.batch import execute
 
 
@@ -848,3 +848,101 @@ class TestHigherOrderDerivatives:
 
         hess = t2.jacobian(grad, params)
         assert hess is None
+
+
+execute_kwargs = [
+    {"gradient_fn": param_shift, "interface": "tensorflow"},
+    {"gradient_fn": finite_diff, "interface": "tensorflow"},
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestHamiltonianWorkflows:
+    """Test that tapes ending with expectations
+    of Hamiltonians provide correct results and gradients"""
+
+    @pytest.fixture
+    def cost_fn(self, execute_kwargs):
+        """Cost function for gradient tests"""
+
+        def _cost_fn(weights, coeffs1, coeffs2, dev=None):
+            obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+            H1 = qml.Hamiltonian(coeffs1, obs1)
+
+            obs2 = [qml.PauliZ(0)]
+            H2 = qml.Hamiltonian(coeffs2, obs2)
+
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H1)
+                qml.expval(H2)
+
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        return _cost_fn
+
+    @staticmethod
+    def cost_fn_expected(weights, coeffs1, coeffs2):
+        """Analytic value of cost_fn above"""
+        a, b, c = coeffs1.numpy()
+        d = coeffs2.numpy()[0]
+        x, y = weights.numpy()
+        return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
+
+    @staticmethod
+    def cost_fn_jacobian(weights, coeffs1, coeffs2):
+        """Analytic jacobian of cost_fn above"""
+        a, b, c = coeffs1.numpy()
+        d = coeffs2.numpy()[0]
+        x, y = weights.numpy()
+        return np.array(
+            [
+                [
+                    -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                    b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                    np.cos(x),
+                    np.cos(x) * np.sin(y),
+                    -(np.sin(x) * np.sin(y)),
+                    0,
+                ],
+                [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)],
+            ]
+        )
+
+    def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64)
+        coeffs2 = tf.constant([0.7], dtype=tf.float64)
+        weights = tf.Variable([0.4, 0.5], tf.float64)
+        dev = qml.device("default.qubit", wires=2)
+
+        with tf.GradientTape() as tape:
+            res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = tape.jacobian(res, [weights, coeffs1, coeffs2])
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0)
+        assert res[1] is None
+        assert res[2] is None
+
+    def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = tf.Variable([0.1, 0.2, 0.3], tf.float64)
+        coeffs2 = tf.Variable([0.7], tf.float64)
+        weights = tf.Variable([0.4, 0.5], tf.float64)
+        dev = qml.device("default.qubit", wires=2)
+
+        with tf.GradientTape() as tape:
+            res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = tape.jacobian(res, [weights, coeffs1, coeffs2])
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[:, 2:5], atol=tol, rtol=0)
+        assert np.allclose(res[2], expected[:, 5:], atol=tol, rtol=0)
diff --git a/tests/interfaces/test_batch_torch.py b/tests/interfaces/test_batch_torch.py
index ecd7ca131..35fd68bbd 100644
--- a/tests/interfaces/test_batch_torch.py
+++ b/tests/interfaces/test_batch_torch.py
@@ -20,7 +20,7 @@ import pytest
 torch = pytest.importorskip("torch")
 
 import pennylane as qml
-from pennylane.gradients import param_shift
+from pennylane.gradients import finite_diff, param_shift
 from pennylane.interfaces.batch import execute
 
 
@@ -997,3 +997,101 @@ class TestHigherOrderDerivatives:
         res = torch.autograd.functional.hessian(cost_fn, params)
         expected = torch.zeros([2, 2], dtype=torch.float64)
         assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0)
+
+
+execute_kwargs = [
+    {"gradient_fn": param_shift, "interface": "torch"},
+    {"gradient_fn": finite_diff, "interface": "torch"},
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestHamiltonianWorkflows:
+    """Test that tapes ending with expectations
+    of Hamiltonians provide correct results and gradients"""
+
+    @pytest.fixture
+    def cost_fn(self, execute_kwargs):
+        """Cost function for gradient tests"""
+
+        def _cost_fn(weights, coeffs1, coeffs2, dev=None):
+            obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+            H1 = qml.Hamiltonian(coeffs1, obs1)
+
+            obs2 = [qml.PauliZ(0)]
+            H2 = qml.Hamiltonian(coeffs2, obs2)
+
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H1)
+                qml.expval(H2)
+
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        return _cost_fn
+
+    @staticmethod
+    def cost_fn_expected(weights, coeffs1, coeffs2):
+        """Analytic value of cost_fn above"""
+        a, b, c = coeffs1.detach().numpy()
+        d = coeffs2.detach().numpy()[0]
+        x, y = weights.detach().numpy()
+        return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
+
+    @staticmethod
+    def cost_fn_jacobian(weights, coeffs1, coeffs2):
+        """Analytic jacobian of cost_fn above"""
+        a, b, c = coeffs1.detach().numpy()
+        d = coeffs2.detach().numpy()[0]
+        x, y = weights.detach().numpy()
+        return np.array(
+            [
+                [
+                    -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                    b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                    np.cos(x),
+                    np.cos(x) * np.sin(y),
+                    -(np.sin(x) * np.sin(y)),
+                    0,
+                ],
+                [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)],
+            ]
+        )
+
+    def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False)
+        coeffs2 = torch.tensor([0.7], requires_grad=False)
+        weights = torch.tensor([0.4, 0.5], requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+        res = torch.hstack(
+            torch.autograd.functional.jacobian(
+                lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)
+            )
+        )
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+    def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
+        coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True)
+        coeffs2 = torch.tensor([0.7], requires_grad=True)
+        weights = torch.tensor([0.4, 0.5], requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+        res = torch.hstack(
+            torch.autograd.functional.jacobian(
+                lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)
+            )
+        )
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)