Repository: PennyLaneAI/pennylane
Pull request #5098 (2024-01-24 12:12:37)

Upgrade qml.commutator to use pauli arithmetic under the hood whenever possible
Upgrade qml.commutator to use pauli artihmetic whenever that is possible (i.e. whenever both operators have a `pauli_rep`. This was initiated in https://github.com/PennyLaneAI/pennylane/pull/5051#discussion_r1463394138 Similar things could be done in qml.dot and other points of the codebase. The grander idea is to use fater pauli arithmetic under the hood whenever that is possible.

The main feature of this update is performance:
```python
op1 = qml.dot(range(30), [X(i) @ X(i+1) for i in range(30)])
op2 = qml.dot(range(30), [Y(i) @ Y(i+1) for i in range(30)])

%timeit qml.commutator(op1, op2) # on main branch
11.9 ms ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit qml.commutator(op1, op2) # this branch
1.28 ms ± 3.42 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```

- [x] Basic functionality
- [x] Handle Zero case (empty PauliSentence) with 0 * Identity()
- [x] Add test ensuring pauli route is taken
- [x] All tests pass
- [x] Test mixed cases (ops and Paulis)
- [x] changelog

[sc-56865]

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/5098?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:
> Comparison is base [(`699701d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/699701d1290d253d414a0e0110affb41c4374d59?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`8f57d73`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5098?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67%.


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


```diff
@@                 Coverage Diff                 @@
##           paulicommutator    #5098      +/-   ##
===================================================
- Coverage            99.68%   99.67%   -0.01%     
===================================================
  Files                  393      393              
  Lines                35639    35647       +8     
===================================================
+ Hits                 35526    35532       +6     
- Misses                 113      115       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5098?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 cd4651f99..5f088b7f4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -125,9 +125,10 @@
     ```
 
   * A new `qml.commutator` function is now available that allows you to compute commutators between
-    operators:
+    `qml.operation.Operator`, `qml.pauli.PauliWord` and `qml.pauli.PauliSentence` instances.
     [(#5051)](https://github.com/PennyLaneAI/pennylane/pull/5051)
     [(#5052)](https://github.com/PennyLaneAI/pennylane/pull/5052)
+    [(#5098)](https://github.com/PennyLaneAI/pennylane/pull/5098)
 
     ```pycon
     >>> qml.commutator(X(0), Y(0))
diff --git a/pennylane/ops/functions/commutator.py b/pennylane/ops/functions/commutator.py
index b3e634c70..b82bdda3c 100644
--- a/pennylane/ops/functions/commutator.py
+++ b/pennylane/ops/functions/commutator.py
@@ -35,44 +35,47 @@ def commutator(op1, op2, pauli=False):
 
     You can compute commutators between operators in PennyLane.
 
-    >>> qml.commutator(qml.PauliX(0), qml.PauliY(0))
-    2j*(PauliZ(wires=[0]))
+    >>> qml.commutator(X(0), Y(0))
+    2j * Z(0)
 
-    >>> op1 = qml.PauliX(0) @ qml.PauliX(1)
-    >>> op2 = qml.PauliY(0) @ qml.PauliY(1)
+    >>> op1 = X(0) @ X(1)
+    >>> op2 = Y(0) @ Y(1)
     >>> qml.commutator(op1, op2)
-    0*(Identity(wires=[0, 1]))
+    0 * I()
 
     We can return a :class:`~PauliSentence` instance by setting ``pauli=True``.
 
-    >>> op1 = qml.PauliX(0) @ qml.PauliX(1)
-    >>> op2 = qml.PauliY(0) + qml.PauliY(1)
-    >>> qml.commutator(op1, op2, pauli=True)
+    >>> op1 = X(0) @ X(1)
+    >>> op2 = Y(0) + Y(1)
+    >>> res = qml.commutator(op1, op2, pauli=True)
+    >>> res
     2j * X(1) @ Z(0)
     + 2j * Z(1) @ X(0)
+    >>> isinstance(res, PauliSentence)
+    True
 
     We can also input :class:`~PauliWord` and :class:`~PauliSentence` instances.
 
     >>> op1 = PauliWord({0:"X", 1:"X"})
     >>> op2 = PauliWord({0:"Y"}) + PauliWord({1:"Y"})
-    >>> qml.commutator(op1, op2, pauli=True)
+    >>> res = qml.commutator(op1, op2, pauli=True)
+    >>> res
     2j * Z(0) @ X(1)
     + 2j * X(0) @ Z(1)
+    >>> isinstance(res, PauliSentence)
+    True
 
     Note that when ``pauli=False``, even if Pauli operators are used
     as inputs, ``qml.commutator`` returns Operators.
 
-    >>> qml.commutator(op1, op2, pauli=True)
-    (2j*(PauliX(wires=[1]) @ PauliZ(wires=[0]))) + (2j*(PauliZ(wires=[1]) @ PauliX(wires=[0])))
-
-
-    It is also worth highlighting that computing commutators with Paulis is typically faster.
-    So whenever all input operators have a pauli representation, we recommend using `pauli=True`.
-    The result can then still be transformed into a PennyLane operator.
-
-    >>> res = qml.commutator(qml.PauliX(0), qml.PauliY(0), pauli=True)
-    >>> res.operation()
+    >>> res = qml.commutator(op1, op2, pauli=False)
+    >>> res
+    2j * (Z(0) @ X(1)) + 2j * (X(0) @ Z(1))
+    >>> isinstance(res, qml.operation.Operator))
+    True
 
+    It is worth noting that computing commutators with Paulis is typically faster.
+    Internally, ``qml.commutator`` uses the ``op.pauli_rep`` whenever that is available for both operators.
 
     .. details::
         :title: Usage Details
@@ -101,13 +104,18 @@ def commutator(op1, op2, pauli=False):
         or insert an extra ``stop_recording()`` context (see :class:`~QueuingManager`).
 
     """
-    if pauli:
+
+    both_have_pauli_rep = op1.pauli_rep is not None and op2.pauli_rep is not None
+
+    if pauli or both_have_pauli_rep:
         if not isinstance(op1, PauliSentence):
             op1 = qml.pauli.pauli_sentence(op1)
         if not isinstance(op2, PauliSentence):
             op2 = qml.pauli.pauli_sentence(op2)
-        return op1.commutator(op2)
+        res = op1.commutator(op2)
+        return res if pauli else res.operation(wire_order=res.wires)
 
+    # If no pauli processing is possible, use operator arithmetic
     with qml.QueuingManager.stop_recording():
         if isinstance(op1, (PauliWord, PauliSentence)):
             op1 = op1.operation()
diff --git a/pennylane/pauli/pauli_arithmetic.py b/pennylane/pauli/pauli_arithmetic.py
index 1b573c3b7..e5ffff944 100644
--- a/pennylane/pauli/pauli_arithmetic.py
+++ b/pennylane/pauli/pauli_arithmetic.py
@@ -192,6 +192,11 @@ class PauliWord(dict):
                 del mapping[wire]
         super().__init__(mapping)
 
+    @property
+    def pauli_rep(self):
+        """Trivial pauli_rep"""
+        return PauliSentence({self: 1.0})
+
     def __reduce__(self):
         """Defines how to pickle and unpickle a PauliWord. Otherwise, un-pickling
         would cause __setitem__ to be called, which is forbidden on PauliWord.
@@ -593,6 +598,11 @@ class PauliSentence(dict):
     # taken from [stackexchange](https://stackoverflow.com/questions/40694380/forcing-multiplication-to-use-rmul-instead-of-numpy-array-mul-or-byp/44634634#44634634)
     __array_priority__ = 1000
 
+    @property
+    def pauli_rep(self):
+        """Trivial pauli_rep"""
+        return self
+
     def __missing__(self, key):
         """If the PauliWord is not in the sentence then the coefficient
         associated with it should be 0."""


Test patch:
diff --git a/tests/ops/functions/test_commutator.py b/tests/ops/functions/test_commutator.py
index 38686bb77..7d6f83cdd 100644
--- a/tests/ops/functions/test_commutator.py
+++ b/tests/ops/functions/test_commutator.py
@@ -175,8 +175,18 @@ class TestcommPauli:
         assert res == true_res
         assert isinstance(res, PauliSentence)
 
+    def test_consistency_with_native_pauli_comms(self):
+        """Test consistent behavior between native comms in PauliWord and PauliSentence and qml.commutor"""
+        op1 = qml.PauliX(0) @ qml.PauliX(1)
+        op2 = qml.PauliY(0) + qml.PauliY(1)
+        res1 = qml.commutator(op1, op2, pauli=True)
+        res2 = PauliWord({0: "X", 1: "X"}).commutator(PauliWord({0: "Y"}) + PauliWord({1: "Y"}))
+        assert isinstance(res1, PauliSentence)
+        assert isinstance(res2, PauliSentence)
+        assert res1 == res2
+
 
-class TestcommPauliFalseSimplify:
+class TestcommPauliFalse:
     """Test qml.comm for pauli=False (default behavior)"""
 
     data_pauli_relations_ops = (
@@ -225,12 +235,12 @@ class TestcommPauliFalseSimplify:
         assert isinstance(res, Operator)
         assert isinstance(res, SProd)
 
-    def test_consistency_with_native_pauli_comms(self):
-        """Test consistent behavior between native comms in PauliWord and PauliSentence and qml.commutor"""
+    def test_paulis_used_when_ever_possible(self, mocker):
+        """Test that pauli_rep is used whenever possible even when ``pauli=False``"""
+        spy = mocker.spy(PauliSentence, "operation")
         op1 = qml.PauliX(0) @ qml.PauliX(1)
         op2 = qml.PauliY(0) + qml.PauliY(1)
-        res1 = qml.commutator(op1, op2, pauli=True)
-        res2 = PauliWord({0: "X", 1: "X"}).commutator(PauliWord({0: "Y"}) + PauliWord({1: "Y"}))
-        assert isinstance(res1, PauliSentence)
-        assert isinstance(res2, PauliSentence)
-        assert res1 == res2
+        res = qml.commutator(op1, op2, pauli=False)
+        spy.assert_called()
+
+        assert isinstance(res, Operator)
diff --git a/tests/pauli/test_pauli_arithmetic.py b/tests/pauli/test_pauli_arithmetic.py
index 4f284be99..ab49412c0 100644
--- a/tests/pauli/test_pauli_arithmetic.py
+++ b/tests/pauli/test_pauli_arithmetic.py
@@ -156,6 +156,11 @@ class TestPauliWord:
         assert 3 not in pw.keys()
         assert pw[3] == I
 
+    def test_pauli_rep(self):
+        """Test trivial pauli_rep property"""
+        pw = PauliWord({0: "I", 1: "X", 2: Y})
+        assert pw.pauli_rep == PauliSentence({pw: 1})
+
     def test_set_items(self):
         """Test that setting items raises an error"""
         pw = PauliWord({0: I, 1: X, 2: Y})
@@ -536,6 +541,11 @@ class TestPauliSentence:
         ps[new_pw] = 3.45
         assert new_pw in ps.keys() and ps[new_pw] == 3.45
 
+    def test_pauli_rep(self):
+        """Test trivial pauli_rep property"""
+        ps = PauliSentence({PauliWord({0: "I", 1: "X", 2: Y}): 1j, X0: 2.0})
+        assert ps.pauli_rep == ps
+
     tup_ps_str = (
         (
             ps1,