Repository: PennyLaneAI/pennylane
Pull request #5070 (2024-01-15 23:35:32)

Add grouping support for arithmetic ops
**Context:**
Adding support for `Prod` and `SProd` to be used with grouping.

**Description of the Change:**
* Update `Tensor.__matmul__` to stop mutating in-place and return a new op instead.
* Added Pauli rep to `Tensor`. Consequently, `Tensor.copy()` and `Tensor.map_wires()` also had to be updated.
* Moved around code inside `Tensor__init__` to fix incorrect validation after the previous changes. The warning about overlapping wires is now raised in `Tensor.queue()` instead of at the beginning of construction, so that raising the warning pre-maturely can be prevented.
* `qml.pauli.group_observables` now returns `Prod`s if the observable list contains new arithmetic ops.
* Update `pauli_to_binary` to rely on the Pauli rep. All operations that can be grouped should now always have a valid Pauli rep, so we don't need to worry about other cases, as `pauli_to_binary` is only a utility function for grouping.
* Update `are_pauli_words_qwc` to correctly handle ops with Pauli reps. This change however means that any new arithmetic ops that have `pauli_rep=None` will still not work with this function. If we eventually start doing more simplification eagerly, this can be addressed.
* `split_non_commuting` now creates `Prod`s instead of `Tensor`s for creating observables for measurements that use wires. Also made a change so that if a measurement without any wires is made, the observable that is created uses `tape.wires`.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


mudit2812: [sc-53504]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `3 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`c0ce3ce`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c0ce3ced1644cc552a3454fec5a09790ee5347fc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`830f1ca`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.57%.
> Report is 3 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/pauli/utils.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3BhdWxpL3V0aWxzLnB5) | 95.12% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | 95.23% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #5070      +/-   ##
==========================================
- Coverage   99.68%   99.57%   -0.12%     
==========================================
  Files         394      395       +1     
  Lines       35720    35676      -44     
==========================================
- Hits        35609    35523      -86     
- Misses        111      153      +42     
```



</details>

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

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-commenter: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`e73213e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e73213edc4dc08182dea69c4c252a76510d307f1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.69% compared to head [(`1567cda`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67%.
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | 92.30% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #5070      +/-   ##
==========================================
- Coverage   99.69%   99.67%   -0.02%     
==========================================
  Files         394      394              
  Lines       36151    35865     -286     
==========================================
- Hits        36040    35749     -291     
- Misses        111      116       +5     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e4ac1f558..4fdeb1b71 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -174,6 +174,9 @@
 
 <h4>Other improvements</h4>
 
+* `qml.pauli.group_observables` now supports grouping `Prod` and `SProd` operators.
+  [(#5070)](https://github.com/PennyLaneAI/pennylane/pull/5070)
+
 * Faster `qml.probs` measurements due to an optimization in `_samples_to_counts`.
   [(#5145)](https://github.com/PennyLaneAI/pennylane/pull/5145)
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index ca1664707..5f5b034ec 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -2054,13 +2054,6 @@ class Tensor(Observable):
         return cls(*data)
 
     def __init__(self, *args):  # pylint: disable=super-init-not-called
-        wires = [op.wires for op in args]
-        if len(wires) != len(set(wires)):
-            warnings.warn(
-                "Tensor object acts on overlapping wires; in some PennyLane functions this will lead to undefined behaviour",
-                UserWarning,
-            )
-
         self._eigvals_cache = None
         self.obs: List[Observable] = []
         self._args = args
@@ -2068,6 +2061,20 @@ class Tensor(Observable):
         self._pauli_rep = None
         self.queue(init=True)
 
+        wires = [op.wires for op in self.obs]
+        if len(wires) != len(set(wires)):
+            warnings.warn(
+                "Tensor object acts on overlapping wires; in some PennyLane functions this will "
+                "lead to undefined behaviour",
+                UserWarning,
+            )
+
+        # Queue before updating pauli_rep because self.queue updates self.obs
+        if all(prs := [o.pauli_rep for o in self.obs]):
+            self._pauli_rep = functools.reduce(lambda a, b: a @ b, prs)
+        else:
+            self._pauli_rep = None
+
     def label(self, decimals=None, base_label=None, cache=None):
         r"""How the operator is represented in diagrams and drawings.
 
@@ -2103,7 +2110,6 @@ class Tensor(Observable):
 
     def queue(self, context=QueuingManager, init=False):  # pylint: disable=arguments-differ
         constituents = self._args if init else self.obs
-
         for o in constituents:
             if init:
                 if isinstance(o, Tensor):
@@ -2124,6 +2130,7 @@ class Tensor(Observable):
         copied_op.obs = self.obs.copy()
         copied_op._eigvals_cache = self._eigvals_cache
         copied_op._batch_size = self._batch_size
+        copied_op._pauli_rep = self._pauli_rep
         return copied_op
 
     def __repr__(self):
@@ -2231,37 +2238,17 @@ class Tensor(Observable):
         if isinstance(other, qml.Hamiltonian):
             return other.__rmatmul__(self)
 
-        if isinstance(other, Tensor):
-            self.obs.extend(other.obs)
-
-        elif isinstance(other, Observable):
-            self.obs.append(other)
+        if isinstance(other, Observable):
+            return Tensor(self, other)
 
-        elif isinstance(other, Operator):
+        if isinstance(other, Operator):
             return qml.prod(*self.obs, other)
 
-        else:
-            return NotImplemented
-
-        wires = [op.wires for op in self.obs]
-        if len(wires) != len(set(wires)):
-            warnings.warn(
-                "Tensor object acts on overlapping wires; in some PennyLane functions this will lead to undefined behaviour",
-                UserWarning,
-            )
-
-        if QueuingManager.recording() and self not in QueuingManager.active_context():
-            QueuingManager.append(self)
-
-        QueuingManager.remove(other)
-
-        return self
+        return NotImplemented
 
     def __rmatmul__(self, other):
         if isinstance(other, Observable):
-            self.obs[:0] = [other]
-            QueuingManager.remove(other)
-            return self
+            return Tensor(other, self)
 
         return NotImplemented
 
@@ -2552,6 +2539,7 @@ class Tensor(Observable):
         new_op.obs = [obs.map_wires(wire_map) for obs in self.obs]
         new_op._eigvals_cache = self._eigvals_cache
         new_op._batch_size = self._batch_size
+        new_op._pauli_rep = self._pauli_rep.map_wires(wire_map)
         return new_op
 
 
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index a1a5579f6..b5fd238a4 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -48,16 +48,12 @@ def _compute_grouping_indices(observables, grouping_type="qwc", method="rlf"):
         indices_this_group = []
         for pauli_word in partition:
             # find index of this pauli word in remaining original observables,
-            for observable in observables:
+            for ind, observable in enumerate(observables):
                 if qml.pauli.are_identical_pauli_words(pauli_word, observable):
-                    for ind, obs in enumerate(observables):
-                        if obs is not observable:
-                            continue
-                        indices_this_group.append(available_indices[ind])
-                        # delete this observable and its index, so it cannot be found again
-                        observables.pop(ind)
-                        available_indices.pop(ind)
-                        break
+                    indices_this_group.append(available_indices[ind])
+                    # delete this observable and its index, so it cannot be found again
+                    observables.pop(ind)
+                    available_indices.pop(ind)
                     break
         indices.append(tuple(indices_this_group))
 
diff --git a/pennylane/pauli/grouping/group_observables.py b/pennylane/pauli/grouping/group_observables.py
index c31ebba54..d6ef018e9 100644
--- a/pennylane/pauli/grouping/group_observables.py
+++ b/pennylane/pauli/grouping/group_observables.py
@@ -20,6 +20,7 @@ from copy import copy
 import numpy as np
 import pennylane as qml
 
+from pennylane.ops import Prod, SProd
 from pennylane.pauli.utils import (
     are_identical_pauli_words,
     binary_to_pauli,
@@ -229,7 +230,18 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
     pauli_grouping = PauliGroupingStrategy(
         observables, grouping_type=grouping_type, graph_colourer=method
     )
-    partitioned_paulis = pauli_grouping.colour_pauli_graph()
+
+    temp_opmath = not qml.operation.active_new_opmath() and any(
+        isinstance(o, (Prod, SProd)) for o in observables
+    )
+    if temp_opmath:
+        qml.operation.enable_new_opmath()
+
+    try:
+        partitioned_paulis = pauli_grouping.colour_pauli_graph()
+    finally:
+        if temp_opmath:
+            qml.operation.disable_new_opmath()
 
     if coefficients is None:
         return partitioned_paulis
@@ -246,15 +258,11 @@ def group_observables(observables, coefficients=None, grouping_type="qwc", metho
         indices = []
         for pauli_word in partition:
             # find index of this pauli word in remaining original observables,
-            for observable in observables:
+            for ind, observable in enumerate(observables):
                 if are_identical_pauli_words(pauli_word, observable):
-                    for ind, obs in enumerate(observables):
-                        if obs is not observable:
-                            continue
-                        indices.append(coeff_indices[ind])
-                        observables.pop(ind)
-                        coeff_indices.pop(ind)
-                        break
+                    indices.append(coeff_indices[ind])
+                    observables.pop(ind)
+                    coeff_indices.pop(ind)
                     break
 
         # add a tensor of coefficients to the grouped coefficients
diff --git a/pennylane/pauli/utils.py b/pennylane/pauli/utils.py
index 5d009e87d..fa5627de7 100644
--- a/pennylane/pauli/utils.py
+++ b/pennylane/pauli/utils.py
@@ -93,7 +93,7 @@ def is_pauli_word(observable):
     >>> is_pauli_word(4 * qml.PauliX(0) @ qml.PauliZ(0))
     True
     """
-    return _is_pauli_word(observable)
+    return _is_pauli_word(observable) or (len(observable.pauli_rep or []) == 1)
 
 
 @singledispatch
@@ -140,24 +140,20 @@ def are_identical_pauli_words(pauli_1, pauli_2):
     # pylint: disable=isinstance-second-argument-not-valid-type
     """Performs a check if two Pauli words have the same ``wires`` and ``name`` attributes.
 
-    This is a convenience function that checks if two given :class:`~.Tensor` instances specify the same
-    Pauli word. This function only checks if both :class:`~.Tensor` instances have the same wires and name
-    attributes, and hence won't perform any simplification to identify if the two Pauli words are
-    algebraically equivalent. For instance, this function will not identify
-    that ``PauliX(0) @ PauliX(0) = Identity(0)``, or ``PauliX(0) @ Identity(1)
-    = PauliX(0)``, or ``Identity(0) = Identity(1)``, etc.
+    This is a convenience function that checks if two given :class:`~.Tensor` or :class:`~.Prod`
+    instances specify the same Pauli word.
 
     Args:
-        pauli_1 (Union[Identity, PauliX, PauliY, PauliZ, Tensor]): the first Pauli word
-        pauli_2 (Union[Identity, PauliX, PauliY, PauliZ, Tensor]): the second Pauli word
+        pauli_1 (Union[Identity, PauliX, PauliY, PauliZ, Tensor, Prod, SProd]): the first Pauli word
+        pauli_2 (Union[Identity, PauliX, PauliY, PauliZ, Tensor, Prod, SProd]): the second Pauli word
 
     Returns:
         bool: whether ``pauli_1`` and ``pauli_2`` have the same wires and name attributes
 
     Raises:
-        TypeError: if ``pauli_1`` or ``pauli_2`` are not :class:`~.Identity`,
-            :class:`~.PauliX`, :class:`~.PauliY`, :class:`~.PauliZ`, or
-            :class:`~.Tensor` instances
+        TypeError: if ``pauli_1`` or ``pauli_2`` are not :class:`~.Identity`, :class:`~.PauliX`,
+            :class:`~.PauliY`, :class:`~.PauliZ`, :class:`~.Tensor`, :class:`~.SProd`, or
+            :class:`~.Prod` instances
 
     **Example**
 
@@ -166,29 +162,13 @@ def are_identical_pauli_words(pauli_1, pauli_2):
     >>> are_identical_pauli_words(qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliX(3))
     False
     """
-
     if not (is_pauli_word(pauli_1) and is_pauli_word(pauli_2)):
         raise TypeError(f"Expected Pauli word observables, instead got {pauli_1} and {pauli_2}.")
 
-    paulis_with_identity = (PauliX, PauliY, PauliZ, Identity)
-
-    # convert tensors of length 1 to plain observables
-    pauli_1 = getattr(pauli_1, "prune", lambda: pauli_1)()
-    pauli_2 = getattr(pauli_2, "prune", lambda: pauli_2)()
-
-    if isinstance(pauli_1, qml.Identity) and isinstance(pauli_2, qml.Identity):
-        return True
-
-    if isinstance(pauli_1, paulis_with_identity) and isinstance(pauli_2, paulis_with_identity):
-        return (pauli_1.wires, pauli_1.name) == (pauli_2.wires, pauli_2.name)
-
-    if isinstance(pauli_1, paulis_with_identity) and isinstance(pauli_2, Tensor):
-        return {(pauli_1.wires, pauli_1.name)} == set(zip(pauli_2.wires, pauli_2.name))
+    if pauli_1.pauli_rep is not None and pauli_2.pauli_rep is not None:
+        return next(iter(pauli_1.pauli_rep)) == next(iter(pauli_2.pauli_rep))
 
-    if isinstance(pauli_1, Tensor) and isinstance(pauli_2, paulis_with_identity):
-        return set(zip(pauli_1.wires, pauli_1.name)) == {(pauli_2.wires, pauli_2.name)}
-
-    return set(zip(pauli_1.wires, pauli_1.name)) == set(zip(pauli_2.wires, pauli_2.name))
+    return False
 
 
 def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_word=True):
@@ -199,7 +179,7 @@ def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_wor
     PauliX placements while the last half specify PauliZ placements.
 
     Args:
-        pauli_word (Union[Identity, PauliX, PauliY, PauliZ, Tensor]): the Pauli word to be
+        pauli_word (Union[Identity, PauliX, PauliY, PauliZ, Tensor, Prod, SProd]): the Pauli word to be
             converted to binary vector representation
         n_qubits (int): number of qubits to specify dimension of binary vector representation
         wire_map (dict): dictionary containing all wire labels used in the Pauli word as keys, and
@@ -273,13 +253,12 @@ def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_wor
     >>> pauli_to_binary(qml.PauliX(0) @ qml.PauliX(5),  wire_map=wire_map)
     array([1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.])
     """
+    wire_map = wire_map or {w: i for i, w in enumerate(pauli_word.wires)}
 
     if check_is_pauli_word and not is_pauli_word(pauli_word):
         raise TypeError(f"Expected a Pauli word Observable instance, instead got {pauli_word}.")
 
-    if wire_map is None:
-        num_wires = len(pauli_word.wires)
-        wire_map = {pauli_word.wires[i]: i for i in range(num_wires)}
+    pw = next(iter(pauli_word.pauli_rep))
 
     n_qubits_min = max(wire_map.values()) + 1
     if n_qubits is None:
@@ -290,28 +269,16 @@ def pauli_to_binary(pauli_word, n_qubits=None, wire_map=None, check_is_pauli_wor
             f" instead got n_qubits={n_qubits}."
         )
 
-    pauli_wires = pauli_word.wires.map(wire_map).tolist()
-
     binary_pauli = np.zeros(2 * n_qubits)
 
-    paulis_with_identity = (PauliX, PauliY, PauliZ, Identity)
-
-    if isinstance(pauli_word, paulis_with_identity):
-        operations_zip = zip(pauli_wires, [pauli_word.name])
-    else:
-        operations_zip = zip(pauli_wires, pauli_word.name)
-
-    for wire, name in operations_zip:
-        if name == "PauliX":
-            binary_pauli[wire] = 1
-
-        elif name == "PauliY":
-            binary_pauli[wire] = 1
-            binary_pauli[n_qubits + wire] = 1
-
-        elif name == "PauliZ":
-            binary_pauli[n_qubits + wire] = 1
-
+    for wire, pauli_type in pw.items():
+        if pauli_type == "X":
+            binary_pauli[wire_map[wire]] = 1
+        elif pauli_type == "Y":
+            binary_pauli[wire_map[wire]] = 1
+            binary_pauli[n_qubits + wire_map[wire]] = 1
+        elif pauli_type == "Z":
+            binary_pauli[n_qubits + wire_map[wire]] = 1
     return binary_pauli
 
 
@@ -328,9 +295,11 @@ def binary_to_pauli(binary_vector, wire_map=None):  # pylint: disable=too-many-b
             unique integer labels as their values
 
     Returns:
-        Tensor: The Pauli word corresponding to the input binary vector. Note
-        that if a zero vector is input, then the resulting Pauli word will be
-        an :class:`~.Identity` instance.
+        Union[Tensor, Prod]: The Pauli word corresponding to the input binary vector.
+        Note that if a zero vector is input, then the resulting Pauli word will be
+        an :class:`~.Identity` instance. If new operator arithmetic is enabled via
+        :func:`~.pennylane.operation.enable_new_opmath`, a :class:`~.Prod` will be
+        returned, else a :class:`~.Tensor` will be returned.
 
     Raises:
         TypeError: if length of binary vector is not even, or if vector does not have strictly
@@ -567,7 +536,7 @@ def pauli_word_to_matrix(pauli_word, wire_map=None):
     product at the correct positions.
 
     Args:
-        pauli_word (Observable): an observable, either a :class:`~.Tensor` instance or
+        pauli_word (Observable): an observable, either a :class:`~.Tensor`, :class:`~.Prod` or
             single-qubit observable representing a Pauli group element.
         wire_map (dict[Union[str, int], int]): dictionary containing all wire labels used in
             the Pauli word as keys, and unique integer labels as their values
@@ -698,8 +667,31 @@ def is_qwc(pauli_vec_1, pauli_vec_2):
     return True
 
 
-def are_pauli_words_qwc(lst_pauli_words):
+def _are_pauli_words_qwc_pauli_rep(lst_pauli_words):
     """Given a list of observables assumed to be valid Pauli words, determine if they are pairwise
+    qubit-wise commuting. This private method is used for operators that have a valid pauli
+    representation"""
+    basis = {}
+    for op in lst_pauli_words:  # iterate over the list of observables
+        if len(pr := op.pauli_rep) > 1:
+            return False
+
+        pw = next(iter(pr))
+
+        for wire, pauli_type in pw.items():  # iterate over wires of the observable,
+            if pauli_type != "I":
+                if wire in basis and pauli_type != basis[wire]:
+                    # Only non-identity paulis are in basis, so if pauli_type doesn't match
+                    # it is guaranteed to not commute
+                    return False
+
+                basis[wire] = pauli_type
+
+    return True  # if we get through all ops, then they are qwc!
+
+
+def are_pauli_words_qwc(lst_pauli_words):
+    """Given a list of observables assumed to be valid Pauli observables, determine if they are pairwise
     qubit-wise commuting.
 
     This implementation has time complexity ~ O(m * n) for m Pauli words and n wires, where n is the
@@ -709,8 +701,12 @@ def are_pauli_words_qwc(lst_pauli_words):
         lst_pauli_words (list[Observable]): List of observables (assumed to be valid Pauli words).
 
     Returns:
-        (bool): True if they are all qubit-wise commuting, false otherwise.
+        (bool): True if they are all qubit-wise commuting, false otherwise. If any of the provided
+        observables are not valid Pauli words, false is returned.
     """
+    if all(op.pauli_rep is not None for op in lst_pauli_words):
+        return _are_pauli_words_qwc_pauli_rep(lst_pauli_words)
+
     latest_op_name_per_wire = {}
 
     for op in lst_pauli_words:  # iterate over the list of observables
@@ -718,15 +714,14 @@ def are_pauli_words_qwc(lst_pauli_words):
         op_wires = op.wires.tolist()
 
         for op_name, wire in zip(op_names, op_wires):  # iterate over wires of the observable,
-            try:
-                if latest_op_name_per_wire[wire] != op_name and (
-                    op_name != "Identity" and latest_op_name_per_wire[wire] != "Identity"
-                ):
-                    return False
-                if latest_op_name_per_wire[wire] == "Identity":
-                    latest_op_name_per_wire[wire] = op_name  # update name
-            except KeyError:
-                latest_op_name_per_wire[wire] = op_name  # add wire and name for the first time
+            latest_op_name = latest_op_name_per_wire.get(wire, "Identity")
+            if latest_op_name != op_name and (
+                op_name != "Identity" and latest_op_name != "Identity"
+            ):
+                return False
+
+            if op_name != "Identity":
+                latest_op_name_per_wire[wire] = op_name
 
     return True  # if we get through all ops, then they are qwc!
 
@@ -739,8 +734,8 @@ def observables_to_binary_matrix(observables, n_qubits=None, wire_map=None):
     non-trivially by the Pauli words in observables.
 
     Args:
-        observables (list[Union[Identity, PauliX, PauliY, PauliZ, Tensor]]): the list of Pauli
-            words
+        observables (list[Union[Identity, PauliX, PauliY, PauliZ, Tensor, Prod, SProd]]): the list
+            of Pauli words
         n_qubits (int): number of qubits to specify dimension of binary vector representation
         wire_map (dict): dictionary containing all wire labels used in the Pauli words as keys, and
             unique integer labels as their values
@@ -1288,7 +1283,10 @@ def diagonalize_qwc_pauli_words(
     if not are_pauli_words_qwc(qwc_grouping):
         raise ValueError("The list of Pauli words are not qubit-wise commuting.")
 
-    if not all(isinstance(op, (Tensor, PauliX, PauliY, PauliZ, Identity)) for op in qwc_grouping):
+    if not all(
+        isinstance(op, (Tensor, PauliX, PauliY, PauliZ, Identity, SProd, Prod))
+        for op in qwc_grouping
+    ):
         raise ValueError("This function only supports Tensor products of pauli ops.")
 
     pauli_operators = []
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index cbc78deac..25d8011df 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -187,15 +187,14 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
         else:
             # create the PauliZ tensor product observable when only wires are provided for the
             # measurements
-            # TODO: Revisit when qml.prod is compatible with qml.pauli.group_observables
-            pauliz_obs = qml.PauliZ(obs.wires[0])
-            for wire in obs.wires[1:]:
-                pauliz_obs = pauliz_obs @ qml.PauliZ(wire)
+            obs_wires = obs.wires if obs.wires else tape.wires
+            pauliz_obs = qml.prod(*(qml.PauliZ(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 = []


Test patch:
diff --git a/tests/pauli/grouping/test_pauli_group_observables.py b/tests/pauli/grouping/test_pauli_group_observables.py
index 69c561173..a2d7c3baf 100644
--- a/tests/pauli/grouping/test_pauli_group_observables.py
+++ b/tests/pauli/grouping/test_pauli_group_observables.py
@@ -18,6 +18,7 @@ import pytest
 import numpy as np
 import pennylane as qml
 from pennylane import Identity, PauliX, PauliY, PauliZ
+from pennylane.operation import Tensor
 from pennylane.pauli import are_identical_pauli_words
 from pennylane.pauli.grouping.group_observables import PauliGroupingStrategy, group_observables
 from pennylane import numpy as pnp
@@ -130,6 +131,13 @@ class TestPauliGroupingStrategy:
 
 observables_list = [
     [PauliX(0) @ PauliZ(1), PauliY(2) @ PauliZ(1), PauliX(1), PauliY(0), PauliZ(1) @ PauliZ(2)],
+    [
+        qml.s_prod(1.5, qml.prod(PauliX(0), PauliZ(1))),
+        qml.prod(PauliY(2), PauliZ(1)),
+        PauliX(1),
+        PauliY(0),
+        qml.prod(PauliZ(1), PauliZ(2)),
+    ],
     [
         Identity(1) @ Identity(0),
         PauliX(1) @ PauliY(0) @ Identity(2),
@@ -138,16 +146,35 @@ observables_list = [
         PauliZ(2) @ Identity(0),
         PauliX(0) @ PauliX(1),
     ],
+    [
+        qml.prod(Identity(1), Identity(0)),
+        qml.prod(PauliX(1), PauliY(0), Identity(2)),
+        PauliZ(2),
+        qml.s_prod(-2.0, Identity(0)),
+        qml.prod(PauliZ(2), Identity(0)),
+        qml.prod(PauliX(0), PauliX(1)),
+    ],
     [
         PauliX("a") @ Identity("b"),
         PauliX("a") @ PauliZ("b"),
         PauliX("b") @ PauliZ("a"),
         PauliZ("a") @ PauliZ("b") @ PauliZ("c"),
     ],
+    [
+        qml.prod(PauliX("a"), Identity("b")),
+        qml.s_prod(0.5, qml.prod(PauliX("a"), PauliZ("b"))),
+        qml.s_prod(0.5, qml.prod(PauliX("b"), PauliZ("a"))),
+        qml.prod(PauliZ("a"), PauliZ("b") @ PauliZ("c")),
+    ],
     [PauliX([(0, 0)]), PauliZ([(0, 0)])],
 ]
 
 qwc_sols = [
+    [
+        [PauliX(wires=[1]), PauliY(wires=[0])],
+        [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
+        [PauliZ(wires=[1]) @ PauliZ(wires=[2])],
+    ],
     [
         [PauliX(wires=[1]), PauliY(wires=[0])],
         [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
@@ -163,6 +190,21 @@ qwc_sols = [
         ],
         [PauliX(wires=[1]) @ PauliX(wires=[0])],
     ],
+    [
+        [
+            Identity(wires=[1]),
+            PauliX(wires=[1]) @ PauliY(wires=[0]),
+            PauliZ(wires=[2]),
+            Identity(wires=[1]),
+            PauliZ(wires=[2]),
+        ],
+        [PauliX(wires=[1]) @ PauliX(wires=[0])],
+    ],
+    [
+        [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
+        [PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"])],
+        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
+    ],
     [
         [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"])],
@@ -177,6 +219,21 @@ commuting_sols = [
         [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
         [PauliZ(wires=[1]) @ PauliZ(wires=[2])],
     ],
+    [
+        [PauliX(wires=[1]), PauliY(wires=[0])],
+        [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliZ(wires=[1]) @ PauliY(wires=[2])],
+        [PauliZ(wires=[1]) @ PauliZ(wires=[2])],
+    ],
+    [
+        [
+            Identity(wires=[1]),
+            PauliX(wires=[1]) @ PauliY(wires=[0]),
+            PauliZ(wires=[2]),
+            Identity(wires=[1]),
+            PauliZ(wires=[2]),
+        ],
+        [PauliX(wires=[1]) @ PauliX(wires=[0])],
+    ],
     [
         [
             Identity(wires=[1]),
@@ -192,6 +249,11 @@ commuting_sols = [
         [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
         [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
     ],
+    [
+        [PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"])],
+        [PauliX(wires=["a"]), PauliX(wires=["a"]) @ PauliZ(wires=["b"])],
+        [PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
+    ],
     [[PauliX([(0, 0)])], [PauliZ([(0, 0)])]],
 ]
 
@@ -204,6 +266,21 @@ anticommuting_sols = [
             PauliZ(wires=[1]) @ PauliZ(wires=[2]),
         ],
     ],
+    [
+        [PauliX(wires=[0]) @ PauliZ(wires=[1]), PauliY(wires=[0])],
+        [
+            PauliZ(wires=[1]) @ PauliY(wires=[2]),
+            PauliX(wires=[1]),
+            PauliZ(wires=[1]) @ PauliZ(wires=[2]),
+        ],
+    ],
+    [
+        [Identity(wires=[1])],
+        [PauliZ(wires=[2])],
+        [Identity(wires=[1])],
+        [PauliZ(wires=[2])],
+        [PauliX(wires=[1]) @ PauliY(wires=[0]), PauliX(wires=[1]) @ PauliX(wires=[0])],
+    ],
     [
         [Identity(wires=[1])],
         [PauliZ(wires=[2])],
@@ -218,6 +295,13 @@ anticommuting_sols = [
         ],
         [PauliX(wires=["a"]), PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
     ],
+    [
+        [
+            PauliX(wires=["a"]) @ PauliZ(wires=["b"]),
+            PauliZ(wires=["a"]) @ PauliZ(wires=["b"]) @ PauliZ(wires=["c"]),
+        ],
+        [PauliX(wires=["a"]), PauliZ(wires=["a"]) @ PauliX(wires=["b"])],
+    ],
     [[PauliX([(0, 0)]), PauliZ([(0, 0)])]],
 ]
 
@@ -309,6 +393,33 @@ class TestGroupObservables:
         _, grouped_coeffs = group_observables(obs, coeffs)
         assert isinstance(grouped_coeffs[0], list)
 
+    def test_return_new_opmath(self):
+        """Test that using new opmath causes grouped observables to have Prods instead of
+        Tensors"""
+        old_observables = [
+            Tensor(PauliX(0), PauliZ(1)),
+            Tensor(PauliY(2), PauliZ(1)),
+            Tensor(PauliZ(1), PauliZ(2)),
+        ]
+        new_observables = [
+            qml.prod(PauliX(0), PauliZ(1)),
+            qml.prod(PauliY(2), PauliZ(1)),
+            qml.s_prod(1.5, qml.prod(PauliZ(1), PauliZ(2))),
+        ]
+        mixed_observables = [
+            Tensor(PauliX(0), PauliZ(1)),
+            qml.prod(PauliY(2), PauliZ(1)),
+            qml.s_prod(1.5, qml.prod(PauliZ(1), PauliZ(2))),
+        ]
+
+        old_groups = group_observables(old_observables)
+        new_groups = group_observables(new_observables)
+        mixed_groups = group_observables(mixed_observables)
+
+        assert all(isinstance(o, Tensor) for g in old_groups for o in g)
+        assert all(isinstance(o, qml.ops.Prod) for g in new_groups for o in g)
+        assert all(isinstance(o, qml.ops.Prod) for g in mixed_groups for o in g)
+
 
 class TestDifferentiable:
     """Tests that grouping observables is differentiable with respect to the coefficients."""
diff --git a/tests/pauli/test_pauli_utils.py b/tests/pauli/test_pauli_utils.py
index 1b465a4e8..c4714f67c 100644
--- a/tests/pauli/test_pauli_utils.py
+++ b/tests/pauli/test_pauli_utils.py
@@ -71,6 +71,9 @@ class TestGroupingUtils:
         (PauliZ(0) @ PauliY(2), np.array([0, 1, 1, 1])),
         (PauliY(1) @ PauliX(2), np.array([1, 1, 1, 0])),
         (Identity(0), np.zeros(2)),
+        (qml.prod(qml.PauliX(0), qml.PauliZ(1), qml.PauliY(2)), np.array([1, 0, 1, 0, 1, 1])),
+        (qml.s_prod(1.5, qml.PauliZ(0)), np.array([0, 1])),
+        (qml.sum(qml.PauliX(0), qml.PauliX(0), qml.PauliX(0)), np.array([1, 0])),
     ]
 
     @pytest.mark.parametrize("op,vec", ops_to_vecs_explicit_wires)
@@ -85,6 +88,15 @@ class TestGroupingUtils:
         (PauliY(6) @ PauliZ("a") @ PauliZ("b"), np.array([0, 0, 0, 1, 1, 1, 0, 1])),
         (PauliX("b") @ PauliY("c"), np.array([0, 1, 1, 0, 0, 0, 1, 0])),
         (Identity("a") @ Identity(6), np.zeros(8)),
+        (
+            qml.prod(qml.PauliX("a"), qml.PauliZ("b"), qml.PauliY(6)),
+            np.array([1, 0, 0, 1, 0, 1, 0, 1]),
+        ),
+        (qml.s_prod(1.5, qml.PauliZ(6)), np.array([0, 0, 0, 0, 0, 0, 0, 1])),
+        (
+            qml.sum(qml.PauliX("b"), qml.PauliX("b"), qml.PauliX("b")),
+            np.array([0, 1, 0, 0, 0, 0, 0, 0]),
+        ),
     ]
 
     @pytest.mark.parametrize("op,vec", ops_to_vecs_abstract_wires)
@@ -227,6 +239,9 @@ class TestGroupingUtils:
             True,
         ),  # multi I
         ([qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliZ(2), qml.PauliX(1) @ qml.PauliY(2)], False),
+        # The following are use cases for `expand_tape`, so should be tested, even though Hadamard is not a Pauli op
+        ([qml.Hadamard(0) @ qml.PauliX(1), qml.Identity(0)], True),
+        ([qml.Hadamard(0) @ qml.PauliX(1), qml.PauliZ(0)], False),
     ]
 
     @pytest.mark.parametrize("obs_lst, expected_qwc", obs_lsts)
@@ -236,6 +251,21 @@ class TestGroupingUtils:
         qwc = are_pauli_words_qwc(obs_lst)
         assert qwc == expected_qwc
 
+    @pytest.mark.parametrize(
+        "ops",
+        (
+            [qml.PauliX(0), qml.sum(qml.PauliX(2), qml.PauliZ(1))],
+            [qml.sum(qml.prod(qml.PauliX(0), qml.PauliY(1)), 2 * qml.PauliZ(0))],
+            [qml.sum(qml.PauliY(0), qml.PauliX(1), qml.PauliZ(2))],
+            [qml.prod(qml.sum(qml.PauliX(0), qml.PauliY(1)), 2 * qml.PauliZ(0)), qml.PauliZ(0)],
+        ),
+    )
+    def test_are_pauli_words_qwc_sum_false(self, ops):
+        """Test that using operators with pauli rep containing more than one term with
+        are_pauli_words_qwc always returns False."""
+
+        assert are_pauli_words_qwc(ops) is False
+
     def test_is_qwc_not_equal_lengths(self):
         """Tests ValueError is raised when input Pauli vectors are not of equal length."""
 
@@ -276,6 +306,10 @@ class TestGroupingUtils:
         (qml.prod(qml.PauliX(0), qml.Hadamard(1)), False),
         (qml.s_prod(5, qml.PauliX(0) @ qml.PauliZ(1)), True),
         (qml.s_prod(5, qml.Hadamard(0)), False),
+        (qml.sum(qml.PauliX(0), qml.PauliY(0)), False),
+        (qml.sum(qml.s_prod(0.5, qml.PauliX(1)), qml.PauliX(1)), True),
+        (qml.sum(qml.s_prod(0.5, qml.Hadamard(1)), qml.PauliX(1)), False),
+        (qml.sum(qml.Hadamard(0), qml.Hadamard(0)), False),
     )
 
     @pytest.mark.parametrize("ob, is_pw", obs_pw_status)
@@ -304,11 +338,25 @@ class TestGroupingUtils:
         pauli_word_2 = PauliY(1) @ PauliX(0)
         pauli_word_3 = Tensor(PauliX(0), PauliY(1))
         pauli_word_4 = PauliX(1) @ PauliZ(2)
+        pauli_word_5 = qml.s_prod(1.5, qml.PauliX(0))
+        pauli_word_6 = qml.sum(qml.s_prod(0.5, qml.PauliX(0)), qml.s_prod(1.0, qml.PauliX(0)))
+        pauli_word_7 = qml.s_prod(2.2, qml.prod(qml.PauliX(0), qml.PauliY(1)))
 
         assert are_identical_pauli_words(pauli_word_1, pauli_word_2)
         assert are_identical_pauli_words(pauli_word_1, pauli_word_3)
         assert not are_identical_pauli_words(pauli_word_1, pauli_word_4)
         assert not are_identical_pauli_words(pauli_word_3, pauli_word_4)
+        assert are_identical_pauli_words(pauli_word_5, pauli_word_6)
+        assert are_identical_pauli_words(pauli_word_7, pauli_word_1)
+        assert not are_identical_pauli_words(pauli_word_7, pauli_word_4)
+        assert not are_identical_pauli_words(pauli_word_6, pauli_word_4)
+
+    def test_are_identical_pauli_words_hamiltonian_unsupported(self):
+        """Test that using Hamiltonians that are valid Pauli words with are_identical_pauli_words
+        always returns False"""
+        pauli_word_1 = qml.Hamiltonian([1.0], [qml.PauliX(0)])
+        pauli_word_2 = qml.PauliX(0)
+        assert not are_identical_pauli_words(pauli_word_1, pauli_word_2)
 
     def test_identities_always_pauli_words(self):
         """Tests that identity terms are always identical."""
@@ -913,7 +961,7 @@ class TestMeasurementTransformations:
         [PauliZ(0) @ Identity(1), PauliZ(0) @ PauliZ(1), PauliX(0) @ Identity(1)],
         [PauliX("a") @ PauliX(0), PauliZ(0) @ PauliZ("a")],
         [PauliZ("a") @ PauliY(1), PauliZ(1) @ PauliY("a")],
-        [Hamiltonian([1.0, 2.0], [PauliX("a"), PauliY("a")])],
+        [qml.prod(qml.PauliZ(0), qml.PauliX(1)), qml.prod(qml.PauliZ(1), qml.PauliX(0))],
     ]
 
     @pytest.mark.parametrize("not_qwc_grouping", not_qwc_groupings)
@@ -923,6 +971,16 @@ class TestMeasurementTransformations:
 
         assert pytest.raises(ValueError, diagonalize_qwc_pauli_words, not_qwc_grouping)
 
+    def test_diagonalize_qwc_pauli_words_catch_invalid_type(self):
+        """Test for ValueError raise when diagonalize_qwc_pauli_words is given a list
+        containing invalid operator types."""
+        invalid_ops = [qml.PauliX(0), qml.Hamiltonian([1.0], [qml.PauliZ(1)])]
+
+        with pytest.raises(
+            ValueError, match="This function only supports Tensor products of pauli ops"
+        ):
+            _ = diagonalize_qwc_pauli_words(invalid_ops)
+
 
 class TestObservableHF:
     @pytest.mark.parametrize(
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 0b5476e3c..8aaa213e0 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -1456,10 +1456,16 @@ class TestTensor:
         assert t.batch_size is None
 
     def test_pauli_rep(self):
-        """Test that the _pauli_rep attribute of the Tensor is initialized as None."""
+        """Test that the _pauli_rep attribute of the Tensor is initialized correctly."""
+        # Pauli rep not None for Pauli observables
         X = qml.PauliX(0)
         Y = qml.PauliY(2)
         t = Tensor(X, Y)
+        assert t.pauli_rep == qml.pauli.PauliSentence({qml.pauli.PauliWord({0: "X", 2: "Y"}): 1.0})
+
+        # Puli rep None if observables not valid Pauli observables
+        H = qml.Hadamard(1)
+        t = Tensor(X, H)
         assert t.pauli_rep is None
 
     def test_has_matrix(self):
@@ -1989,6 +1995,15 @@ class TestTensor:
         for obs1, obs2 in zip(mapped_tensor.obs, final_obs):
             assert qml.equal(obs1, obs2)
 
+    def test_matmul_not_implemented(self):
+        """Test that matrix multiplication raises TypeError if unsupported
+        object is used."""
+
+        op = Tensor(qml.PauliX(0), qml.PauliZ(1))
+
+        with pytest.raises(TypeError, match="unsupported operand type"):
+            _ = op @ 1.0
+
 
 equal_obs = [
     (qml.PauliZ(0), qml.PauliZ(0), True),
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index 9c68c0e27..820089fee 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -94,6 +94,7 @@ class TestUnittestSplitNonCommuting:
             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])
@@ -279,6 +280,7 @@ class TestUnittestSplitNonCommuting:
             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])
 
@@ -288,7 +290,8 @@ class TestUnittestSplitNonCommuting:
         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(wires=[0, 1]))
+        assert qml.equal(split[1].measurements[0], meas_type_2())
+        assert qml.equal(split[1].measurements[1], meas_type_2(wires=[0, 1]))
 
 
 # measurements that require shots=True