Speed up GroverOperator and add GlobalPhase to its decomposition

doc/releases/changelog-dev.md

@@ -9,6 +9,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)
+
 * `AmplitudeEmbedding` now also supports batching when used with Tensorflow.
   [(#4818)](https://github.com/PennyLaneAI/pennylane/pull/4818)
 
@@ -75,6 +79,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)
+
 * Fixes a bug where the adjoint method differentiation would fail if
   an operation with `grad_method=None` that has a parameter is present.
   [(#4820)](https://github.com/PennyLaneAI/pennylane/pull/4820)

pennylane/devices/qubit/apply_operation.py

@@ -263,6 +263,40 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
     return math.stack([state[sl_0], state_x], axis=control_axes)
 
 
+@apply_operation.register
+def apply_grover(op: qml.GroverOperator, state, is_state_batched: bool = False, debugger=None):
+    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.
+    """
+    if (
+        len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD
+        and math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
+    ):
+        return apply_operation_einsum(op, state, is_state_batched)
+    if len(op.wires) < 9:
+        return apply_operation_tensordot(op, state, is_state_batched)
+
+    # 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]
+    n_wires = len(sum_axes)
+    collapsed = math.sum(state, axis=tuple(sum_axes))
+
+    # 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 - n_wires)
+    all_plus = math.cast_like(math.full([2] * n_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"""

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 @@ def compute_decomposition(
         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)

tests/devices/qubit/test_apply_operation.py

@@ -206,6 +206,30 @@ def test_cnot(self, method, wire, ml_framework):
         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 a cnot gate 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."""
 
@@ -824,11 +848,38 @@ def test_double_excitation(self, method):
 
         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 = np.copy(self.state)
+        for _op in op.decomposition():
+            expected_state = method(_op, expected_state)
+
+        assert qml.math.allclose(expected_state, new_state)
+
+
+@pytest.mark.parametrize(
+    "num_wires, einsum_called, tensordot_called",
+    [(2, True, False), (3, False, True), (9, False, False)],
+)
+def test_grover_dispatching(num_wires, einsum_called, tensordot_called, mocker):
+    """Test that apply_grover dispatches to einsum correctly for small numbers of wires."""
+    op = qml.GroverOperator(list(range(num_wires)))
+    state = np.zeros([2] * num_wires, dtype=complex)
+    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.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
 

tests/templates/test_subroutines/test_grover.py

@@ -246,3 +246,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), f"{mat1}\n{mat2}"