Repository: PennyLaneAI/pennylane
Pull request #5729 (2024-05-23 15:03:40)

A unified `split_non_commuting` transform for any tape
**Context:**

Right now `qml.Device` choose between applying `hamiltonian_expand`, `sum_expand`, and `split_non_commuting` based on the contents of the measurements and selectively applies them. Under the new device API, we cannot simply stick these three transforms into the transform program, as they currently raise errors when applied to ordinary circuits.

```
>>> tape = qml.tape.QuantumScript([], [qml.probs(wires=(0,1))])
>>> qml.transforms.hamiltonian_expand(tape)
ValueError: Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`
>>> qml.transforms.split_non_commuting(tape)
NotImplementedError: When non-commuting observables are used, only `qml.expval` and `qml.var` are supported.
>>> qml.transforms.sum_expand(tape)
([<QuantumScript: wires=[0, 1], params=0>],
 <function pennylane.transforms.hamiltonian_expand.sum_expand.<locals>.processing_fn(expanded_results)>)
```

While `sum_expand` has sensible default behaviour that will not raise errors on standard functions, it is not as efficient for large hamiltonians and commuting measurements. We therefore need a new transform that selectively applies each of the above transforms based on the contents of the circuit. This transform is necessary for the development of devices that do not natively support hamiltonians or non-commuting measurements.

**Description of the Change:**

- `split_non_commuting` now incorporates the capabilities of `sum_expand` and `hamiltonian_expand`
- Tests from `test_hamiltonian_expand.py` are copied to `test_split_non_commuting.py`
- `grouping_strategy` keyword argument added to `split_non_commuting` to allow users to choose between grouping based on overlapping wires (`"wires"`), or using the graph colouring algorithm (`"qwc"`) implemented in `pauli.group_observables`.

**Follow-up TODOs**

- Clean up `Device.batch_transform` to use the new `split_non_commuting`
- Option to add diagonalizing gates to each tape
- A transforms transform `shots_only` to wrap around `split_non_commuting`

**Related GitHub Issues:**
[sc-50536]
[sc-59046]

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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5729?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`e600bc9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e600bc958bc7c7c205a6ed9bde745199e62000ec?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`6535be9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5729?dropdown=coverage&src=pr&el=desc&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                @@
##           sum-expand-bug    #5729    +/-   ##
================================================
  Coverage           99.67%   99.67%            
================================================
  Files                 416      416            
  Lines               38741    38878   +137     
================================================
+ Hits                38616    38753   +137     
  Misses                125      125            
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5729?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).

astralcai: @albi3ro @trbromley Should I rename `"naive"` and `"pauli"` to `"wires"` and `"qwc"`? Seems that the latter more accurately describes what the two grouping strategies are. 
astralcai: New tests for `split_non_commuting` also merged into this PR [sc-64069]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4fa1c09a2..6409d734c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -24,6 +24,9 @@
 * The sorting order of parameter-shift terms is now guaranteed to resolve ties in the absolute value with the sign of the shifts.
   [(#5582)](https://github.com/PennyLaneAI/pennylane/pull/5582)
 
+* `qml.transforms.split_non_commuting` can now handle circuits containing measurements of multi-term observables.
+  [(#5729)](https://github.com/PennyLaneAI/pennylane/pull/5729)
+
 <h4>Mid-circuit measurements and dynamic circuits</h4>
 
 * The `dynamic_one_shot` transform uses a single auxiliary tape with a shot vector and `default.qubit` implements the loop over shots with `jax.vmap`.
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 0d6a1d968..3f7ec2923 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -775,8 +775,9 @@ class Hamiltonian(Observable):
         ops1 = self.ops.copy()
 
         if isinstance(H, (Tensor, Observable)):
+            qml.QueuingManager.remove(H)
+            qml.QueuingManager.remove(self)
             terms = [copy(H) @ op for op in ops1]
-
             return qml.simplify(Hamiltonian(coeffs1, terms))
 
         return NotImplemented
@@ -790,11 +791,15 @@ class Hamiltonian(Observable):
             return self
 
         if isinstance(H, Hamiltonian):
+            qml.QueuingManager.remove(H)
+            qml.QueuingManager.remove(self)
             coeffs = qml.math.concatenate([self_coeffs, copy(H.coeffs)], axis=0)
             ops.extend(H.ops.copy())
             return qml.simplify(Hamiltonian(coeffs, ops))
 
         if isinstance(H, (Tensor, Observable)):
+            qml.QueuingManager.remove(H)
+            qml.QueuingManager.remove(self)
             coeffs = qml.math.concatenate(
                 [self_coeffs, qml.math.cast_like([1.0], self_coeffs)], axis=0
             )
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index 6697273f1..fa161a5eb 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -1,4 +1,4 @@
-# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+# Copyright 2018-2024 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.
@@ -11,42 +11,43 @@
 # 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.
+
 """
-Contains the tape transform that splits non-commuting terms
+Contains the tape transform that splits a tape into tapes measuring commuting observables.
 """
-from functools import reduce
 
-# pylint: disable=protected-access
-from typing import Callable, Sequence
+# pylint: disable=too-many-arguments
+
+from functools import partial
+from typing import Callable, Dict, List, Optional, Sequence, Tuple
 
 import pennylane as qml
+from pennylane.measurements import ExpectationMP, MeasurementProcess, Shots, StateMP
+from pennylane.ops import Hamiltonian, LinearCombination, Prod, SProd, Sum
 from pennylane.transforms import transform
-
-
-def null_postprocessing(results):
-    """A postprocesing function returned by a transform that only converts the batch of results
-    into a result for a single ``QuantumTape``.
-    """
-    return results[0]
+from pennylane.typing import Result, ResultBatch
 
 
 @transform
-def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
-    r"""
-    Splits a qnode measuring non-commuting observables into groups of commuting observables.
+def split_non_commuting(
+    tape: qml.tape.QuantumScript,
+    grouping_strategy: Optional[str] = "default",
+) -> Tuple[Sequence[qml.tape.QuantumTape], Callable]:
+    r"""Splits a circuit into tapes measuring groups of commuting observables.
 
     Args:
-        tape (QNode or QuantumTape or Callable): A circuit that contains a list of
-            non-commuting observables to measure.
+        tape (QNode or QuantumScript or Callable): The quantum circuit to be split.
+        grouping_strategy (str): The strategy to use for computing disjoint groups of
+            commuting observables, can be ``"default"``, ``"wires"``, ``"qwc"``,
+            or ``None`` to disable grouping.
 
     Returns:
-        qnode (QNode) or tuple[List[QuantumTape], function]: The transformed circuit as described in
-        :func:`qml.transform <pennylane.transform>`.
+        qnode (QNode) or tuple[List[QuantumScript], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
 
-    **Example**
+    **Examples:**
 
-    This transform allows us to transform a QNode that measures non-commuting observables to
-    *multiple* circuit executions with qubit-wise commuting groups:
+    This transform allows us to transform a QNode measuring multiple observables into multiple
+    circuit executions, each measuring a group of commuting observables.
 
     .. code-block:: python3
 
@@ -55,84 +56,143 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
         @qml.transforms.split_non_commuting
         @qml.qnode(dev)
         def circuit(x):
-            qml.RX(x,wires=0)
-            return [qml.expval(qml.X(0)), qml.expval(qml.Z(0))]
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=1)
+            return [
+                qml.expval(qml.X(0)),
+                qml.expval(qml.Y(1)),
+                qml.expval(qml.Z(0) @ qml.Z(1)),
+                qml.expval(qml.X(0) @ qml.Z(1) + 0.5 * qml.Y(1) + qml.Z(0)),
+            ]
 
-    Instead of decorating the QNode, we can also create a new function that yields the same result
-    in the following way:
+    Instead of decorating the QNode, we can also create a new function that yields the same
+    result in the following way:
 
     .. code-block:: python3
 
         @qml.qnode(dev)
         def circuit(x):
-            qml.RX(x,wires=0)
-            return [qml.expval(qml.X(0)), qml.expval(qml.Z(0))]
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=1)
+            return [
+                qml.expval(qml.X(0)),
+                qml.expval(qml.Y(1)),
+                qml.expval(qml.Z(0) @ qml.Z(1)),
+                qml.expval(qml.X(0) @ qml.Z(1) + 0.5 * qml.Y(1) + qml.Z(0)),
+            ]
 
         circuit = qml.transforms.split_non_commuting(circuit)
 
-    Internally, the QNode is split into groups of commuting observables when executed:
-
-    >>> print(qml.draw(circuit)(0.5))
-    0: ──RX(0.50)─┤  <X>
-    \
-    0: ──RX(0.50)─┤  <Z>
-
-    Note that while internally multiple QNodes are created, the end result has the same ordering as
-    the user provides in the return statement.
-    Here is a more involved example where we can see the different ordering at the execution level
-    but restoring the original ordering in the output:
+    Internally, the QNode is split into multiple circuits when executed:
+
+    >>> print(qml.draw(circuit)([np.pi/4, np.pi/4]))
+    0: ──RY(0.79)─┤ ╭<Z@Z>  <Z>
+    1: ──RX(0.79)─┤ ╰<Z@Z>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤  <X>
+    1: ──RX(0.79)─┤  <Y>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤ ╭<X@Z>
+    1: ──RX(0.79)─┤ ╰<X@Z>
+
+    Note that the observable ``Y(1)`` occurs twice in the original QNode, but only once in the
+    transformed circuits. When there are multiple expecatation value measurements that rely on
+    the same observable, this observable is measured only once, and the result is copied to each
+    original measurement.
+
+    While internally multiple tapes are created, the end result has the same ordering as the user
+    provides in the return statement. Executing the above QNode returns the original ordering of
+    the expectation values.
+
+    >>> circuit([np.pi/4, np.pi/4])
+    [0.7071067811865475, -0.7071067811865475, 0.5, 0.5]
+
+    There are two algorithms used to compute disjoint groups of commuting observables: ``"qwc"``
+    grouping uses :func:`~pennylane.pauli.group_observables` which computes groups of qubit-wise
+    commuting observables, producing the fewest number of circuit executions, but can be expensive
+    to compute for large multi-term Hamiltonians, while ``"wires"`` grouping simply ensures
+    that no circuit contains two measurements with overlapping wires, disregarding commutativity
+    between the observables being measured.
+
+    The ``grouping_strategy`` keyword argument can be used to specify the grouping strategy. By
+    default, qwc grouping is used whenever possible, except when the circuit contains multiple
+    measurements that includes an expectation value of a ``qml.Hamiltonian``, in which case wires
+    grouping is used in case the Hamiltonian is very large, to save on classical runtime. To force
+    qwc grouping in all cases, set ``grouping_strategy="qwc"``. Similarly, to force wires grouping,
+    set ``grouping_strategy="wires"``:
 
     .. code-block:: python3
 
-        @qml.transforms.split_non_commuting
+        @functools.partial(qml.transforms.split_non_commuting, grouping="wires")
         @qml.qnode(dev)
-        def circuit0(x):
+        def circuit(x):
             qml.RY(x[0], wires=0)
-            qml.RX(x[1], wires=0)
-            return [qml.expval(qml.X(0)),
-                    qml.expval(qml.Z(0)),
-                    qml.expval(qml.Y(1)),
-                    qml.expval(qml.Z(0) @ qml.Z(1)),
-                    ]
+            qml.RX(x[1], wires=1)
+            return [
+                qml.expval(qml.X(0)),
+                qml.expval(qml.Y(1)),
+                qml.expval(qml.Z(0) @ qml.Z(1)),
+                qml.expval(qml.X(0) @ qml.Z(1) + 0.5 * qml.Y(1) + qml.Z(0)),
+            ]
 
-    Drawing this QNode unveils the separate executions in the background
+    In this case, four circuits are created as follows:
 
-    >>> print(qml.draw(circuit0)([np.pi/4, np.pi/4]))
-    0: ──RY(0.79)──RX(0.79)─┤  <X>
-    1: ─────────────────────┤  <Y>
-    \
-    0: ──RY(0.79)──RX(0.79)─┤  <Z> ╭<Z@Z>
-    1: ─────────────────────┤      ╰<Z@Z>
+    >>> print(qml.draw(circuit)([np.pi/4, np.pi/4]))
+    0: ──RY(0.79)─┤  <X>
+    1: ──RX(0.79)─┤  <Y>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤ ╭<Z@Z>
+    1: ──RX(0.79)─┤ ╰<Z@Z>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤ ╭<X@Z>
+    1: ──RX(0.79)─┤ ╰<X@Z>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤  <Z>
+    1: ──RX(0.79)─┤
 
-    Yet, executing it returns the original ordering of the expectation values. The outputs
-    correspond to
-    :math:`(\langle \sigma_x^0 \rangle, \langle \sigma_z^0 \rangle, \langle \sigma_y^1 \rangle,
-    \langle \sigma_z^0\sigma_z^1 \rangle)`.
+    Alternatively, to disable grouping completely, set ``grouping_strategy=None``:
 
-    >>> circuit0([np.pi/4, np.pi/4])
-    [0.7071067811865475, 0.49999999999999994, 0.0, 0.49999999999999994]
+    .. code-block:: python3
 
+        @functools.partial(qml.transforms.split_non_commuting, grouping_strategy=None)
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=1)
+            return [
+                qml.expval(qml.X(0)),
+                qml.expval(qml.Y(1)),
+                qml.expval(qml.Z(0) @ qml.Z(1)),
+                qml.expval(qml.X(0) @ qml.Z(1) + 0.5 * qml.Y(1) + qml.Z(0)),
+            ]
+
+    In this case, each observable is measured in a separate circuit execution.
+
+    >>> print(qml.draw(circuit)([np.pi/4, np.pi/4]))
+    0: ──RY(0.79)─┤  <X>
+    1: ──RX(0.79)─┤
+    <BLANKLINE>
+    0: ──RY(0.79)─┤
+    1: ──RX(0.79)─┤  <Y>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤ ╭<Z@Z>
+    1: ──RX(0.79)─┤ ╰<Z@Z>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤ ╭<X@Z>
+    1: ──RX(0.79)─┤ ╰<X@Z>
+    <BLANKLINE>
+    0: ──RY(0.79)─┤  <Z>
+    1: ──RX(0.79)─┤
+
+    Note that there is an exception to the above rules: if the circuit only contains a single
+    expectation value measurement of a ``Hamiltonian`` or ``Sum`` with pre-computed grouping
+    indices, the grouping information will be used regardless of the requested ``grouping_strategy``
 
     .. details::
         :title: Usage Details
 
-        Internally, this function works with tapes. We can create a tape with non-commuting
-        observables:
-
-        .. code-block:: python3
-
-            measurements = [qml.expval(qml.Z(0)), qml.expval(qml.Y(0))]
-            tape = qml.tape.QuantumTape(measurements=measurements)
-
-            tapes, processing_fn = qml.transforms.split_non_commuting(tape)
-
-        Now ``tapes`` is a list of two tapes, each for one of the non-commuting terms:
-
-        >>> [t.observables for t in tapes]
-        [[expval(Z(0))], [expval(Y(0))]]
-
-        The processing function becomes important when creating the commuting groups as the order
-        of the inputs has been modified:
+        Internally, this function works with tapes. We can create a tape with multiple
+        measurements of non-commuting observables:
 
         .. code-block:: python3
 
@@ -142,18 +202,25 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
                 qml.expval(qml.Z(0)),
                 qml.expval(qml.X(0))
             ]
-            tape = qml.tape.QuantumTape(measurements=measurements)
-
+            tape = qml.tape.QuantumScript(measurements=measurements)
             tapes, processing_fn = qml.transforms.split_non_commuting(tape)
 
-        In this example, the groupings are ``group_coeffs = [[0,2], [1,3]]`` and ``processing_fn``
-        makes sure that the final output is of the same shape and ordering:
+        Now ``tapes`` is a list of two tapes, each contains a group of commuting observables:
+
+        >>> [t.measurements for t in tapes]
+        [[expval(Z(0) @ Z(1)), expval(Z(0))], [expval(X(0) @ X(1)), expval(X(0))]]
+
+        The processing function becomes important as the order of the inputs has been modified.
 
-        >>> processing_fn([t.measurements for t in tapes])
-        (expval(Z(0) @ Z(1)),
-        expval(X(0) @ X(1)),
-        expval(Z(0)),
-        expval(X(0)))
+        >>> dev = qml.device("default.qubit", wires=2)
+        >>> result_batch = [dev.execute(t) for t in tapes]
+        >>> result_batch
+        [(1.0, 1.0), (0.0, 0.0)]
+
+        The processing function can be used to reorganize the results:
+
+        >>> processing_fn(result_batch)
+        (1.0, 0.0, 1.0, 0.0)
 
         Measurements that accept both observables and ``wires`` so that e.g. ``qml.counts``,
         ``qml.probs`` and ``qml.sample`` can also be used. When initialized using only ``wires``,
@@ -167,75 +234,486 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
                 qml.probs(wires=[1]),
                 qml.probs(wires=[0, 1])
             ]
-            tape = qml.tape.QuantumTape(measurements=measurements)
-
+            tape = qml.tape.QuantumScript(measurements=measurements)
             tapes, processing_fn = qml.transforms.split_non_commuting(tape)
 
         This results in two tapes, each with commuting measurements:
 
         >>> [t.measurements for t in tapes]
         [[expval(X(0)), probs(wires=[1])], [probs(wires=[0, 1])]]
+
     """
 
-    # Construct a list of observables to group based on the measurements in the tape
-    obs_list = []
-    for obs in tape.observables:
-        # observable provided for a measurement
-        if isinstance(obs, qml.operation.Operator):
-            obs_list.append(obs)
-        # measurements using wires instead of observables
-        else:
-            # create the PauliZ tensor product observable when only wires are provided for the
-            # measurements
-            obs_wires = obs.wires if obs.wires else tape.wires
-            pauliz_obs = qml.prod(*(qml.Z(wire) for wire in obs_wires))
-
-            obs_list.append(pauliz_obs)
-
-    # If there is more than one group of commuting observables, split tapes
-    _, group_coeffs = qml.pauli.group_observables(obs_list, range(len(obs_list)))
-
-    if len(group_coeffs) > 1:
-        # make one tape per commuting group
-        tapes = []
-        for indices in group_coeffs:
-            new_tape = tape.__class__(
-                tape.operations, (tape.measurements[i] for i in indices), shots=tape.shots
+    # Special case for a single measurement of a Sum or Hamiltonian, in which case
+    # the grouping information can be computed and cached in the observable.
+    if (
+        len(tape.measurements) == 1
+        and isinstance(tape.measurements[0], ExpectationMP)
+        and isinstance(tape.measurements[0].obs, (Hamiltonian, Sum))
+        and (
+            grouping_strategy in ("default", "qwc")
+            or tape.measurements[0].obs.grouping_indices is not None
+        )
+    ):
+        return _split_ham_with_grouping(tape)
+
+    single_term_obs_mps, offsets = _split_all_multi_term_obs_mps(tape)
+
+    if grouping_strategy is None:
+        measurements = list(single_term_obs_mps.keys())
+        tapes = [tape.__class__(tape.operations, [m], shots=tape.shots) for m in measurements]
+        return tapes, partial(
+            _processing_fn_no_grouping,
+            single_term_obs_mps=single_term_obs_mps,
+            offsets=offsets,
+            shots=tape.shots,
+            batch_size=tape.batch_size,
+        )
+
+    if (
+        grouping_strategy == "wires"
+        or grouping_strategy == "default"
+        and any(
+            isinstance(m, ExpectationMP) and isinstance(m.obs, (LinearCombination, Hamiltonian))
+            for m in tape.measurements
+        )
+    ):
+        # This is a loose check to see whether wires grouping or qwc grouping should be used,
+        # which does not necessarily make perfect sense but is consistent with the old decision
+        # logic in `Device.batch_transform`. The premise is that qwc grouping is classically
+        # expensive but produces fewer tapes, whereas wires grouping is classically faster to
+        # compute, but inefficient quantum-wise. If this transform is to be added to a device's
+        # `preprocess`, it will be performed for every circuit execution, which can get very
+        # expensive if there is a large number of observables. The reasoning here is, large
+        # Hamiltonians typically come in the form of a `LinearCombination` or `Hamiltonian`, so
+        # if we see one of those, use wires grouping to be safe. Otherwise, use qwc grouping.
+        return _split_using_wires_grouping(tape, single_term_obs_mps, offsets)
+
+    return _split_using_qwc_grouping(tape, single_term_obs_mps, offsets)
+
+
+def _split_ham_with_grouping(tape: qml.tape.QuantumScript):
+    """Splits a tape measuring a single Hamiltonian or Sum and group commuting observables."""
+
+    obs = tape.measurements[0].obs
+    if obs.grouping_indices is None:
+        obs.compute_grouping()
+
+    coeffs, obs_list = obs.terms()
+
+    # The constant offset of the Hamiltonian, typically arising from Identity terms.
+    offset = 0
+
+    # A dictionary for measurements of each unique single-term observable, mapped to the
+    # indices of the original measurements it belongs to, its coefficients, the index of
+    # the group it belongs to, and the index of the measurement in the group.
+    single_term_obs_mps = {}
+
+    # A list of lists for each group of commuting measurement processes.
+    mp_groups = []
+
+    # The number of measurements in each group
+    group_sizes = []
+
+    # obs.grouping_indices is a list of lists, where each list contains the indices of
+    # observables that belong in each group.
+    for group_idx, obs_indices in enumerate(obs.grouping_indices):
+        mp_group = []
+        group_size = 0
+        for obs_idx in obs_indices:
+            # Do not measure Identity terms, but track their contribution with the offset.
+            if isinstance(obs_list[obs_idx], qml.Identity):
+                offset += coeffs[obs_idx]
+            else:
+                new_mp = qml.expval(obs_list[obs_idx])
+                if new_mp in single_term_obs_mps:
+                    # If the Hamiltonian contains duplicate observables, it can be reused,
+                    # and the coefficients for each duplicate should be combined.
+                    single_term_obs_mps[new_mp] = (
+                        single_term_obs_mps[new_mp][0],
+                        [single_term_obs_mps[new_mp][1][0] + coeffs[obs_idx]],
+                        single_term_obs_mps[new_mp][2],
+                        single_term_obs_mps[new_mp][3],
+                    )
+                else:
+                    mp_group.append(new_mp)
+                    single_term_obs_mps[new_mp] = (
+                        [0],
+                        [coeffs[obs_idx]],
+                        group_idx,
+                        group_size,  # the index of this measurement in the group
+                    )
+                    group_size += 1
+
+        if group_size > 0:
+            mp_groups.append(mp_group)
+            group_sizes.append(group_size)
+
+    tapes = [tape.__class__(tape.operations, mps, shots=tape.shots) for mps in mp_groups]
+    return tapes, partial(
+        _processing_fn_with_grouping,
+        single_term_obs_mps=single_term_obs_mps,
+        offsets=[offset],
+        group_sizes=group_sizes,
+        shots=tape.shots,
+        batch_size=tape.batch_size,
+    )
+
+
+def _split_using_qwc_grouping(
+    tape: qml.tape.QuantumScript,
+    single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]],
+    offsets: List[float],
+):
+    """Split tapes using group_observables in the Pauli module.
+
+    Args:
+        tape (~qml.tape.QuantumScript): The tape to be split.
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+            of measurements of each unique single-term observable, mapped to the indices of the
+            original measurements it belongs to, and its coefficients.
+        offsets (List[float]): Offsets associated with each original measurement in the tape.
+
+    """
+
+    # The legacy device does not support state measurements combined with any other
+    # measurement, so each state measurement must be in its own tape.
+    state_measurements = [m for m in single_term_obs_mps if isinstance(m, StateMP)]
+
+    measurements = [m for m in single_term_obs_mps if not isinstance(m, StateMP)]
+    obs_list = [_mp_to_obs(m, tape) for m in measurements]
+    index_groups = []
+    if len(obs_list) > 0:
+        _, index_groups = qml.pauli.group_observables(obs_list, range(len(obs_list)))
+
+    # A dictionary for measurements of each unique single-term observable, mapped to the
+    # indices of the original measurements it belongs to, its coefficients, the index of
+    # the group it belongs to, and the index of the measurement in the group.
+    single_term_obs_mps_grouped = {}
+
+    mp_groups = [[] for _ in index_groups]
+    group_sizes = []
+    for group_idx, obs_indices in enumerate(index_groups):
+        group_size = 0
+        for obs_idx in obs_indices:
+            new_mp = measurements[obs_idx]
+            mp_groups[group_idx].append(new_mp)
+            single_term_obs_mps_grouped[new_mp] = (
+                *single_term_obs_mps[new_mp],
+                group_idx,
+                group_size,
             )
+            group_size += 1
+        group_sizes.append(group_size)
+
+    for state_mp in state_measurements:
+        mp_groups.append([state_mp])
+        single_term_obs_mps_grouped[state_mp] = (
+            *single_term_obs_mps[state_mp],
+            len(mp_groups) - 1,
+            0,
+        )
+        group_sizes.append(1)
+
+    tapes = [tape.__class__(tape.operations, mps, shots=tape.shots) for mps in mp_groups]
+    return tapes, partial(
+        _processing_fn_with_grouping,
+        single_term_obs_mps=single_term_obs_mps_grouped,
+        offsets=offsets,
+        group_sizes=group_sizes,
+        shots=tape.shots,
+        batch_size=tape.batch_size,
+    )
+
+
+def _split_using_wires_grouping(
+    tape: qml.tape.QuantumScript,
+    single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]],
+    offsets: List[float],
+):
+    """Split tapes by grouping observables based on overlapping wires.
 
-            tapes.append(new_tape)
+    Args:
+        tape (~qml.tape.QuantumScript): The tape to be split.
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+            of measurements of each unique single-term observable, mapped to the indices of the
+            original measurements it belongs to, and its coefficients.
+        offsets (List[float]): Offsets associated with each original measurement in the tape.
 
-        def reorder_fn(res):
-            """re-order the output to the original shape and order"""
-            # determine if shot vector is used
-            if len(tapes[0].measurements) == 1:
-                shot_vector_defined = isinstance(res[0], tuple)
+    """
+
+    mp_groups = []
+    wires_for_each_group = []
+    group_sizes = []
+
+    # A dictionary for measurements of each unique single-term observable, mapped to the
+    # indices of the original measurements it belongs to, its coefficient, the index of
+    # the group it belongs to, and the index of the measurement in the group.
+    single_term_obs_mps_grouped = {}
+    num_groups = 0
+
+    for smp, (mp_indices, coeffs) in single_term_obs_mps.items():
+
+        if len(smp.wires) == 0:  # measurement acting on all wires
+            mp_groups.append([smp])
+            wires_for_each_group.append(tape.wires)
+            group_sizes.append(1)
+            single_term_obs_mps_grouped[smp] = (mp_indices, coeffs, num_groups, 0)
+            num_groups += 1
+            continue
+
+        group_idx = 0
+        added_to_existing_group = False
+        while not added_to_existing_group and group_idx < num_groups:
+            wires = wires_for_each_group[group_idx]
+            if len(wires) != 0 and len(qml.wires.Wires.shared_wires([wires, smp.wires])) == 0:
+                mp_groups[group_idx].append(smp)
+                wires_for_each_group[group_idx] += smp.wires
+                single_term_obs_mps_grouped[smp] = (
+                    mp_indices,
+                    coeffs,
+                    group_idx,
+                    group_sizes[group_idx],
+                )
+                group_sizes[group_idx] += 1
+                added_to_existing_group = True
+            group_idx += 1
+
+        if not added_to_existing_group:
+            mp_groups.append([smp])
+            wires_for_each_group.append(smp.wires)
+            group_sizes.append(1)
+            single_term_obs_mps_grouped[smp] = (mp_indices, coeffs, num_groups, 0)
+            num_groups += 1
+
+    tapes = [tape.__class__(tape.operations, mps, shots=tape.shots) for mps in mp_groups]
+    return tapes, partial(
+        _processing_fn_with_grouping,
+        single_term_obs_mps=single_term_obs_mps_grouped,
+        offsets=offsets,
+        group_sizes=group_sizes,
+        shots=tape.shots,
+        batch_size=tape.batch_size,
+    )
+
+
+def _split_all_multi_term_obs_mps(tape: qml.tape.QuantumScript):
+    """Splits all multi-term observables in a tape to measurements of single-term observables.
+
+    Args:
+        tape (~qml.tape.QuantumScript): The tape with measurements to split.
+
+    Returns:
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A
+            dictionary for measurements of each unique single-term observable, mapped to the
+            indices of the original measurements it belongs to, and its coefficients.
+        offsets (List[float]): Offsets associated with each original measurement in the tape.
+
+    """
+
+    # The dictionary for measurements of each unique single-term observable, mapped the indices
+    # of the original measurements it belongs to, and its coefficients.
+    single_term_obs_mps = {}
+
+    # Offsets associated with each original measurement in the tape (from Identity)
+    offsets = []
+
+    for mp_idx, mp in enumerate(tape.measurements):
+        obs = mp.obs
+        offset = 0
+        if isinstance(mp, ExpectationMP) and isinstance(obs, (Hamiltonian, Sum, Prod, SProd)):
+            if isinstance(obs, SProd):
+                # This is necessary because SProd currently does not flatten into
+                # multiple terms if the base is a sum, which is needed here.
+                obs = obs.simplify()
+            # Break the observable into terms, and construct an ExpectationMP with each term.
+            for c, o in zip(*obs.terms()):
+                # If the observable is an identity, track it with a constant offset
+                if isinstance(o, qml.Identity):
+                    offset += c
+                # If the single-term measurement already exists, it can be reused by all original
+                # measurements. In this case, add the existing single-term measurement to the list
+                # corresponding to this original measurement.
+                # pylint: disable=superfluous-parens
+                elif (sm := qml.expval(o)) in single_term_obs_mps:
+                    single_term_obs_mps[sm][0].append(mp_idx)
+                    single_term_obs_mps[sm][1].append(c)
+                # Otherwise, add this new measurement to the list of single-term measurements.
+                else:
+                    single_term_obs_mps[sm] = ([mp_idx], [c])
+        else:
+            # For all other measurement types, simply add them to the list of measurements.
+            if mp not in single_term_obs_mps:
+                single_term_obs_mps[mp] = ([mp_idx], [1])
             else:
-                shot_vector_defined = isinstance(res[0][0], tuple)
+                single_term_obs_mps[mp][0].append(mp_idx)
+                single_term_obs_mps[mp][1].append(1)
 
-            res = list(zip(*res)) if shot_vector_defined else [res]
+        offsets.append(offset)
 
-            reorder_indxs = qml.math.concatenate(group_coeffs)
+    return single_term_obs_mps, offsets
 
-            res_ordered = []
-            for shot_res in res:
-                # flatten the results
-                shot_res = reduce(
-                    lambda x, y: x + list(y) if isinstance(y, (tuple, list)) else x + [y],
-                    shot_res,
-                    [],
-                )
 
-                # reorder the tape results to match the user-provided order
-                shot_res = list(zip(range(len(shot_res)), shot_res))
-                shot_res = sorted(shot_res, key=lambda r: reorder_indxs[r[0]])
-                shot_res = [r[1] for r in shot_res]
+def _processing_fn_no_grouping(
+    res: ResultBatch,
+    single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]],
+    offsets: List[float],
+    shots: Shots,
+    batch_size: int,
+):
+    """Postprocessing function for the split_non_commuting transform without grouping.
+
+    Args:
+        res (ResultBatch): The results from executing the tapes. Assumed to have a shape
+            of (n_groups [,n_shots] [,n_mps] [,batch_size])
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+            of measurements of each unique single-term observable, mapped to the indices of the
+            original measurements it belongs to, and its coefficients.
+        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        shots (Shots): The shots settings of the original tape.
+
+    """
+
+    res_batch_for_each_mp = [[] for _ in offsets]
+    coeffs_for_each_mp = [[] for _ in offsets]
+
+    for smp_idx, (_, (mp_indices, coeffs)) in enumerate(single_term_obs_mps.items()):
+
+        for mp_idx, coeff in zip(mp_indices, coeffs):
+            res_batch_for_each_mp[mp_idx].append(res[smp_idx])
+            coeffs_for_each_mp[mp_idx].append(coeff)
+
+    result_shape = _infer_result_shape(shots, batch_size)
+
+    # Sum up the results for each original measurement
+    res_for_each_mp = [
+        _sum_terms(_sub_res, coeffs, offset, result_shape)
+        for _sub_res, coeffs, offset in zip(res_batch_for_each_mp, coeffs_for_each_mp, offsets)
+    ]
+
+    # res_for_each_mp should have shape (n_mps, [,n_shots] [,batch_size])
+    if len(res_for_each_mp) == 1:
+        return res_for_each_mp[0]
+
+    if shots.has_partitioned_shots:
+        # If the shot vector dimension exists, it should be moved to the first axis
+        # Basically, the shape becomes (n_shots, n_mps, [,batch_size])
+        res_for_each_mp = [
+            tuple(res_for_each_mp[j][i] for j in range(len(res_for_each_mp)))
+            for i in range(shots.num_copies)
+        ]
+
+    return tuple(res_for_each_mp)
+
+
+def _processing_fn_with_grouping(
+    res: ResultBatch,
+    single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float], int, int]],
+    offsets: List[float],
+    group_sizes: List[int],
+    shots: Shots,
+    batch_size: int,
+):
+    """Postprocessing function for the split_non_commuting transform with grouping.
+
+    Args:
+        res (ResultBatch): The results from executing the tapes. Assumed to have a shape
+            of (n_groups [,n_shots] [,n_mps_in_group] [,batch_size])
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float], int, int]]):
+            A dictionary of measurements of each unique single-term observable, mapped to the
+            indices of the original measurements it belongs to, its coefficients, its group
+            index, and the index of the measurement within the group.
+        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        group_sizes (List[int]): The number of tapes in each group.
+        shots (Shots): The shots setting of the original tape.
+
+    Returns:
+        The results combined into a single result for each original measurement.
+
+    """
+
+    res_batch_for_each_mp = [[] for _ in offsets]  # ([n_mps] [,n_shots] [,batch_size])
+    coeffs_for_each_mp = [[] for _ in offsets]
+
+    for _, (mp_indices, coeffs, group_idx, mp_idx_in_group) in single_term_obs_mps.items():
+
+        res_group = res[group_idx]  # ([n_shots] [,n_mps] [,batch_size])
+        group_size = group_sizes[group_idx]
+
+        if group_size > 1 and shots.has_partitioned_shots:
+            # Each result should have shape ([n_shots] [,batch_size])
+            sub_res = [_res[mp_idx_in_group] for _res in res_group]
+        else:
+            # If there is only one term in the group, the n_mps dimension would have
+            # been squeezed out, use the entire result directly.
+            sub_res = res_group if group_size == 1 else res_group[mp_idx_in_group]
+
+        # Add this result to the result batch for the corresponding original measurement
+        for mp_idx, coeff in zip(mp_indices, coeffs):
+            res_batch_for_each_mp[mp_idx].append(sub_res)
+            coeffs_for_each_mp[mp_idx].append(coeff)
+
+    result_shape = _infer_result_shape(shots, batch_size)
+
+    # Sum up the results for each original measurement
+    res_for_each_mp = [
+        _sum_terms(_sub_res, coeffs, offset, result_shape)
+        for _sub_res, coeffs, offset in zip(res_batch_for_each_mp, coeffs_for_each_mp, offsets)
+    ]
+
+    # res_for_each_mp should have shape (n_mps, [,n_shots] [,batch_size])
+    if len(res_for_each_mp) == 1:
+        return res_for_each_mp[0]
+
+    if shots.has_partitioned_shots:
+        # If the shot vector dimension exists, it should be moved to the first axis
+        # Basically, the shape becomes (n_shots, n_mps, [,batch_size])
+        res_for_each_mp = [
+            tuple(res_for_each_mp[j][i] for j in range(len(res_for_each_mp)))
+            for i in range(shots.num_copies)
+        ]
+
+    return tuple(res_for_each_mp)
+
+
+def _sum_terms(res: ResultBatch, coeffs: List[float], offset: float, shape: Tuple) -> Result:
+    """Sum results from measurements of multiple terms in a multi-term observable."""
+
+    # Trivially return the original result
+    if coeffs == [1] and offset == 0:
+        return res[0]
+
+    # The shape of res at this point is (n_terms, [,n_shots] [,batch_size])
+    dot_products = []
+    for c, r in zip(coeffs, res):
+        dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
+    if len(dot_products) == 0:
+        return qml.math.ones(shape) * offset
+    summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
+    return qml.math.convert_like(summed_dot_products + offset, res[0])
+
+
+def _mp_to_obs(mp: MeasurementProcess, tape: qml.tape.QuantumScript) -> qml.operation.Operator:
+    """Extract the observable from a measurement process.
+
+    If the measurement process has an observable, return it. Otherwise, return a dummy
+    observable that is a tensor product of Z gates on every wire.
+
+    """
+
+    if mp.obs is not None:
+        return mp.obs
 
-                res_ordered.append(tuple(shot_res))
+    obs_wires = mp.wires if mp.wires else tape.wires
+    return qml.prod(*(qml.Z(wire) for wire in obs_wires))
 
-            return tuple(res_ordered) if shot_vector_defined else res_ordered[0]
 
-        return tapes, reorder_fn
+def _infer_result_shape(shots: Shots, batch_size: int) -> Tuple:
+    """Based on the result, infer the ([,n_shots] [,batch_size]) shape of the result."""
 
-    # if the group is already commuting, no need to do anything
-    return [tape], null_postprocessing
+    shape = ()
+    if shots.has_partitioned_shots:
+        shape += (shots.num_copies,)
+    if batch_size and batch_size > 1:
+        shape += (batch_size,)
+    return shape


Test patch:
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index e02c4d89d..c7ff19c5f 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -1,4 +1,4 @@
-# Copyright 2018-2020 Xanadu Quantum Technologies Inc.
+# Copyright 2018-2024 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.
@@ -11,962 +11,1030 @@
 # 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.
-""" Tests for the transform ``qml.transform.split_non_commuting()`` """
-# pylint: disable=no-self-use, import-outside-toplevel, no-member, import-error
+
+"""Tests for the transform ``qml.transform.split_non_commuting`` """
+
+# pylint: disable=import-outside-toplevel,unnecessary-lambda
+
 import itertools
+from functools import partial
 
 import numpy as np
 import pytest
 
 import pennylane as qml
-import pennylane.numpy as pnp
 from pennylane.transforms import split_non_commuting
 
-# list of observables with 2 commuting groups [[1, 3], [0, 2, 4]]
-obs_list_2 = [
-    qml.prod(qml.Z(0), qml.Z(1)),
-    qml.prod(qml.PauliX(0), qml.PauliX(1)),
-    qml.PauliZ(0),
-    qml.PauliX(0),
-    qml.PauliZ(1),
+# Two commuting groups: [[0, 3], [1, 2, 4]]
+# Four groups based on wire overlaps: [[0, 2], [1], [3], [4]]
+single_term_obs_list = [
+    qml.X(0),
+    qml.Y(0),
+    qml.Z(1),
+    qml.X(0) @ qml.Y(1),
+    qml.Y(0) @ qml.Z(1),
 ]
 
-# list of observables with 3 commuting groups [[0,3], [1,4], [2,5]]
-obs_list_3 = [
-    qml.prod(qml.PauliZ(0), qml.PauliZ(1)),
-    qml.prod(qml.PauliX(0), qml.PauliX(1)),
-    qml.prod(qml.PauliY(0), qml.PauliY(1)),
-    qml.PauliZ(0),
-    qml.PauliX(0),
-    qml.PauliY(0),
+single_term_qwc_groups = [
+    [qml.X(0), qml.X(0) @ qml.Y(1)],
+    [qml.Y(0), qml.Z(1), qml.Y(0) @ qml.Z(1)],
 ]
 
-# measurements that accept observables as arguments
-obs_meas_fn = [qml.expval, qml.var, qml.probs, qml.counts, qml.sample]
+single_term_wires_groups = [
+    [qml.X(0), qml.Z(1)],
+    [qml.Y(0)],
+    [qml.X(0) @ qml.Y(1)],
+    [qml.Y(0) @ qml.Z(1)],
+]
 
-# measurements that accept wires as arguments
-wire_meas_fn = [qml.probs, qml.counts, qml.sample]
+# contains the following observables: X(0), Y(0), Y(0) @ Z(1), X(1), Z(1), X(0) @ Y(1)
+# qwc groups: [[0, 5], [1, 3], [2, 4]]
+# wires groups: [[0, 3], [1, 4], [2], [5]]
+complex_obs_list = [
+    qml.X(0),  # single observable
+    0.5 * qml.Y(0),  # scalar product
+    qml.X(0) + qml.Y(0) @ qml.Z(1) + 2.0 * qml.X(1) + qml.I(),  # sum
+    qml.Hamiltonian(
+        [0.1, 0.2, 0.3, 0.4], [qml.Z(1), qml.X(0) @ qml.Y(1), qml.Y(0) @ qml.Z(1), qml.I()]
+    ),
+    1.5 * qml.I(),  # identity
+]
 
+complex_no_grouping_obs = [
+    qml.X(0),
+    qml.Y(0),
+    qml.Y(0) @ qml.Z(1),
+    qml.X(1),
+    qml.Z(1),
+    qml.X(0) @ qml.Y(1),
+]
 
-class TestUnittestSplitNonCommuting:
-    """Unit tests on ``qml.transforms.split_non_commuting()``"""
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", obs_meas_fn)
-    def test_commuting_group_no_split(self, mocker, meas_type, convert_to_opmath):
-        """Testing that commuting groups are not split for all supported measurement types"""
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.PauliZ(0)
-            qml.Hadamard(0)
-            qml.CNOT((0, 1))
-            meas_type(op=qml.PauliZ(0))
-            meas_type(op=qml.PauliZ(0))
-            meas_type(op=qml.PauliX(1))
-            meas_type(op=qml.PauliZ(2))
-            if convert_to_opmath:
-                meas_type(op=qml.prod(qml.Z(0), qml.Z(3)))
-            else:
-                meas_type(op=qml.operation.Tensor(qml.Z(0), qml.Z(3)))
+def _convert_obs_to_legacy_opmath(obs):
+    """Convert single-term observables to legacy opmath"""
 
-        # test transform on tape
-        tape = qml.tape.QuantumScript.from_queue(q, shots=100)
-        split, fn = split_non_commuting(tape)
-        for t in split:
-            assert t.shots == tape.shots
+    if isinstance(obs, qml.ops.Prod):
+        return qml.operation.Tensor(*obs.operands)
 
-        spy = mocker.spy(qml.math, "concatenate")
+    if isinstance(obs, list):
+        return [_convert_obs_to_legacy_opmath(o) for o in obs]
 
-        assert len(split) == 1
-        assert all(isinstance(t, qml.tape.QuantumScript) for t in split)
-        assert fn([0.5]) == 0.5
+    return obs
 
-        # test transform on qscript
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements, shots=50)
-        split, fn = split_non_commuting(qs)
-        for t in split:
-            assert t.shots == qs.shots
 
-        assert len(split) == 1
-        assert all(isinstance(i_qs, qml.tape.QuantumScript) for i_qs in split)
-        assert fn([0.5]) == 0.5
+def complex_no_grouping_processing_fn(results):
+    """The expected processing function without grouping of complex_obs_list"""
 
-        spy.assert_not_called()
+    return (
+        results[0],
+        0.5 * results[1],
+        results[0] + results[2] + 2.0 * results[3] + 1.0,
+        0.1 * results[4] + 0.2 * results[5] + 0.3 * results[2] + 0.4,
+        1.5,
+    )
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", wire_meas_fn)
-    def test_wire_commuting_group_no_split(self, mocker, meas_type, convert_to_opmath):
-        """Testing that commuting MPs initialized using wires or observables are not split"""
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.PauliZ(0)
-            qml.Hadamard(0)
-            qml.CNOT((0, 1))
-            meas_type()  # Included to check splitting with all-wire measurements
-            meas_type(wires=[0])
-            meas_type(wires=[1])
-            meas_type(wires=[0, 1])
-            meas_type(op=qml.PauliZ(0))
-            if convert_to_opmath:
-                meas_type(op=qml.prod(qml.PauliZ(0), qml.PauliZ(2)))
-            else:
-                meas_type(op=qml.operation.Tensor(qml.PauliZ(0), qml.PauliZ(2)))
 
-        # test transform on tape
-        tape = qml.tape.QuantumScript.from_queue(q)
-        split, fn = split_non_commuting(tape)
+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)],
+]
 
-        spy = mocker.spy(qml.math, "concatenate")
 
-        assert len(split) == 1
-        assert all(isinstance(t, qml.tape.QuantumScript) for t in split)
-        assert fn([0.5]) == 0.5
+def complex_qwc_processing_fn(results):
+    """The expected processing function for qwc grouping of complex_obs_list"""
+    group0, group1, group2 = 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,
+        1.5,
+    )
 
-        # test transform on qscript
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
-        split, fn = split_non_commuting(qs)
 
-        assert len(split) == 1
-        assert all(isinstance(i_qs, qml.tape.QuantumScript) for i_qs in split)
-        assert fn([0.5]) == 0.5
+complex_wires_groups = [
+    [qml.X(0), qml.X(1)],
+    [qml.Y(0), qml.Z(1)],
+    [qml.Y(0) @ qml.Z(1)],
+    [qml.X(0) @ qml.Y(1)],
+]
 
-        spy.assert_not_called()
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", obs_meas_fn)
-    @pytest.mark.parametrize("obs_list, expected", [(obs_list_2, 2), (obs_list_3, 3)])
-    def test_non_commuting_group_right_number(
-        self, meas_type, obs_list, expected, convert_to_opmath
-    ):
-        """Test that the no. of tapes after splitting into commuting groups is of the right size"""
+def complex_wires_processing_fn(results):
+    """The expected processing function for wires grouping of complex_obs_list"""
 
-        if not convert_to_opmath:
-            obs_list = [
-                qml.operation.Tensor(*o) if isinstance(o, qml.ops.Prod) else o for o in obs_list
-            ]
-        # create a queue with several measurements of same type but with differnent non-commuting
-        # observables
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.PauliZ(0)
-            qml.Hadamard(0)
-            qml.CNOT((0, 1))
-            for obs in obs_list:
-                meas_type(op=obs)
+    group0, group1, group2, group3 = results
+    return (
+        group0[0],
+        0.5 * group1[0],
+        group0[0] + group2 + 2.0 * group0[1] + 1.0,
+        0.1 * group1[1] + 0.2 * group3 + 0.3 * group2 + 0.4,
+        1.5,
+    )
 
-            # if MP type can accept wires, then add two extra measurements using wires and test no.
-            # of tapes after splitting commuting groups
-            if meas_type in wire_meas_fn:
-                meas_type(wires=[0])
-                meas_type(wires=[0, 1])
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        split, _ = split_non_commuting(tape)
-        assert len(split) == expected
+# Measurements that accept observables as arguments
+obs_measurements = [qml.expval, qml.var, qml.probs, qml.counts, qml.sample]
 
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
-        split, _ = split_non_commuting(qs)
-        assert len(split) == expected
+# measurements that accept wires as arguments
+wire_measurements = [qml.probs, qml.counts, qml.sample]
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", obs_meas_fn)
+
+class TestUnits:
+    """Unit tests for components of the ``split_non_commuting`` transform"""
+
+    @pytest.mark.parametrize("measure_fn", obs_measurements)
     @pytest.mark.parametrize(
-        "obs_list, group_coeffs",
-        [(obs_list_2, [[1, 3], [0, 2, 4]]), (obs_list_3, [[0, 3], [1, 4], [2, 5]])],
+        "grouping_strategy, n_tapes", [(None, 5), ("default", 2), ("qwc", 2), ("wires", 4)]
     )
-    def test_non_commuting_group_right_reorder(
-        self, meas_type, obs_list, convert_to_opmath, group_coeffs
-    ):
-        """Test that the output is of the correct order"""
-        # create a queue with several measurements of same type but with differnent non-commuting
-        # observables
+    def test_number_of_tapes(self, measure_fn, grouping_strategy, n_tapes):
+        """Tests that the correct number of tapes is returned"""
 
-        if not convert_to_opmath:
-            obs_list = [
-                qml.operation.Tensor(*o) if isinstance(o, qml.ops.Prod) else o for o in obs_list
-            ]
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.PauliZ(0)
-            qml.Hadamard(0)
-            qml.CNOT((0, 1))
-            for obs in obs_list:
-                meas_type(op=obs)
+        measurements = [measure_fn(op=o) for o in single_term_obs_list]
+        tape = qml.tape.QuantumScript([qml.X(0), qml.CNOT([0, 1])], measurements, shots=100)
+        tapes, _ = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == n_tapes
+        assert all(t.operations == [qml.X(0), qml.CNOT([0, 1])] for t in tapes)
+        assert all(t.shots == tape.shots for t in tapes)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        _, fn = split_non_commuting(tape)
-        assert all(np.array(fn(group_coeffs)) == np.arange(len(tape.measurements)))
+    @pytest.mark.parametrize(
+        "grouping_strategy, n_tapes", [(None, 5), ("default", 2), ("qwc", 2), ("wires", 4)]
+    )
+    @pytest.mark.parametrize(
+        "make_H",
+        [
+            lambda coeffs, obs: qml.Hamiltonian(coeffs, obs),
+            lambda coeffs, obs: qml.sum(*(qml.s_prod(c, o) for c, o in zip(coeffs, obs))),
+        ],
+    )
+    def test_number_of_tapes_single_hamiltonian(self, grouping_strategy, n_tapes, make_H):
+        """Tests that the correct number of tapes is returned for a single Hamiltonian"""
+
+        obs_list = single_term_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = _convert_obs_to_legacy_opmath(obs_list)
 
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
-        _, fn = split_non_commuting(qs)
-        assert all(np.array(fn(group_coeffs)) == np.arange(len(tape.measurements)))
+        obs_list = obs_list + [qml.Y(0), qml.X(0) @ qml.Y(1)]  # add duplicate terms
+        coeffs = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+
+        H = make_H(coeffs, obs_list)
+        tape = qml.tape.QuantumScript([qml.X(0), qml.CNOT([0, 1])], [qml.expval(H)], shots=100)
+        tapes, _ = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == n_tapes
+        assert all(t.operations == [qml.X(0), qml.CNOT([0, 1])] for t in tapes)
+        assert all(t.shots == tape.shots for t in tapes)
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", wire_meas_fn)
     @pytest.mark.parametrize(
-        "obs_list, group_coeffs",
-        [(obs_list_2, [[1, 3], [0, 2, 4, 5]]), (obs_list_3, [[1, 4], [2, 5], [0, 3, 6]])],
+        "grouping_strategy, n_tapes", [(None, 6), ("default", 4), ("qwc", 3), ("wires", 4)]
     )
-    def test_wire_non_commuting_group_right_reorder(
-        self, meas_type, obs_list, convert_to_opmath, group_coeffs
-    ):
-        """Test that the output is of the correct order with wire MPs initialized using a
-        combination of wires and observables"""
-        # create a queue with several measurements of same type but with differnent non-commuting
-        # observables
-        if not convert_to_opmath:
-            obs_list = [
-                qml.operation.Tensor(*o) if isinstance(o, qml.ops.Prod) else o for o in obs_list
-            ]
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.PauliZ(0)
-            qml.Hadamard(0)
-            qml.CNOT((0, 1))
-            for obs in obs_list:
-                meas_type(op=obs)
+    def test_number_of_tapes_complex_obs(self, grouping_strategy, n_tapes):
+        """Tests number of tapes with mixed types of observables"""
+
+        measurements = [qml.expval(o) for o in complex_obs_list]
+        tape = qml.tape.QuantumScript([qml.X(0), qml.CNOT([0, 1])], measurements, shots=100)
+        tapes, _ = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == n_tapes
+        assert all(t.operations == [qml.X(0), qml.CNOT([0, 1])] for t in tapes)
+        assert all(t.shots == tape.shots for t in tapes)
+
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_state_measurement_in_separate_tape(self, grouping_strategy):
+        """Tests that a state measurement is in a separate tape"""
+
+        measurements = [qml.expval(qml.Z(0) @ qml.Z(1)), qml.state()]
+        tape = qml.tape.QuantumScript([qml.X(0), qml.CNOT([0, 1])], measurements, shots=100)
+        tapes, _ = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == 2
+        assert all(t.operations == [qml.X(0), qml.CNOT([0, 1])] for t in tapes)
+        assert all(t.shots == tape.shots for t in tapes)
+
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    @pytest.mark.parametrize(
+        "make_H",
+        [
+            lambda obs_list: qml.Hamiltonian([0.1, 0.2, 0.3, 0.4, 0.5], obs_list),
+            lambda obs_list: qml.sum(
+                *(qml.s_prod(c, o) for c, o in zip([0.1, 0.2, 0.3, 0.4, 0.5], obs_list))
+            ),
+        ],
+    )
+    def test_existing_grouping_used_for_single_hamiltonian(self, grouping_strategy, make_H):
+        """Tests that if a Hamiltonian has an existing grouping, it is used regardless of
+        what is requested through the ``grouping_strategy`` argument."""
 
-            # initialize measurements using wires
-            meas_type(wires=[0])
+        obs_list = single_term_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = _convert_obs_to_legacy_opmath(obs_list)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        _, fn = split_non_commuting(tape)
-        assert all(np.array(fn(group_coeffs)) == np.arange(len(tape.measurements)))
+        H = make_H(obs_list)
+        H.compute_grouping()
 
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
-        _, fn = split_non_commuting(qs)
-        assert all(np.array(fn(group_coeffs)) == np.arange(len(tape.measurements)))
+        tape = qml.tape.QuantumScript([qml.X(0), qml.CNOT([0, 1])], [qml.expval(H)], shots=100)
+        tapes, _ = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == 2
+        assert all(t.operations == [qml.X(0), qml.CNOT([0, 1])] for t in tapes)
+        assert all(t.shots == tape.shots for t in tapes)
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    @pytest.mark.parametrize("meas_type", obs_meas_fn)
-    def test_different_measurement_types(self, meas_type, convert_to_opmath):
-        """Test that the measurements types of the split tapes are correct"""
+    @pytest.mark.parametrize("measure_fn", obs_measurements)
+    def test_single_group(self, measure_fn):
+        """Tests when all measurements can be taken at the same time"""
 
-        prod_type = qml.prod if convert_to_opmath else qml.operation.Tensor
         with qml.queuing.AnnotatedQueue() as q:
             qml.PauliZ(0)
             qml.Hadamard(0)
             qml.CNOT((0, 1))
-            meas_type(op=prod_type(qml.Z(0), qml.Z(1)))
-            meas_type(op=prod_type(qml.X(0), qml.X(1)))
-            meas_type(op=qml.PauliZ(0))
-            meas_type(op=qml.PauliX(0))
+            measure_fn(op=qml.X(0))
+            measure_fn(op=qml.Y(1))
+            measure_fn(op=qml.Z(2))
+            measure_fn(op=qml.X(0) @ qml.Y(1))
+            measure_fn(op=qml.Y(1) @ qml.Z(2))
 
-            # if the MP can also accept wires as arguments, add extra measurements to test
-            if meas_type in wire_meas_fn:
-                meas_type(wires=[0])
-                meas_type(wires=[0, 1])
+        tape = qml.tape.QuantumScript.from_queue(q, shots=100)
+        tapes, fn = split_non_commuting(tape)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        the_return_type = tape.measurements[0].return_type
-        split, _ = split_non_commuting(tape)
-        for new_tape in split:
-            for meas in new_tape.measurements:
-                assert meas.return_type == the_return_type
-
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
-        split, _ = split_non_commuting(qs)
-        for new_tape in split:
-            for meas in new_tape.measurements:
-                assert meas.return_type == the_return_type
-
-    @pytest.mark.parametrize("meas_type_1, meas_type_2", itertools.combinations(obs_meas_fn, 2))
-    def test_mixed_measurement_types(self, meas_type_1, meas_type_2):
-        """Test that mixing different combintations of MPs initialized using obs works correctly."""
+        assert len(tapes) == 1
+        assert fn([[0.1, 0.2, 0.3, 0.4, 0.5]]) == (0.1, 0.2, 0.3, 0.4, 0.5)
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.Hadamard(0)
-            qml.Hadamard(1)
-            meas_type_1(op=qml.PauliX(0))
-            meas_type_1(op=qml.PauliZ(1))
-            meas_type_2(op=qml.PauliZ(0))
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_single_observable(self, grouping_strategy):
+        """Tests a circuit that contains a single observable"""
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        split, _ = split_non_commuting(tape)
+        tape = qml.tape.QuantumScript([], [qml.expval(qml.X(0))])
+        tapes, fn = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == 1
+        assert fn([0.1]) == 0.1
+
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_single_hamiltonian_single_observable(self, grouping_strategy):
+        """Tests a circuit that contains a single observable"""
 
-        assert len(split) == 2
-        assert qml.equal(split[0].measurements[0], meas_type_1(op=qml.PauliX(0)))
-        assert qml.equal(split[0].measurements[1], meas_type_1(op=qml.PauliZ(1)))
-        assert qml.equal(split[1].measurements[0], meas_type_2(op=qml.PauliZ(0)))
+        tape = qml.tape.QuantumScript([], [qml.expval(qml.Hamiltonian([0.1], [qml.X(0)]))])
+        tapes, fn = split_non_commuting(tape, grouping_strategy=grouping_strategy)
+        assert len(tapes) == 1
+        assert qml.math.allclose(fn([0.1]), 0.01)
 
-    @pytest.mark.parametrize("meas_type_1, meas_type_2", itertools.combinations(wire_meas_fn, 2))
-    def test_mixed_wires_measurement_types(self, meas_type_1, meas_type_2):
-        """Test that mixing different combinations of MPs initialized using wires works correctly"""
+    @pytest.mark.parametrize("measure_fn", wire_measurements)
+    def test_all_wire_measurements(self, measure_fn):
+        """Tests that measurements based on wires don't need to be split"""
 
         with qml.queuing.AnnotatedQueue() as q:
+            qml.PauliZ(0)
             qml.Hadamard(0)
-            qml.Hadamard(1)
-            meas_type_1(op=qml.PauliX(0))
-            meas_type_1(wires=[1])
-            meas_type_2(wires=[0])
+            qml.CNOT((0, 1))
+            measure_fn()
+            measure_fn(wires=[0])
+            measure_fn(wires=[1])
+            measure_fn(wires=[0, 1])
+            measure_fn(op=qml.PauliZ(0))
+            measure_fn(op=qml.PauliZ(0) @ qml.PauliZ(2))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        split, _ = split_non_commuting(tape)
+        tapes, fn = split_non_commuting(tape)
 
-        assert len(split) == 2
-        assert qml.equal(split[0].measurements[0], meas_type_1(op=qml.PauliX(0)))
-        assert qml.equal(split[0].measurements[1], meas_type_1(wires=[1]))
-        assert qml.equal(split[1].measurements[0], meas_type_2(wires=[0]))
+        assert len(tapes) == 1
+        assert fn([[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]]) == (0.1, 0.2, 0.3, 0.4, 0.5, 0.6)
 
+    @pytest.mark.parametrize("obs_meas_1, obs_meas_2", itertools.combinations(obs_measurements, 2))
     @pytest.mark.parametrize(
-        "meas_type_1, meas_type_2", itertools.product(obs_meas_fn, wire_meas_fn)
+        "wire_meas_1, wire_meas_2", itertools.combinations(wire_measurements, 2)
     )
-    def test_mixed_wires_obs_measurement_types(self, meas_type_1, meas_type_2):
-        """Test that mixing different combinations of MPs initialized using wires and obs works
-        correctly"""
+    def test_mix_measurement_types(self, obs_meas_1, obs_meas_2, wire_meas_1, wire_meas_2):
+        """Tests that tapes mixing different measurement types is handled correctly"""
 
         with qml.queuing.AnnotatedQueue() as q:
-            qml.Hadamard(0)
-            qml.Hadamard(1)
-            meas_type_1(op=qml.PauliX(0))
-            meas_type_2()
-            meas_type_2(wires=[1])
-            meas_type_2(wires=[0, 1])
+            obs_meas_1(op=qml.PauliX(0))
+            obs_meas_2(op=qml.PauliZ(1))
+            obs_meas_1(op=qml.PauliZ(0))
+            wire_meas_1(wires=[0])
+            wire_meas_2(wires=[1])
+            wire_meas_1(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        split, _ = split_non_commuting(tape)
+        tapes, _ = split_non_commuting(tape)
+        assert len(tapes) == 2
+        assert tapes[0].measurements == [
+            obs_meas_1(op=qml.PauliX(0)),
+            obs_meas_2(op=qml.PauliZ(1)),
+            wire_meas_2(wires=[1]),
+        ]
+        assert tapes[1].measurements == [
+            obs_meas_1(op=qml.PauliZ(0)),
+            wire_meas_1(wires=[0]),
+            wire_meas_1(wires=[0, 1]),
+        ]
+
+    def test_grouping_strategies(self):
+        """Tests that the tape is split correctly for different grouping strategies"""
+
+        measurements = [
+            qml.expval(c * o) for c, o in zip([0.1, 0.2, 0.3, 0.4, 0.5], single_term_obs_list)
+        ]
+        tape = qml.tape.QuantumScript([], measurements, shots=100)
+
+        expected_tapes_no_grouping = [
+            qml.tape.QuantumScript([], [qml.expval(o)], shots=100) for o in single_term_obs_list
+        ]
+
+        # qwc grouping produces [[0, 3], [1, 2, 4]]
+        expected_tapes_qwc_grouping = [
+            qml.tape.QuantumScript([], [qml.expval(o) for o in group], shots=100)
+            for group in single_term_qwc_groups
+        ]
+
+        # wires grouping produces [[0, 2], [1], [3], [4]]
+        expected_tapes_wires_grouping = [
+            qml.tape.QuantumScript([], [qml.expval(o) for o in group], shots=100)
+            for group in single_term_wires_groups
+        ]
+
+        tapes, fn = split_non_commuting(tape, grouping_strategy=None)
+        for actual_tape, expected_tape in zip(tapes, expected_tapes_no_grouping):
+            assert qml.equal(actual_tape, expected_tape)
+        assert qml.math.allclose(fn([0.1, 0.2, 0.3, 0.4, 0.5]), [0.01, 0.04, 0.09, 0.16, 0.25])
+
+        tapes, fn = split_non_commuting(tape, grouping_strategy="default")
+        # When new opmath is disabled, c * o gives Hamiltonians, which leads to wires grouping
+        if qml.operation.active_new_opmath():
+            for actual_tape, expected_tape in zip(tapes, expected_tapes_qwc_grouping):
+                assert qml.equal(actual_tape, expected_tape)
+            assert qml.math.allclose(
+                fn([[0.1, 0.2], [0.3, 0.4, 0.5]]), [0.01, 0.06, 0.12, 0.08, 0.25]
+            )
+        else:
+            for actual_tape, expected_tape in zip(tapes, expected_tapes_wires_grouping):
+                assert qml.equal(actual_tape, expected_tape)
+            assert qml.math.allclose(
+                fn([[0.1, 0.2], 0.3, 0.4, 0.5]), [0.01, 0.06, 0.06, 0.16, 0.25]
+            )
 
-        assert len(split) == 2
-        assert qml.equal(split[0].measurements[0], meas_type_1(op=qml.PauliX(0)))
-        assert qml.equal(split[0].measurements[1], meas_type_2(wires=[1]))
-        assert qml.equal(split[1].measurements[0], meas_type_2())
-        assert qml.equal(split[1].measurements[1], meas_type_2(wires=[0, 1]))
+        tapes, fn = split_non_commuting(tape, grouping_strategy="qwc")
+        for actual_tape, expected_tape in zip(tapes, expected_tapes_qwc_grouping):
+            assert qml.equal(actual_tape, expected_tape)
+        assert qml.math.allclose(fn([[0.1, 0.2], [0.3, 0.4, 0.5]]), [0.01, 0.06, 0.12, 0.08, 0.25])
 
-    @pytest.mark.parametrize("batch_type", (tuple, list))
-    def test_batch_of_tapes(self, batch_type):
-        """Test that `split_non_commuting` can transform a batch of tapes"""
+        tapes, fn = split_non_commuting(tape, grouping_strategy="wires")
+        for actual_tape, expected_tape in zip(tapes, expected_tapes_wires_grouping):
+            assert qml.equal(actual_tape, expected_tape)
+        assert qml.math.allclose(fn([[0.1, 0.2], 0.3, 0.4, 0.5]), [0.01, 0.06, 0.06, 0.16, 0.25])
 
-        # create a batch with two simple tapes
-        tape1 = qml.tape.QuantumScript(
-            [qml.RX(1.2, 0)], [qml.expval(qml.X(0)), qml.expval(qml.Y(0)), qml.expval(qml.X(1))]
-        )
-        tape2 = qml.tape.QuantumScript(
-            [qml.RY(0.5, 0)], [qml.expval(qml.Z(0)), qml.expval(qml.Y(0))]
-        )
-        batch = batch_type([tape1, tape2])
+    @pytest.mark.parametrize(
+        "make_H",
+        [
+            lambda coeffs, obs_list: qml.Hamiltonian(coeffs, obs_list),
+            lambda coeffs, obs_list: qml.sum(*(qml.s_prod(c, o) for c, o in zip(coeffs, obs_list))),
+        ],
+    )
+    def test_grouping_strategies_single_hamiltonian(self, make_H):
+        """Tests that a single Hamiltonian or Sum is split correctly"""
 
-        # test transform on the batch
-        new_batch, post_proc_fn = split_non_commuting(batch)
+        coeffs, obs_list = [0.1, 0.2, 0.3, 0.4, 0.5], single_term_obs_list
+        qwc_groups = single_term_qwc_groups
 
-        # test that transform has been applied correctly on the batch by explicitly comparing with splitted tapes
-        tp1 = qml.tape.QuantumScript([qml.RX(1.2, 0)], [qml.expval(qml.X(0)), qml.expval(qml.X(1))])
-        tp2 = qml.tape.QuantumScript([qml.RX(1.2, 0)], [qml.expval(qml.Y(0))])
-        tp3 = qml.tape.QuantumScript([qml.RY(0.5, 0)], [qml.expval(qml.Z(0))])
-        tp4 = qml.tape.QuantumScript([qml.RY(0.5, 0)], [qml.expval(qml.Y(0))])
+        if not qml.operation.active_new_opmath():
+            obs_list = _convert_obs_to_legacy_opmath(obs_list)
+            qwc_groups = _convert_obs_to_legacy_opmath(single_term_qwc_groups)
 
-        assert all(qml.equal(tapeA, tapeB) for tapeA, tapeB in zip(new_batch, [tp1, tp2, tp3, tp4]))
+        expected_tapes_no_grouping = [
+            qml.tape.QuantumScript([], [qml.expval(o)], shots=100) for o in obs_list
+        ]
 
-        # test postprocessing function applied to the transformed batch
-        assert all(
-            qml.equal(tapeA, tapeB)
-            for sublist1, sublist2 in zip(post_proc_fn(new_batch), ((tp1, tp2), (tp3, tp4)))
-            for tapeA, tapeB in zip(sublist1, sublist2)
-        )
+        expected_tapes_qwc_grouping = [
+            qml.tape.QuantumScript([], [qml.expval(o) for o in group], shots=100)
+            for group in qwc_groups
+        ]
 
-        # final (double) check: test postprocessing function on a fictitious results
-        result = ("tp1", "tp2", "tp3", "tp4")
-        assert post_proc_fn(result) == (("tp1", "tp2"), ("tp3", "tp4"))
+        if qml.operation.active_new_opmath():
+            coeffs, obs_list = coeffs + [0.6], obs_list + [qml.I()]
 
-    def test_sprod_support(self):
-        """Test that split_non_commuting works with scalar product pauli words."""
+        H = make_H(coeffs, obs_list)  # Tests that constant offset is handled
 
-        ob1 = 2.0 * qml.prod(qml.X(0), qml.X(1))
-        ob2 = 3.0 * qml.prod(qml.Y(0), qml.Y(1))
-        ob3 = qml.s_prod(4.0, qml.X(1))
+        if not qml.operation.active_new_opmath() and isinstance(H, qml.ops.Sum):
+            pytest.skip("Sum is not part of legacy opmath")
 
-        tape = qml.tape.QuantumScript([], [qml.expval(o) for o in [ob1, ob2, ob3]])
-        batch, fn = qml.transforms.split_non_commuting(tape)
+        tape = qml.tape.QuantumScript([], [qml.expval(H)], shots=100)
 
-        tape0 = qml.tape.QuantumScript([], [qml.expval(ob2)])
-        assert qml.equal(tape0, batch[0])
-        tape1 = qml.tape.QuantumScript([], [qml.expval(ob1), qml.expval(ob3)])
-        assert qml.equal(tape1, batch[1])
+        tapes, fn = split_non_commuting(tape, grouping_strategy=None)
+        for actual_tape, expected_tape in zip(tapes, expected_tapes_no_grouping):
+            assert qml.equal(actual_tape, expected_tape)
+        expected = 0.55 if not qml.operation.active_new_opmath() else 1.15
+        assert qml.math.allclose(fn([0.1, 0.2, 0.3, 0.4, 0.5]), expected)
 
-        in_res = (1.0, (2.0, 3.0))
-        assert fn(in_res) == (2.0, 1.0, 3.0)
+        tapes, fn = split_non_commuting(tape, grouping_strategy="default")
+        for actual_tape, expected_tape in zip(tapes, expected_tapes_qwc_grouping):
+            assert qml.equal(actual_tape, expected_tape)
+        expected = 0.52 if not qml.operation.active_new_opmath() else 1.12
+        assert qml.math.allclose(fn([[0.1, 0.2], [0.3, 0.4, 0.5]]), expected)
 
+    @pytest.mark.parametrize(
+        "grouping_strategy, expected_tapes, processing_fn, mock_results",
+        [
+            (
+                None,
+                [
+                    qml.tape.QuantumScript([], [qml.expval(o)], shots=100)
+                    for o in complex_no_grouping_obs
+                ],
+                complex_no_grouping_processing_fn,
+                [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
+            ),
+            (
+                "wires",
+                [
+                    qml.tape.QuantumScript([], [qml.expval(o) for o in group], shots=100)
+                    for group in complex_wires_groups
+                ],
+                complex_wires_processing_fn,
+                [[0.1, 0.2], [0.3, 0.4], 0.5, 0.6],
+            ),
+            (
+                "qwc",
+                [
+                    qml.tape.QuantumScript([], [qml.expval(o) for o in group], shots=100)
+                    for group in complex_qwc_groups
+                ],
+                complex_qwc_processing_fn,
+                [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
+            ),
+        ],
+    )
+    def test_grouping_strategies_complex(
+        self, grouping_strategy, expected_tapes, processing_fn, mock_results
+    ):
+        """Tests that the tape is split correctly when containing more complex observables"""
 
-# measurements that require shots=True
-required_shot_meas_fn = [qml.sample, qml.counts]
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
-# measurements that can optionally have shots=True
-optional_shot_meas_fn = [qml.probs, qml.expval, qml.var]
+        measurements = [qml.expval(o) for o in obs_list]
+        tape = qml.tape.QuantumScript([], measurements, shots=100)
+        tapes, fn = split_non_commuting(tape, grouping_strategy=grouping_strategy)
 
+        for actual_tape, expected_tape in zip(tapes, expected_tapes):
+            assert qml.equal(actual_tape, expected_tape)
 
-class TestIntegration:
-    """Integration tests for ``qml.transforms.split_non_commuting()``"""
+        expected = processing_fn(mock_results)
+        if not qml.operation.active_new_opmath():
+            expected = expected[:-1]  # exclude the identity term
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    def test_expval_non_commuting_observables(self, convert_to_opmath):
-        """Test expval with multiple non-commuting operators"""
-        dev = qml.device("default.qubit", wires=6)
+        assert qml.math.allclose(fn(mock_results), expected)
 
-        prod_type = qml.prod if convert_to_opmath else qml.operation.Tensor
+    @pytest.mark.parametrize("batch_type", (tuple, list))
+    def test_batch_of_tapes(self, batch_type):
+        """Test that `split_non_commuting` can transform a batch of tapes"""
 
-        @qml.qnode(dev)
-        def circuit():
-            qml.Hadamard(1)
-            qml.Hadamard(0)
-            qml.PauliZ(0)
-            qml.Hadamard(3)
-            qml.Hadamard(5)
-            qml.T(5)
-            return (
-                qml.expval(prod_type(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliX(0)),
-                qml.expval(qml.PauliZ(1)),
-                qml.expval(prod_type(qml.X(1), qml.X(4))),
-                qml.expval(qml.PauliX(3)),
-                qml.expval(qml.PauliY(5)),
-            )
+        tape_batch = batch_type(
+            [
+                qml.tape.QuantumScript(
+                    [qml.RX(1.2, 0)],
+                    [qml.expval(qml.X(0)), qml.expval(qml.Y(0)), qml.expval(qml.X(1))],
+                ),
+                qml.tape.QuantumScript(
+                    [qml.RY(0.5, 0)], [qml.expval(qml.Z(0)), qml.expval(qml.Y(0))]
+                ),
+            ]
+        )
+        tapes, fn = split_non_commuting(tape_batch)
 
-        res = circuit()
-        assert isinstance(res, tuple)
-        assert len(res) == 6
-        assert all(isinstance(r, np.ndarray) for r in res)
-        assert all(r.shape == () for r in res)
+        expected_tapes = [
+            qml.tape.QuantumScript([qml.RX(1.2, 0)], [qml.expval(qml.X(0)), qml.expval(qml.X(1))]),
+            qml.tape.QuantumScript([qml.RX(1.2, 0)], [qml.expval(qml.Y(0))]),
+            qml.tape.QuantumScript([qml.RY(0.5, 0)], [qml.expval(qml.Z(0))]),
+            qml.tape.QuantumScript([qml.RY(0.5, 0)], [qml.expval(qml.Y(0))]),
+        ]
+        for actual_tape, expected_tape in zip(tapes, expected_tapes):
+            assert qml.equal(actual_tape, expected_tape)
 
-        res = qml.math.stack(res)
+        result = ([0.1, 0.2], 0.2, 0.3, 0.4)
+        assert fn(result) == ((0.1, 0.2, 0.2), (0.3, 0.4))
 
-        assert all(np.isclose(res, np.array([0.0, -1.0, 0.0, 0.0, 1.0, 1 / np.sqrt(2)])))
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    def test_expval_non_commuting_observables_qnode(self, convert_to_opmath):
-        """Test expval with multiple non-commuting operators as a transform program on the qnode."""
-        dev = qml.device("default.qubit", wires=6)
+class TestIntegration:
+    """Tests the ``split_non_commuting`` transform performed on a QNode"""
 
-        prod_type = qml.prod if convert_to_opmath else qml.operation.Tensor
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    @pytest.mark.parametrize("shots", [None, 20000, [20000, 30000, 40000]])
+    @pytest.mark.parametrize(
+        "params, expected_results",
+        [
+            (
+                [np.pi / 4, 3 * np.pi / 4],
+                [
+                    0.5,
+                    -np.cos(np.pi / 4),
+                    -0.5,
+                    -0.5 * np.cos(np.pi / 4),
+                    0.5 * np.cos(np.pi / 4),
+                ],
+            ),
+            (
+                [[np.pi / 4, 3 * np.pi / 4], [3 * np.pi / 4, 3 * np.pi / 4]],
+                [
+                    [0.5, -0.5],
+                    [-np.cos(np.pi / 4), -np.cos(np.pi / 4)],
+                    [-0.5, 0.5],
+                    [-0.5 * np.cos(np.pi / 4), 0.5 * np.cos(np.pi / 4)],
+                    [0.5 * np.cos(np.pi / 4), -0.5 * np.cos(np.pi / 4)],
+                ],
+            ),
+        ],
+    )
+    def test_single_expval(self, grouping_strategy, shots, params, expected_results):
+        """Tests that a QNode with a single expval measurement is executed correctly"""
 
-        @qml.qnode(dev)
-        def circuit():
-            qml.Hadamard(1)
-            qml.Hadamard(0)
-            qml.PauliZ(0)
-            qml.Hadamard(3)
-            qml.Hadamard(5)
-            qml.T(5)
-            return (
-                qml.expval(prod_type(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliX(0)),
-                qml.expval(qml.PauliZ(1)),
-                qml.expval(prod_type(qml.X(1), qml.X(4))),
-                qml.expval(qml.PauliX(3)),
-                qml.expval(qml.PauliY(5)),
-            )
+        coeffs, obs = [0.1, 0.2, 0.3, 0.4, 0.5], single_term_obs_list
 
-        res = split_non_commuting(circuit)()
+        if not qml.operation.active_new_opmath():
+            obs = _convert_obs_to_legacy_opmath(obs)
 
-        assert isinstance(res, tuple)
-        assert len(res) == 6
-        assert all(isinstance(r, np.ndarray) for r in res)
-        assert all(r.shape == () for r in res)
+        if qml.operation.active_new_opmath():
+            # test constant offset with new opmath
+            coeffs, obs = coeffs + [0.6], obs + [qml.I()]
 
-        res = qml.math.stack(res)
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        assert all(np.isclose(res, np.array([0.0, -1.0, 0.0, 0.0, 1.0, 1 / np.sqrt(2)])))
+        @qml.qnode(dev)
+        def circuit(angles):
+            qml.RX(angles[0], wires=0)
+            qml.RY(angles[1], wires=0)
+            qml.RX(angles[0], wires=1)
+            qml.RY(angles[1], wires=1)
+            return qml.expval(qml.Hamiltonian(coeffs, obs))
+
+        circuit = split_non_commuting(circuit, grouping_strategy=grouping_strategy)
+        res = circuit(params)
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    def test_expval_probs_non_commuting_observables_qnode(self, convert_to_opmath):
-        """Test expval with multiple non-commuting operators and probs with non-commuting wires as a
-        transform program on the qnode."""
-        dev = qml.device("default.qubit", wires=6)
+        if qml.operation.active_new_opmath():
+            identity_results = [1] if len(np.shape(params)) == 1 else [[1, 1]]
+            expected_results = expected_results + identity_results
 
-        prod_type = qml.prod if convert_to_opmath else qml.operation.Tensor
+        expected = np.dot(coeffs, expected_results)
 
-        @qml.qnode(dev)
-        def circuit():
-            qml.Hadamard(1)
-            qml.Hadamard(0)
-            qml.PauliZ(0)
-            qml.Hadamard(3)
-            qml.Hadamard(5)
-            qml.T(5)
-            return (
-                qml.probs(wires=[0, 1]),
-                qml.probs(wires=[1]),
-                qml.expval(qml.PauliZ(0)),
-                qml.expval(prod_type(qml.X(1), qml.X(4))),
-                qml.expval(qml.PauliX(3)),
-                qml.expval(qml.PauliY(5)),
-            )
+        if isinstance(shots, list):
+            assert qml.math.shape(res) == (3,) if len(np.shape(res)) == 1 else (3, 2)
+            for i in range(3):
+                assert qml.math.allclose(res[i], expected, atol=0.05)
+        else:
+            assert qml.math.allclose(res, expected, atol=0.05)
 
-        res = split_non_commuting(circuit)()
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    @pytest.mark.parametrize("shots", [None, 20000, [20000, 30000, 40000]])
+    @pytest.mark.parametrize(
+        "params, expected_results",
+        [
+            (
+                [np.pi / 4, 3 * np.pi / 4],
+                [
+                    0.5,
+                    -0.5 * np.cos(np.pi / 4),
+                    0.5 + np.cos(np.pi / 4) * 0.5 + 2.0 * 0.5 + 1,
+                    np.dot(
+                        [0.1, 0.2, 0.3, 0.4],
+                        [-0.5, -0.5 * np.cos(np.pi / 4), 0.5 * np.cos(np.pi / 4), 1],
+                    ),
+                    1.5,
+                ],
+            ),
+            (
+                [[np.pi / 4, 3 * np.pi / 4], [3 * np.pi / 4, 3 * np.pi / 4]],
+                [
+                    [0.5, -0.5],
+                    [-0.5 * np.cos(np.pi / 4), -0.5 * np.cos(np.pi / 4)],
+                    [
+                        0.5 + np.cos(np.pi / 4) * 0.5 + 2.0 * 0.5 + 1,
+                        -0.5 - np.cos(np.pi / 4) * 0.5 - 2.0 * 0.5 + 1,
+                    ],
+                    [
+                        np.dot(
+                            [0.1, 0.2, 0.3, 0.4],
+                            [-0.5, -0.5 * np.cos(np.pi / 4), 0.5 * np.cos(np.pi / 4), 1],
+                        ),
+                        np.dot(
+                            [0.1, 0.2, 0.3, 0.4],
+                            [0.5, 0.5 * np.cos(np.pi / 4), -0.5 * np.cos(np.pi / 4), 1],
+                        ),
+                    ],
+                    [1.5, 1.5],
+                ],
+            ),
+        ],
+    )
+    def test_multiple_expval(self, grouping_strategy, shots, params, expected_results):
+        """Tests that a QNode with multiple expval measurements is executed correctly"""
 
-        assert isinstance(res, tuple)
-        assert len(res) == 6
-        assert all(isinstance(r, np.ndarray) for r in res)
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        res_probs = qml.math.concatenate(res[:2], axis=0)
-        res_expval = qml.math.stack(res[2:])
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
-        assert all(np.isclose(res_probs, np.array([0.25, 0.25, 0.25, 0.25, 0.5, 0.5])))
+        @qml.qnode(dev)
+        def circuit(angles):
+            qml.RX(angles[0], wires=0)
+            qml.RY(angles[1], wires=0)
+            qml.RX(angles[0], wires=1)
+            qml.RY(angles[1], wires=1)
+            return [qml.expval(obs) for obs in obs_list]
+
+        circuit = split_non_commuting(circuit, grouping_strategy=grouping_strategy)
+        res = circuit(params)
 
-        assert all(np.isclose(res_expval, np.array([0.0, 0.0, 1.0, 1 / np.sqrt(2)])))
+        if not qml.operation.active_new_opmath():
+            expected_results = expected_results[:-1]  # exclude the identity term
 
-    @pytest.mark.parametrize("convert_to_opmath", (True, False))
-    def test_shot_vector_support(self, convert_to_opmath):
-        """Test output is correct when using shot vectors"""
+        if isinstance(shots, list):
+            assert qml.math.shape(res) == (3, *np.shape(expected_results))
+            for i in range(3):
+                assert qml.math.allclose(res[i], expected_results, atol=0.05)
+        else:
+            assert qml.math.allclose(res, expected_results, atol=0.05)
+
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    @pytest.mark.parametrize("shots", [20000, [20000, 30000, 40000]])
+    @pytest.mark.parametrize(
+        "params, expected_results",
+        [
+            (
+                [np.pi / 4, 3 * np.pi / 4],
+                [
+                    0.5,
+                    -0.5 * np.cos(np.pi / 4),
+                    0.5 + np.cos(np.pi / 4) * 0.5 + 2.0 * 0.5 + 1,
+                    np.dot(
+                        [0.1, 0.2, 0.3, 0.4],
+                        [-0.5, -0.5 * np.cos(np.pi / 4), 0.5 * np.cos(np.pi / 4), 1],
+                    ),
+                    1.5,
+                ],
+            ),
+            (
+                [[np.pi / 4, 3 * np.pi / 4], [3 * np.pi / 4, 3 * np.pi / 4]],
+                [
+                    [0.5, -0.5],
+                    [-0.5 * np.cos(np.pi / 4), -0.5 * np.cos(np.pi / 4)],
+                    [
+                        0.5 + np.cos(np.pi / 4) * 0.5 + 2.0 * 0.5 + 1,
+                        -0.5 - np.cos(np.pi / 4) * 0.5 - 2.0 * 0.5 + 1,
+                    ],
+                    [
+                        np.dot(
+                            [0.1, 0.2, 0.3, 0.4],
+                            [-0.5, -0.5 * np.cos(np.pi / 4), 0.5 * np.cos(np.pi / 4), 1],
+                        ),
+                        np.dot(
+                            [0.1, 0.2, 0.3, 0.4],
+                            [0.5, 0.5 * np.cos(np.pi / 4), -0.5 * np.cos(np.pi / 4), 1],
+                        ),
+                    ],
+                    [1.5, 1.5],
+                ],
+            ),
+        ],
+    )
+    def test_mixed_measurement_types(self, grouping_strategy, shots, params, expected_results):
+        """Tests that a QNode with mixed measurement types is executed correctly"""
 
-        dev = qml.device("default.qubit", wires=6, shots=(10000, (20000, 2), 30000))
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        prod_type = qml.prod if convert_to_opmath else qml.operation.Tensor
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
         @qml.qnode(dev)
-        def circuit():
-            qml.Hadamard(1)
-            qml.Hadamard(0)
-            qml.PauliZ(0)
-            qml.Hadamard(3)
-            qml.Hadamard(5)
-            qml.T(5)
+        def circuit(angles):
+            qml.RX(angles[0], wires=0)
+            qml.RY(angles[1], wires=0)
+            qml.RX(angles[0], wires=1)
+            qml.RY(angles[1], wires=1)
             return (
-                qml.expval(prod_type(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliX(0)),
-                qml.expval(qml.PauliZ(1)),
-                qml.expval(prod_type(qml.Y(0), qml.Y(1), qml.Z(3), qml.Y(4), qml.X(5))),
-                qml.expval(prod_type(qml.X(1), qml.X(4))),
-                qml.expval(qml.PauliX(3)),
-                qml.expval(qml.PauliY(5)),
+                qml.probs(wires=0),
+                qml.probs(wires=[0, 1]),
+                qml.counts(wires=0),
+                qml.sample(wires=0),
+                *[qml.expval(obs) for obs in obs_list],
             )
 
-        res = circuit()
-        assert isinstance(res, tuple)
-        assert len(res) == 4
-        assert all(isinstance(shot_res, tuple) for shot_res in res)
-        assert all(len(shot_res) == 7 for shot_res in res)
-        # pylint:disable=not-an-iterable
-        assert all(
-            all(list(isinstance(r, np.ndarray) and r.shape == () for r in shot_res))
-            for shot_res in res
-        )
+        circuit = split_non_commuting(circuit, grouping_strategy=grouping_strategy)
+        res = circuit(params)
 
-        res = qml.math.stack([qml.math.stack(r) for r in res])
+        if not qml.operation.active_new_opmath():
+            expected_results = expected_results[:-1]  # exclude the identity term
+
+        if isinstance(shots, list):
+            assert len(res) == 3
+            for i in range(3):
+
+                prob_res_0 = res[i][0]
+                prob_res_1 = res[i][1]
+                counts_res = res[i][2]
+                sample_res = res[i][3]
+                if len(qml.math.shape(params)) == 1:
+                    assert qml.math.shape(prob_res_0) == (2,)
+                    assert qml.math.shape(prob_res_1) == (4,)
+                    assert isinstance(counts_res, dict)
+                    assert qml.math.shape(sample_res) == (shots[i],)
+                else:
+                    assert qml.math.shape(prob_res_0) == (2, 2)
+                    assert qml.math.shape(prob_res_1) == (2, 4)
+                    assert all(isinstance(_res, dict) for _res in counts_res)
+                    assert qml.math.shape(sample_res) == (2, shots[i])
+
+                expval_res = res[i][4:]
+                assert qml.math.allclose(expval_res, expected_results, atol=0.05)
+        else:
+            prob_res_0 = res[0]
+            prob_res_1 = res[1]
+            counts_res = res[2]
+            sample_res = res[3]
+            if len(qml.math.shape(params)) == 1:
+                assert qml.math.shape(prob_res_0) == (2,)
+                assert qml.math.shape(prob_res_1) == (4,)
+                assert isinstance(counts_res, dict)
+                assert qml.math.shape(sample_res) == (shots,)
+            else:
+                assert qml.math.shape(prob_res_0) == (2, 2)
+                assert qml.math.shape(prob_res_1) == (2, 4)
+                assert all(isinstance(_res, dict) for _res in counts_res)
+                assert qml.math.shape(sample_res) == (2, shots)
 
-        assert np.allclose(
-            res, np.array([0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 1 / np.sqrt(2)]), atol=0.05
-        )
+            expval_res = res[4:]
+            assert qml.math.allclose(expval_res, expected_results, atol=0.05)
 
-    def test_shot_vector_support_sample(self):
-        """Test output is correct when using shots and sample and expval measurements"""
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_single_hamiltonian_only_constant_offset(self, grouping_strategy):
+        """Tests that split_non_commuting can handle a single Identity observable"""
 
-        dev = qml.device("default.qubit", wires=2, shots=(10, 20))
+        dev = qml.device("default.qubit", wires=2)
+        H = qml.Hamiltonian([1.5, 2.5], [qml.I(), qml.I()])
 
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
         @qml.qnode(dev)
         def circuit():
-            qml.PauliZ(0)
-            return (qml.sample(wires=[0, 1]), qml.expval(qml.PauliZ(0)))
-
-        res = split_non_commuting(circuit)()
-        assert isinstance(res, tuple)
-        assert len(res) == 2
-        assert all(isinstance(shot_res, tuple) for shot_res in res)
-        assert all(len(shot_res) == 2 for shot_res in res)
-        # pylint:disable=not-an-iterable
-        assert all(all(list(isinstance(r, np.ndarray) for r in shot_res)) for shot_res in res)
-
-        assert all(
-            shot_res[0].shape in [(10, 2), (20, 2)] and shot_res[1].shape == () for shot_res in res
-        )
+            return qml.expval(H)
 
-        # check all the wire samples are as expected
-        sample_res = qml.math.concatenate(
-            [qml.math.concatenate(shot_res[0], axis=0) for shot_res in res], axis=0
-        )
-        assert np.allclose(sample_res, 0.0, atol=0.05)
-
-        expval_res = qml.math.stack([shot_res[1] for shot_res in res])
-        assert np.allclose(expval_res, np.array([1.0, 1.0]), atol=0.05)
+        with dev.tracker:
+            res = circuit()
+        assert dev.tracker.totals == {}
+        assert qml.math.allclose(res, 4.0)
 
-    def test_shot_vector_support_counts(self):
-        """Test output is correct when using shots, counts and expval measurements"""
+    @pytest.mark.usefixtures("new_opmath_only")
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_no_obs_tape(self, grouping_strategy):
+        """Tests tapes with only constant offsets (only measurements on Identity)"""
 
-        dev = qml.device("default.qubit", wires=2, shots=(10, 20))
+        _dev = qml.device("default.qubit", wires=1)
 
-        @qml.qnode(dev)
+        @qml.qnode(_dev)
         def circuit():
-            qml.PauliZ(0)
-            return (qml.counts(wires=[0, 1]), qml.expval(qml.PauliZ(0)))
-
-        res = split_non_commuting(circuit)()
-        assert isinstance(res, tuple)
-        assert len(res) == 2
-        assert all(isinstance(shot_res, tuple) for shot_res in res)
-        assert all(len(shot_res) == 2 for shot_res in res)
-        # pylint:disable=not-an-iterable
-        assert all(
-            isinstance(shot_res[0], dict) and isinstance(shot_res[1], np.ndarray)
-            for shot_res in res
-        )
+            return qml.expval(1.5 * qml.I(0))
 
-        assert all(shot_res[1].shape == () for shot_res in res)
+        circuit = split_non_commuting(circuit, grouping_strategy=grouping_strategy)
 
-        # check all the wire counts are as expected
-        assert all(shot_res[0]["00"] in [10, 20] for shot_res in res)
+        with _dev.tracker:
+            res = circuit()
 
-        expval_res = qml.math.stack([shot_res[1] for shot_res in res])
-        assert np.allclose(expval_res, np.array([1.0, 1.0]), atol=0.05)
+        assert _dev.tracker.totals == {}
+        assert qml.math.allclose(res, 1.5)
 
+    @pytest.mark.usefixtures("new_opmath_only")
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_no_obs_tape_multi_measurement(self, grouping_strategy):
+        """Tests tapes with only constant offsets (only measurements on Identity)"""
 
-# Autodiff tests
-exp_res = np.array([0.77015115, -0.47942554, 0.87758256])
-exp_grad = np.array(
-    [[-4.20735492e-01, -4.20735492e-01], [-8.77582562e-01, 0.0], [-4.79425539e-01, 0.0]]
-)
+        _dev = qml.device("default.qubit", wires=1)
 
-exp_res_probs = np.array([0.88132907, 0.05746221, 0.05746221, 0.00374651, 0.0])
-exp_grad_probs = np.array(
-    [
-        [-0.22504026, -0.22504026],
-        [-0.01467251, 0.22504026],
-        [0.22504026, -0.01467251],
-        [0.01467251, 0.01467251],
-        [0.0, 0.0],
-    ]
-)
+        @qml.qnode(_dev)
+        def circuit():
+            return qml.expval(1.5 * qml.I()), qml.expval(2.5 * qml.I())
 
+        circuit = split_non_commuting(circuit, grouping_strategy=grouping_strategy)
 
-class TestAutodiffSplitNonCommuting:
-    """Autodiff tests for all frameworks"""
+        with _dev.tracker:
+            res = circuit()
 
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    @pytest.mark.autograd
-    def test_split_with_autograd(self, prod):
-        """Test that results after splitting are still differentiable with autograd"""
-        dev = qml.device("default.qubit", wires=3)
+        assert _dev.tracker.totals == {}
+        assert qml.math.allclose(res, [1.5, 2.5])
 
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (
-                qml.expval(prod(qml.PauliZ(0), qml.PauliZ(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
 
-        def cost(params):
-            res = circuit(params)
-            return qml.math.stack(res)
-
-        params = pnp.array([0.5, 0.5])
-        res = cost(params)
-        grad = qml.jacobian(cost)(params)
-        assert all(np.isclose(res, exp_res))
-        assert all(np.isclose(grad, exp_grad).flatten())
+expected_grad_param_0 = [
+    0.125,
+    0.125,
+    0.125,
+    -0.375,
+    -0.5,
+    0.5 * -np.cos(np.pi / 4),
+    -0.5 - 2.0 * 0.5,
+    0.1 * 0.5,
+    0,
+]
 
-    @pytest.mark.autograd
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_autograd_probs(self, prod):
-        """Test resulting after splitting non-commuting tapes with expval and probs measurements
-        are still differentiable with autograd"""
-        dev = qml.device("default.qubit", wires=2)
+expected_grad_param_1 = [
+    -0.125,
+    -0.125,
+    -0.125,
+    0.375,
+    -0.5,
+    0,
+    -0.5 + np.cos(np.pi / 4) / 2 - 2.0 * 0.5,
+    np.dot([0.1, 0.2, 0.3], [-0.5, np.cos(np.pi / 4) / 2, np.cos(np.pi / 4) / 2]),
+    0,
+]
 
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return qml.probs(wires=[0, 1]), qml.expval(prod(qml.PauliX(0), qml.PauliX(1)))
 
-        def cost(params):
-            res = split_non_commuting(circuit)(params)
-            return qml.math.concatenate([res[0]] + [qml.math.stack(res[1:])], axis=0)
+class TestDifferentiability:
+    """Tests the differentiability of the ``split_non_commuting`` transform"""
 
-        params = pnp.array([0.5, 0.5])
-        res = cost(params)
-        grad = qml.jacobian(cost)(params)
-        assert all(np.isclose(res, exp_res_probs))
-        assert all(np.isclose(grad, exp_grad_probs).flatten())
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_autograd(self, grouping_strategy):
+        """Tests that the output of ``split_non_commuting`` is differentiable with autograd"""
 
-    @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax(self, prod):
-        """Test that results after splitting are still differentiable with jax"""
+        import pennylane.numpy as pnp
 
-        import jax
-        import jax.numpy as jnp
+        dev = qml.device("default.qubit", wires=2)
 
-        dev = qml.device("default.qubit.jax", wires=3)
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
         @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (
-                qml.expval(prod(qml.PauliZ(0), qml.PauliZ(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
+        def circuit(theta, phi):
+            qml.RX(theta, wires=0)
+            qml.RY(phi, wires=0)
+            qml.RX(theta, wires=1)
+            qml.RY(phi, wires=1)
+            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
-        params = jnp.array([0.5, 0.5])
-        res = split_non_commuting(circuit)(params)
-        grad = jax.jacobian(split_non_commuting(circuit))(params)
-        assert all(np.isclose(res, exp_res, atol=0.05))
-        assert all(np.isclose(grad, exp_grad, atol=0.05).flatten())
+        def cost(theta, phi):
+            res = circuit(theta, phi)
+            return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
 
-    @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax_probs(self, prod):
-        """Test resulting after splitting non-commuting tapes with expval and probs measurements
-        are still differentiable with jax"""
-        import jax
-        import jax.numpy as jnp
+        params = pnp.array(pnp.pi / 4), pnp.array(3 * pnp.pi / 4)
+        grad1, grad2 = qml.jacobian(cost)(*params)
 
-        dev = qml.device("default.qubit.jax", wires=2)
-
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (qml.probs(wires=[0, 1]), qml.expval(prod(qml.X(0), qml.X(1))))
+        expected_grad_1 = expected_grad_param_0
+        expected_grad_2 = expected_grad_param_1
+        if not qml.operation.active_new_opmath():
+            expected_grad_1 = expected_grad_param_0[:-1]
+            expected_grad_2 = expected_grad_param_1[:-1]
 
-        params = jnp.array([0.5, 0.5])
-        res = split_non_commuting(circuit)(params)
-        res = jnp.concatenate([res[0]] + [jnp.stack(res[1:])], axis=0)
-
-        grad = jax.jacobian(split_non_commuting(circuit))(params)
-        grad = jnp.concatenate([grad[0]] + [jnp.stack(grad[1:])], axis=0)
-
-        assert all(np.isclose(res, exp_res_probs, atol=0.05))
-        assert all(np.isclose(grad, exp_grad_probs, atol=0.05).flatten())
+        assert qml.math.allclose(grad1, expected_grad_1)
+        assert qml.math.allclose(grad2, expected_grad_2)
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax_multi_params(self, prod):
-        """Test that results after splitting are still differentiable with jax
-        with multiple parameters"""
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_jax(self, grouping_strategy):
+        """Tests that the output of ``split_non_commuting`` is differentiable with jax"""
 
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.jax", wires=3)
-
-        @qml.qnode(dev)
-        def circuit(x, y):
-            qml.RX(x, wires=0)
-            qml.RY(y, wires=1)
-            return (
-                qml.expval(prod(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
-
-        x = jnp.array(0.5)
-        y = jnp.array(0.5)
-
-        res = split_non_commuting(circuit)(x, y)
-        grad = jax.jacobian(split_non_commuting(circuit), argnums=[0, 1])(x, y)
-
-        assert all(np.isclose(res, exp_res))
-
-        assert isinstance(grad, tuple)
-        assert len(grad) == 3
-
-        for i, meas_grad in enumerate(grad):
-            assert isinstance(meas_grad, tuple)
-            assert len(meas_grad) == 2
-            assert all(isinstance(g, jnp.ndarray) and g.shape == () for g in meas_grad)
-
-            assert np.allclose(meas_grad, exp_grad[i], atol=1e-5)
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax_multi_params_probs(self, prod):
-        """Test that results after splitting are still differentiable with jax
-        with multiple parameters"""
-
-        import jax
-        import jax.numpy as jnp
+        dev = qml.device("default.qubit", wires=2)
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
         @qml.qnode(dev)
-        def circuit(x, y):
-            qml.RX(x, wires=0)
-            qml.RY(y, wires=1)
-            return (qml.probs(wires=[0, 1]), qml.expval(prod(qml.X(0), qml.X(1))))
+        def circuit(theta, phi):
+            qml.RX(theta, wires=0)
+            qml.RY(phi, wires=0)
+            qml.RX(theta, wires=1)
+            qml.RY(phi, wires=1)
+            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
-        x = jnp.array(0.5)
-        y = jnp.array(0.5)
+        def cost(theta, phi):
+            res = circuit(theta, phi)
+            return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
 
-        res = split_non_commuting(circuit)(x, y)
-        res = jnp.concatenate([res[0]] + [jnp.stack(res[1:])], axis=0)
-        assert all(np.isclose(res, exp_res_probs))
+        params = jnp.array(jnp.pi / 4), jnp.array(3 * jnp.pi / 4)
+        grad1, grad2 = jax.jacobian(cost, argnums=[0, 1])(*params)
 
-        grad = jax.jacobian(split_non_commuting(circuit), argnums=[0, 1])(x, y)
+        expected_grad_1 = expected_grad_param_0
+        expected_grad_2 = expected_grad_param_1
+        if not qml.operation.active_new_opmath():
+            expected_grad_1 = expected_grad_param_0[:-1]
+            expected_grad_2 = expected_grad_param_1[:-1]
 
-        assert isinstance(grad, tuple)
-        assert len(grad) == 2
-
-        for meas_grad in grad:
-            assert isinstance(meas_grad, tuple)
-            assert len(meas_grad) == 2
-            assert all(isinstance(g, jnp.ndarray) for g in meas_grad)
-
-        # reshape the returned gradient to the right shape
-        grad = jnp.concatenate(
-            [
-                jnp.concatenate([grad[0][0].reshape(-1, 1), grad[0][1].reshape(-1, 1)], axis=1),
-                jnp.concatenate([grad[1][0].reshape(-1, 1), grad[1][1].reshape(-1, 1)], axis=1),
-            ],
-            axis=0,
-        )
-        assert all(np.isclose(grad, exp_grad_probs, atol=0.05).flatten())
+        assert qml.math.allclose(grad1, expected_grad_1)
+        assert qml.math.allclose(grad2, expected_grad_2)
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax_jit(self, prod):
-        """Test that results after splitting are still differentiable with jax-jit"""
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_jax_jit(self, grouping_strategy):
+        """Tests that the output of ``split_non_commuting`` is differentiable with jax and jit"""
 
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit", wires=2)
+
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
         @jax.jit
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
         @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (
-                qml.expval(prod(qml.Z(0), qml.PauliZ(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
-
-        params = jnp.array([0.5, 0.5])
-        res = circuit(params)
-        grad = jax.jacobian(circuit)(params)
-        assert all(np.isclose(res, exp_res))
-        assert all(np.isclose(grad, exp_grad, atol=1e-5).flatten())
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_jax_jit_probs(self, prod):
-        """Test resulting after splitting non-commuting tapes with expval and probs measurements
-        are still differentiable with jax"""
+        def circuit(theta, phi):
+            qml.RX(theta, wires=0)
+            qml.RY(phi, wires=0)
+            qml.RX(theta, wires=1)
+            qml.RY(phi, wires=1)
+            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
-        import jax
-        import jax.numpy as jnp
-
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (qml.probs(wires=[0, 1]), qml.expval(prod(qml.X(0), qml.PauliX(1))))
+        def cost(theta, phi):
+            res = circuit(theta, phi)
+            return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
 
-        params = jnp.array([0.5, 0.5])
-        res = split_non_commuting(circuit)(params)
-        res = jnp.concatenate([res[0]] + [jnp.stack(res[1:])], axis=0)
+        params = jnp.array(jnp.pi / 4), jnp.array(3 * jnp.pi / 4)
+        grad1, grad2 = jax.jacobian(cost, argnums=[0, 1])(*params)
 
-        grad = jax.jacobian(split_non_commuting(circuit))(params)
-        grad = jnp.concatenate([grad[0]] + [jnp.stack(grad[1:])], axis=0)
+        expected_grad_1 = expected_grad_param_0
+        expected_grad_2 = expected_grad_param_1
+        if not qml.operation.active_new_opmath():
+            expected_grad_1 = expected_grad_param_0[:-1]
+            expected_grad_2 = expected_grad_param_1[:-1]
 
-        assert all(np.isclose(res, exp_res_probs, atol=0.05))
-        assert all(np.isclose(grad, exp_grad_probs, atol=0.05).flatten())
+        assert qml.math.allclose(grad1, expected_grad_1)
+        assert qml.math.allclose(grad2, expected_grad_2)
 
     @pytest.mark.torch
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_torch(self, prod):
-        """Test that results after splitting are still differentiable with torch"""
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_torch(self, grouping_strategy):
+        """Tests that the output of ``split_non_commuting`` is differentiable with torch"""
 
         import torch
         from torch.autograd.functional import jacobian
 
-        dev = qml.device("default.qubit.torch", wires=3)
-
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (
-                qml.expval(prod(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
-
-        def cost(params):
-            res = split_non_commuting(circuit)(params)
-            return qml.math.stack(res)
-
-        params = torch.tensor([0.5, 0.5], requires_grad=True)
-        res = cost(params)
-        grad = jacobian(cost, (params))
-        assert all(np.isclose(res.detach().numpy(), exp_res))
-        assert all(np.isclose(grad.detach().numpy(), exp_grad, atol=1e-5).flatten())
+        dev = qml.device("default.qubit", wires=2)
 
-    @pytest.mark.torch
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_torch_probs(self, prod):
-        """Test resulting after splitting non-commuting tapes with expval and probs measurements
-        are still differentiable with torch"""
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
-        import torch
-        from torch.autograd.functional import jacobian
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
+        @qml.qnode(dev)
+        def circuit(theta, phi):
+            qml.RX(theta, wires=0)
+            qml.RY(phi, wires=0)
+            qml.RX(theta, wires=1)
+            qml.RY(phi, wires=1)
+            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
-        dev = qml.device("default.qubit", wires=2)
+        def cost(theta, phi):
+            res = circuit(theta, phi)
+            return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
 
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (qml.probs(wires=[0, 1]), qml.expval(prod(qml.X(0), qml.X(1))))
+        params = torch.tensor(np.pi / 4), torch.tensor(3 * np.pi / 4)
+        grad1, grad2 = jacobian(cost, params)
 
-        def cost(params):
-            res = split_non_commuting(circuit)(params)
-            return qml.math.concatenate([res[0]] + [qml.math.stack(res[1:])], axis=0)
+        expected_grad_1 = expected_grad_param_0
+        expected_grad_2 = expected_grad_param_1
+        if not qml.operation.active_new_opmath():
+            expected_grad_1 = expected_grad_param_0[:-1]
+            expected_grad_2 = expected_grad_param_1[:-1]
 
-        params = torch.tensor([0.5, 0.5], requires_grad=True)
-        res = cost(params)
-        grad = jacobian(cost, (params))
-        assert all(np.isclose(res.detach().numpy(), exp_res_probs))
-        assert all(np.isclose(grad.detach().numpy(), exp_grad_probs, atol=1e-5).flatten())
+        assert qml.math.allclose(grad1, expected_grad_1, atol=1e-5)
+        assert qml.math.allclose(grad2, expected_grad_2, atol=1e-5)
 
     @pytest.mark.tf
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    def test_split_with_tf(self, prod):
-        """Test that results after splitting are still differentiable with tf"""
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_tensorflow(self, grouping_strategy):
+        """Tests that the output of ``split_non_commuting`` is differentiable with torch"""
 
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=3)
-
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (
-                qml.expval(prod(qml.Z(0), qml.Z(1))),
-                qml.expval(qml.PauliY(0)),
-                qml.expval(qml.PauliZ(0)),
-            )
-
-        params = tf.Variable([0.5, 0.5])
-        res = circuit(params)
-        with tf.GradientTape() as tape:
-            loss = split_non_commuting(circuit)(params)
-            loss = tf.stack(loss)
+        dev = qml.device("default.qubit", wires=2)
 
-        grad = tape.jacobian(loss, params)
-        assert all(np.isclose(res, exp_res))
-        assert all(np.isclose(grad, exp_grad, atol=1e-5).flatten())
+        obs_list = complex_obs_list
+        if not qml.operation.active_new_opmath():
+            obs_list = obs_list[:-1]  # exclude the identity term
 
-    @pytest.mark.parametrize("prod", (qml.prod, qml.operation.Tensor))
-    @pytest.mark.tf
-    def test_split_with_tf_probs(self, prod):
-        """Test that results after splitting are still differentiable with tf"""
+        @qml.qnode(dev)
+        def circuit(theta, phi):
+            qml.RX(theta, wires=0)
+            qml.RY(phi, wires=0)
+            qml.RX(theta, wires=1)
+            qml.RY(phi, wires=1)
+            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
-        import tensorflow as tf
+        params = tf.Variable(np.pi / 4), tf.Variable(3 * np.pi / 4)
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        with tf.GradientTape() as tape:
+            res = split_non_commuting(circuit, grouping_strategy=grouping_strategy)(*params)
+            cost = qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
 
-        @qml.qnode(dev)
-        def circuit(params):
-            qml.RX(params[0], wires=0)
-            qml.RY(params[1], wires=1)
-            return (qml.probs(wires=[0, 1]), qml.expval(prod(qml.X(0), qml.X(1))))
+        grad1, grad2 = tape.jacobian(cost, params)
 
-        params = tf.Variable([0.5, 0.5])
-        with tf.GradientTape() as tape:
-            res = split_non_commuting(circuit)(params)
-            res = tf.concat([res[0]] + [tf.stack(res[1:])], axis=0)
+        expected_grad_1 = expected_grad_param_0
+        expected_grad_2 = expected_grad_param_1
+        if not qml.operation.active_new_opmath():
+            expected_grad_1 = expected_grad_param_0[:-1]
+            expected_grad_2 = expected_grad_param_1[:-1]
 
-        grad = tape.jacobian(res, params)
-        assert all(np.isclose(res, exp_res_probs))
-        assert all(np.isclose(grad, exp_grad_probs, atol=1e-5).flatten())
+        assert qml.math.allclose(grad1, expected_grad_1, atol=1e-5)
+        assert qml.math.allclose(grad2, expected_grad_2, atol=1e-5)