Repository: PennyLaneAI/pennylane
Pull request #4666 (2023-10-12 11:05:06)

Speed up `GroverOperator` and add GlobalPhase to its decomposition
**Context:**
1. The `GroverOperator` is applied by `DefaultQubit` by applying its matrix for up to 12 wires (incl).
However, `GroverOperator` has a lot of structure as it is `2*P-1` for a projector `P`, and a custom application rule is faster than constructing and applying the matrix.
2. The `compute_matrix` method of `GroverOperator` in addition uses a rather contrived way to construct the matrix `np.full((dim, dim), normalization_constant)`. Instead, we can simply use the broadcasting behaviour of numpy and do `normalization_constant - np.eye(dim)` (timings below).
3. The `GroverOperator` decomposition is only correct up to a minus sign (or phase of $\pi$).


**Description of the Change:**
1. Adds a dispatch registration for `apply_operation` in `/qubit/apply_operation.py` for `GroverOperator` that is not based on the matrix representation of `GroverOperator`.
2. Modifies the `compute_matrix` method of `GroverOperator` to skip repetitive calling of `np.kron` and `np.outer` to construct a trivial matrix.
3. Introduces a `GlobalPhase($\pi$)` gate in the decomposition.

**Benefits:**
Speedups and correct decomposition.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**


Timing code for the matrix construction:
```python
def f(n_wires):
    s1 = np.array([1, 1]) / np.sqrt(2)
    # uniform superposition state |s>
    s = reduce(np.kron, list(it.repeat(s1, n_wires)))
    # Grover diffusion operator
    return 2 * np.outer(s, s) - np.identity(2**n_wires)

for N in [3, 6, 9, 12]:
    dim = 2**N

    print(f"N={N}")
    %timeit f(N)
    %timeit 2 / dim - np.eye(dim)
    print()
```
![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/a1b477d3-30ec-4aa5-b655-0cbdcc1caf0a)

Timings for the different `apply_operation` options, for three scenarios: Number of wires for operation and state is equal, number of wires for state is that for operation +4, and number of wires for state is double of that for operation.

Given the results, we could do einsum for `len(op.wires)<8` and the dispatch method for bigger wire counts.

![Screenshot from 2023-10-12 14-03-56](https://github.com/PennyLaneAI/pennylane/assets/20772922/5a24dc6c-9732-4875-ab74-36b800eafd44)
![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/8e4ce6b7-f72a-46f2-8371-7a577a02bd75)
![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/91442c01-cc2e-4549-9316-576c242b7720)



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4666?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`ea7e1d4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ea7e1d4e4c5039aec9d4b9c059f1829d7dd10630?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63% compared to head [(`b6eccbb`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4666?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

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


```diff
@@           Coverage Diff           @@
##           master    #4666   +/-   ##
=======================================
  Coverage   99.63%   99.63%           
=======================================
  Files         376      376           
  Lines       33672    33677    +5     
=======================================
+ Hits        33548    33553    +5     
  Misses        124      124           
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4666?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/apply\_operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4666?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYXBwbHlfb3BlcmF0aW9uLnB5) | `99.27% <100.00%> (+0.05%)` | :arrow_up: |
| [pennylane/templates/subroutines/grover.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4666?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9ncm92ZXIucHk=) | `100.00% <100.00%> (ø)` | |


</details>

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

dwierichs: >You may need to tweak "use matrix versus decomposition" threshold in default qubit as well.

Good call. I actually removed it now, because it is supposed to address the large wire counts, which is the regime where the dispatched `apply_grover` is particularly fast. Moreover, as the decomposition also creates a dense matrix (via `MultiControlledX`), this PR makes `GroverOperator` accessible for qubit counts (`>=15` on my laptop) that were inaccessible before :)

> But I have an open PR that will be moving it. We can resolve that merge conflict when we get to it.

:+1:
dwierichs: Updated timings for all interfaces (line style gives interface, color gives target wires): As there is an outlier in Tensorflow which leads to `num_wires=16`, `num_op_wires=8` having a factor 2 slowdown, still, I will move the threshold from `num_op_wires<8` to `num_op_wires<9`.

![image](https://github.com/PennyLaneAI/pennylane/assets/20772922/57baf063-3109-480b-a337-220a5ac3e75d)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index df60fc60f..bbc53b8dd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -42,6 +42,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* `default.qubit` now applies `GroverOperator` faster by not using its matrix representation but a
+  custom rule for `apply_operation`. Also, the matrix representation of `GroverOperator` runs faster now.
+  [(#4666)](https://github.com/PennyLaneAI/pennylane/pull/4666)
+
 * `default.qubit` now can evolve already batched states with `ParametrizedEvolution`
   [(#4863)](https://github.com/PennyLaneAI/pennylane/pull/4863)
 
@@ -97,6 +101,9 @@
 
 <h3>Breaking changes 💔</h3>
 
+* The decomposition of `GroverOperator` now has an additional global phase operation.
+  [(#4666)](https://github.com/PennyLaneAI/pennylane/pull/4666)
+
 * Moved `qml.cond` and the `Conditional` operation from the `transforms` folder to the `ops/op_math` folder.
   `qml.transforms.Conditional` will now be available as `qml.ops.Conditional`.
   [(#4860)](https://github.com/PennyLaneAI/pennylane/pull/4860)
@@ -148,6 +155,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* The decomposition of `GroverOperator` now has the same global phase as its matrix.
+  [(#4666)](https://github.com/PennyLaneAI/pennylane/pull/4666)
+
 * Jax can now differeniate a batch of circuits where one tape does not have trainable parameters.
   [(#4837)](https://github.com/PennyLaneAI/pennylane/pull/4837)
 
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 40c086c3a..75b4f68ec 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -315,6 +315,48 @@ def apply_multicontrolledx(
     return state
 
 
+@apply_operation.register
+def apply_grover(op: qml.GroverOperator, state, is_state_batched: bool = False, debugger=None):
+    """Apply GroverOperator either via a custom matrix-free method (more than 8 operation
+    wires) or via standard matrix based methods (else)."""
+    if len(op.wires) < 9:
+        return _apply_operation_default(op, state, is_state_batched, debugger)
+    return _apply_grover_without_matrix(state, op.wires, is_state_batched)
+
+
+def _apply_grover_without_matrix(state, op_wires, is_state_batched):
+    r"""Apply GroverOperator to state. This method uses that this operator
+    is :math:`2*P-\mathbb{I}`, where :math:`P` is the projector onto the
+    all-plus state. This allows us to compute the new state by replacing summing
+    over all axes on which the operation acts, and "filling in" the all-plus state
+    in the resulting lower-dimensional state via a Kronecker product.
+    """
+    num_wires = len(op_wires)
+    # 2 * Squared normalization of the all-plus state on the op wires
+    # (squared, because we skipped the normalization when summing, 2* because of the def of Grover)
+    prefactor = 2 ** (1 - num_wires)
+
+    # The axes to sum over in order to obtain <+|\psi>, where <+| only acts on the op wires.
+    sum_axes = [w + is_state_batched for w in op_wires]
+    collapsed = math.sum(state, axis=tuple(sum_axes))
+
+    if num_wires == (len(qml.math.shape(state)) - is_state_batched):
+        # If the operation acts on all wires, we can skip the tensor product with all-ones state
+        new_shape = (-1,) + (1,) * num_wires if is_state_batched else (1,) * num_wires
+        return prefactor * math.reshape(collapsed, new_shape) - state
+        # [todo]: Once Tensorflow support expand_dims with multiple axes in the second argument,
+        # use the following line instead of the two above.
+        # return prefactor * math.expand_dims(collapsed, sum_axes) - state
+
+    all_plus = math.cast_like(math.full([2] * num_wires, prefactor), state)
+    # After the Kronecker product (realized with tensordot with axes=0), we need to move
+    # the new axes to the summed-away axes' positions. Finally, subtract the original state.
+    source = list(range(math.ndim(collapsed), math.ndim(state)))
+    # Probably it will be better to use math.full or math.tile to create the outer product
+    # here computed with math.tensordot. However, Tensorflow and Torch do not have full support
+    return math.moveaxis(math.tensordot(collapsed, all_plus, axes=0), source, sum_axes) - state
+
+
 @apply_operation.register
 def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debugger=None):
     """Take a snapshot of the state"""
diff --git a/pennylane/templates/subroutines/grover.py b/pennylane/templates/subroutines/grover.py
index 5f3cbf094..c6d936437 100644
--- a/pennylane/templates/subroutines/grover.py
+++ b/pennylane/templates/subroutines/grover.py
@@ -14,11 +14,9 @@
 """
 Contains the Grover Operation template.
 """
-import itertools
-import functools
 import numpy as np
 from pennylane.operation import AnyWires, Operation
-from pennylane.ops import Hadamard, PauliZ, MultiControlledX
+from pennylane.ops import Hadamard, PauliZ, MultiControlledX, GlobalPhase
 from pennylane.wires import Wires
 
 
@@ -165,17 +163,34 @@ class GroverOperator(Operation):
         for wire in wires[:-1]:
             op_list.append(Hadamard(wire))
 
+        op_list.append(GlobalPhase(np.pi, wires))
+
         return op_list
 
     @staticmethod
-    @functools.lru_cache()
     def compute_matrix(n_wires, work_wires):  # pylint: disable=arguments-differ,unused-argument
-        # s1 = H|0>, Hadamard on a single qubit in the ground state
-        s1 = np.array([1, 1]) / np.sqrt(2)
+        r"""Representation of the operator as a canonical matrix in the computational basis
+        (static method).
+
+        The canonical matrix is the textbook matrix representation that does not consider wires.
+        Implicitly, this assumes that the wires of the operator correspond to the global wire order.
 
-        # uniform superposition state |s>
-        s = functools.reduce(np.kron, list(itertools.repeat(s1, n_wires)))
+        .. seealso:: :meth:`.GroverOperator.matrix` and :func:`qml.matrix() <pennylane.matrix>`
+
+        Args:
+            n_wires (int): Number of wires the ``GroverOperator`` acts on
+            work_wires (Any or Iterable[Any]): optional auxiliary wires to assist decompositions.
+                *Unused argument*.
 
-        # Grover diffusion operator
-        G = 2 * np.outer(s, s) - np.identity(2**n_wires)
-        return G
+        Returns:
+            tensor_like: matrix representation
+
+        The Grover diffusion operator is :math:`2|+\rangle\langle +| - \mathbb{I}`.
+        The first term is an all-ones matrix multiplied with two times the squared
+        normalization factor of the all-plus state, i.e. all entries of the first term are
+        :math:`2^{1-N}` for :math:`N` wires.
+        """
+        dim = 2**n_wires
+        # Grover diffusion operator. Realize the all-ones entry via broadcasting when subtracting
+        # the second term.
+        return 2 / dim - np.eye(dim)


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 56fa51069..7ed66f5e7 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -206,6 +206,30 @@ class TestTwoQubitStateSpecialCases:
         assert qml.math.allclose(initial1[1], new1[0])
         assert qml.math.allclose(initial1[0], new1[1])
 
+    def test_grover(self, method, wire, ml_framework):
+        """Test the application of GroverOperator on a two qubit state."""
+
+        initial_state = np.array(
+            [
+                [0.04624539 + 0.3895457j, 0.22399401 + 0.53870339j],
+                [-0.483054 + 0.2468498j, -0.02772249 - 0.45901669j],
+            ]
+        )
+        initial_state = qml.math.asarray(initial_state, like=ml_framework)
+
+        wires = [wire, 1 - wire]
+        op = qml.GroverOperator(wires)
+        new_state = method(op, initial_state)
+
+        overlap = qml.math.sum(initial_state) / 2
+        ones_state = qml.math.ones_like(initial_state) / 2
+        expected_state = 2 * ones_state * overlap - initial_state
+        assert qml.math.allclose(new_state, expected_state)
+        state_via_mat = qml.math.tensordot(
+            op.matrix().reshape([2] * 4), initial_state, axes=[[2, 3], [0, 1]]
+        )
+        assert qml.math.allclose(new_state, state_via_mat)
+
     def test_identity(self, method, wire, ml_framework):
         """Test the application of a GlobalPhase gate on a two qubit state."""
 
@@ -841,6 +865,218 @@ class TestLargerOperations:
 
         assert qml.math.allclose(state_v1, state_v2)
 
+    @pytest.mark.parametrize("apply_wires", ([0, 3], [0, 1, 3, 2], [2, 1], [1, 3]))
+    def test_grover(self, method, apply_wires):
+        """Tests a four qubit GroverOperator."""
+        op = qml.GroverOperator(apply_wires)
+        new_state = method(op, self.state)
+
+        expected_state = self.state
+        for _op in op.decomposition():
+            expected_state = method(_op, expected_state)
+
+        assert qml.math.allclose(expected_state, new_state)
+
+
+class TestApplyGroverOperator:
+    """Test that GroverOperator is applied correctly."""
+
+    def grover_kernel_full_wires(self, state, op_wires, batched):
+        """Additional kernel to apply GroverOperator to all state wires."""
+        prefactor = 2 ** (1 - len(op_wires))
+        sum_axes = tuple(range(batched, np.ndim(state)))
+        collapsed = np.sum(state, axis=sum_axes)
+        return prefactor * np.expand_dims(collapsed, sum_axes) - state
+
+    def grover_kernel_partial_wires(self, state, op_wires, batched):
+        """Additional kernel to apply GroverOperator to some of all state wires."""
+        num_wires = len(op_wires)
+        sum_axes = [w + batched for w in op_wires]
+        collapsed = np.sum(state, tuple(sum_axes))
+        prefactor = 2 ** (1 - num_wires)
+        bcast_shape = [2] * num_wires + list(state.shape[:-num_wires])
+        expanded = np.broadcast_to(prefactor * collapsed, bcast_shape)
+        source = list(range(num_wires))
+        expanded = np.moveaxis(expanded, source, sum_axes)
+        return expanded - state
+
+    @pytest.mark.parametrize(
+        "op_wires, state_wires, einsum_called, tensordot_called",
+        [
+            (2, 2, True, False),
+            (3, 3, False, True),
+            (9, 9, False, False),
+            (2, 13, False, True),
+            (3, 9, False, True),
+            (9, 13, False, False),
+        ],
+    )
+    def test_dispatching(self, op_wires, state_wires, einsum_called, tensordot_called, mocker):
+        """Test that apply_operation dispatches to einsum, tensordot and the kernel correctly."""
+        # pylint: disable=too-many-arguments
+        np.random.seed(752)
+        state = np.random.random([2] * state_wires) + 1j * np.random.random([2] * state_wires)
+
+        op = qml.GroverOperator(list(range(op_wires)))
+        spy_einsum = mocker.spy(qml.math, "einsum")
+        spy_tensordot = mocker.spy(qml.math, "argsort")
+        apply_operation(op, state, is_state_batched=False, debugger=None)
+        assert spy_einsum.call_count == int(einsum_called)
+        assert spy_tensordot.call_count == int(tensordot_called)
+
+    @pytest.mark.parametrize("op_wires, state_wires", [(2, 2), (3, 3), (9, 9)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_full_wires(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        when applying it to all wires of a state."""
+        np.random.seed(752)
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        flat_shape = (batch_dim, 2**state_wires) if batched else (2**state_wires,)
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        op = qml.GroverOperator(list(range(op_wires)))
+        out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+        # Double transpose to accomodate for batching
+        expected_via_mat = (op.matrix() @ state.reshape(flat_shape).T).T.reshape(shape)
+        expected_via_kernel = self.grover_kernel_full_wires(state, op.wires, batched)
+        assert np.allclose(out, expected_via_mat)
+        assert np.allclose(out, expected_via_kernel)
+
+    @pytest.mark.parametrize("op_wires, state_wires", [(3, 5), (9, 13)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_partial_wires(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        but einsum (because Grover can't act on a single wire)
+        when applying it only to some of the wires of a state."""
+        np.random.seed(752)
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        for start_wire in [0, 1, state_wires - op_wires]:
+            wires = list(range(start_wire, start_wire + op_wires))
+            op = qml.GroverOperator(wires)
+            out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+            expected_via_mat = apply_operation_tensordot(op, state, batched)
+            expected_via_kernel = self.grover_kernel_partial_wires(state, wires, batched)
+            assert np.allclose(out, expected_via_mat)
+            assert np.allclose(out, expected_via_kernel)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("op_wires, state_wires", [(2, 2), (3, 3), (9, 9), (3, 5), (9, 13)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_autograd(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        when applying it to an Autograd state."""
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        # Input state
+        np.random.seed(752)
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        wires = list(range(op_wires))
+        op = qml.GroverOperator(wires)
+        expected_via_mat = apply_operation_tensordot(op, state, batched)
+        if op_wires == state_wires:
+            expected_via_kernel = self.grover_kernel_full_wires(state, wires, batched)
+        else:
+            expected_via_kernel = self.grover_kernel_partial_wires(state, wires, batched)
+
+        # Cast to interface and apply operation
+        state = qml.numpy.array(state)
+        out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+
+        assert qml.math.allclose(out, expected_via_mat)
+        assert qml.math.allclose(out, expected_via_kernel)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("op_wires, state_wires", [(2, 2), (3, 3), (9, 9), (3, 5), (9, 13)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_tf(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        when applying it to a Tensorflow state."""
+        import tensorflow as tf
+
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        # Input state
+        np.random.seed(752)
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        wires = list(range(op_wires))
+        op = qml.GroverOperator(wires)
+        expected_via_mat = apply_operation_tensordot(op, state, batched)
+        if op_wires == state_wires:
+            expected_via_kernel = self.grover_kernel_full_wires(state, wires, batched)
+        else:
+            expected_via_kernel = self.grover_kernel_partial_wires(state, wires, batched)
+
+        # Cast to interface and apply operation
+        state = tf.Variable(state)
+        out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+
+        assert qml.math.allclose(out, expected_via_mat)
+        assert qml.math.allclose(out, expected_via_kernel)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("op_wires, state_wires", [(2, 2), (3, 3), (9, 9), (3, 5), (9, 13)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_jax(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        when applying it to a Jax state."""
+        import jax
+
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        # Input state
+        np.random.seed(752)
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        wires = list(range(op_wires))
+        op = qml.GroverOperator(wires)
+        expected_via_mat = apply_operation_tensordot(op, state, batched)
+        if op_wires == state_wires:
+            expected_via_kernel = self.grover_kernel_full_wires(state, wires, batched)
+        else:
+            expected_via_kernel = self.grover_kernel_partial_wires(state, wires, batched)
+
+        # Cast to interface and apply operation
+        state = jax.numpy.array(state)
+        out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+
+        assert qml.math.allclose(out, expected_via_mat)
+        assert qml.math.allclose(out, expected_via_kernel)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("op_wires, state_wires", [(2, 2), (3, 3), (9, 9), (3, 5), (9, 13)])
+    @pytest.mark.parametrize("batch_dim", [None, 1, 3])
+    def test_correctness_torch(self, op_wires, state_wires, batch_dim):
+        """Test that apply_operation is correct for GroverOperator for all dispatch branches
+        when applying it to a Torch state."""
+        import torch
+
+        batched = batch_dim is not None
+        shape = [batch_dim] + [2] * state_wires if batched else [2] * state_wires
+        # Input state
+        np.random.seed(752)
+        state = np.random.random(shape) + 1j * np.random.random(shape)
+
+        wires = list(range(op_wires))
+        op = qml.GroverOperator(wires)
+        expected_via_mat = apply_operation_tensordot(op, state, batched)
+        if op_wires == state_wires:
+            expected_via_kernel = self.grover_kernel_full_wires(state, wires, batched)
+        else:
+            expected_via_kernel = self.grover_kernel_partial_wires(state, wires, batched)
+
+        # Cast to interface and apply operation
+        state = torch.tensor(state, requires_grad=True)
+        out = apply_operation(op, state, is_state_batched=batched, debugger=None)
+
+        assert qml.math.allclose(out, expected_via_mat)
+        assert qml.math.allclose(out, expected_via_kernel)
+
 
 class TestMultiControlledXKernel:
     """Test the specialized kernel for MultiControlledX and its dispatching."""
@@ -953,7 +1189,7 @@ class TestMultiControlledXKernel:
 @pytest.mark.tf
 @pytest.mark.parametrize("op", (qml.PauliZ(8), qml.CNOT((5, 6))))
 def test_tf_large_state(op):
-    """ "Tests that custom kernels that use slicing fall back to a different method when
+    """Tests that custom kernels that use slicing fall back to a different method when
     the state has a large number of wires."""
     import tensorflow as tf
 
diff --git a/tests/templates/test_subroutines/test_grover.py b/tests/templates/test_subroutines/test_grover.py
index 987ef2d5d..4a21aea5c 100644
--- a/tests/templates/test_subroutines/test_grover.py
+++ b/tests/templates/test_subroutines/test_grover.py
@@ -92,6 +92,7 @@ decomp_3wires = [
     qml.PauliZ,
     qml.Hadamard,
     qml.Hadamard,
+    qml.GlobalPhase,
 ]
 
 
@@ -104,6 +105,7 @@ def decomposition_wires(wires):
         wires[2],
         wires[0],
         wires[1],
+        wires,
     ]
     return wire_order
 
@@ -149,7 +151,9 @@ def test_grover_diffusion_matrix(n_wires):
 
 
 def test_grover_diffusion_matrix_results():
-    """Test that the matrix gives the same result as when running the example in the documentation `here <https://pennylane.readthedocs.io/en/stable/code/api/pennylane.templates.subroutines.GroverOperator.html>`_"""
+    """Test that the matrix gives the same result as when running the example in the documentation
+    `here <https://pennylane.readthedocs.io/en/stable/code/api/pennylane.templates.subroutines.GroverOperator.html>`_
+    """
     n_wires = 3
     wires = list(range(n_wires))
 
@@ -201,9 +205,11 @@ def test_expand(wires):
 
     expected_wires = decomposition_wires(wires)
 
-    for actual_op, expected_class, expected_wires in zip(decomp, decomp_3wires, expected_wires):
+    assert len(decomp) == len(decomp_3wires) == len(expected_wires)
+
+    for actual_op, expected_class, expected_wire in zip(decomp, decomp_3wires, expected_wires):
         assert isinstance(actual_op, expected_class)
-        assert actual_op.wires == qml.wires.Wires(expected_wires)
+        assert actual_op.wires == qml.wires.Wires(expected_wire)
 
 
 @pytest.mark.parametrize("n_wires", [6, 13])
@@ -246,3 +252,13 @@ def test_matrix(tol):
     assert np.allclose(res_dynamic, expected, atol=tol, rtol=0)
     # reordering should not affect this particular matrix
     assert np.allclose(res_reordered, expected, atol=tol, rtol=0)
+
+
+@pytest.mark.parametrize("n_wires", [2, 3, 5])
+def test_decomposition_matrix(n_wires):
+    """Test that the decomposition and the matrix match."""
+    wires = list(range(n_wires))
+    op = qml.GroverOperator(wires)
+    mat1 = op.matrix()
+    mat2 = qml.matrix(qml.tape.QuantumScript(op.decomposition()))
+    assert np.allclose(mat1, mat2)