Speed up default.qubit by using permutations to apply some gates

.github/CHANGELOG.md

@@ -15,6 +15,11 @@
 
 <h3>Improvements</h3>
 
+* Sped up the application of certain gates in ``default.qubit`` by using array/tensor
+  manipulation tricks. The following gates are affected: ``PauliX``, ``PauliY``, ``PauliZ``,
+  ``Hadamard``, ``SWAP``, ``S``, ``T``, ``CNOT``, ``CZ``.
+  [(#772)](https://github.com/PennyLaneAI/pennylane/pull/772)
+
 * Adds arithmetic operations (addition, tensor product, 
   subtraction, and scalar multiplication) between ``Hamiltonian``, 
   ``Tensor``, and ``Observable`` objects, and inline arithmetic 

pennylane/_qubit_device.py

@@ -84,6 +84,8 @@ class QubitDevice(Device):
     _tensordot = staticmethod(np.tensordot)
     _conj = staticmethod(np.conj)
     _imag = staticmethod(np.imag)
+    _roll = staticmethod(np.roll)
+    _stack = staticmethod(np.stack)
 
     @staticmethod
     def _scatter(indices, array, new_dimensions):

pennylane/devices/default_qubit.py

@@ -31,8 +31,40 @@
 
 # tolerance for numerical errors
 tolerance = 1e-10
+SQRT2INV = 1 / np.sqrt(2)
+TPHASE = np.exp(1j * np.pi / 4)
 
 
+def _get_slice(index, axis, num_axes):
+    """Allows slicing along an arbitrary axis of an array or tensor.
+
+    Args:
+        index (int): the index to access
+        axis (int): the axis to slice into
+        num_axes (int): total number of axes
+
+    Returns:
+        tuple[slice or int]: a tuple that can be used to slice into an array or tensor
+
+    **Example:**
+
+    Accessing the 2 index along axis 1 of a 3-axis array:
+
+    >>> sl = _get_slice(2, 1, 3)
+    >>> sl
+    (slice(None, None, None), 2, slice(None, None, None))
+    >>> a = np.arange(27).reshape((3, 3, 3))
+    >>> a[sl]
+    array([[ 6,  7,  8],
+           [15, 16, 17],
+           [24, 25, 26]])
+    """
+    idx = [slice(None)] * num_axes
+    idx[axis] = index
+    return tuple(idx)
+
+
+# pylint: disable=unused-argument
 class DefaultQubit(QubitDevice):
     """Default qubit device for PennyLane.
 
@@ -95,6 +127,18 @@ def __init__(self, wires, *, shots=1000, analytic=True):
         self._state = self._create_basis_state(0)
         self._pre_rotated_state = self._state
 
+        self._apply_ops = {
+            "PauliX": self._apply_x,
+            "PauliY": self._apply_y,
+            "PauliZ": self._apply_z,
+            "Hadamard": self._apply_hadamard,
+            "S": self._apply_s,
+            "T": self._apply_t,
+            "CNOT": self._apply_cnot,
+            "SWAP": self._apply_swap,
+            "CZ": self._apply_cz,
+        }
+
     def apply(self, operations, rotations=None, **kwargs):
         rotations = rotations or []
 
@@ -132,6 +176,13 @@ def _apply_operation(self, operation):
             self._apply_basis_state(operation.parameters[0], wires)
             return
 
+        if operation.name in self._apply_ops:
+            axes = self.wires.indices(wires)
+            self._state = self._apply_ops[operation.name](
+                self._state, axes, inverse=operation.inverse
+            )
+            return
+
         matrix = self._get_unitary_matrix(operation)
 
         if isinstance(operation, DiagonalOperation):
@@ -142,6 +193,157 @@ def _apply_operation(self, operation):
         else:
             self._apply_unitary(matrix, wires)
 
+    def _apply_x(self, state, axes, **kwargs):
+        """Applies a PauliX gate by rolling 1 unit along the axis specified in ``axes``.
+
+        Rolling by 1 unit along the axis means that the :math:`|0 \rangle` state with index ``0`` is
+        shifted to the :math:`|1 \rangle` state with index ``1``. Likewise, since rolling beyond
+        the last index loops back to the first, :math:`|1 \rangle` is transformed to
+        :math:`|0\rangle`.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return self._roll(state, 1, axes[0])
+
+    def _apply_y(self, state, axes, **kwargs):
+        """Applies a PauliY gate by adding a negative sign to the 1 index along the axis specified
+        in ``axes``, rolling one unit along the same axis, and multiplying the result by 1j.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return 1j * self._apply_x(self._apply_z(state, axes), axes)
+
+    def _apply_z(self, state, axes, **kwargs):
+        """Applies a PauliZ gate by adding a negative sign to the 1 index along the axis specified
+        in ``axes``.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return self._apply_phase(state, axes, -1)
+
+    def _apply_hadamard(self, state, axes, **kwargs):
+        """Apply the Hadamard gate by combining the results of applying the PauliX and PauliZ gates.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        state_x = self._apply_x(state, axes)
+        state_z = self._apply_z(state, axes)
+        return SQRT2INV * (state_x + state_z)
+
+    def _apply_s(self, state, axes, inverse=False):
+        return self._apply_phase(state, axes, 1j, inverse)
+
+    def _apply_t(self, state, axes, inverse=False):
+        return self._apply_phase(state, axes, TPHASE, inverse)
+
+    def _apply_cnot(self, state, axes, **kwargs):
+        """Applies a CNOT gate by slicing along the first axis specified in ``axes`` and then
+        applying an X transformation along the second axis.
+
+        By slicing along the first axis, we are able to select all of the amplitudes with a
+        corresponding :math:`|1\rangle` for the control qubit. This means we then just need to apply
+        a :class:`~.PauliX` (NOT) gate to the result.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        sl_0 = _get_slice(0, axes[0], self.num_wires)
+        sl_1 = _get_slice(1, axes[0], self.num_wires)
+
+        # We will be slicing into the state according to state[sl_1], giving us all of the
+        # amplitudes with a |1> for the control qubit. The resulting array has lost an axis
+        # relative to state and we need to be careful about the axis we apply the PauliX rotation
+        # to. If axes[1] is larger than axes[0], then we need to shift the target axis down by
+        # one, otherwise we can leave as-is. For example: a state has [0, 1, 2, 3], control=1,
+        # target=3. Then, state[sl_1] has 3 axes and target=3 now corresponds to the second axis.
+        if axes[1] > axes[0]:
+            target_axes = [axes[1] - 1]
+        else:
+            target_axes = [axes[1]]
+
+        state_x = self._apply_x(state[sl_1], axes=target_axes)
+        return self._stack([state[sl_0], state_x], axis=axes[0])
+
+    def _apply_swap(self, state, axes, **kwargs):
+        """Applies a SWAP gate by performing a partial transposition along the specified axes.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        all_axes = list(range(len(state.shape)))
+        all_axes[axes[0]] = axes[1]
+        all_axes[axes[1]] = axes[0]
+        return self._transpose(state, all_axes)
+
+    def _apply_cz(self, state, axes, **kwargs):
+        """Applies a CZ gate by slicing along the first axis specified in ``axes`` and then
+        applying a Z transformation along the second axis.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        sl_0 = _get_slice(0, axes[0], self.num_wires)
+        sl_1 = _get_slice(1, axes[0], self.num_wires)
+
+        if axes[1] > axes[0]:
+            target_axes = [axes[1] - 1]
+        else:
+            target_axes = [axes[1]]
+
+        state_z = self._apply_z(state[sl_1], axes=target_axes)
+        return self._stack([state[sl_0], state_z], axis=axes[0])
+
+    def _apply_phase(self, state, axes, parameters, inverse=False):
+        """Applies a phase onto the 1 index along the axis specified in ``axes``.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+            parameters (float): phase to apply
+            inverse (bool): whether to apply the inverse phase
+
+        Returns:
+            array[complex]: output state
+        """
+        num_wires = len(state.shape)
+        sl_0 = _get_slice(0, axes[0], num_wires)
+        sl_1 = _get_slice(1, axes[0], num_wires)
+
+        phase = self._conj(parameters) if inverse else parameters
+        return self._stack([state[sl_0], phase * state[sl_1]], axis=axes[0])
+
     def _get_unitary_matrix(self, unitary):  # pylint: disable=no-self-use
         """Return the matrix representing a unitary operation.
 

pennylane/devices/default_qubit_autograd.py

@@ -106,6 +106,17 @@ class DefaultQubitAutograd(DefaultQubit):
     _tensordot = staticmethod(np.tensordot)
     _conj = staticmethod(np.conj)
     _imag = staticmethod(np.imag)
+    _roll = staticmethod(np.roll)
+    _stack = staticmethod(np.stack)
+
+    def __init__(self, wires, *, shots=1000, analytic=True):
+        super().__init__(wires, shots=shots, analytic=analytic)
+
+        # prevent using special apply methods for these gates due to slowdown in Autograd
+        # implementation
+        del self._apply_ops["PauliY"]
+        del self._apply_ops["Hadamard"]
+        del self._apply_ops["CZ"]
 
     @staticmethod
     def _scatter(indices, array, new_dimensions):

pennylane/devices/default_qubit_tf.py

@@ -156,6 +156,14 @@ class DefaultQubitTF(DefaultQubit):
     _tensordot = staticmethod(tf.tensordot)
     _conj = staticmethod(tf.math.conj)
     _imag = staticmethod(tf.math.conj)
+    _roll = staticmethod(tf.roll)
+    _stack = staticmethod(tf.stack)
+
+    def __init__(self, wires, *, shots=1000, analytic=True):
+        super().__init__(wires, shots=shots, analytic=analytic)
+
+        # prevent using special apply method for this gate due to slowdown in TF implementation
+        del self._apply_ops["CZ"]
 
     @staticmethod
     def _scatter(indices, array, new_dimensions):

tests/devices/test_default_qubit.py

@@ -21,6 +21,7 @@
 import pytest
 import pennylane as qml
 from pennylane import numpy as np, DeviceError
+from pennylane.devices.default_qubit import _get_slice
 from pennylane.operation import Operation
 
 U = np.array(
@@ -1733,3 +1734,103 @@ def test_wires_expval(self, wires1, wires2, tol):
         circuit2 = self.make_circuit_expval(wires2)
 
         assert np.allclose(circuit1(), circuit2(), tol)
+
+
+class TestGetSlice:
+    """Tests for the _get_slice function."""
+
+    def test_get_slice(self):
+        """Test that the _get_slice function returns the expected slice and allows us to slice
+        correctly into an array."""
+
+        sl = _get_slice(1, 1, 3)
+        array = np.arange(27).reshape((3, 3, 3))
+        target = array[:, 1, :]
+
+        assert sl == (slice(None, None, None), 1, slice(None, None, None))
+        assert np.allclose(array[sl], target)
+
+    def test_get_slice_first(self):
+        """Test that the _get_slice function returns the expected slice when accessing the first
+        axis of an array."""
+
+        sl = _get_slice(2, 0, 3)
+        array = np.arange(27).reshape((3, 3, 3))
+        target = array[2]
+
+        assert sl == (2, slice(None, None, None), slice(None, None, None))
+        assert np.allclose(array[sl], target)
+
+    def test_get_slice_last(self):
+        """Test that the _get_slice function returns the expected slice when accessing the last
+        axis of an array."""
+
+        sl = _get_slice(0, 2, 3)
+        array = np.arange(27).reshape((3, 3, 3))
+        target = array[:, :, 0]
+
+        assert sl == (slice(None, None, None), slice(None, None, None), 0)
+        assert np.allclose(array[sl], target)
+
+    def test_get_slice_1d(self):
+        """Test that the _get_slice function returns the expected slice when accessing a
+        1-dimensional array."""
+
+        sl = _get_slice(2, 0, 1)
+        array = np.arange(27)
+        target = array[2]
+
+        assert sl == (2,)
+        assert np.allclose(array[sl], target)
+
+
+@pytest.mark.parametrize("inverse", [True, False])
+class TestApplyOps:
+    """Tests for special methods listed in _apply_ops that use array manipulation tricks to apply
+    gates in DefaultQubit."""
+
+    state = np.arange(2 ** 3).reshape((2, 2, 2))
+    dev = qml.device("default.qubit", wires=3)
+    single_qubit_ops = [
+        (qml.PauliX, dev._apply_x),
+        (qml.PauliY, dev._apply_y),
+        (qml.PauliZ, dev._apply_z),
+        (qml.Hadamard, dev._apply_hadamard),
+        (qml.S, dev._apply_s),
+        (qml.T, dev._apply_t),
+    ]
+    two_qubit_ops = [
+        (qml.CNOT, dev._apply_cnot),
+        (qml.SWAP, dev._apply_swap),
+        (qml.CZ, dev._apply_cz),
+    ]
+
+    @pytest.mark.parametrize("op, method", single_qubit_ops)
+    def test_apply_single_qubit_op(self, op, method, inverse):
+        """Test if the application of single qubit operations is correct."""
+        state_out = method(self.state, axes=[1], inverse=inverse)
+        op = op(wires=[1])
+        matrix = op.inv().matrix if inverse else op.matrix
+        state_out_einsum = np.einsum("ab,ibk->iak", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)
+
+    @pytest.mark.parametrize("op, method", two_qubit_ops)
+    def test_apply_two_qubit_op(self, op, method, inverse):
+        """Test if the application of two qubit operations is correct."""
+        state_out = method(self.state, axes=[0, 1])
+        op = op(wires=[0, 1])
+        matrix = op.inv().matrix if inverse else op.matrix
+        matrix = matrix.reshape((2, 2, 2, 2))
+        state_out_einsum = np.einsum("abcd,cdk->abk", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)
+
+    @pytest.mark.parametrize("op, method", two_qubit_ops)
+    def test_apply_two_qubit_op_reverse(self, op, method, inverse):
+        """Test if the application of two qubit operations is correct when the applied wires are
+        reversed."""
+        state_out = method(self.state, axes=[2, 1])
+        op = op(wires=[2, 1])
+        matrix = op.inv().matrix if inverse else op.matrix
+        matrix = matrix.reshape((2, 2, 2, 2))
+        state_out_einsum = np.einsum("abcd,idc->iba", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)