Repository: PennyLaneAI/pennylane
Pull request #1515 (2021-08-11 15:00:08)

Store grouping in Hamiltonian
**Context:**

With the new VQE workflow, we need a way for the user to define whether PennyLane and devices are supposed to use grouping. Grouping observables can take a long time and should therefore only be computed once. A way to meet both requirements is to allow users to store grouping information in the Hamiltonian, which - if found - is used by PennyLane:

  ``` python
  obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
  coeffs = np.array([1., 2., 3.])
  H = qml.Hamiltonian(coeffs, obs, compute_groupings=True)
  ```
  Initialization with ``compute_groupings=True`` stores the indices required to make groups of 
  commuting observables and their coefficients. These are used 
  (or, if not found, computed) by the ``get_groupings()`` method, which returns the 
  actual groupings.
  
  ``` pycon
  >>> H.grouping_indices
  [[0, 1], [2]]
  grouped_coeffs, grouped_obs = H.get_groupings()
  >>> grouped_coeffs
  [np.array([1., 2.]), np.array(3.)]
  >>> grouped_obs
  [[qml.PauliX(0), qml.PauliX(1)], [qml.PauliZ(0)]]
  ```

PennyLane's qnode, which currently splits a tape with `expval(H)` into tapes with a single Pauli measurement each, looks for the `grouping_indices` attribute, and if found, uses groups of Pauli measurements per tape. Also devices that support Hamiltonians can use this strategy if desired. 

**Description of the Change:**

Added grouping functionality to the Hamiltonian class, and uses it in `hamiltonian_expand`, which in turn is used by the QNode to split `expval(H)` into Pauli expectations.

**Benefits:**

One can now control the use of grouping behaviour in the observable.

**Possible Drawbacks:**

I am quite unsure about the design details. Feedback welcome!



codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1515](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5edcc03) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/2cf9d2998229b7ee8cfc0b219c7997053c8eaadf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2cf9d29) will **decrease** coverage by `0.10%`.
> The diff coverage is `97.19%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1515      +/-   ##
==========================================
- Coverage   98.21%   98.11%   -0.11%     
==========================================
  Files         184      184              
  Lines       13201    13250      +49     
==========================================
+ Hits        12966    13000      +34     
- Misses        235      250      +15     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515?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/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `98.08% <ø> (ø)` | |
| [pennylane/vqe/vqe.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3ZxZS92cWUucHk=) | `93.03% <96.20%> (+0.79%)` | :arrow_up: |
| [pennylane/grouping/group\_observables.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyb3VwaW5nL2dyb3VwX29ic2VydmFibGVzLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.47% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `95.78% <100.00%> (ø)` | |
| [pennylane/transforms/hamiltonian\_expand.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/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%> (ø)` | |
| [pennylane/transforms/measurement\_grouping.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWVhc3VyZW1lbnRfZ3JvdXBpbmcucHk=) | `17.64% <0.00%> (-82.36%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515?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/1515?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [2cf9d29...5edcc03](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1515?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).

mariaschuld: Locally, codecov has 100% coverage for all files I changed, except from one line that I never changed (and which is a NotImplementedError). Also the `measurement_grouping` file it mentions has 100% coverage. I really don't understand why it complains.
mariaschuld: Thanks @glassnotes I think I implemented all your suggestions and also checked the docs once more...

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index c73d3bc09..15500bc8a 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,24 @@
 
 <h3>New features since last release</h3>
 
+* The Hamiltonian can now store grouping information, which can be accessed by a device to 
+  speed up computations of the expectation value of a Hamiltonian. 
+  [(#1515)](https://github.com/PennyLaneAI/pennylane/pull/1515)
+
+  ```python
+  obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
+  coeffs = np.array([1., 2., 3.])
+  H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+  ```
+  
+  Initialization with a ``grouping_type`` other than ``None`` stores the indices 
+  required to make groups of commuting observables and their coefficients. 
+  
+  ``` pycon
+  >>> H.grouping_indices
+  [[0, 1], [2]]
+  ```
+
 * Hamiltonians are now trainable with respect to their coefficients.
   [(#1483)](https://github.com/PennyLaneAI/pennylane/pull/1483)
 
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 590c8dd53..0bd9104a8 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -26,7 +26,8 @@ def hamiltonian_expand(tape, group=True):
     Args:
         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
+        group (bool): Whether to compute disjoint groups of commuting Pauli observables, leading to fewer tapes.
+            If grouping information can be found in the Hamiltonian, it will be used even if group=False.
 
     Returns:
         tuple[list[.QuantumTape], function]: Returns a tuple containing a list of
@@ -67,18 +68,42 @@ def hamiltonian_expand(tape, group=True):
     >>> fn(res)
     -0.5
 
-    .. Warning::
+    Fewer tapes can be constructed by grouping commuting observables. This can be achieved
+    by the ``group`` keyword argument:
 
-         Note that defining Hamiltonians inside of QNodes using arithmetic can lead to errors.
-         See :class:`~pennylane.Hamiltonian` for more information.
+    .. 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)
+
+    With grouping, the Hamiltonian gets split into two groups of observables (here ``[qml.PauliZ(0)]`` and
+    ``[qml.PauliX(1), qml.PauliX(0)]``):
+
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape)
+    >>> len(tapes)
+    2
+
+    Without grouping it gets split into three groups (``[qml.PauliZ(0)]``, ``[qml.PauliX(1)]`` and ``[qml.PauliX(0)]``):
+
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+    >>> len(tapes)
+    3
 
-    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:
+    Alternatively, if the Hamiltonian has already computed groups, they are used even if ``group=False``:
 
     .. code-block:: python3
 
-        H = qml.Hamiltonian([1., 2., 3.], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)])
+        obs = [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)]
+        coeffs = [1., 2., 3.]
+        H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+
+        # the initialisation already computes grouping information and stores it in the Hamiltonian
+        assert H.grouping_indices is not None
 
         with qml.tape.QuantumTape() as tape:
             qml.Hadamard(wires=0)
@@ -86,13 +111,11 @@ def hamiltonian_expand(tape, group=True):
             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
+    Grouping information has been used to reduce the number of tapes from 3 to 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
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+    >>> len(tapes)
+    2
     """
 
     hamiltonian = tape.measurements[0].obs
@@ -106,27 +129,48 @@ def hamiltonian_expand(tape, group=True):
             "Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`"
         )
 
-    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:
-        # we need to create a new tape here, because
-        # updating metadata of a copied tape is error-prone
-        # when the observables were changed
-        with tape.__class__() as new_tape:
-            for op in tape.operations:
-                qml.apply(op)
-            qml.expval(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))
+    if group or hamiltonian.grouping_indices is not None:
+
+        if hamiltonian.grouping_indices is None:
+            hamiltonian.compute_grouping()
+
+        # use groups of observables if available or explicitly requested
+        coeffs = [
+            qml.math.squeeze(qml.math.take(hamiltonian.coeffs, indices, axis=0))
+            for indices in hamiltonian.grouping_indices
+        ]
+        obs_groupings = [
+            [hamiltonian.ops[i] for i in indices] for indices in hamiltonian.grouping_indices
         ]
+
+        tapes = []
+        for obs in obs_groupings:
+
+            with tape.__class__() as new_tape:
+                for op in tape.operations:
+                    op.queue()
+
+                for o in obs:
+                    qml.expval(o)
+
+            new_tape = new_tape.expand(stop_at=lambda obj: True)
+            tapes.append(new_tape)
+    else:
+        coeffs = hamiltonian.coeffs
+
+        tapes = []
+        for o in hamiltonian.ops:
+            with tape.__class__() as new_tape:
+                for op in tape.operations:
+                    op.queue()
+                qml.expval(o)
+
+            tapes.append(new_tape)
+
+    def processing_fn(res):
+        # note: res could have an extra dimension here if a shots_distribution
+        # is used for evaluation
+        dot_products = [qml.math.dot(qml.math.squeeze(r), c) for c, r in zip(coeffs, 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 069e68155..961975133 100644
--- a/pennylane/vqe/vqe.py
+++ b/pennylane/vqe/vqe.py
@@ -30,6 +30,35 @@ from pennylane.wires import Wires
 OBS_MAP = {"PauliX": "X", "PauliY": "Y", "PauliZ": "Z", "Hadamard": "H", "Identity": "I"}
 
 
+def _compute_grouping_indices(observables, grouping_type="qwc", method="rlf"):
+
+    # todo: directly compute the
+    # indices, instead of extracting groups of observables first
+    observable_groups = qml.grouping.group_observables(
+        observables, coefficients=None, grouping_type=grouping_type, method=method
+    )
+
+    observables = copy(observables)
+
+    indices = []
+    available_indices = list(range(len(observables)))
+    for partition in observable_groups:
+        indices_this_group = []
+        for pauli_word in partition:
+            # find index of this pauli word in remaining original observables,
+            for observable in observables:
+                if qml.grouping.utils.are_identical_pauli_words(pauli_word, observable):
+                    ind = observables.index(observable)
+                    indices_this_group.append(available_indices[ind])
+                    # delete this observable and its index, so it cannot be found again
+                    observables.pop(ind)
+                    available_indices.pop(ind)
+                    break
+        indices.append(indices_this_group)
+
+    return indices
+
+
 class Hamiltonian(qml.operation.Observable):
     r"""Operator representing a Hamiltonian.
 
@@ -41,6 +70,13 @@ class Hamiltonian(qml.operation.Observable):
         observables (Iterable[Observable]): observables in the Hamiltonian expression, of same length as coeffs
         simplify (bool): Specifies whether the Hamiltonian is simplified upon initialization
                          (like-terms are combined). The default value is `False`.
+        grouping_type (str): If not None, compute and store information on how to group commuting
+            observables upon initialization. This information may be accessed when QNodes containing this
+            Hamiltonian are executed on devices. The string refers to the type of binary relation between Pauli words.
+            Can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or ``'anticommuting'``.
+        method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+            can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
+        id (str): name to be assigned to this Hamiltonian instance
 
     **Example:**
 
@@ -93,9 +129,34 @@ class Hamiltonian(qml.operation.Observable):
 
     >>> H1 = qml.Hamiltonian(torch.tensor([1.]), [qml.PauliX(0)])
     >>> H2 = qml.Hamiltonian(torch.tensor([2., 3.]), [qml.PauliY(0), qml.PauliX(1)])
-    >>> H3 = qml.Hamiltonian(torch.tensor([1., 2., 3.]), [qml.PauliX(0), qml.PauliY(0), qml.PauliX(1)])
+    >>> obs3 = [qml.PauliX(0), qml.PauliY(0), qml.PauliX(1)]
+    >>> H3 = qml.Hamiltonian(torch.tensor([1., 2., 3.]), obs3)
     >>> H3.compare(H1 + H2)
     True
+
+    A Hamiltonian can store information on which commuting observables should be measured together in
+    a circuit:
+
+    >>> obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
+    >>> coeffs = np.array([1., 2., 3.])
+    >>> H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+    >>> H.grouping_indices
+    [[0, 1], [2]]
+
+    This attribute can be used to compute groups of coefficients and observables:
+
+    >>> grouped_coeffs = [coeffs[indices] for indices in H.grouping_indices]
+    >>> grouped_obs = [[H.ops[i] for i in indices] for indices in H.grouping_indices]
+    >>> grouped_coeffs
+    [tensor([1., 2.], requires_grad=True), tensor([3.], requires_grad=True)]
+    >>> grouped_obs
+    [[qml.PauliX(0), qml.PauliX(1)], [qml.PauliZ(0)]]
+
+    Devices that evaluate a Hamiltonian expectation by splitting it into its local observables can
+    use this information to reduce the number of circuits evaluated.
+
+    Note that one can compute the ``grouping_indices`` for an already initialized Hamiltonian by
+    using the :func:`compute_grouping <pennylane.Hamiltonian.compute_grouping>` method.
     """
 
     num_wires = qml.operation.AnyWires
@@ -103,7 +164,16 @@ class Hamiltonian(qml.operation.Observable):
     par_domain = "A"
     grad_method = "A"  # supports analytic gradients
 
-    def __init__(self, coeffs, observables, simplify=False, id=None, do_queue=True):
+    def __init__(
+        self,
+        coeffs,
+        observables,
+        simplify=False,
+        grouping_type=None,
+        method="rlf",
+        id=None,
+        do_queue=True,
+    ):
 
         if qml.math.shape(coeffs)[0] != len(observables):
             raise ValueError(
@@ -123,8 +193,16 @@ class Hamiltonian(qml.operation.Observable):
 
         self.return_type = None
 
+        # attribute to store indices used to form groups of
+        # commuting observables, since recomputation is costly
+        self._grouping_indices = None
+
         if simplify:
             self.simplify()
+        if grouping_type is not None:
+            self._grouping_indices = qml.transforms.invisible(_compute_grouping_indices)(
+                self.ops, grouping_type=grouping_type, method=method
+            )
 
         coeffs_flat = [self._coeffs[i] for i in range(qml.math.shape(self._coeffs)[0])]
         # overwrite this attribute, now that we have the correct info
@@ -175,6 +253,31 @@ class Hamiltonian(qml.operation.Observable):
     def name(self):
         return "Hamiltonian"
 
+    @property
+    def grouping_indices(self):
+        """Return the grouping indices attribute.
+
+        Returns:
+            list[list[int]]: indices needed to form groups of commuting observables
+        """
+        return self._grouping_indices
+
+    def compute_grouping(self, grouping_type="qwc", method="rlf"):
+        """
+        Compute groups of indices corresponding to commuting observables of this
+        Hamiltonian, and store it in the ``grouping_indices`` attribute.
+
+        Args:
+            grouping_type (str): The type of binary relation between Pauli words used to compute the grouping.
+                Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
+            method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+                can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
+        """
+
+        self._grouping_indices = qml.transforms.invisible(_compute_grouping_indices)(
+            self.ops, grouping_type=grouping_type, method=method
+        )
+
     def simplify(self):
         r"""Simplifies the Hamiltonian by combining like-terms.
 
@@ -186,6 +289,11 @@ class Hamiltonian(qml.operation.Observable):
         >>> print(H)
           (-1) [X0]
         + (1) [Y2]
+
+        .. warning::
+
+            Calling this method will reset ``grouping_indices`` to None, since
+            the observables it refers to are updated.
         """
         data = []
         ops = []
@@ -213,6 +321,8 @@ class Hamiltonian(qml.operation.Observable):
         self._coeffs = qml.math.stack(data) if data else []
         self.data = data
         self._ops = ops
+        # reset grouping, since the indices refer to the old observables and coefficients
+        self._grouping_indices = None
 
     def __str__(self):
         # Lambda function that formats the wires
@@ -277,7 +387,7 @@ class Hamiltonian(qml.operation.Observable):
 
         return data
 
-    def compare(self, H):
+    def compare(self, other):
         r"""Compares with another :class:`~Hamiltonian`, :class:`~.Observable`, or :class:`~.Tensor`,
         to determine if they are equivalent.
 
@@ -317,15 +427,15 @@ class Hamiltonian(qml.operation.Observable):
         >>> ob1.compare(ob2)
         False
         """
-        if isinstance(H, Hamiltonian):
+        if isinstance(other, Hamiltonian):
             self.simplify()
-            H.simplify()
-            return self._obs_data() == H._obs_data()  # pylint: disable=protected-access
+            other.simplify()
+            return self._obs_data() == other._obs_data()  # pylint: disable=protected-access
 
-        if isinstance(H, (Tensor, Observable)):
+        if isinstance(other, (Tensor, Observable)):
             self.simplify()
             return self._obs_data() == {
-                (1, frozenset(H._obs_data()))  # pylint: disable=protected-access
+                (1, frozenset(other._obs_data()))  # pylint: disable=protected-access
             }
 
         raise ValueError("Can only compare a Hamiltonian, and a Hamiltonian/Observable/Tensor.")


Test patch:
diff --git a/tests/ops/test_hamiltonian.py b/tests/ops/test_hamiltonian.py
index 2d429b9dc..541405587 100644
--- a/tests/ops/test_hamiltonian.py
+++ b/tests/ops/test_hamiltonian.py
@@ -404,6 +404,86 @@ class TestHamiltonianArithmeticJax:
         assert H.compare(H1 @ H2)
 
 
+class TestGrouping:
+    """Tests for the grouping functionality"""
+
+    def test_grouping_is_correct_kwarg(self):
+        """Basic test checking that grouping with a kwarg works as expected"""
+        a = qml.PauliX(0)
+        b = qml.PauliX(1)
+        c = qml.PauliZ(0)
+        obs = [a, b, c]
+        coeffs = [1.0, 2.0, 3.0]
+
+        H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc")
+        assert H.grouping_indices == [[0, 1], [2]]
+
+    def test_grouping_is_correct_compute_grouping(self):
+        """Basic test checking that grouping with compute_grouping works as expected"""
+        a = qml.PauliX(0)
+        b = qml.PauliX(1)
+        c = qml.PauliZ(0)
+        obs = [a, b, c]
+        coeffs = [1.0, 2.0, 3.0]
+
+        H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc")
+        H.compute_grouping()
+        assert H.grouping_indices == [[0, 1], [2]]
+
+    def test_grouping_for_non_groupable_hamiltonians(self):
+        """Test that grouping is computed correctly, even if no observables commute"""
+        a = qml.PauliX(0)
+        b = qml.PauliY(0)
+        c = qml.PauliZ(0)
+        obs = [a, b, c]
+        coeffs = [1.0, 2.0, 3.0]
+
+        H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc")
+        assert H.grouping_indices == [[0], [1], [2]]
+
+    def test_grouping_is_reset_when_simplifying(self):
+        """Tests that calling simplify() resets the grouping"""
+        obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
+        coeffs = [1.0, 2.0, 3.0]
+
+        H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc")
+        assert H.grouping_indices is not None
+
+        H.simplify()
+        assert H.grouping_indices is None
+
+    def test_grouping_does_not_alter_queue(self):
+        """Tests that grouping is invisible to the queue."""
+        a = qml.PauliX(0)
+        b = qml.PauliX(1)
+        c = qml.PauliZ(0)
+        obs = [a, b, c]
+        coeffs = [1.0, 2.0, 3.0]
+
+        with qml.tape.QuantumTape() as tape:
+            H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc")
+
+        assert tape.queue == [a, b, c, H]
+
+    def test_grouping_method_can_be_set(self):
+        r"""Tests that the grouping method can be controlled by kwargs.
+        This is done by changing from default to 'rlf' and checking the result."""
+        a = qml.PauliX(0)
+        b = qml.PauliX(1)
+        c = qml.PauliZ(0)
+        obs = [a, b, c]
+        coeffs = [1.0, 2.0, 3.0]
+
+        # compute grouping during construction
+        H2 = qml.Hamiltonian(coeffs, obs, grouping_type="qwc", method="lf")
+        assert H2.grouping_indices == [[2, 1], [0]]
+
+        # compute grouping separately
+        H3 = qml.Hamiltonian(coeffs, obs, grouping_type=None)
+        H3.compute_grouping(method="lf")
+        assert H3.grouping_indices == [[2, 1], [0]]
+
+
 class TestHamiltonianEvaluation:
     """Test the usage of a Hamiltonian as an observable"""
 
@@ -453,11 +533,12 @@ class TestHamiltonianEvaluation:
         assert pars == [0.1, 3.0]
 
 
+@pytest.mark.parametrize("simplify", [True, False])
+@pytest.mark.parametrize("group", [None, "qwc"])
 class TestHamiltonianDifferentiation:
     """Test that the Hamiltonian coefficients are differentiable"""
 
-    @pytest.mark.parametrize("simplify", [True, False])
-    def test_vqe_differentiation_paramshift(self, simplify):
+    def test_vqe_differentiation_paramshift(self, simplify, group):
         """Test the parameter-shift method by comparing the differentiation of linearly combined subcircuits
         with the differentiation of a Hamiltonian expectation"""
         coeffs = np.array([-0.05, 0.17])
@@ -469,7 +550,12 @@ class TestHamiltonianDifferentiation:
             qml.RX(param, wires=0)
             qml.RY(param, wires=0)
             return qml.expval(
-                qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify)
+                qml.Hamiltonian(
+                    coeffs,
+                    [qml.PauliX(0), qml.PauliZ(0)],
+                    simplify=simplify,
+                    grouping_type=group,
+                )
             )
 
         grad_fn = qml.grad(circuit)
@@ -489,8 +575,7 @@ class TestHamiltonianDifferentiation:
         assert np.allclose(grad[0], grad_expected[0])
         assert np.allclose(grad[1], grad_expected[1])
 
-    @pytest.mark.parametrize("simplify", [True, False])
-    def test_vqe_differentiation_autograd(self, simplify):
+    def test_vqe_differentiation_autograd(self, simplify, group):
         """Test the autograd interface by comparing the differentiation of linearly combined subcircuits
         with the differentiation of a Hamiltonian expectation"""
         coeffs = pnp.array([-0.05, 0.17], requires_grad=True)
@@ -502,7 +587,12 @@ class TestHamiltonianDifferentiation:
             qml.RX(param, wires=0)
             qml.RY(param, wires=0)
             return qml.expval(
-                qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify)
+                qml.Hamiltonian(
+                    coeffs,
+                    [qml.PauliX(0), qml.PauliZ(0)],
+                    simplify=simplify,
+                    grouping_type=group,
+                )
             )
 
         grad_fn = qml.grad(circuit)
@@ -522,8 +612,7 @@ class TestHamiltonianDifferentiation:
         assert np.allclose(grad[0], grad_expected[0])
         assert np.allclose(grad[1], grad_expected[1])
 
-    @pytest.mark.parametrize("simplify", [True, False])
-    def test_vqe_differentiation_jax(self, simplify):
+    def test_vqe_differentiation_jax(self, simplify, group):
         """Test the jax interface by comparing the differentiation of linearly combined subcircuits
         with the differentiation of a Hamiltonian expectation"""
 
@@ -538,7 +627,12 @@ class TestHamiltonianDifferentiation:
             qml.RX(param, wires=0)
             qml.RY(param, wires=0)
             return qml.expval(
-                qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify)
+                qml.Hamiltonian(
+                    coeffs,
+                    [qml.PauliX(0), qml.PauliZ(0)],
+                    simplify=simplify,
+                    grouping_type=group,
+                )
             )
 
         grad_fn = jax.grad(circuit)
@@ -558,8 +652,7 @@ class TestHamiltonianDifferentiation:
         assert np.allclose(grad[0], grad_expected[0])
         assert np.allclose(grad[1], grad_expected[1])
 
-    @pytest.mark.parametrize("simplify", [True, False])
-    def test_vqe_differentiation_torch(self, simplify):
+    def test_vqe_differentiation_torch(self, simplify, group):
         """Test the torch interface by comparing the differentiation of linearly combined subcircuits
         with the differentiation of a Hamiltonian expectation"""
 
@@ -573,7 +666,12 @@ class TestHamiltonianDifferentiation:
             qml.RX(param, wires=0)
             qml.RY(param, wires=0)
             return qml.expval(
-                qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify)
+                qml.Hamiltonian(
+                    coeffs,
+                    [qml.PauliX(0), qml.PauliZ(0)],
+                    simplify=simplify,
+                    grouping_type=group,
+                )
             )
 
         res = circuit(coeffs, param)
@@ -600,8 +698,7 @@ class TestHamiltonianDifferentiation:
         assert np.allclose(grad[0], grad_expected[0])
         assert np.allclose(grad[1], grad_expected[1])
 
-    @pytest.mark.parametrize("simplify", [True, False])
-    def test_vqe_differentiation_tf(self, simplify):
+    def test_vqe_differentiation_tf(self, simplify, group):
         """Test the tf interface by comparing the differentiation of linearly combined subcircuits
         with the differentiation of a Hamiltonian expectation"""
 
@@ -615,7 +712,12 @@ class TestHamiltonianDifferentiation:
             qml.RX(param, wires=0)
             qml.RY(param, wires=0)
             return qml.expval(
-                qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify)
+                qml.Hamiltonian(
+                    coeffs,
+                    [qml.PauliX(0), qml.PauliZ(0)],
+                    simplify=simplify,
+                    grouping_type=group,
+                )
             )
 
         with tf.GradientTape() as tape:
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 8754804ab..3c74455cc 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -1253,7 +1253,7 @@ class TestNewVQE:
 
             res_expected.append(circuit())
 
-        res_expected = np.sum(c * r for c, r in zip(coeffs, res_expected))
+        res_expected = np.sum([c * r for c, r in zip(coeffs, res_expected)])
 
         assert np.isclose(res, res_expected, atol=tol)
 
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index e1cde7c5d..cbfcd758b 100644
--- a/tests/transforms/test_hamiltonian_expand.py
+++ b/tests/transforms/test_hamiltonian_expand.py
@@ -98,6 +98,22 @@ class TestHamiltonianExpval:
 
         assert np.isclose(output, expval)
 
+    def test_grouping_is_used(self):
+        """Test that the grouping in a Hamiltonian is used"""
+        H = qml.Hamiltonian(
+            [1.0, 2.0, 3.0], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)], grouping_type="qwc"
+        )
+        assert H.grouping_indices is not None
+
+        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) == 2
+
     def test_number_of_tapes(self):
         """Tests that the the correct number of tapes is produced"""
 
diff --git a/tests/transforms/test_measurement_grouping.py b/tests/transforms/test_measurement_grouping.py
new file mode 100644
index 000000000..527b1695e
--- /dev/null
+++ b/tests/transforms/test_measurement_grouping.py
@@ -0,0 +1,37 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import numpy as np
+import pennylane as qml
+
+
+def test_measurement_grouping():
+    """Test that measurement grouping works as expected."""
+
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(0.1, wires=0)
+        qml.RX(0.2, wires=1)
+        qml.CNOT(wires=[0, 1])
+        qml.CNOT(wires=[1, 2])
+
+    obs = [qml.PauliZ(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliX(2)]
+    coeffs = [2.0, -0.54, 0.1]
+
+    tapes, fn = qml.transforms.measurement_grouping(tape, obs, coeffs)
+    assert len(tapes) == 2
+
+    dev = qml.device("default.qubit", wires=3)
+    res = fn(dev.batch_execute(tapes))
+    assert np.isclose(res, 2.0007186)