Repository: PennyLaneAI/pennylane
Pull request #1474 (2021-07-26 08:53:40)

Support expval(H) syntax
**Context:**

This PR is the first in a series of PRs that simplify, extend and speedup VQE problems in PennyLane.

**Description of the Change:**

Adds an if statement to the qnode call, which, if encountering a Hamiltonian observable on the tape which is not supported on the device, splits the tape, evaluates the subtapes, and recombines them.

Basic tests for differentiation are added.

**Benefits:**

This allows users to compute and differentiate through VQE problems using the syntax `expval(H)`. The splitting and recombination is tracked by the classical interface, and gradients are automatically supported (since quantum gradients of each subtape are provided, and the recombination taken care of by autodifferentiation).

**Possible drawbacks:** 

Are there any obvious tests one could add?

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/1474?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1474](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e7be706) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/76c9b5747f8b1cacf5b1033d3ead255277299c4a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (76c9b57) will **decrease** coverage by `0.01%`.
> The diff coverage is `92.30%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474/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/1474?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1474      +/-   ##
==========================================
- Coverage   98.30%   98.28%   -0.02%     
==========================================
  Files         170      170              
  Lines       12483    12505      +22     
==========================================
+ Hits        12271    12291      +20     
- Misses        212      214       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `97.31% <80.00%> (-0.70%)` | :arrow_down: |
| [pennylane/transforms/hamiltonian\_expand.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvaGFtaWx0b25pYW5fZXhwYW5kLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474?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/1474?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [76c9b57...e7be706](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1474?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: @mariaschuld had a sudden thought - what happens when you try and `qml.draw(qnode)(params)` for a QNode that ends in `expval(H)`?
mariaschuld: > qml.draw(qnode)(params)

Excellent thought. It does not work, because `Hamiltonian` has no atribute `num_params`. Let me see what I can do.
mariaschuld: Some crude benchmarking (since the benchmarking suite still uses the old VQE setup) shows that the new VQE design is a little bit faster (around 10%):

``` python
import pennylane as qml
from pennylane import numpy as np
from functools import partial
from pennylane.templates.subroutines import UCCSD
from pennylane import Identity, PauliX, PauliY, PauliZ
from time import time

s_wires = [[0, 1, 2], [1, 2, 3]]
d_wires = [[[0, 1], [2, 3]]]

hf_state = np.array([1, 1, 0, 0])

ansatz = partial(UCCSD, init_state=hf_state, s_wires=s_wires, d_wires=d_wires)
params = np.array([3.14545258, 3.13766988, -0.21446816])

h2_coeffs = np.array(
    [
        -0.05963862,
        0.17575739,
        0.17575739,
        -0.23666489,
        -0.23666489,
        0.17001485,
        0.04491735,
        -0.04491735,
        -0.04491735,
        0.04491735,
        0.12222641,
        0.16714376,
        0.16714376,
        0.12222641,
        0.17570278,
    ]
)

h2_ops = [
    Identity(wires=[0]),
    PauliZ(wires=[0]),
    PauliZ(wires=[1]),
    PauliZ(wires=[2]),
    PauliZ(wires=[3]),
    PauliZ(wires=[0]) @ PauliZ(wires=[1]),
    PauliY(wires=[0]) @ PauliX(wires=[1]) @ PauliX(wires=[2]) @ PauliY(wires=[3]),
    PauliY(wires=[0]) @ PauliY(wires=[1]) @ PauliX(wires=[2]) @ PauliX(wires=[3]),
    PauliX(wires=[0]) @ PauliX(wires=[1]) @ PauliY(wires=[2]) @ PauliY(wires=[3]),
    PauliX(wires=[0]) @ PauliY(wires=[1]) @ PauliY(wires=[2]) @ PauliX(wires=[3]),
    PauliZ(wires=[0]) @ PauliZ(wires=[2]),
    PauliZ(wires=[0]) @ PauliZ(wires=[3]),
    PauliZ(wires=[1]) @ PauliZ(wires=[2]),
    PauliZ(wires=[1]) @ PauliZ(wires=[3]),
    PauliZ(wires=[2]) @ PauliZ(wires=[3]),
]

ham = qml.Hamiltonian(h2_coeffs, h2_ops)
n_steps = 3
device = qml.device("default.qubit", wires=len(hf_state), shots=None)
opt = qml.GradientDescentOptimizer(stepsize=0.4)


def benchmark_old_vqe():
    params_temp = params
    cost_fn = qml.ExpvalCost(ansatz, ham, device, optimize=False)
    for _ in range(n_steps):
        params_temp, energy = opt.step_and_cost(cost_fn, params_temp)


def benchmark_new_vqe():

    @qml.qnode(device)
    def cost_fn(weights):
        ansatz(weights, wires=range(4))
        return qml.expval(ham)

    params_temp = params
    for _ in range(n_steps):
        params_temp, energy = opt.step_and_cost(cost_fn, params_temp)


start = time()
benchmark_old_vqe()
print(time() - start)  # 2.8592017209994083

start = time()
benchmark_new_vqe()
print(time() - start) # 2.3986086410004646
```

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index c5bc8f887..e2a406b36 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,31 @@
 
 <h3>New features since last release</h3>
 
+* VQE problems can now intuitively been set up by passing the Hamiltonian 
+  as an observable. [(#1474)](https://github.com/PennyLaneAI/pennylane/pull/1474)
+
+  ``` python
+  dev = qml.device("default.qubit", wires=2)
+  H = qml.Hamiltonian([1., 2., 3.],  [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(1)])
+  w = qml.init.strong_ent_layers_uniform(1, 2, seed=1967)
+  
+  @qml.qnode(dev)
+  def circuit(w):
+      qml.templates.StronglyEntanglingLayers(w, wires=range(2))
+      return qml.expval(H)
+  ```
+  
+  ```pycon
+  >>> print(circuit(w))
+  -1.5133943637878295
+  >>> print(qml.grad(circuit)(w))
+  [[[-8.32667268e-17  1.39122955e+00 -9.12462052e-02]
+  [ 1.02348685e-16 -7.77143238e-01 -1.74708049e-01]]]
+  ```
+  
+  Note that other measurement types like `var(H)` or `sample(H)`, as well 
+  as multiple expectations like `expval(H1), expval(H2)` are not supported. 
+
 * A new gradients module `qml.gradients` has been added, which provides
   differentiable quantum gradient transforms.
   [(#1476)](https://github.com/PennyLaneAI/pennylane/pull/1476)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 857e34e05..29e6b3425 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -605,8 +605,25 @@ class QNode:
             # construct the tape
             self.construct(args, kwargs)
 
-        # execute the tape
-        res = self.qtape.execute(device=self.device)
+        # Execute the tape.
+        # If the observable contains a Hamiltonian and the device does not
+        # support Hamiltonians, split tape into multiple tapes and execute
+        # sequentially. In future, this logic should be moved to the device
+        # to allow for more efficient batch execution.
+        supports_hamiltonian = self.device.supports_observable("Hamiltonian")
+        hamiltonian_in_obs = "Hamiltonian" in [obs.name for obs in self.qtape.observables]
+        if hamiltonian_in_obs and not supports_hamiltonian:
+            try:
+                tapes, fn = qml.transforms.hamiltonian_expand(self.qtape, group=False)
+            except ValueError as e:
+                raise ValueError(
+                    "At the moment, only single expectations of Hamiltonian observables can be returned"
+                    "on the {} device.".format(self.device.name)
+                ) from e
+            results = [tape.execute(device=self.device) for tape in tapes]
+            res = fn(results)
+        else:
+            res = self.qtape.execute(device=self.device)
 
         # if shots was changed
         if original_shots != -1:
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 305258f9e..927375fef 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -14,17 +14,19 @@
 """
 Contains the hamiltonian expand tape transform
 """
+# pylint: disable=protected-access
 import pennylane as qml
 
 
-def hamiltonian_expand(tape):
+def hamiltonian_expand(tape, group=True):
     r"""
-    Returns a list of tapes, and a classical processing function, for computing the expectation
-    value of a Hamiltonian.
+    Splits a tape measuring a Hamiltonian expectation into mutliple tapes of Pauli expectations,
+    and provides a function to recombine the results.
 
     Args:
-        tape (.QuantumTape) the tape used when calculating the expectation value
-            of the Hamiltonian.
+        tape (.QuantumTape): the tape used when calculating the expectation value
+            of the Hamiltonian
+        group (bool): whether to compute groups of non-commuting Pauli observables, leading to fewer tapes
 
     Returns:
         tuple[list[.QuantumTape], function]: Returns a tuple containing a list of
@@ -69,6 +71,28 @@ def hamiltonian_expand(tape):
 
          Note that defining Hamiltonians inside of QNodes using arithmetic can lead to errors.
          See :class:`~pennylane.Hamiltonian` for more information.
+
+    The ``group`` keyword argument toggles between the creation of one tape per Pauli observable, or
+    one tape per group of non-commuting Pauli observables computed by the :func:`.measurement_grouping`
+    transform:
+
+    .. code-block:: python3
+
+        H = qml.Hamiltonian([1., 2., 3.], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliX(wires=2)
+            qml.expval(H)
+
+        # split H into observable groups [qml.PauliZ(0)] and [qml.PauliX(1), qml.PauliX(0)]
+        tapes, fn = qml.transforms.hamiltonian_expand(tape)
+        print(len(tapes)) # 2
+
+        # split H into observables [qml.PauliZ(0)], [qml.PauliX(1)] and [qml.PauliX(0)]
+        tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+        print(len(tapes)) # 3
     """
 
     hamiltonian = tape.measurements[0].obs
@@ -77,6 +101,24 @@ def hamiltonian_expand(tape):
         raise ValueError(
             "Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`"
         )
-
-    hamiltonian.simplify()
-    return qml.transforms.measurement_grouping(tape, hamiltonian.ops, hamiltonian.coeffs)
+    if group:
+        hamiltonian.simplify()
+        return qml.transforms.measurement_grouping(tape, hamiltonian.ops, hamiltonian.coeffs)
+
+    # create tapes that measure the Pauli-words in the Hamiltonian
+    tapes = []
+    for ob in hamiltonian.ops:
+        new_tape = tape.copy()
+        new_tape._measurements = [
+            qml.measure.MeasurementProcess(return_type=qml.operation.Expectation, obs=ob)
+        ]
+        tapes.append(new_tape)
+
+    # create processing function that performs linear recombination
+    def processing_fn(res):
+        dot_products = [
+            qml.math.dot(qml.math.squeeze(res[i]), hamiltonian.coeffs[i]) for i in range(len(res))
+        ]
+        return qml.math.sum(qml.math.stack(dot_products), axis=0)
+
+    return tapes, processing_fn
diff --git a/pennylane/vqe/vqe.py b/pennylane/vqe/vqe.py
index 365ea0aae..fef96c5e2 100644
--- a/pennylane/vqe/vqe.py
+++ b/pennylane/vqe/vqe.py
@@ -89,6 +89,8 @@ class Hamiltonian:
 
         Note that this issue also arises when calling the ``simplify()`` method.
     """
+    # Todo: this is a temporary solution to make the circuit drawer work
+    num_params = 0
 
     def __init__(self, coeffs, observables, simplify=False):
 


Test patch:
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 2243fb08c..d84c9ae0c 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -18,6 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
+from pennylane import numpy as pnp
 from pennylane.wires import Wires
 from pennylane.devices import DefaultQubit, DefaultMixed
 
@@ -1212,6 +1213,232 @@ class TestVQE:
             qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, hamiltonian, dev, measure="var")
 
 
+class TestNewVQE:
+    """Test the new VQE syntax of passing the Hamiltonian as an observable."""
+
+    @pytest.mark.parametrize("ansatz, params", CIRCUITS)
+    @pytest.mark.parametrize("observables", OBSERVABLES)
+    def test_circuits_evaluate(self, ansatz, observables, params, tol):
+        """Tests simple VQE evaluations."""
+        coeffs = [1.0] * len(observables)
+        dev = qml.device("default.qubit", wires=3)
+        H = qml.Hamiltonian(coeffs, observables)
+
+        # pass H directly
+        @qml.qnode(dev)
+        def circuit():
+            ansatz(params, wires=range(3))
+            return qml.expval(H)
+
+        res = circuit()
+
+        res_expected = []
+        for obs in observables:
+
+            @qml.qnode(dev)
+            def circuit():
+                ansatz(params, wires=range(3))
+                return qml.expval(obs)
+
+            res_expected.append(circuit())
+
+        res_expected = np.sum(c * r for c, r in zip(coeffs, res_expected))
+
+        assert np.isclose(res, res_expected, atol=tol)
+
+    def test_acting_on_subcircuit(self, tol):
+        """Tests a VQE circuit where the observable does not act on all wires."""
+        dev = qml.device("default.qubit", wires=3)
+        coeffs = [1.0, 1.0, 1.0]
+        w = qml.init.strong_ent_layers_uniform(1, 2, seed=1967)
+
+        observables1 = [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(1)]
+        H1 = qml.Hamiltonian(coeffs, observables1)
+
+        @qml.qnode(dev)
+        def circuit1():
+            qml.templates.StronglyEntanglingLayers(w, wires=range(2))
+            return qml.expval(H1)
+
+        observables2 = [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(1) @ qml.Identity(2)]
+        H2 = qml.Hamiltonian(coeffs, observables2)
+
+        @qml.qnode(dev)
+        def circuit2():
+            qml.templates.StronglyEntanglingLayers(w, wires=range(2))
+            return qml.expval(H2)
+
+        res1 = circuit1()
+        res2 = circuit2()
+
+        assert np.allclose(res1, res2, atol=tol)
+
+    @pytest.mark.parametrize("shots, dim", [([(1000, 2)], 2), ([30, 30], 2), ([2, 3, 4], 3)])
+    def test_shot_distribution(self, shots, dim):
+        """Tests that distributed shots work with the new VQE design."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(weights, coeffs):
+            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
+            H = qml.Hamiltonian(coeffs, obs)
+            return qml.expval(H)
+
+        obs = [qml.PauliZ(0), qml.PauliX(0) @ qml.PauliZ(1)]
+        coeffs = np.array([0.1, 0.2])
+        weights = pnp.random.random([2, 2, 3], requires_grad=True)
+
+        res = circuit(weights, coeffs)
+        grad = qml.jacobian(circuit, argnum=1)(weights, coeffs)
+        assert len(res) == dim
+        assert grad.shape == (dim, 2)
+
+    def test_circuit_drawer(self):
+        """Test that the circuit drawer displays Hamiltonians well."""
+        dev = qml.device("default.qubit", wires=3)
+        coeffs = [1.0, 1.0, 1.0]
+        observables1 = [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(2)]
+        H1 = qml.Hamiltonian(coeffs, observables1)
+
+        @qml.qnode(dev)
+        def circuit1():
+            qml.Hadamard(wires=0)
+            return qml.expval(H1)
+
+        res = qml.draw(circuit1)()
+        expected = " 0: ──H──╭Hamiltonian──┤  \n" + " 2: ─────╰Hamiltonian──┤  \n"
+        assert res == expected
+
+    def test_error_multiple_expvals(self):
+        """Tests that error is thrown if more than one expval is evaluated."""
+        observables = [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(1)]
+        coeffs = [1.0] * len(observables)
+        dev = qml.device("default.qubit", wires=3)
+        H = qml.Hamiltonian(coeffs, observables)
+        w = qml.init.strong_ent_layers_uniform(2, 4, seed=1967)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H), qml.expval(qml.PauliX(3))
+
+        with pytest.raises(ValueError, match="At the moment"):
+            circuit()
+
+    def test_error_non_expval_measurement(self):
+        """Tests that error is thrown if sample() or var() is used."""
+        observables = [qml.PauliZ(0), qml.PauliY(0), qml.PauliZ(1)]
+        coeffs = [1.0] * len(observables)
+        dev = qml.device("default.qubit", wires=3)
+        H = qml.Hamiltonian(coeffs, observables)
+        w = qml.init.strong_ent_layers_uniform(2, 4, seed=1967)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.sample(H)
+
+        with pytest.raises(qml.QuantumFunctionError, match="Hamiltonian is not an observable"):
+            circuit()
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.var(H)
+
+        with pytest.raises(qml.QuantumFunctionError, match="Hamiltonian is not an observable"):
+            circuit()
+
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "best"])
+    def test_grad_autograd(self, diff_method, tol):
+        """Tests the VQE gradient in the autograd interface."""
+        dev = qml.device("default.qubit", wires=4)
+        H = big_hamiltonian
+        w = pnp.array(qml.init.strong_ent_layers_uniform(2, 4, seed=1967), requires_grad=True)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(w):
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H)
+
+        dc = qml.grad(circuit)(w)
+        assert np.allclose(dc, big_hamiltonian_grad, atol=tol)
+
+    def test_grad_zero_hamiltonian(self, tol):
+        """Tests the VQE gradient for a "zero" Hamiltonian."""
+        dev = qml.device("default.qubit", wires=4)
+        H = qml.Hamiltonian([0], [qml.PauliX(0)])
+        w = qml.init.strong_ent_layers_uniform(2, 4, seed=1967)
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit(w):
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H)
+
+        dc = qml.grad(circuit)(w)
+        assert np.allclose(dc, 0, atol=tol)
+
+    def test_grad_torch(self, torch_support, tol):
+        """Tests VQE gradients in the torch interface."""
+        if not torch_support:
+            pytest.skip("This test requires Torch")
+
+        dev = qml.device("default.qubit", wires=4)
+        H = big_hamiltonian
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(w):
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H)
+
+        w = torch.tensor(qml.init.strong_ent_layers_uniform(2, 4, seed=1967), requires_grad=True)
+
+        res = circuit(w)
+        res.backward()
+        dc = w.grad.detach().numpy()
+
+        assert np.allclose(dc, big_hamiltonian_grad, atol=tol)
+
+    def test_grad_tf(self, tf_support, tol):
+        """Tests VQE gradients in the tf interface."""
+        if not tf_support:
+            pytest.skip("This test requires TensorFlow")
+
+        dev = qml.device("default.qubit", wires=4)
+        H = big_hamiltonian
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(w):
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H)
+
+        w = tf.Variable(qml.init.strong_ent_layers_uniform(2, 4, seed=1967), dtype=tf.double)
+
+        with tf.GradientTape() as tape:
+            res = circuit(w)
+
+        dc = tape.gradient(res, w).numpy()
+
+        assert np.allclose(dc, big_hamiltonian_grad, atol=tol)
+
+    def test_grad_jax(self, tol):
+        """Tests VQE gradients in the jax interface."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        dev = qml.device("default.qubit", wires=4)
+        H = big_hamiltonian
+        w = jnp.array(qml.init.strong_ent_layers_uniform(2, 4, seed=1967))
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(w):
+            qml.templates.StronglyEntanglingLayers(w, wires=range(4))
+            return qml.expval(H)
+
+        dc = jax.grad(circuit)(w)
+        assert np.allclose(dc, big_hamiltonian_grad, atol=tol)
+
+
 class TestAutogradInterface:
     """Tests for the Autograd interface (and the NumPy interface for backward compatibility)"""
 
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index 2b8d080f6..a2533a8d2 100644
--- a/tests/transforms/test_hamiltonian_expand.py
+++ b/tests/transforms/test_hamiltonian_expand.py
@@ -88,6 +88,34 @@ class TestHamiltonianExpval:
 
         assert np.isclose(output, expval)
 
+    @pytest.mark.parametrize(("tape", "output"), zip(TAPES, OUTPUTS))
+    def test_hamiltonians_no_grouping(self, tape, output):
+        """Tests that the hamiltonian_expand transform returns the correct value
+        if we switch grouping off"""
+
+        tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+        results = dev.batch_execute(tapes)
+        expval = fn(results)
+
+        assert np.isclose(output, expval)
+
+    def test_number_of_tapes(self):
+        """Tests that the the correct number of tapes is produced"""
+
+        H = qml.Hamiltonian([1.0, 2.0, 3.0], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliX(wires=2)
+            qml.expval(H)
+
+        tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+        assert len(tapes) == 3
+
+        tapes, fn = qml.transforms.hamiltonian_expand(tape, group=True)
+        assert len(tapes) == 2
+
     def test_hamiltonian_error(self):
 
         with pennylane.tape.QuantumTape() as tape: