Repository: PennyLaneAI/pennylane
Pull request #6043 (2024-07-26 18:03:25)

Improve implementation of `group_observables`
**Context:**

`qml.pauli.group_observables` has several areas of improvement that make the calculation slow. This PR attempts to address some of them.

**Description of the Change:**

1. Solve the Graph colouring algorithm using `Rustworkx` instead of custom implementation. This adds 2 additional methods to solve the Minimum Clique cover problem: DSATUR (`'dsatur'`) and IndependentSet (`'gis'`)
2. Use the symplectic representation of the Pauli observables to construct the adjacency matrix of the complement graph.
3. Introduce a new function to obtain the indices of the partitions directly: `qml.pauli.compute_partition_indices`

**Benefits:**

Improves of orders of magnitude compared to the current implementation of the `qwc` grouping type strategy. The current implementation is based on an $\mathcal{O}(m^2)$ computation, with $m$ the number of observables to be grouped. Instead, taking advantage of the symplectic inner product, and its relation with the commutator and anti commutator of Pauli observables [1], we can obtain a significant improvement. 

![image](https://github.com/user-attachments/assets/3f0ec70b-2082-44a3-acf9-9cdf4238da08)

In addition, by introducing `qml.pauli.compute_partition_indices`, we now avoid the (rather often) inefficient computation of partitions of observables, just to then calculate the indices from them - see https://github.com/PennyLaneAI/pennylane/blob/2892a9a0aa6797912ef4a03f3c4759eaff01d8ef/pennylane/ops/qubit/hamiltonian.py#L37

Instead, we obtain the indices of the partitions directly from the graph colouring algorithm. Putting all together, the improvement in performance when calculating the partitions of  large Hamiltonians is evident.

![image](https://github.com/user-attachments/assets/7bd047a9-a628-46c8-9d3e-58818fd7acc0)


This graph was obtained by running 
``` python
%timeit hamiltonian.compute_grouping(grouping_type="qwc", method='lf')
```

**Notes:**

Profiling the implementation from `linear_combination` demonstrates that the post processing to obtain the indices remains a significant burden in the calculations, and therefore should be addressed for the remaining functions - see reference code above.

![image](https://github.com/user-attachments/assets/09cf4ca6-601a-4e91-8865-8bf727dd3f1d)

[1] Andrew Jena (2019). Partitioning Pauli Operators: in Theory and in Practice. UWSpace. http://hdl.handle.net/10012/15017
[2] "Fast Partitioning of Pauli Strings into Commuting Families for Optimal Expectation Value Measurements of Dense Operators". https://arxiv.org/abs/2305.11847


[sc-64594]

## Performance comparisons between colouring algorithms

Comparison for `LinearCombination.compute_grouping()`

![image](https://github.com/user-attachments/assets/042d040a-f75b-4642-b6e4-0e903c5d58c5)

![image](https://github.com/user-attachments/assets/a89197ee-9c3d-424d-9bba-36829ee33388)

Comparison for `colour_paui_graph` (only between `Rustworkx` algorithms)

![image](https://github.com/user-attachments/assets/52ae4d0f-acaa-4d45-aab5-d3c40c37863a)





github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.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.
albi3ro: Thanks for opening this PR 🚀 

I was just running some initial on some large problems:

```
dataset = qml.data.load("qchem", molname="BeH2", bondlength=0.78, basis="STO-3G")
H = dataset[0].hamiltonian
obs = H.terms()[1]
%timeit qml.pauli.group_observables(obs, method="lf")
```
![Screenshot 2024-07-26 at 2 57 10 PM](https://github.com/user-attachments/assets/541bcaaf-e92b-45a9-a824-7c9dee384eac)

Since that is now the main bottleneck, it might be worth looking into whether or not we need to calculate the rustworkx graph that way.

Also, when looking at:
```
%snakeviz H.compute_grouping(method="lf")
```

It seems like 46% of the time is in `are_identical_pauli_words`. This leads credence to the fact that we need a different design than the one provided by `group_observables`.  
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `73.52941%` with `27 lines` in your changes missing coverage. Please review.
> Project coverage is 99.59%. Comparing base [(`3ad15d7`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3ad15d7ee01b1c07b74d3d69f82c6c1699967076?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`eb37d6f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/eb37d6f15b90b1ee0eea21f8f4e357c8ff38e039?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 37 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/ops/op\_math/linear\_combination.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?src=pr&el=tree&filepath=pennylane%2Fops%2Fop_math%2Flinear_combination.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2xpbmVhcl9jb21iaW5hdGlvbi5weQ==) | 15.78% | [16 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/pauli/grouping/group\_observables.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?src=pr&el=tree&filepath=pennylane%2Fpauli%2Fgrouping%2Fgroup_observables.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3BhdWxpL2dyb3VwaW5nL2dyb3VwX29ic2VydmFibGVzLnB5) | 86.58% | [11 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6043      +/-   ##
==========================================
- Coverage   99.66%   99.59%   -0.08%     
==========================================
  Files         427      430       +3     
  Lines       41057    41333     +276     
==========================================
+ Hits        40919    41165     +246     
- Misses        138      168      +30     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6043?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5b4765931..e120047dd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -52,6 +52,14 @@
 * A new method `to_mat` has been added to the `FermiWord` and `FermiSentence` classes, which allows
   computing the matrix representation of these Fermi operators.
   [(#5920)](https://github.com/PennyLaneAI/pennylane/pull/5920)
+  
+* `qml.pauli.group_observables` now uses `Rustworkx` colouring algorithms to solve the Minimum Clique Cover problem.
+  This adds two new options for the `method` argument: `dsatur` and `gis`. In addition, the creation of the adjancecy matrix 
+  now takes advantage of the symplectic representation of the Pauli observables. An additional function `qml.pauli.compute_partition_indices` 
+  is added to calculate the indices from the partitioned observables more efficiently. `qml.pauli.grouping.PauliGroupingStrategy.idx_partitions_from_graph` 
+  can be used to compute partitions of custom indices. These changes improve the wall time of `qml.LinearCombination.compute_grouping` 
+  and the `grouping_type='qwc'` by orders of magnitude. 
+  [(#6043)](https://github.com/PennyLaneAI/pennylane/pull/6043)
 
 <h4>Improvements to operators</h4>
 
diff --git a/doc/requirements.txt b/doc/requirements.txt
index 79daaa35a..4449a5399 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -31,7 +31,7 @@ torch==1.9.0; sys_platform == "darwin" and python_version < "3.10"
 torch==1.13.1; sys_platform != "darwin" and python_version == "3.10"
 torch==1.13.1; sys_platform == "darwin" and python_version == "3.10"
 jinja2==3.0.3
-rustworkx==0.12.1
+rustworkx>=0.14.0
 networkx==2.6
 requests~=2.28.1
 # we do not pin the sphinx theme, to allow the
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 2d87d3e50..ee43cf07e 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -63,7 +63,6 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
     # measurements with no observables
     mp_no_obs = []
     mp_no_obs_indices = []
-
     for i, mp in enumerate(mps):
         if isinstance(mp.obs, (Sum, SProd, Prod)):
             mps[i].obs = qml.simplify(mp.obs)
@@ -78,13 +77,11 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
         else:
             mp_other_obs.append([mp])
             mp_other_obs_indices.append([i])
-
     if mp_pauli_obs:
         i_to_pauli_mp = dict(mp_pauli_obs)
         _, group_indices = qml.pauli.group_observables(
             [mp.obs for mp in i_to_pauli_mp.values()], list(i_to_pauli_mp.keys())
         )
-
         mp_pauli_groups = []
         for indices in group_indices:
             mp_group = [i_to_pauli_mp[i] for i in indices]
@@ -94,7 +91,6 @@ def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, Sh
 
     mp_no_obs_indices = [mp_no_obs_indices] if mp_no_obs else []
     mp_no_obs = [mp_no_obs] if mp_no_obs else []
-
     all_mp_groups = mp_pauli_groups + mp_no_obs + mp_other_obs
     all_indices = group_indices + mp_no_obs_indices + mp_other_obs_indices
 
@@ -240,7 +236,6 @@ def measure_with_samples(
     mps = measurements[0 : -len(mid_measurements)] if mid_measurements else measurements
 
     groups, indices = _group_measurements(mps)
-
     all_res = []
     for group in groups:
         if isinstance(group[0], ExpectationMP) and isinstance(
diff --git a/pennylane/ops/functions/dot.py b/pennylane/ops/functions/dot.py
index a49dd668c..dd5b85cb5 100644
--- a/pennylane/ops/functions/dot.py
+++ b/pennylane/ops/functions/dot.py
@@ -49,7 +49,7 @@ def dot(
         grouping_type (str): The type of binary relation between Pauli words used to compute
             the grouping. Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``. Note that if
             ``pauli=True``, the grouping will be ignored.
-        method (str): The graph coloring heuristic to use in solving minimum clique cover for
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
 
diff --git a/pennylane/ops/op_math/linear_combination.py b/pennylane/ops/op_math/linear_combination.py
index 722843a92..853461ac3 100644
--- a/pennylane/ops/op_math/linear_combination.py
+++ b/pennylane/ops/op_math/linear_combination.py
@@ -38,16 +38,20 @@ class LinearCombination(Sum):
         coeffs (tensor_like): coefficients of the ``LinearCombination`` expression
         observables (Iterable[Observable]): observables in the ``LinearCombination`` expression, of same length as ``coeffs``
         simplify (bool): Specifies whether the ``LinearCombination`` is simplified upon initialization
-                         (like-terms are combined). The default value is `False`. Note that ``coeffs`` cannot
-                         be differentiated when using the ``'torch'`` interface and ``simplify=True``. Use of this argument is deprecated.
+                        (like-terms are combined). The default value is `False`. Note that ``coeffs`` cannot
+                        be differentiated when using the ``'torch'`` interface and ``simplify=True``. Use of this argument is deprecated.
         grouping_type (str): If not ``None``, compute and store information on how to group commuting
             observables upon initialization. This information may be accessed when a :class:`~.QNode` containing this
             ``LinearCombination`` is 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). Ignored if ``grouping_type=None``.
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
+            can be ``'lf'`` (Largest First), ``'rlf'`` (Recursive Largest First), ``'dsatur'`` (Degree of Saturation), or ``'gis'`` (IndependentSet).
+            Defaults to ``'lf'``. Ignored if ``grouping_type=None``.
         id (str): name to be assigned to this ``LinearCombination`` instance
 
+    .. seealso:: `rustworkx.ColoringStrategy <https://www.rustworkx.org/apiref/rustworkx.ColoringStrategy.html#coloringstrategy>`_
+        for more information on the ``('lf', 'dsatur', 'gis')`` strategies.
+
     .. warning::
         The ``simplify`` argument is deprecated and will be removed in a future release.
         Instead, you can call ``qml.simplify`` on the constructed operator.
@@ -122,7 +126,7 @@ class LinearCombination(Sum):
         observables: list[Operator],
         simplify=False,
         grouping_type=None,
-        method="rlf",
+        method="lf",
         _grouping_indices=None,
         _pauli_rep=None,
         id=None,
@@ -229,7 +233,7 @@ class LinearCombination(Sum):
         """
         return self.coeffs, self.ops
 
-    def compute_grouping(self, grouping_type="qwc", method="rlf"):
+    def compute_grouping(self, grouping_type="qwc", method="lf"):
         """
         Compute groups of operators and coefficients corresponding to commuting
         observables of this ``LinearCombination``.
@@ -242,9 +246,10 @@ class LinearCombination(Sum):
         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
+                Defaults to ``'qwc'``.
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for
                 grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
-                First).
+                First). Defaults to ``'lf'``.
 
         **Example**
 
@@ -271,27 +276,9 @@ class LinearCombination(Sum):
 
         _, ops = self.terms()
 
-        with qml.QueuingManager.stop_recording():
-            op_groups = qml.pauli.group_observables(ops, grouping_type=grouping_type, method=method)
-
-        ops = copy(ops)
-
-        indices = []
-        available_indices = list(range(len(ops)))
-        for partition in op_groups:  # pylint:disable=too-many-nested-blocks
-            indices_this_group = []
-            for pauli_word in partition:
-                # find index of this pauli word in remaining original observables,
-                for ind, observable in enumerate(ops):
-                    if qml.pauli.are_identical_pauli_words(pauli_word, observable):
-                        indices_this_group.append(available_indices[ind])
-                        # delete this observable and its index, so it cannot be found again
-                        ops.pop(ind)
-                        available_indices.pop(ind)
-                        break
-            indices.append(tuple(indices_this_group))
-
-        self._grouping_indices = tuple(indices)
+        self._grouping_indices = qml.pauli.compute_partition_indices(
+            ops, grouping_type=grouping_type, method=method
+        )
 
     @property
     def wires(self):
diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py
index 62c705e82..422fdc128 100644
--- a/pennylane/ops/op_math/sum.py
+++ b/pennylane/ops/op_math/sum.py
@@ -43,7 +43,7 @@ def sum(*summands, grouping_type=None, method="rlf", id=None, lazy=True):
             of the operators is already a sum operator, its operands (summands) will be used instead.
         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
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
 
@@ -129,7 +129,7 @@ class Sum(CompositeOp):
     Keyword 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
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for
             grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
             First). This keyword argument is ignored if ``grouping_type`` is ``None``.
         id (str or None): id for the sum operator. Default is None.
@@ -480,7 +480,7 @@ class Sum(CompositeOp):
         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
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for
                 grouping, which can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest
                 First).
 
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index d9cd3ddf5..239060107 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -87,7 +87,7 @@ class Hamiltonian(Observable):
             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
+        method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
             can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First). Ignored if ``grouping_type=None``.
         id (str): name to be assigned to this Hamiltonian instance
 
@@ -476,7 +476,7 @@ class Hamiltonian(Observable):
         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
+            method (str): The graph colouring heuristic to use in solving minimum clique cover for grouping, which
                 can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
         """
 
diff --git a/pennylane/pauli/__init__.py b/pennylane/pauli/__init__.py
index 87decb1cb..ab90d27df 100644
--- a/pennylane/pauli/__init__.py
+++ b/pennylane/pauli/__init__.py
@@ -48,6 +48,7 @@ from .grouping import (
     PauliGroupingStrategy,
     optimize_measurements,
     graph_colouring,
+    compute_partition_indices,
 )
 
 from .dla import PauliVSpace, lie_closure, structure_constants, center
diff --git a/pennylane/pauli/grouping/__init__.py b/pennylane/pauli/grouping/__init__.py
index 134906ea1..84ab08ab4 100644
--- a/pennylane/pauli/grouping/__init__.py
+++ b/pennylane/pauli/grouping/__init__.py
@@ -17,5 +17,5 @@ functionality used in measurement optimization.
 """
 
 from . import graph_colouring
-from .group_observables import PauliGroupingStrategy, group_observables
+from .group_observables import PauliGroupingStrategy, group_observables, compute_partition_indices
 from .optimize_measurements import optimize_measurements
diff --git a/pennylane/pauli/grouping/group_observables.py b/pennylane/pauli/grouping/group_observables.py
index bfd93358c..b8549e67d 100644
--- a/pennylane/pauli/grouping/group_observables.py
+++ b/pennylane/pauli/grouping/group_observables.py
@@ -15,10 +15,14 @@
 This module contains the high-level Pauli-word-partitioning functionality used in measurement optimization.
 """
 
+from collections import defaultdict
 from copy import copy
-from typing import Literal, Optional
+from functools import cached_property
+from operator import itemgetter
+from typing import Literal, Optional, Sequence
 
 import numpy as np
+import rustworkx as rx
 
 import pennylane as qml
 from pennylane.ops import Prod, SProd
@@ -26,15 +30,27 @@ from pennylane.pauli.utils import (
     are_identical_pauli_words,
     binary_to_pauli,
     observables_to_binary_matrix,
-    qwc_complement_adj_matrix,
 )
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
-from .graph_colouring import largest_first, recursive_largest_first
+from .graph_colouring import recursive_largest_first
 
 GROUPING_TYPES = frozenset(["qwc", "commuting", "anticommuting"])
-GRAPH_COLOURING_METHODS = {"lf": largest_first, "rlf": recursive_largest_first}
+
+# ColoringStrategy is only available from version 0.15.0
+new_rx = True
+try:
+    RX_STRATEGIES = {
+        "lf": rx.ColoringStrategy.Degree,
+        "dsatur": rx.ColoringStrategy.Saturation,
+        "gis": rx.ColoringStrategy.IndependentSet,
+    }
+except AttributeError:  # pragma: no cover
+    new_rx = False  # pragma: no cover. # This error is raised for versions lower than 0.15.0
+    RX_STRATEGIES = {"lf": None}  # pragma: no cover # Only "lf" can be used without a strategy
+
+GRAPH_COLOURING_METHODS = frozenset(RX_STRATEGIES.keys()).union({"rlf"})
 
 
 class PauliGroupingStrategy:  # pylint: disable=too-many-instance-attributes
@@ -55,49 +71,65 @@ class PauliGroupingStrategy:  # pylint: disable=too-many-instance-attributes
     find approximate solutions in polynomial time.
 
     Args:
-        observables (list[Observable]): a list of Pauli words to be partitioned according to a
-            grouping strategy
-        grouping_type (str): the binary relation used to define partitions of
+        observables (list[Operator]): A list of Pauli words to be partitioned according to a
+            grouping strategy.
+        grouping_type (str): The binary relation used to define partitions of
             the Pauli words, can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or
             ``'anticommuting'``.
-        graph_colourer (str): the heuristic algorithm to employ for graph
-            colouring, can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive
-            Largest First)
+        graph_colourer (str): The heuristic algorithm to employ for graph
+                colouring, can be ``'lf'`` (Largest First), ``'rlf'`` (Recursive
+                Largest First), ``'dsatur'`` (Degree of Saturation), or ``'gis'`` (IndependentSet). Defaults to ``'lf'``.
 
     Raises:
-        ValueError: if arguments specified for ``grouping_type`` or
-            ``graph_colourer`` are not recognized
+        ValueError: If arguments specified for ``grouping_type`` or ``graph_colourer``
+            are not recognized.
+
+    .. seealso:: `rustworkx.ColoringStrategy <https://www.rustworkx.org/apiref/rustworkx.ColoringStrategy.html#coloringstrategy>`_
+        for more information on the ``('lf', 'dsatur', 'gis')`` strategies.
     """
 
     def __init__(
         self,
         observables,
         grouping_type: Literal["qwc", "commuting", "anticommuting"] = "qwc",
-        graph_colourer: Literal["lf", "rlf"] = "rlf",
+        graph_colourer: Literal["lf", "rlf", "dsatur", "gis"] = "lf",
     ):
-        if grouping_type.lower() not in GROUPING_TYPES:
+        self.graph_colourer = graph_colourer.lower()
+        self.grouping_type = grouping_type.lower()
+
+        if self.grouping_type not in GROUPING_TYPES:
             raise ValueError(
                 f"Grouping type must be one of: {GROUPING_TYPES}, instead got {grouping_type}."
             )
 
-        self.grouping_type = grouping_type.lower()
+        if self.graph_colourer in ["dsatur", "gis"] and not new_rx:
+            raise ValueError(
+                f"The strategy '{graph_colourer}' is not supported in this version of Rustworkx. "
+                "Please install rustworkx>=0.15.0 to access the 'dsatur' and 'gis' colouring strategies."
+            )
 
-        if graph_colourer.lower() not in GRAPH_COLOURING_METHODS:
+        if self.graph_colourer not in GRAPH_COLOURING_METHODS:
             raise ValueError(
-                f"Graph colouring method must be one of: {list(GRAPH_COLOURING_METHODS)}, "
-                f"instead got {graph_colourer}."
+                f"Graph colouring method must be one of: {GRAPH_COLOURING_METHODS}, "
+                f"instead got '{graph_colourer}'."
             )
 
-        self.graph_colourer = GRAPH_COLOURING_METHODS[graph_colourer.lower()]
         self.observables = observables
         self._wire_map = None
-        self._n_qubits = None
-        self.binary_observables = None
-        self.adj_matrix = None
-        self.grouped_paulis = None
+
+    @cached_property
+    def binary_observables(self):
+        """Binary Matrix corresponding to the symplectic representation of ``self.observables``.
+
+        It is an m x n matrix where each row is the symplectic (binary) representation of
+        ``self.observables``, with ``m = len(self.observables)`` and n the
+        number of qubits acted on by the observables.
+        """
+        return self.binary_repr()
 
     def binary_repr(self, n_qubits=None, wire_map=None):
-        """Converts the list of Pauli words to a binary matrix.
+        """Converts the list of Pauli words to a binary matrix,
+        i.e. a matrix where row m is the symplectic representation of ``self.observables[m]``.
 
         Args:
             n_qubits (int): number of qubits to specify dimension of binary vector representation
@@ -119,75 +151,315 @@ class PauliGroupingStrategy:  # pylint: disable=too-many-instance-attributes
         else:
             self._wire_map = wire_map
 
-        self._n_qubits = n_qubits
-
         return observables_to_binary_matrix(self.observables, n_qubits, self._wire_map)
 
-    def complement_adj_matrix_for_operator(self):
-        """Constructs the adjacency matrix for the complement of the Pauli graph.
+    @cached_property
+    def adj_matrix(self) -> np.ndarray:
+        """Adjacency matrix for the complement of the Pauli graph determined by the ``grouping_type``.
+
+        The adjacency matrix for an undirected graph of N nodes is an N x N symmetric binary
+        matrix, where matrix elements of 1 denote an edge (grouping strategy is **not** satisfied), and matrix elements of 0 denote no edge (grouping strategy is satisfied).
+        """
+        return _adj_matrix_from_symplectic(
+            self.binary_observables, grouping_type=self.grouping_type
+        )
+
+    @property
+    def complement_graph(self) -> rx.PyGraph:
+        """
+        Complement graph of the (anti-)commutation graph constructed from the Pauli observables.
+
+        Edge ``(i,j)`` is present in the graph if ``observable[i]`` and ``observable[j]`` do **not** satisfy
+        the ``grouping_type`` strategy.
+
+        The nodes are the observables (can only be accesssed through their integer index).
+        """
+        # Use upper triangle since adjacency matrix is symmetric and we have an undirected graph
+        edges = list(zip(*np.where(np.triu(self.adj_matrix, k=1))))
+        # Create complement graph
+        graph = rx.PyGraph(
+            node_count_hint=len(self.observables),
+            edge_count_hint=len(edges),
+        )
+
+        graph.add_nodes_from(self.observables)
+        graph.add_edges_from_no_data(edges)
+
+        return graph
+
+    def partition_observables(self) -> list[list]:
+        """
+        Partition the observables into groups of observables mutually satisfying the binary relation determined by ``self.grouping_type``.
+
+        Returns:
+            list[list[Operator]]: List of partitions of the Pauli observables made up of mutually (anti-)commuting observables.
+        """
+        if self.graph_colourer != "rlf":
+            return self.pauli_partitions_from_graph()
+        coloured_binary_paulis = recursive_largest_first(self.binary_observables, self.adj_matrix)
+
+        # Need to convert back from the symplectic representation
+        return [
+            [binary_to_pauli(pauli_word, wire_map=self._wire_map) for pauli_word in grouping]
+            for grouping in coloured_binary_paulis.values()
+        ]
+
+    @cached_property
+    def _idx_partitions_dict_from_graph(self) -> dict[int, list[int]]:
+        """Dictionary containing the solution to the graph colouring problem of ``self.complement_graph``.
 
-        The adjacency matrix for an undirected graph of N vertices is an N by N symmetric binary
-        matrix, where matrix elements of 1 denote an edge, and matrix elements of 0 denote no edge.
+        Colours the complement graph using a greedy colouring algorithm and groups indices by colour.
+
+        Uses the ``graph_greedy_color`` function from ``Rustworkx`` to colour the graph defined by
+        ``self.complement_graph`` using a specified strategy from ``RX_STRATEGIES``. It then groups the indices
+        (nodes) of the graph by their assigned colours.
 
         Returns:
-            array[int]: the square and symmetric adjacency matrix
+            dict[int, list[int]]: A dictionary where the keys are colours (integers) and the values are lists
+                of indices (nodes) that have been assigned that colour.
         """
+        # A dictionary where keys are node indices and the value is the colour
+        if new_rx:
+            # 'strategy' kwarg was implemented in Rustworkx 0.15
+            colouring_dict = rx.graph_greedy_color(
+                self.complement_graph, strategy=RX_STRATEGIES[self.graph_colourer]
+            )
+        else:
+            # Default value for <0.15.0 was 'lf'.
+            colouring_dict = rx.graph_greedy_color(self.complement_graph)
 
-        if self.binary_observables is None:
-            self.binary_observables = self.binary_repr()
+        # group together indices (values) of the same colour (keys)
+        groups = defaultdict(list)
+        for idx, colour in sorted(colouring_dict.items()):
+            groups[colour].append(idx)
 
-        n_qubits = int(np.shape(self.binary_observables)[1] / 2)
+        return groups
+
+    def idx_partitions_from_graph(self, observables_indices=None) -> tuple[tuple[int, ...], ...]:
+        """Use ``Rustworkx`` graph colouring algorithms to partition the indices of the Pauli observables into
+        tuples containing the indices of observables satisying the binary relation determined by ``self.grouping_type``.
+
+        Args:
+            observables_indices (tensor_like, optional): A tensor or list of indices associated to each observable.
+                This argument is helpful when the observables used in the graph colouring are part of a bigger set of observables.
+                Defaults to None. If ``None``, the partitions are made up of the relative indices, i.e. assuming ``self.observables``
+                have indices in [0, len(observables)-1].
 
-        if self.grouping_type == "qwc":
-            adj = qwc_complement_adj_matrix(self.binary_observables)
+        Raises:
+            IndexError: When the tensor_like of observables_indices is not of the same length as the observables.
 
-        elif self.grouping_type in frozenset(["commuting", "anticommuting"]):
-            symplectic_form = np.block(
-                [
-                    [np.zeros((n_qubits, n_qubits)), np.eye(n_qubits)],
-                    [np.eye(n_qubits), np.zeros((n_qubits, n_qubits))],
-                ]
+        Returns:
+            tuple[tuple[int]]: Tuple of tuples containing the indices of the partitioned observables.
+        """
+        if observables_indices is None:
+            return tuple(
+                tuple(indices) for indices in self._idx_partitions_dict_from_graph.values()
             )
-            mat_prod = (
-                self.binary_observables @ symplectic_form @ np.transpose(self.binary_observables)
+        if len(observables_indices) != len(self.observables):
+            raise ValueError(
+                f"The length of the list of indices: {len(observables_indices)} does not "
+                f"match the length of the list of observables: {len(self.observables)}. "
             )
+        return self._partition_custom_indices(observables_indices)
 
-            if self.grouping_type == "commuting":
-                adj = mat_prod % 2
+    def _partition_custom_indices(self, observables_indices) -> list[list]:
+        """Compute the indices of the partititions of the observables when these have custom indices.
 
-            elif self.grouping_type == "anticommuting":
-                adj = (mat_prod + 1) % 2
-                np.fill_diagonal(adj, 0)
+        TODO: Use this function to calculate custom indices instead of calculating observables first.
 
-        return adj
+        Args:
+            observables_indices (tensor_like, optional): A tensor or list of indices associated to each observable.
+                This argument is helpful when the observables used in the graph colouring are part of a bigger set of observables.
+                Defaults to None.
 
-    def colour_pauli_graph(self):
+        Returns:
+            tuple[tuple[int]]: Tuple of tuples containing the indices of the observables on each partition.
         """
-        Runs the graph colouring heuristic algorithm to obtain the partitioned Pauli words.
+        partition_indices = items_partitions_from_idx_partitions(
+            observables_indices, self._idx_partitions_dict_from_graph.values(), return_tuples=True
+        )
+
+        return partition_indices
+
+    def pauli_partitions_from_graph(self) -> list[list]:
+        """Partition Pauli observables into lists of (anti-)commuting observables
+        using ``Rustworkx`` graph colouring algorithms based on binary relation determined by  ``self.grouping_type``.
 
         Returns:
-            list[list[Observable]]: a list of the obtained groupings. Each grouping is itself a
-            list of Pauli word ``Observable`` instances
+            list[list[Observable]]: List of partitions of the Pauli observables made up of mutually (anti-)commuting terms.
         """
+        # Get the observables from the indices. itemgetter outperforms list comprehension
+        pauli_partitions = items_partitions_from_idx_partitions(
+            self.observables, self._idx_partitions_dict_from_graph.values()
+        )
+        return pauli_partitions
 
-        if self.adj_matrix is None:
-            self.adj_matrix = self.complement_adj_matrix_for_operator()
 
-        coloured_binary_paulis = self.graph_colourer(self.binary_observables, self.adj_matrix)
+def items_partitions_from_idx_partitions(
+    items: Sequence, idx_partitions: Sequence[Sequence[int]], return_tuples: bool = False
+) -> Sequence[Sequence]:
+    """Get the partitions of the items corresponding to the partitions of the indices.
 
-        self.grouped_paulis = [
-            [binary_to_pauli(pauli_word, wire_map=self._wire_map) for pauli_word in grouping]
-            for grouping in coloured_binary_paulis.values()
+    Args:
+        items (Sequence): A Sequence of items to be partitioned according to the partition of the indices.
+        idx_partitions (Sequence[Sequence[int]]): Sequence of sequences containing the indices of the partitioned items.
+        return_tuples (bool): Whether to return tuples of tuples or list of lists.
+            Useful when dealing with indices or observables.
+    Returns:
+        Sequence[Sequence]: Sequence of partitions of the items according to the partition of the indices.
+    """
+    if return_tuples:
+        items_partitioned = tuple(
+            (
+                tuple(itemgetter(*indices)(items))
+                if len(indices) > 1
+                else (itemgetter(*indices)(items),)
+            )
+            for indices in idx_partitions
+        )
+    else:
+        items_partitioned = [
+            (
+                list(itemgetter(*indices)(items))
+                if len(indices) > 1
+                else [itemgetter(*indices)(items)]
+            )
+            for indices in idx_partitions
         ]
 
-        return self.grouped_paulis
+    return items_partitioned
+
+
+def _adj_matrix_from_symplectic(symplectic_matrix: np.ndarray, grouping_type: str):
+    """Get adjacency matrix of (anti-)commuting graph based on grouping type.
+
+    This is the adjacency matrix of the complement graph. Based on symplectic representations and inner product of [1].
+
+    [1] Andrew Jena (2019). Partitioning Pauli Operators: in Theory and in Practice.
+    UWSpace. http://hdl.handle.net/10012/15017
+
+    Args:
+        symplectic_matrix (np.ndarray): 2D symplectic matrix. Each row corresponds to the
+        symplectic representation of the Pauli observables.
+        grouping_type (str): the binary relation used to define partitions of
+            the Pauli words, can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or
+            ``'anticommuting'``.
+
+    Returns:
+        np.ndarray: Adjacency matrix. Binary matrix such that adj_matrix[i,j] = 1 if observables[i]
+        observables[j] do **not** (anti-)commute, as determined by the ``grouping_type``.
+    """
+
+    n_qubits = symplectic_matrix.shape[1] // 2
+
+    # Convert symplectic representation to integer format.
+    # This is equivalent to the map: {0: I, 1: X, 2:Y, Z:3}
+    pauli_matrix_int = 2 * symplectic_matrix[:, :n_qubits] + symplectic_matrix[:, n_qubits:]
+    pauli_matrix_int = pauli_matrix_int.astype(np.int8)
+    # Broadcast the second dimension, sucht that pauli_matrix_broad.shape = (m, 1, n_qubits)
+    # with m = len(observables). This allows for calculation of all possible combinations of Pauli observable pairs (Pi, Pj).
+    # Something like: result[i, j, k] = pauli_matrix_int[i, k] * pauli_matrix_int[j, k]
+    pauli_matrix_broad = pauli_matrix_int[:, None]
+    # Calculate the symplectic inner product in [1], using the integer representation - hence the difference in the equation form.
+    # qubit_anticommutation_mat[i,j, k] is k=0 if Pi and Pj commute, else k!=0 if they anticommute.
+    qubit_anticommutation_mat = (pauli_matrix_int * pauli_matrix_broad) * (
+        pauli_matrix_int - pauli_matrix_broad
+    )
+    # 'adjacency_mat[i, j]' is True iff Paulis 'i' and 'j' do not (anti-)commute under given grouping_type.
+    if grouping_type == "qwc":
+        # True if any term anti commutes
+        adj_matrix = np.logical_or.reduce(qubit_anticommutation_mat, axis=2)
+    elif grouping_type == "commuting":
+        # True if the number of anti commuting terms is odd (anti commte)
+        adj_matrix = np.logical_xor.reduce(qubit_anticommutation_mat, axis=2)
+    else:
+        # True if the number of anti commuting terms is even (commute)
+        adj_matrix = ~np.logical_xor.reduce(qubit_anticommutation_mat, axis=2)
+
+    return adj_matrix
+
+
+def compute_partition_indices(
+    observables: list, grouping_type: str = "qwc", method: str = "lf"
+) -> tuple[tuple[int]]:
+    """
+    Computes the partition indices of a list of observables using a specified grouping type
+    and graph colouring method.
+
+    Args:
+        observables (list[Observable]): A list of Pauli Observables to be partitioned.
+        grouping_type (str): The type of binary relation between Pauli observables.
+            Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``. Defaults to ``'qwc'``.
+        method (str): The graph colouring heuristic to use in solving minimum clique cover.
+            Can be ``'lf'`` (Largest First), ``'dsatur'`` (Degree of Saturation), or ``'gis'`` (Greedy Independent Set).
+            Defaults to ``'lf'``.
+
+    Returns:
+        tuple[tuple[int]]: A tuple of tuples where each inner tuple contains the indices of
+        observables that are grouped together according to the specified grouping type and
+        graph colouring method.
+
+    .. seealso:: `rustworkx.ColoringStrategy <https://www.rustworkx.org/apiref/rustworkx.ColoringStrategy.html#coloringstrategy>`_
+        for more information on the ``('lf', 'dsatur', 'gis')`` strategies.
+
+    **Example**
+
+    >>> observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+    >>> compute_partition_indices(observables, grouping_type="qwc", method="lf")
+    ((0,), (1, 2))
+    """
+    if method != "rlf":
+
+        idx_no_wires = [idx for idx, obs in enumerate(observables) if len(obs.wires) == 0]
+
+        if len(idx_no_wires) == len(observables):
+            return (tuple(idx_no_wires),)
+
+        pauli_groupper = PauliGroupingStrategy(
+            observables, grouping_type=grouping_type, graph_colourer=method
+        )
+
+        return pauli_groupper.idx_partitions_from_graph()
+    # 'rlf' method is not compatible with the rx implementation.
+    return _compute_partition_indices_rlf(observables, grouping_type=grouping_type)
+
+
+def _compute_partition_indices_rlf(observables: list, grouping_type: str):
+    """Computes the partition indices of a list of observables using a specified grouping type and 'rlf' method.
+
+    This option is much less efficient so should be avoided.
+    """
+
+    with qml.QueuingManager.stop_recording():
+        obs_groups = qml.pauli.group_observables(
+            observables, grouping_type=grouping_type, method="rlf"
+        )
+
+    observables = copy(observables)
+
+    indices = []
+    available_indices = list(range(len(observables)))
+    for partition in obs_groups:  # pylint:disable=too-many-nested-blocks
+        indices_this_group = []
+        for pauli_word in partition:
+            # find index of this pauli word in remaining original observables,
+            for ind, observable in enumerate(observables):
+                if qml.pauli.are_identical_pauli_words(pauli_word, 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(tuple(indices_this_group))
+
+    return tuple(indices)
 
 
 def group_observables(
     observables: list["qml.operation.Operator"],
     coefficients: Optional[TensorLike] = None,
     grouping_type: Literal["qwc", "commuting", "anticommuting"] = "qwc",
-    method: Literal["lf", "rlf"] = "rlf",
+    method: Literal["lf", "rlf", "dsatur", "gis"] = "lf",
 ):
     """Partitions a list of observables (Pauli operations and tensor products thereof) into
     groupings according to a binary relation (qubit-wise commuting, fully-commuting, or
@@ -204,8 +476,9 @@ def group_observables(
             output ``partitioned_coeffs`` is not returned.
         grouping_type (str): The type of binary relation between Pauli words.
             Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
-        method (str): the graph coloring heuristic to use in solving minimum clique cover, which
-            can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First)
+        method (str): The graph colouring heuristic to use in solving minimum clique cover, which
+            can be ``'lf'`` (Largest First), ``'rlf'`` (Recursive Largest First),
+            ``'dsatur'`` (Degree of Saturation), or ``'gis'`` (IndependentSet). Defaults to ``'lf'``.
 
     Returns:
        tuple:
@@ -220,6 +493,9 @@ def group_observables(
         IndexError: if the input list of coefficients is not of the same length as the input list
             of Pauli words
 
+    .. seealso:: `rustworkx.ColoringStrategy <https://www.rustworkx.org/apiref/rustworkx.ColoringStrategy.html#coloringstrategy>`_
+        for more information on the ``('lf', 'dsatur', 'gis')`` strategies.
+
     **Example**
 
     >>> obs = [qml.Y(0), qml.X(0) @ qml.X(1), qml.Z(1)]
@@ -232,28 +508,30 @@ def group_observables(
     [[0.97, 4.21], [1.43]]
     """
 
-    if coefficients is not None:
-        if qml.math.shape(coefficients)[0] != len(observables):
-            raise IndexError(
-                "The coefficients list must be the same length as the observables list."
-            )
+    if coefficients is not None and qml.math.shape(coefficients)[0] != len(observables):
+        raise IndexError("The coefficients list must be the same length as the observables list.")
+
+    # Separate observables based on whether they have wires or not.
+    no_wires_obs, wires_obs = [], []
 
-    no_wires_obs = []
-    wires_obs = []
     for ob in observables:
         if len(ob.wires) == 0:
             no_wires_obs.append(ob)
         else:
             wires_obs.append(ob)
+
+    # Handle case where all observables have no wires
     if not wires_obs:
         if coefficients is None:
-            return [no_wires_obs]
-        return [no_wires_obs], [coefficients]
+            return [observables]
+        return [observables], [coefficients]
 
-    pauli_grouping = PauliGroupingStrategy(
+    # Initialize PauliGroupingStrategy
+    pauli_groupper = PauliGroupingStrategy(
         wires_obs, grouping_type=grouping_type, graph_colourer=method
     )
 
+    # Handles legacy op_math
     temp_opmath = not qml.operation.active_new_opmath() and any(
         isinstance(o, (Prod, SProd)) for o in observables
     )
@@ -261,11 +539,12 @@ def group_observables(
         qml.operation.enable_new_opmath(warn=False)
 
     try:
-        partitioned_paulis = pauli_grouping.colour_pauli_graph()
+        partitioned_paulis = pauli_groupper.partition_observables()
     finally:
         if temp_opmath:
             qml.operation.disable_new_opmath(warn=False)
 
+    # Add observables without wires back to the first partition
     partitioned_paulis[0].extend(no_wires_obs)
 
     if coefficients is None:
@@ -277,7 +556,12 @@ def group_observables(
 
 
 def _partition_coeffs(partitioned_paulis, observables, coefficients):
-    """Partition the coefficients according to the Pauli word groupings."""
+    """Partition the coefficients according to the Pauli word groupings.
+
+    This function is necessary in the cases where the coefficients are not the trivial
+    integers range(0, len(observables)). In the latter case, using `compute_partition_indices`
+    is recommended.
+    """
 
     partitioned_coeffs = [
         qml.math.cast_like([0] * len(g), coefficients) for g in partitioned_paulis
@@ -293,11 +577,14 @@ def _partition_coeffs(partitioned_paulis, observables, coefficients):
             # find index of this pauli word in remaining original observables,
             for ind, observable in enumerate(observables):
                 if isinstance(observable, qml.ops.Hamiltonian):
-                    # Converts single-term Hamiltonian to SProd because
                     # are_identical_pauli_words cannot handle Hamiltonian
                     coeffs, ops = observable.terms()
                     # Assuming the Hamiltonian has only one term
                     observable = qml.s_prod(coeffs[0], ops[0])
+                if isinstance(pauli_word, qml.ops.Hamiltonian):
+                    # Need to add this case because rx methods do not change type of observables.
+                    coeffs, ops = pauli_word.terms()
+                    pauli_word = qml.s_prod(coeffs[0], ops[0])
                 if are_identical_pauli_words(pauli_word, observable):
                     indices.append(coeff_indices[ind])
                     observables.pop(ind)
diff --git a/pennylane/pauli/utils.py b/pennylane/pauli/utils.py
index 7c5315727..cfe81dff3 100644
--- a/pennylane/pauli/utils.py
+++ b/pennylane/pauli/utils.py
@@ -165,7 +165,9 @@ def are_identical_pauli_words(pauli_1, pauli_2):
 
     **Example**
 
-    >>> are_identical_pauli_words(qml.Z(0) @ qml.Z(1), qml.Z(0) @ qml.Z(1))
+    >>> are_identical_pauli_words(qml.Z(0) @ qml.Z(1), qml.Z(1) @ qml.Z(0))
+    True
+    >>> are_identical_pauli_words(qml.I(0) @ qml.X(1), qml.X(1))
     True
     >>> are_identical_pauli_words(qml.Z(0) @ qml.Z(1), qml.Z(0) @ qml.X(3))
     False
@@ -183,7 +185,7 @@ def are_identical_pauli_words(pauli_1, pauli_2):
 
 def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_word=True):
     # pylint: disable=isinstance-second-argument-not-valid-type
-    """Converts a Pauli word to the binary vector representation.
+    """Converts a Pauli word to the binary vector (symplectic) representation.
 
     This functions follows convention that the first half of binary vector components specify
     PauliX placements while the last half specify PauliZ placements.
@@ -195,7 +197,7 @@ def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_wor
         wire_map (dict): dictionary containing all wire labels used in the Pauli word as keys, and
             unique integer labels as their values
         check_is_pauli_word (bool): If True (default) then a check is run to verify that pauli_word
-            is infact a Pauli word
+            is in fact a Pauli word.
 
     Returns:
         array: the ``2*n_qubits`` dimensional binary vector representation of the input Pauli word
@@ -731,11 +733,11 @@ def are_pauli_words_qwc(lst_pauli_words):
 
 
 def observables_to_binary_matrix(observables, n_qubits=None, wire_map=None):
-    """Converts a list of Pauli words to the binary vector representation and yields a row matrix
-    of the binary vectors.
+    """Converts a list of Pauli words into a matrix where each row is the binary vector (symplectic)
+    representation of the ``observables``.
 
-    The dimension of the binary vectors will be implied from the highest wire being acted on
-    non-trivially by the Pauli words in observables.
+    The dimension of the binary vectors (the number of columns) will be implied from the highest wire
+    being acted on non-trivially by the Pauli words in observables.
 
     Args:
         observables (list[Union[Identity, PauliX, PauliY, PauliZ, Tensor, Prod, SProd]]): the list
@@ -746,7 +748,7 @@ def observables_to_binary_matrix(observables, n_qubits=None, wire_map=None):
 
 
     Returns:
-        array[array[int]]: a matrix whose rows are Pauli words in binary vector representation
+        array[array[int]]: a matrix whose rows are Pauli words in binary vector (symplectic) representation.
 
     **Example**
 
diff --git a/requirements.txt b/requirements.txt
index 56163fd86..b5ab79f1a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,7 +4,7 @@ cvxpy~=1.2
 cvxopt~=1.3.0
 cachetools~=5.0.0
 networkx~=2.8
-rustworkx~=0.12.1
+rustworkx~=0.15.1
 autograd~=1.4
 toml~=0.10
 appdirs~=1.4
diff --git a/setup.py b/setup.py
index 2bcf05540..6c0bcab70 100644
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,7 @@ requirements = [
     "numpy<2.0",
     "scipy",
     "networkx",
-    "rustworkx",
+    "rustworkx>=0.14.0",
     "autograd",
     "toml",
     "appdirs",


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 6fcd69a30..74d9628d1 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1615,7 +1615,7 @@ class TestHamiltonianSamples:
         qs_exp = qml.tape.QuantumScript(ops, [qml.expval(H)])
         expected = dev.execute(qs_exp)
 
-        assert np.allclose(res, expected, atol=0.002)
+        assert np.allclose(res, expected, atol=0.02)
 
 
 class TestClassicalShadows:
diff --git a/tests/ops/op_math/test_linear_combination.py b/tests/ops/op_math/test_linear_combination.py
index 409062fda..553427ce7 100644
--- a/tests/ops/op_math/test_linear_combination.py
+++ b/tests/ops/op_math/test_linear_combination.py
@@ -1583,21 +1583,22 @@ class TestGrouping:
 
     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."""
+        This is done by changing from default to 'lf' and checking the result."""
+        # Create a graph with unique solution so that test does not depend on solver/implementation
         a = X(0)
-        b = X(1)
+        b = X(0)
         c = Z(0)
         obs = [a, b, c]
         coeffs = [1.0, 2.0, 3.0]
 
         # compute grouping during construction
         H2 = qml.ops.LinearCombination(coeffs, obs, grouping_type="qwc", method="lf")
-        assert H2.grouping_indices == ((2, 1), (0,))
+        assert set(H2.grouping_indices) == set(((0, 1), (2,)))
 
         # compute grouping separately
         H3 = qml.ops.LinearCombination(coeffs, obs, grouping_type=None)
         H3.compute_grouping(method="lf")
-        assert H3.grouping_indices == ((2, 1), (0,))
+        assert set(H3.grouping_indices) == set(((0, 1), (2,)))
 
     def test_grouping_with_duplicate_terms(self):
         """Test that the grouping indices are correct when the LinearCombination has duplicate
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index f027976ed..72681f156 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -1437,7 +1437,7 @@ class TestGrouping:
 
     def test_grouping_method_can_be_set(self):
         """Tests that the grouping method can be controlled by kwargs.
-        This is done by changing from default to 'rlf' and checking the result."""
+        This is done by changing from default to 'lf' and checking the result."""
         a = qml.PauliX(0)
         b = qml.PauliX(1)
         c = qml.PauliZ(0)
@@ -1446,21 +1446,21 @@ class TestGrouping:
 
         # compute grouping during construction with qml.dot
         op1 = qml.dot(coeffs, obs, grouping_type="qwc", method="lf")
-        assert op1.grouping_indices == ((2, 1), (0,))
+        assert set(op1.grouping_indices) == set(((0, 1), (2,)))
 
         # compute grouping during construction with qml.sum
         sprods = [qml.s_prod(c, o) for c, o in zip(coeffs, obs)]
         op2 = qml.sum(*sprods, grouping_type="qwc", method="lf")
-        assert op2.grouping_indices == ((2, 1), (0,))
+        assert set(op2.grouping_indices) == set(((0, 1), (2,)))
 
         # compute grouping during construction with Sum
         op3 = Sum(*sprods, grouping_type="qwc", method="lf")
-        assert op3.grouping_indices == ((2, 1), (0,))
+        assert set(op3.grouping_indices) == set(((0, 1), (2,)))
 
         # compute grouping separately
         op4 = qml.dot(coeffs, obs, grouping_type=None)
         op4.compute_grouping(method="lf")
-        assert op4.grouping_indices == ((2, 1), (0,))
+        assert set(op4.grouping_indices) == set(((0, 1), (2,)))
 
     @pytest.mark.parametrize(
         "grouping_type, grouping_indices",
diff --git a/tests/ops/qubit/test_hamiltonian.py b/tests/ops/qubit/test_hamiltonian.py
index 13fbf76ac..169988f8c 100644
--- a/tests/ops/qubit/test_hamiltonian.py
+++ b/tests/ops/qubit/test_hamiltonian.py
@@ -1691,21 +1691,22 @@ class TestGrouping:
 
     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."""
+        This is done by changing from default to 'lf' and checking the result."""
+        # Create a graph with unique solution so that test does not depend on solver/implementation
         a = qml.PauliX(0)
-        b = qml.PauliX(1)
+        b = qml.PauliX(0)
         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,))
+        assert set(H2.grouping_indices) == set(((0, 1), (2,)))
 
         # compute grouping separately
         H3 = qml.Hamiltonian(coeffs, obs, grouping_type=None)
         H3.compute_grouping(method="lf")
-        assert H3.grouping_indices == ((2, 1), (0,))
+        assert set(H3.grouping_indices) == set(((0, 1), (2,)))
 
 
 @pytest.mark.usefixtures("use_legacy_opmath")
diff --git a/tests/pauli/grouping/test_pauli_group_observables.py b/tests/pauli/grouping/test_pauli_group_observables.py
index df8663d8c..b70f84250 100644
--- a/tests/pauli/grouping/test_pauli_group_observables.py
+++ b/tests/pauli/grouping/test_pauli_group_observables.py
@@ -14,6 +14,8 @@
 """
 Unit tests for ``PauliGroupingStrategy`` and ``group_observables`` in ``/pauli/grouping/group_observables.py``.
 """
+import sys
+
 import numpy as np
 import pytest
 
@@ -21,8 +23,43 @@ import pennylane as qml
 from pennylane import Identity, PauliX, PauliY, PauliZ
 from pennylane import numpy as pnp
 from pennylane.operation import Tensor
-from pennylane.pauli import are_identical_pauli_words
-from pennylane.pauli.grouping.group_observables import PauliGroupingStrategy, group_observables
+from pennylane.pauli import are_identical_pauli_words, are_pauli_words_qwc
+from pennylane.pauli.grouping.group_observables import (
+    PauliGroupingStrategy,
+    compute_partition_indices,
+    group_observables,
+    items_partitions_from_idx_partitions,
+)
+
+
+class TestOldRX:
+    """Test PauliGroupingStrategy behaves correctly when versions of rx older than 0.15 are used"""
+
+    @pytest.mark.parametrize("new_colourer", ["dsatur", "gis"])
+    def test_new_strategies_with_old_rx_raise_error(self, monkeypatch, new_colourer):
+        """Test that an error is raised if a new strategy is used with old rx"""
+        # Monkey patch the new_rx variable to False
+        grouping = sys.modules["pennylane.pauli.grouping.group_observables"]
+        monkeypatch.setattr(grouping, "new_rx", False)
+
+        observables = [qml.X(0)]
+        with pytest.raises(ValueError, match="not supported in this version of Rustworkx"):
+            PauliGroupingStrategy(observables, graph_colourer=new_colourer)
+
+    def test_old_rx_produces_right_results(self, monkeypatch):
+        """Test that the results produced with an older rx version is the same as with lf"""
+        observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+
+        new_groupper = PauliGroupingStrategy(observables, graph_colourer="lf")
+        new_partitions = new_groupper.partition_observables()
+
+        grouping = sys.modules["pennylane.pauli.grouping.group_observables"]
+        monkeypatch.setattr(grouping, "new_rx", False)
+
+        old_groupper = PauliGroupingStrategy(observables, graph_colourer="lf")
+        old_partitions = old_groupper.partition_observables()
+
+        assert new_partitions == old_partitions
 
 
 class TestPauliGroupingStrategy:
@@ -46,7 +83,7 @@ class TestPauliGroupingStrategy:
             ValueError, PauliGroupingStrategy, observables, graph_colourer="invalid"
         )
 
-    def test_construct_qwc_complement_adj_matrix_for_operators(self):
+    def test_construct_qwc_adj_matrix(self):
         """Constructing the complement graph adjacency matrix for a list of Pauli words according
         to qubit-wise commutativity."""
 
@@ -54,12 +91,9 @@ class TestPauliGroupingStrategy:
         qwc_complement_adjacency_matrix = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]])
 
         grouping_instance = PauliGroupingStrategy(observables, "qwc")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == qwc_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == qwc_complement_adjacency_matrix).all()
 
-    def test_construct_commuting_complement_adj_matrix_for_operators(self):
+    def test_construct_commuting_adj_matrix(self):
         """Constructing the complement graph adjacency matrix for a list of Pauli words according
         to general commutativity."""
 
@@ -67,23 +101,17 @@ class TestPauliGroupingStrategy:
         commuting_complement_adjacency_matrix = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 0]])
 
         grouping_instance = PauliGroupingStrategy(observables, "commuting")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == commuting_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == commuting_complement_adjacency_matrix).all()
 
-    def test_construct_anticommuting_complement_adj_matrix_for_operators(self):
+    def test_construct_anticommuting_adj_matrix(self):
         """Constructing the complement graph adjacency matrix for a list of Pauli words according
         to anticommutativity."""
 
         observables = [PauliY(0), PauliZ(0) @ PauliZ(1), PauliY(0) @ PauliX(1)]
-        anticommuting_complement_adjacency_matrix = np.array([[0, 0, 1], [0, 0, 1], [1, 1, 0]])
+        anticommuting_complement_adjacency_matrix = np.array([[1, 0, 1], [0, 1, 1], [1, 1, 1]])
 
         grouping_instance = PauliGroupingStrategy(observables, "anticommuting")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == anticommuting_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == anticommuting_complement_adjacency_matrix).all()
 
     trivial_ops = [
         [Identity(0), Identity(0), Identity(7)],
@@ -98,10 +126,7 @@ class TestPauliGroupingStrategy:
         qwc_complement_adjacency_matrix = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
 
         grouping_instance = PauliGroupingStrategy(observables, "qwc")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == qwc_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == qwc_complement_adjacency_matrix).all()
 
     @pytest.mark.parametrize("observables", trivial_ops)
     def test_construct_complement_commuting_adj_matrix_for_trivial_operators(self, observables):
@@ -111,23 +136,51 @@ class TestPauliGroupingStrategy:
         commuting_complement_adjacency_matrix = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
 
         grouping_instance = PauliGroupingStrategy(observables, "commuting")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == commuting_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == commuting_complement_adjacency_matrix).all()
 
     @pytest.mark.parametrize("observables", trivial_ops)
     def test_construct_complement_anticommuting_adj_matrix_for_trivial_operators(self, observables):
         """Constructing the complement of anticommutativity graph's adjacency matrix for a list of
         identity operations and various symmetric binary relations"""
 
-        anticommuting_complement_adjacency_matrix = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
+        anticommuting_complement_adjacency_matrix = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
 
         grouping_instance = PauliGroupingStrategy(observables, "anticommuting")
-        assert (
-            grouping_instance.complement_adj_matrix_for_operator()
-            == anticommuting_complement_adjacency_matrix
-        ).all()
+        assert (grouping_instance.adj_matrix == anticommuting_complement_adjacency_matrix).all()
+
+    def test_wrong_length_of_custom_indices(self):
+        """Test that an error is raised if the length of indices does not match the length of observables"""
+        observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+        groupper = PauliGroupingStrategy(observables=observables)
+
+        custom_indices = [1, 3, 5, 7]
+
+        with pytest.raises(ValueError, match="The length of the list of indices:"):
+            groupper.idx_partitions_from_graph(observables_indices=custom_indices)
+
+    def test_custom_indices_partition(self):
+        """Test that a custom list indices is partitioned according to the observables they correspond to."""
+        observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+        groupper = PauliGroupingStrategy(observables=observables)
+
+        custom_indices = [1, 3, 5]
+        # the indices rustworkx assigns to each observable
+        standard_indices = list(range(len(observables)))
+        # map between custom and standard indices
+        map_indices = dict(zip(custom_indices, standard_indices))
+
+        # compute observable and custom indices partitions
+        observables_partitioned = groupper.partition_observables()
+        # pylint: disable=protected-access
+        indices_partitioned = groupper.idx_partitions_from_graph(observables_indices=custom_indices)
+        for group_obs, group_custom_indices in zip(observables_partitioned, indices_partitioned):
+            for i, custom_idx in enumerate(group_custom_indices):
+                standard_idx = map_indices[custom_idx]
+                # observable corresponding to the custom index
+                observable_from_idx_partition = observables[standard_idx]
+                # observable in partition in the position of the custom idx
+                observable_from_partition = group_obs[i]
+                assert observable_from_idx_partition == observable_from_partition
 
 
 observables_list = [
@@ -172,13 +225,13 @@ observables_list = [
 
 qwc_sols = [
     [
-        [PauliX(wires=[1]), PauliY(wires=[0])],
         [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
+        [PauliX(wires=[1]), PauliY(wires=[0])],
         [PauliZ(wires=[1]) @ PauliZ(wires=[2])],
     ],
     [
-        [PauliX(wires=[1]), PauliY(wires=[0])],
         [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
+        [PauliX(wires=[1]), PauliY(wires=[0])],
         [PauliZ(wires=[1]) @ PauliZ(wires=[2])],
     ],
     [
@@ -202,14 +255,14 @@ qwc_sols = [
         [PauliX(wires=[1]) @ PauliX(wires=[0])],
     ],
     [
+        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"])],
-        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
     ],
     [
+        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"])],
-        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
     ],
     [[PauliX([(0, 0)])], [PauliZ([(0, 0)])]],
 ]
@@ -306,6 +359,29 @@ anticommuting_sols = [
     [[PauliX([(0, 0)]), PauliZ([(0, 0)])]],
 ]
 
+com_tuples = list(zip(observables_list, commuting_sols))
+
+anticom_tuples = list(zip(observables_list, anticommuting_sols))
+
+
+def are_partitions_equal(partition_1: list, partition_2: list) -> bool:
+    """Checks whether two partitions are the same, i.e. contain the same Pauli terms.
+
+    We check this way since the partitions might vary in the order of the elements
+
+    Args:
+        partition_1 (list[Observable]): list of Pauli word ``Observable`` instances corresponding to a partition.
+        partition_2 (list[Observable]): list of Pauli word ``Observable`` instances corresponding to a partition.
+
+    """
+    partition_3 = set(
+        partition_2
+    )  # to improve the lookup time for similar obs in the second partition
+    for pauli in partition_1:
+        if not any(are_identical_pauli_words(pauli, other) for other in partition_3):
+            return False
+    return True
+
 
 class TestGroupObservables:
     """
@@ -319,50 +395,39 @@ class TestGroupObservables:
         qwc_partitions = group_observables(observables, grouping_type="qwc")
 
         # assert the correct number of partitions:
-        n_partitions = len(qwc_partitions_sol)
-        assert len(qwc_partitions) == n_partitions
-        # assert each partition is of the correct length:
-        assert all(
-            len(part) == len(part_sol) for part, part_sol in zip(qwc_partitions, qwc_partitions_sol)
-        )
-        # assert each partition contains the same Pauli terms as the solution partition:
-        for i, partition in enumerate(qwc_partitions):
-            for j, pauli in enumerate(partition):
-                assert are_identical_pauli_words(pauli, qwc_partitions_sol[i][j])
+        assert len(qwc_partitions) == len(qwc_partitions_sol)
 
-    com_tuples = list(zip(observables_list, commuting_sols))
+        # assert each computed partition contains appears in the computed solution.
+        for comp_partition in qwc_partitions:
+            assert any(
+                are_partitions_equal(exp_partition, comp_partition)
+                for exp_partition in qwc_partitions_sol
+            )
 
     @pytest.mark.parametrize("observables,com_partitions_sol", com_tuples)
     def test_commuting_partitioning(self, observables, com_partitions_sol):
         com_partitions = group_observables(observables, grouping_type="commuting")
 
-        # assert the correct number of partitions:
-        n_partitions = len(com_partitions_sol)
-        assert len(com_partitions) == n_partitions
-        # assert each partition is of the correct length:
-        assert all(len(p) == len(p_sol) for p, p_sol in zip(com_partitions, com_partitions_sol))
-        # assert each partition contains the same Pauli terms as the solution partition:
-        for i, partition in enumerate(com_partitions):
-            for j, pauli in enumerate(partition):
-                assert are_identical_pauli_words(pauli, com_partitions_sol[i][j])
-
-    anticom_tuples = list(zip(observables_list, anticommuting_sols))
+        assert len(com_partitions) == len(com_partitions_sol)
+        # assert each computed partition contains appears in the computed solution.
+        for comp_partition in com_partitions:
+            assert any(
+                are_partitions_equal(exp_partition, comp_partition)
+                for exp_partition in com_partitions_sol
+            )
 
     @pytest.mark.parametrize("observables,anticom_partitions_sol", anticom_tuples)
     def test_anticommuting_partitioning(self, observables, anticom_partitions_sol):
         anticom_partitions = group_observables(observables, grouping_type="anticommuting")
 
         # assert the correct number of partitions:
-        n_partitions = len(anticom_partitions_sol)
-        assert len(anticom_partitions) == n_partitions
-        # assert each partition is of the correct length:
-        assert all(
-            len(p) == len(p_sol) for p, p_sol in zip(anticom_partitions, anticom_partitions_sol)
-        )
-        # assert each partition contains the same Pauli terms as the solution partition:
-        for i, partition in enumerate(anticom_partitions):
-            for j, pauli in enumerate(partition):
-                assert are_identical_pauli_words(pauli, anticom_partitions_sol[i][j])
+        assert len(anticom_partitions) == len(anticom_partitions_sol)
+        # assert each computed partition contains appears in the computed solution.
+        for comp_partition in anticom_partitions:
+            assert any(
+                are_partitions_equal(exp_partition, comp_partition)
+                for exp_partition in anticom_partitions_sol
+            )
 
     def test_group_observables_exception(self):
         """Tests that the ``group_observables`` function raises an exception if
@@ -394,26 +459,6 @@ class TestGroupObservables:
         _, grouped_coeffs = group_observables(obs, coeffs)
         assert isinstance(grouped_coeffs[0], list)
 
-    def test_return_new_opmath(self):
-        """Test that using new opmath causes grouped observables to have Prods instead of
-        Tensors"""
-        new_observables = [
-            qml.prod(PauliX(0), PauliZ(1)),
-            qml.prod(PauliY(2), PauliZ(1)),
-            qml.s_prod(1.5, qml.prod(PauliZ(1), PauliZ(2))),
-        ]
-        mixed_observables = [
-            Tensor(PauliX(0), PauliZ(1)),
-            qml.prod(PauliY(2), PauliZ(1)),
-            qml.s_prod(1.5, qml.prod(PauliZ(1), PauliZ(2))),
-        ]
-
-        new_groups = group_observables(new_observables)
-        mixed_groups = group_observables(mixed_observables)
-
-        assert all(isinstance(o, qml.ops.Prod) for g in new_groups for o in g)
-        assert all(isinstance(o, qml.ops.Prod) for g in mixed_groups for o in g)
-
     @pytest.mark.usefixtures("use_legacy_opmath")
     def test_return_new_opmath_legacy_opmath(self):
         """Test that using new opmath causes grouped observables to have Prods instead of
@@ -483,6 +528,96 @@ class TestGroupObservables:
         assert out_coeffs == [[1, 3, 4], [2]]
 
 
+class TestComputePartitionIndices:
+    """Tests for ``compute_partition_indices``"""
+
+    OBS_IDX_PARTITIONS = [
+        (
+            [qml.I(), qml.X(0) @ qml.X(1), qml.Z(0) @ qml.Z(1), 2 * qml.I(), 2 * qml.Z(0)],
+            ((0, 1, 3), (2, 4)),
+            [[qml.I(), qml.X(0) @ qml.X(1), 2 * qml.I()], [qml.Z(0) @ qml.Z(1), 2 * qml.Z(0)]],
+        ),
+    ]
+
+    def test_invalid_colouring_method(self):
+        """Test that passing an invalid colouring method raises an error"""
+        observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+        with pytest.raises(ValueError, match="Graph colouring method must be one of"):
+            compute_partition_indices(observables=observables, method="recursive")
+
+    def test_only_observables_without_wires(self):
+        """Test that if none of the observables has wires, they are all in one single partition."""
+
+        observables = [qml.I(), 2 * qml.I()]
+        partition_indices = compute_partition_indices(observables=observables)
+        assert partition_indices == (tuple(range(len(observables))),)
+
+    @pytest.mark.parametrize("observables, indices, obs_partitions", OBS_IDX_PARTITIONS)
+    def test_obs_from_indices_partitions(self, observables, indices, obs_partitions):
+        """Test that obs_partition_from_idx_partitions returns the correct observables"""
+
+        partition_obs = items_partitions_from_idx_partitions(observables, indices)
+        assert partition_obs == obs_partitions
+
+    def test_mixed_observables_qwc(self):
+        """Test that if both observables with wires and without wires are present,
+        the latter are appended on the first element of the former and the partitions are qwc."""
+        observables = [qml.I(), qml.X(0), qml.Z(0), 2 * qml.I(), 2 * qml.Z(0)]
+        partition_indices = compute_partition_indices(observables=observables, grouping_type="qwc")
+        indices_no_wires = (0, 3)
+        assert set(indices_no_wires) < set(partition_indices[0])
+
+        partition_obs = items_partitions_from_idx_partitions(observables, partition_indices)
+        for partition in partition_obs:
+            assert are_pauli_words_qwc(partition)
+
+    @pytest.mark.parametrize("observables,com_partitions_sol", com_tuples)
+    def test_commuting_partitioning(self, observables, com_partitions_sol):
+        """Test that using the commuting grouping type returns the correct solutions."""
+
+        partition_indices = compute_partition_indices(
+            observables=observables, grouping_type="commuting"
+        )
+
+        com_partitions = items_partitions_from_idx_partitions(observables, partition_indices)
+
+        assert len(com_partitions) == len(com_partitions_sol)
+        # assert each computed partition contains appears in the computed solution.
+        for comp_partition in com_partitions:
+            assert any(
+                are_partitions_equal(exp_partition, comp_partition)
+                for exp_partition in com_partitions_sol
+            )
+
+    @pytest.mark.parametrize("observables,anticom_partitions_sol", anticom_tuples)
+    def test_anticommuting_partitioning(self, observables, anticom_partitions_sol):
+        """Test that using the anticommuting grouping type returns the correct solutions."""
+
+        partition_indices = compute_partition_indices(
+            observables=observables, grouping_type="anticommuting"
+        )
+
+        anticom_partitions = items_partitions_from_idx_partitions(observables, partition_indices)
+
+        # assert the correct number of partitions:
+        assert len(anticom_partitions) == len(anticom_partitions_sol)
+        # assert each computed partition contains appears in the computed solution.
+        for comp_partition in anticom_partitions:
+            assert any(
+                are_partitions_equal(exp_partition, comp_partition)
+                for exp_partition in anticom_partitions_sol
+            )
+
+    @pytest.mark.parametrize("method", ("rlf", "lf", "dsatur", "gis"))
+    def test_colouring_methods(self, method):
+        """Test that all colouring methods return the correct results."""
+        observables = [qml.X(0) @ qml.Z(1), qml.Z(0), qml.X(1)]
+        partition_indices = compute_partition_indices(
+            observables, grouping_type="qwc", method=method
+        )
+        assert set(partition_indices) == set(((0,), (1, 2)))
+
+
 class TestDifferentiable:
     """Tests that grouping observables is differentiable with respect to the coefficients."""
 
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index 0f8a80f6a..ea71a4ff8 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -96,8 +96,8 @@ def complex_no_grouping_processing_fn(results):
 
 complex_qwc_groups = [
     [qml.X(0), qml.X(0) @ qml.Y(1)],
-    [qml.Y(0), qml.X(1)],
-    [qml.Y(0) @ qml.Z(1), qml.Z(1)],
+    [qml.Y(0), qml.Y(0) @ qml.Z(1), qml.Z(1)],
+    [qml.X(1)],
 ]
 
 
@@ -107,8 +107,8 @@ def complex_qwc_processing_fn(results):
     return (
         group0[0],
         0.5 * group1[0],
-        group0[0] + group2[0] + 2.0 * group1[1] + 1.0,
-        0.1 * group2[1] + 0.2 * group0[1] + 0.3 * group2[0] + 0.4,
+        group0[0] + group1[1] + 2.0 * group2[0] + 1.0,
+        0.1 * group1[2] + 0.2 * group0[1] + 0.3 * group1[1] + 0.4,
         1.5,
     )
 
@@ -488,7 +488,7 @@ class TestUnits:
                     for group in complex_qwc_groups
                 ],
                 complex_qwc_processing_fn,
-                [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
+                [[0.1, 0.2], [0.3, 0.5, 0.6], [0.4]],
             ),
         ],
     )