Added a custom gate application for Toffoli in default.qubit (#1249)

* first draft implementation of _apply_toffoli

* second draft _apply_toffoli

* testing first pass with output

* _apply_toffoli is working with forward test case

* passing all test cases

* deleted comments, working for full pytest. still need documentation

* documentation done

* fixed typo

* Update CHANGELOG.md

* replaced reshaping with stacking

* converting to 4 qubit tests

* still writing test cases and working on implementation, not ready for review yet

* passing 21/24 4-qubit 3-permutations

* passing all test cases, documented & ran black

* ran black on test_default_qubit.py

* wire ordering typo

* added `cntrl_min = cntrl_max ^ 1` for readability

Co-authored-by: antalszava <antalszava@gmail.com>

.github/CHANGELOG.md

@@ -37,6 +37,9 @@
 
 <h3>Improvements</h3>
 
+* Added custom gate application for Toffoli in `default.qubit`.
+  [(#1249)](https://github.com/PennyLaneAI/pennylane/pull/1249)
+
 * The device test suite now provides test cases for checking gates by comparing
   expectation values.
   [(#1212)](https://github.com/PennyLaneAI/pennylane/pull/1212)
@@ -95,6 +98,7 @@ Thomas Bromley, Olivia Di Matteo, Diego Guala, Anthony Hayes, Josh Izaac, Antal
 
 This release contains contributions from (in alphabetical order):
 
+Thomas Bromley, Olivia Di Matteo, Diego Guala, Anthony Hayes, Ryan Hill,
 Josh Izaac, Maria Schuld, Antal Száva.
 
 # Release 0.15.0

pennylane/devices/default_qubit.py

@@ -151,6 +151,7 @@ def __init__(self, wires, *, shots=None, cache=0, analytic=None):
             "CNOT": self._apply_cnot,
             "SWAP": self._apply_swap,
             "CZ": self._apply_cz,
+            "Toffoli": self._apply_toffoli,
         }
 
     def map_wires(self, wires):
@@ -336,6 +337,44 @@ def _apply_cnot(self, state, axes, **kwargs):
         state_x = self._apply_x(state[sl_1], axes=target_axes)
         return self._stack([state[sl_0], state_x], axis=axes[0])
 
+    def _apply_toffoli(self, state, axes, **kwargs):
+        """Applies a Toffoli gate by slicing along the axis of the greater control qubit, slicing
+        each of the resulting sub-arrays along the axis of the smaller control qubit, and then applying
+        an X transformation along the axis of the target qubit of the fourth sub-sub-array.
+
+        By performing two consecutive slices in this way, we are able to select all of the amplitudes with
+        a corresponding :math:`|11\rangle` for the two control qubits. 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
+        """
+        cntrl_max = np.argmax(axes[:2])
+        cntrl_min = cntrl_max ^ 1
+        sl_a0 = _get_slice(0, axes[cntrl_max], self.num_wires)
+        sl_a1 = _get_slice(1, axes[cntrl_max], self.num_wires)
+        sl_b0 = _get_slice(0, axes[cntrl_min], self.num_wires - 1)
+        sl_b1 = _get_slice(1, axes[cntrl_min], self.num_wires - 1)
+
+        # If both controls are smaller than the target, shift the target axis down by two. If one
+        # control is greater and one control is smaller than the target, shift the target axis
+        # down by one. If both controls are greater than the target, leave the target axis as-is.
+        if axes[cntrl_min] > axes[2]:
+            target_axes = [axes[2]]
+        elif axes[cntrl_max] > axes[2]:
+            target_axes = [axes[2] - 1]
+        else:
+            target_axes = [axes[2] - 2]
+
+        # state[sl_a1][sl_b1] gives us all of the amplitudes with a |11> for the two control qubits.
+        state_x = self._apply_x(state[sl_a1][sl_b1], axes=target_axes)
+        state_stacked_a1 = self._stack([state[sl_a1][sl_b0], state_x], axis=axes[cntrl_min])
+        return self._stack([state[sl_a0], state_stacked_a1], axis=axes[cntrl_max])
+
     def _apply_swap(self, state, axes, **kwargs):
         """Applies a SWAP gate by performing a partial transposition along the specified axes.
 

tests/devices/test_default_qubit.py

@@ -1909,8 +1909,8 @@ 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)
+    state = np.arange(2 ** 4).reshape((2, 2, 2, 2))
+    dev = qml.device("default.qubit", wires=4)
     single_qubit_ops = [
         (qml.PauliX, dev._apply_x),
         (qml.PauliY, dev._apply_y),
@@ -1925,14 +1925,17 @@ class TestApplyOps:
         (qml.SWAP, dev._apply_swap),
         (qml.CZ, dev._apply_cz),
     ]
+    three_qubit_ops = [
+        (qml.Toffoli, dev._apply_toffoli),
+    ]
 
     @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)
+        state_out_einsum = np.einsum("ab,ibjk->iajk", matrix, self.state)
         assert np.allclose(state_out, state_out_einsum)
 
     @pytest.mark.parametrize("op, method", two_qubit_ops)
@@ -1942,7 +1945,7 @@ def test_apply_two_qubit_op(self, op, method, inverse):
         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)
+        state_out_einsum = np.einsum("abcd,cdjk->abjk", matrix, self.state)
         assert np.allclose(state_out, state_out_einsum)
 
     @pytest.mark.parametrize("op, method", two_qubit_ops)
@@ -1953,7 +1956,40 @@ def test_apply_two_qubit_op_reverse(self, op, method, inverse):
         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)
+        state_out_einsum = np.einsum("abcd,idck->ibak", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)
+
+    @pytest.mark.parametrize("op, method", three_qubit_ops)
+    def test_apply_three_qubit_op_controls_smaller(self, op, method, inverse):
+        """Test if the application of three qubit operations is correct when both control wires are
+        smaller than the target wire."""
+        state_out = method(self.state, axes=[0, 2, 3])
+        op = op(wires=[0, 2, 3])
+        matrix = op.inv().matrix if inverse else op.matrix
+        matrix = matrix.reshape((2, 2) * 3)
+        state_out_einsum = np.einsum("abcdef,dkef->akbc", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)
+
+    @pytest.mark.parametrize("op, method", three_qubit_ops)
+    def test_apply_three_qubit_op_controls_greater(self, op, method, inverse):
+        """Test if the application of three qubit operations is correct when both control wires are
+        greater than the target wire."""
+        state_out = method(self.state, axes=[2, 1, 0])
+        op = op(wires=[2, 1, 0])
+        matrix = op.inv().matrix if inverse else op.matrix
+        matrix = matrix.reshape((2, 2) * 3)
+        state_out_einsum = np.einsum("abcdef,fedk->cbak", matrix, self.state)
+        assert np.allclose(state_out, state_out_einsum)
+
+    @pytest.mark.parametrize("op, method", three_qubit_ops)
+    def test_apply_three_qubit_op_controls_split(self, op, method, inverse):
+        """Test if the application of three qubit operations is correct when one control wire is smaller
+        and one control wire is greater than the target wire."""
+        state_out = method(self.state, axes=[3, 1, 2])
+        op = op(wires=[3, 1, 2])
+        matrix = op.inv().matrix if inverse else op.matrix
+        matrix = matrix.reshape((2, 2) * 3)
+        state_out_einsum = np.einsum("abcdef,kdfe->kacb", matrix, self.state)
         assert np.allclose(state_out, state_out_einsum)