Repository: PennyLaneAI/pennylane (python)

Issue:
All controlled operations should belong to the class `Controlled`
#### Issue description

Most of controlled gates are not in the `ControlledOperation` class, therefore you cannot easily check if an operation is controlled and you cannot access the control wires directly if you don't know the operation. Because of this behavior, it is complicated to write algorithms that use the fact that certain gates are controlled or not. Every controlled operation should be built from a target operation and a set of control wires within the class `ControlledOperation`.

* *Expected behavior:* (What you expect to happen)
A small example of what we should expect from any controlled operation:
```
op = qml.Toffoli(wires=[1,2,0])
print(isinstance(operation, ControlledOperation))
>>> True
print(op.control_wires())
>>> Wires(1,2)
print(op.target_wires())
>>> Wires(0)
print(op.target_operation)
>>> PauliX
```

* *Actual behavior:* (What actually happens)

```
op= qml.Toffoli(wires=[1,2,0])
print(isinstance(operation, ControlledOperation))
>>> False
```

I have started to work on this issue in the commutation DAG PR: https://github.com/PennyLaneAI/pennylane/pull/1446


Pull request:
Implement transform to push commuting operations through controlled gates
**Context:** Addition of compilation transforms to PennyLane.

**Description of the Change:** Adds a new transform, `commute_controlled` that pushes commuting gates through controls/targets of controlled operations. All diagonal gates are pushed ahead (or behind, if specified) of control qubits, and `X`/`Y`/`Z`-type gates are pushed through targets when appropriate. Additional compilation attributes and properties are also added to support the transform. 

**Benefits:** Moving as many gates as possible through controlled operations will enable more single-qubit gate simplifications to be performed, as well as fusion and cancellation of two-qubit controlled rotations once those in-between gates are out of the way.

**Possible Drawbacks:** Rather than adding attributes, it would be nice if we could instead resolve #1447 and have things inherit from `ControlledOperation`. However after a preliminary look, it seems that would be a pretty extensive refactor.

**Related GitHub Issues:** Uses some of the code and ideas from #1446 :tada:


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index e848655ed..746f13a01 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,33 @@
 
 <h3>New features since last release</h3>
 
+* A new quantum function transform has been added to push commuting
+  single-qubit gates through controlled operations.
+  [(#1464)](https://github.com/PennyLaneAI/pennylane/pull/1464)
+
+  The `commute_controlled` transform works as follows:
+
+  ```python
+  def circuit(theta):
+      qml.PauliX(wires=2)
+      qml.S(wires=0)
+      qml.CNOT(wires=[0, 1])
+      qml.PhaseShift(theta/2, wires=0)
+      qml.T(wires=0)
+      qml.Toffoli(wires=[0, 1, 2])
+      return qml.expval(qml.PauliZ(0))
+  ```
+
+  ```pycon
+  >>> optimized_circuit = qml.transforms.commute_controlled(direction="right")(circuit)
+  >>> dev = qml.device('default.qubit', wires=3)
+  >>> qnode = qml.QNode(optimized_circuit, dev)
+  >>> print(qml.draw(qnode)(0.5))
+   0: ──╭C──╭C──S──Rϕ(0.25)──T──┤ ⟨Z⟩
+   1: ──╰X──├C──────────────────┤
+   2: ──────╰X──X───────────────┤
+  ```
+
 * Grover Diffusion Operator template added.
   [(#1442)](https://github.com/PennyLaneAI/pennylane/pull/1442)
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 1332de2c8..99168eac5 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -582,6 +582,24 @@ class Operation(Operator):
     transforms that merge adjacent rotations.
     """
 
+    basis = None
+    """str or None: The basis of an operation, or for controlled gates, of the
+    target operation. If not ``None``, should take a value of ``"X"``, ``"Y"``,
+    or ``"Z"``.
+
+    For example, ``X`` and ``CNOT`` have ``basis = "X"``, whereas
+    ``ControlledPhaseShift`` and ``RZ`` have ``basis = "Z"``.
+    """
+
+    @property
+    def control_wires(self):  # pragma: no cover
+        r"""For operations that are controlled, returns the set of control wires.
+
+        Returns:
+            Wires: The set of control wires of the operation.
+        """
+        raise NotImplementedError
+
     @property
     def single_qubit_rot_angles(self):
         r"""The parameters required to implement a single-qubit gate as an
diff --git a/pennylane/ops/qubit.py b/pennylane/ops/qubit.py
index 9434ddccc..03827a30f 100644
--- a/pennylane/ops/qubit.py
+++ b/pennylane/ops/qubit.py
@@ -121,6 +121,7 @@ class PauliX(Observable, Operation):
     num_wires = 1
     par_domain = None
     is_self_inverse = True
+    basis = "X"
     eigvals = pauli_eigs(1)
     matrix = np.array([[0, 1], [1, 0]])
 
@@ -184,6 +185,7 @@ class PauliY(Observable, Operation):
     num_wires = 1
     par_domain = None
     is_self_inverse = True
+    basis = "Y"
     eigvals = pauli_eigs(1)
     matrix = np.array([[0, -1j], [1j, 0]])
 
@@ -249,6 +251,7 @@ class PauliZ(Observable, DiagonalOperation):
     num_wires = 1
     par_domain = None
     is_self_inverse = True
+    basis = "Z"
     eigvals = pauli_eigs(1)
     matrix = np.array([[1, 0], [0, -1]])
 
@@ -299,6 +302,7 @@ class S(DiagonalOperation):
     num_params = 0
     num_wires = 1
     par_domain = None
+    basis = "Z"
 
     @classmethod
     def _matrix(cls, *params):
@@ -341,6 +345,7 @@ class T(DiagonalOperation):
     num_params = 0
     num_wires = 1
     par_domain = None
+    basis = "Z"
 
     @classmethod
     def _matrix(cls, *params):
@@ -383,6 +388,7 @@ class SX(Operation):
     num_params = 0
     num_wires = 1
     par_domain = None
+    basis = "X"
 
     @classmethod
     def _matrix(cls, *params):
@@ -435,6 +441,7 @@ class CNOT(Operation):
     num_wires = 2
     par_domain = None
     is_self_inverse = True
+    basis = "X"
     matrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]])
 
     @classmethod
@@ -447,6 +454,10 @@ class CNOT(Operation):
     def _controlled(self, wire):
         Toffoli(wires=Wires(wire) + self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class CZ(DiagonalOperation):
     r"""CZ(wires)
@@ -474,6 +485,7 @@ class CZ(DiagonalOperation):
     par_domain = None
     is_self_inverse = True
     is_symmetric_over_all_wires = True
+    basis = "Z"
     eigvals = np.array([1, 1, 1, -1])
     matrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])
 
@@ -488,6 +500,10 @@ class CZ(DiagonalOperation):
     def adjoint(self):
         return CZ(wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class CY(Operation):
     r"""CY(wires)
@@ -514,6 +530,7 @@ class CY(Operation):
     num_wires = 2
     par_domain = None
     is_self_inverse = True
+    basis = "Y"
     matrix = np.array(
         [
             [1, 0, 0, 0],
@@ -535,6 +552,10 @@ class CY(Operation):
     def adjoint(self):
         return CY(wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class SWAP(Operation):
     r"""SWAP(wires)
@@ -718,6 +739,7 @@ class Toffoli(Operation):
     par_domain = None
     is_self_inverse = True
     is_symmetric_over_control_wires = True
+    basis = "X"
     matrix = np.array(
         [
             [1, 0, 0, 0, 0, 0, 0, 0],
@@ -759,6 +781,10 @@ class Toffoli(Operation):
     def adjoint(self):
         return Toffoli(wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[:2])
+
 
 class RX(Operation):
     r"""RX(phi, wires)
@@ -784,6 +810,7 @@ class RX(Operation):
     num_wires = 1
     par_domain = "R"
     is_composable_rotation = True
+    basis = "X"
     grad_method = "A"
     generator = [PauliX, -1 / 2]
 
@@ -830,6 +857,7 @@ class RY(Operation):
     num_wires = 1
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Y"
     grad_method = "A"
     generator = [PauliY, -1 / 2]
 
@@ -876,6 +904,7 @@ class RZ(DiagonalOperation):
     num_wires = 1
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Z"
     grad_method = "A"
     generator = [PauliZ, -1 / 2]
 
@@ -928,6 +957,7 @@ class PhaseShift(DiagonalOperation):
     num_wires = 1
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Z"
     grad_method = "A"
     generator = [np.array([[0, 0], [0, 1]]), 1]
 
@@ -985,6 +1015,7 @@ class ControlledPhaseShift(DiagonalOperation):
     num_wires = 2
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Z"
     grad_method = "A"
     generator = [np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1]]), 1]
 
@@ -1012,6 +1043,10 @@ class ControlledPhaseShift(DiagonalOperation):
     def adjoint(self):
         return ControlledPhaseShift(-self.data[0], wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 CPhase = ControlledPhaseShift
 
@@ -1407,6 +1442,7 @@ class CRX(Operation):
     num_wires = 2
     par_domain = "R"
     is_composable_rotation = True
+    basis = "X"
     grad_method = "A"
     grad_recipe = four_term_grad_recipe
 
@@ -1435,6 +1471,10 @@ class CRX(Operation):
     def adjoint(self):
         return CRX(-self.data[0], wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class CRY(Operation):
     r"""CRY(phi, wires)
@@ -1477,6 +1517,7 @@ class CRY(Operation):
     num_wires = 2
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Y"
     grad_method = "A"
     grad_recipe = four_term_grad_recipe
 
@@ -1503,6 +1544,10 @@ class CRY(Operation):
     def adjoint(self):
         return CRY(-self.data[0], wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class CRZ(DiagonalOperation):
     r"""CRZ(phi, wires)
@@ -1548,6 +1593,7 @@ class CRZ(DiagonalOperation):
     num_wires = 2
     par_domain = "R"
     is_composable_rotation = True
+    basis = "Z"
     grad_method = "A"
     grad_recipe = four_term_grad_recipe
 
@@ -1590,6 +1636,10 @@ class CRZ(DiagonalOperation):
     def adjoint(self):
         return CRZ(-self.data[0], wires=self.wires)
 
+    @property
+    def control_wires(self):
+        return Wires(self.wires[0])
+
 
 class CRot(Operation):
     r"""CRot(phi, theta, omega, wires)
@@ -2334,7 +2384,7 @@ class ControlledQubitUnitary(QubitUnitary):
 
         # Saving for the circuit drawer
         self._target_wires = wires
-        self.control_wires = control_wires
+        self._control_wires = control_wires
         self.U = U
 
         wires = control_wires + wires
@@ -2368,6 +2418,10 @@ class ControlledQubitUnitary(QubitUnitary):
         params[0] = self._CU
         return super()._matrix(*params)
 
+    @property
+    def control_wires(self):
+        return self._control_wires
+
     @staticmethod
     def _parse_control_values(control_wires, control_values):
         """Ensure any user-specified control strings have the right format."""
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 6817a1f48..3058e054f 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -45,6 +45,7 @@ containing quantum operations) that are used to construct QNodes.
 
     ~adjoint
     ~transforms.cancel_inverses
+    ~transforms.commute_controlled
     ~ctrl
     ~transforms.merge_rotations
     ~transforms.single_qubit_fusion
@@ -98,7 +99,12 @@ from .hamiltonian_expand import hamiltonian_expand
 from .invisible import invisible
 from .measurement_grouping import measurement_grouping
 from .metric_tensor import metric_tensor, metric_tensor_tape
-from .optimization import cancel_inverses, merge_rotations, single_qubit_fusion
+from .optimization import (
+    cancel_inverses,
+    commute_controlled,
+    merge_rotations,
+    single_qubit_fusion,
+)
 from .specs import specs
 from .qmc import apply_controlled_Q, quantum_monte_carlo
 from .unitary_to_rot import unitary_to_rot
diff --git a/pennylane/transforms/control.py b/pennylane/transforms/control.py
index 3a583b131..44884246d 100644
--- a/pennylane/transforms/control.py
+++ b/pennylane/transforms/control.py
@@ -95,12 +95,16 @@ class ControlledOperation(Operation):
         self.subtape = tape
         """QuantumTape: The tape that defines the underlying operation."""
 
-        self.control_wires = Wires(control_wires)
+        self._control_wires = Wires(control_wires)
         """Wires: The control wires."""
 
         wires = self.control_wires + tape.wires
         super().__init__(*tape.get_parameters(), wires=wires, do_queue=do_queue)
 
+    @property
+    def control_wires(self):
+        return self._control_wires
+
     def expand(self):
         tape = self.subtape
         for wire in self.control_wires:
diff --git a/pennylane/transforms/optimization/__init__.py b/pennylane/transforms/optimization/__init__.py
index 67d4989ce..7a3bbc697 100644
--- a/pennylane/transforms/optimization/__init__.py
+++ b/pennylane/transforms/optimization/__init__.py
@@ -16,5 +16,6 @@ This subpackage contains quantum function transforms for optimizing quantum circ
 """
 
 from .cancel_inverses import cancel_inverses
+from .commute_controlled import commute_controlled
 from .merge_rotations import merge_rotations
 from .single_qubit_fusion import single_qubit_fusion
diff --git a/pennylane/transforms/optimization/cancel_inverses.py b/pennylane/transforms/optimization/cancel_inverses.py
index 8e1536dd4..1fd553a71 100644
--- a/pennylane/transforms/optimization/cancel_inverses.py
+++ b/pennylane/transforms/optimization/cancel_inverses.py
@@ -28,6 +28,9 @@ def cancel_inverses(tape):
     Args:
         qfunc (function): A quantum function.
 
+    Returns:
+        function: the transformed quantum function
+
     **Example**
 
     Consider the following quantum function:
diff --git a/pennylane/transforms/optimization/commute_controlled.py b/pennylane/transforms/optimization/commute_controlled.py
new file mode 100644
index 000000000..e879f8e17
--- /dev/null
+++ b/pennylane/transforms/optimization/commute_controlled.py
@@ -0,0 +1,227 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transforms for pushing commuting gates through targets/control qubits."""
+
+from pennylane import apply
+from pennylane.wires import Wires
+from pennylane.transforms import qfunc_transform
+
+from .optimization_utils import find_next_gate
+
+
+def _commute_controlled_right(op_list):
+    """Push commuting single qubit gates to the right of controlled gates.
+
+    Args:
+        op_list (list[Operation]): The initial list of operations.
+
+    Returns:
+        list[Operation]: The modified list of operations with all single-qubit
+        gates as far right as possible.
+    """
+    # We will go through the list backwards; whenever we find a single-qubit
+    # gate, we will extract it and push it through 2-qubit gates as far as
+    # possible to the right.
+    current_location = len(op_list) - 1
+
+    while current_location >= 0:
+
+        current_gate = op_list[current_location]
+
+        # We are looking only at the gates that can be pushed through
+        # controls/targets; these are single-qubit gates with the basis
+        # property specified.
+        if current_gate.basis is None or len(current_gate.wires) != 1:
+            current_location -= 1
+            continue
+
+        # Find the next gate that contains an overlapping wire
+        next_gate_idx = find_next_gate(current_gate.wires, op_list[current_location + 1 :])
+
+        new_location = current_location
+
+        # Loop as long as a valid next gate exists
+        while next_gate_idx is not None:
+            next_gate = op_list[new_location + next_gate_idx + 1]
+
+            # Only go ahead if information is available
+            if next_gate.basis is None:
+                break
+
+            # If the next gate does not have control_wires defined, it is not
+            # controlled so we can't push through.
+            try:
+                shared_controls = Wires.shared_wires(
+                    [Wires(current_gate.wires), next_gate.control_wires]
+                )
+            except (NotImplementedError, AttributeError):
+                break
+
+            # Case 1: overlap is on the control wires. Only Z-type gates go through
+            if len(shared_controls) > 0:
+                if current_gate.basis == "Z":
+                    new_location += next_gate_idx + 1
+                else:
+                    break
+
+            # Case 2: since we know the gates overlap somewhere, and it's a
+            # single-qubit gate, if it wasn't on a control it's the target.
+            else:
+                if current_gate.basis == next_gate.basis:
+                    new_location += next_gate_idx + 1
+                else:
+                    break
+
+            next_gate_idx = find_next_gate(current_gate.wires, op_list[new_location + 1 :])
+
+        # After we have gone as far as possible, move the gate to new location
+        op_list.insert(new_location + 1, current_gate)
+        op_list.pop(current_location)
+        current_location -= 1
+
+    return op_list
+
+
+def _commute_controlled_left(op_list):
+    """Push commuting single qubit gates to the left of controlled gates.
+
+    Args:
+        op_list (list[Operation]): The initial list of operations.
+
+    Returns:
+        list[Operation]: The modified list of operations with all single-qubit
+        gates as far left as possible.
+    """
+    # We will go through the list forwards; whenever we find a single-qubit
+    # gate, we will extract it and push it through 2-qubit gates as far as
+    # possible back to the left.
+    current_location = 0
+
+    while current_location < len(op_list):
+
+        current_gate = op_list[current_location]
+
+        if current_gate.basis is None or len(current_gate.wires) != 1:
+            current_location += 1
+            continue
+
+        # Pass a backwards copy of the list
+        prev_gate_idx = find_next_gate(current_gate.wires, op_list[:current_location][::-1])
+
+        new_location = current_location
+
+        while prev_gate_idx is not None:
+            prev_gate = op_list[new_location - prev_gate_idx - 1]
+
+            if prev_gate.basis is None:
+                break
+
+            try:
+                shared_controls = Wires.shared_wires(
+                    [Wires(current_gate.wires), prev_gate.control_wires]
+                )
+            except (NotImplementedError, AttributeError):
+                break
+
+            if len(shared_controls) > 0:
+                if current_gate.basis == "Z":
+                    new_location = new_location - prev_gate_idx - 1
+                else:
+                    break
+
+            else:
+                if current_gate.basis == prev_gate.basis:
+                    new_location = new_location - prev_gate_idx - 1
+                else:
+                    break
+
+            prev_gate_idx = find_next_gate(current_gate.wires, op_list[:new_location][::-1])
+
+        op_list.pop(current_location)
+        op_list.insert(new_location, current_gate)
+        current_location += 1
+
+    return op_list
+
+
+@qfunc_transform
+def commute_controlled(tape, direction="right"):
+    """Quantum function transform to move commuting gates past
+    control and target qubits of controlled operations.
+
+    Args:
+        qfunc (function): A quantum function.
+        direction (str): The direction in which to move single-qubit gates.
+            Options are "right" (default), or "left". Single-qubit gates will
+            be pushed through controlled operations as far as possible in the
+            specified direction.
+
+    Returns:
+        function: the transformed quantum function
+
+    **Example**
+
+    Consider the following quantum function.
+
+    .. code-block:: python
+
+        def qfunc(theta):
+            qml.CZ(wires=[0, 2])
+            qml.PauliX(wires=2)
+            qml.S(wires=0)
+
+            qml.CNOT(wires=[0, 1])
+
+            qml.PauliY(wires=1)
+            qml.CRY(theta, wires=[0, 1])
+            qml.PhaseShift(theta/2, wires=0)
+
+            qml.Toffoli(wires=[0, 1, 2])
+            qml.T(wires=0)
+            qml.RZ(theta/2, wires=1)
+
+            return qml.expval(qml.PauliZ(0))
+
+    >>> dev = qml.device('default.qubit', wires=3)
+    >>> qnode = qml.QNode(qfunc, dev)
+    >>> print(qml.draw(qnode)(0.5))
+     0: ──╭C──S──╭C─────╭C────────Rϕ(0.25)──╭C──T─────────┤ ⟨Z⟩
+     1: ──│──────╰X──Y──╰RY(0.5)────────────├C──RZ(0.25)──┤
+     2: ──╰Z──X─────────────────────────────╰X────────────┤
+
+    Diagonal gates on either side of control qubits do not affect the outcome
+    of controlled gates; thus we can push all the single-qubit gates on the
+    first qubit together on the right (and fuse them if desired). Similarly, X
+    gates commute with the target of ``CNOT`` and ``Toffoli`` (and ``PauliY``
+    with ``CRY``). We can use the transform to push single-qubit gates as
+    far as possible through the controlled operations:
+
+    >>> optimized_qfunc = commute_controlled(direction="right")(qfunc)
+    >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
+    >>> print(qml.draw(optimized_qnode)(0.5))
+     0: ──╭C──╭C──╭C───────────╭C──S─────────Rϕ(0.25)──T──┤ ⟨Z⟩
+     1: ──│───╰X──╰RY(0.5)──Y──├C──RZ(0.25)───────────────┤
+     2: ──╰Z───────────────────╰X──X──────────────────────┤
+    """
+    if direction not in ("left", "right"):
+        raise ValueError("Direction for commute_controlled must be 'left' or 'right'")
+
+    if direction == "right":
+        op_list = _commute_controlled_right(tape.operations)
+    else:
+        op_list = _commute_controlled_left(tape.operations)
+
+    # Once the list is rearranged, queue all the operations
+    for op in op_list + tape.measurements:
+        apply(op)
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index f78a9d5ef..bf239cc7c 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -37,6 +37,9 @@ def merge_rotations(tape, atol=1e-8, include_gates=None):
             attribute set to ``True`` will be merged. Otherwise, only the operations whose
             names match those in the list will undergo merging.
 
+    Returns:
+        function: the transformed quantum function
+
     **Example**
 
     Consider the following quantum function.
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 8592fc725..de6482009 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -39,6 +39,9 @@ def single_qubit_fusion(tape, atol=1e-8, exclude_gates=None):
             from full fusion. If set to ``None``, all single-qubit gates that can
             be fused will be fused.
 
+    Returns:
+        function: the transformed quantum function
+
     **Example**
 
     Consider the following quantum function.


Test patch:
diff --git a/tests/transforms/test_optimization/test_commute_controlled.py b/tests/transforms/test_optimization/test_commute_controlled.py
new file mode 100644
index 000000000..f02a2306c
--- /dev/null
+++ b/tests/transforms/test_optimization/test_commute_controlled.py
@@ -0,0 +1,440 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from pennylane import numpy as np
+import pennylane as qml
+from pennylane.wires import Wires
+
+from pennylane.transforms.optimization import commute_controlled
+from utils import (
+    compare_operation_lists,
+    compute_matrix_from_ops_two_qubit,
+    check_matrix_equivalence,
+)
+
+
+class TestCommuteControlled:
+    """Tests for single-qubit gates being pushed through controlled gates."""
+
+    def test_invalid_direction(self):
+        """Test that any direction other than 'left' or 'right' raises an error."""
+
+        def qfunc():
+            qml.PauliX(wires=2)
+            qml.CNOT(wires=[0, 2])
+            qml.RX(0.2, wires=2)
+
+        transformed_qfunc = commute_controlled(direction="sideways")(qfunc)
+
+        with pytest.raises(ValueError, match="must be 'left' or 'right'"):
+            ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+    @pytest.mark.parametrize("direction", [("left"), ("right")])
+    def test_gate_with_no_basis(self, direction):
+        """Test that gates with no basis specified are ignored."""
+
+        def qfunc():
+            qml.PauliX(wires=2)
+            qml.ControlledQubitUnitary(np.array([[0, 1], [1, 0]]), control_wires=0, wires=2)
+            qml.PauliX(wires=2)
+
+        transformed_qfunc = commute_controlled(direction=direction)(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+        print(ops)
+
+        names_expected = ["PauliX", "ControlledQubitUnitary", "PauliX"]
+        wires_expected = [Wires(2), Wires([0, 2]), Wires(2)]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    @pytest.mark.parametrize("direction", [("left"), ("right")])
+    def test_gate_blocked_different_basis(self, direction):
+        """Test that gates do not get pushed through controlled gates whose target bases don't match."""
+
+        def qfunc():
+            qml.PauliZ(wires="b")
+            qml.CNOT(wires=[2, "b"])
+            qml.PauliY(wires="b")
+
+        transformed_qfunc = commute_controlled(direction=direction)(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+        print(ops)
+
+        names_expected = ["PauliZ", "CNOT", "PauliY"]
+        wires_expected = [Wires("b"), Wires([2, "b"]), Wires("b")]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_x_gates_right(self):
+        """Test that X-basis gates before controlled-X-type gates on targets get pushed ahead."""
+
+        def qfunc():
+            qml.PauliX(wires=2)
+            qml.CNOT(wires=[0, 2])
+            qml.RX(0.2, wires=2)
+            qml.Toffoli(wires=[0, 1, 2])
+            qml.SX(wires=1)
+            qml.PauliX(wires=1)
+            qml.CRX(0.1, wires=[0, 1])
+
+        transformed_qfunc = commute_controlled()(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["CNOT", "Toffoli", "PauliX", "RX", "CRX", "SX", "PauliX"]
+        wires_expected = [
+            Wires([0, 2]),
+            Wires([0, 1, 2]),
+            Wires(2),
+            Wires(2),
+            Wires([0, 1]),
+            Wires(1),
+            Wires(1),
+        ]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_x_gates_left(self):
+        """Test that X-basis gates after controlled-X-type gates on targets get pushed back."""
+
+        def qfunc():
+            qml.CNOT(wires=[0, 2])
+            qml.PauliX(wires=2)
+            qml.RX(0.2, wires=2)
+            qml.Toffoli(wires=[0, 1, 2])
+            qml.CRX(0.1, wires=[0, 1])
+            qml.SX(wires=1)
+            qml.PauliX(wires=1)
+
+        transformed_qfunc = commute_controlled(direction="left")(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = [
+            "PauliX",
+            "RX",
+            "CNOT",
+            "Toffoli",
+            "SX",
+            "PauliX",
+            "CRX",
+        ]
+        wires_expected = [
+            Wires(2),
+            Wires(2),
+            Wires([0, 2]),
+            Wires([0, 1, 2]),
+            Wires(1),
+            Wires(1),
+            Wires([0, 1]),
+        ]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    @pytest.mark.parametrize("direction", [("left"), ("right")])
+    def test_dont_push_x_gates(self, direction):
+        """Test that X-basis gates before controlled-X-type gates on controls don't get pushed."""
+
+        def qfunc():
+            qml.PauliX(wires="a")
+            qml.CNOT(wires=["a", "c"])
+            qml.RX(0.2, wires="a")
+            qml.Toffoli(wires=["c", "a", "b"])
+
+        transformed_qfunc = commute_controlled(direction=direction)(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["PauliX", "CNOT", "RX", "Toffoli"]
+        wires_expected = [Wires("a"), Wires(["a", "c"]), Wires("a"), Wires(["c", "a", "b"])]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_y_gates_right(self):
+        """Test that Y-basis gates before controlled-Y-type gates on targets get pushed ahead."""
+
+        def qfunc():
+            qml.PauliY(wires=2)
+            qml.CRY(-0.5, wires=["a", 2])
+            qml.CNOT(wires=[1, 2])
+            qml.RY(0.3, wires=1)
+            qml.CY(wires=["a", 1])
+
+        transformed_qfunc = commute_controlled()(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["CRY", "PauliY", "CNOT", "CY", "RY"]
+        wires_expected = [Wires(["a", 2]), Wires(2), Wires([1, 2]), Wires(["a", 1]), Wires(1)]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_y_gates_left(self):
+        """Test that Y-basis gates after controlled-Y-type gates on targets get pushed behind."""
+
+        def qfunc():
+            qml.CRY(-0.5, wires=["a", 2])
+            qml.PauliY(wires=2)
+            qml.CNOT(wires=[1, 2])
+            qml.CY(wires=["a", 1])
+            qml.RY(0.3, wires=1)
+
+        transformed_qfunc = commute_controlled(direction="left")(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["PauliY", "CRY", "CNOT", "RY", "CY"]
+        wires_expected = [Wires(2), Wires(["a", 2]), Wires([1, 2]), Wires(1), Wires(["a", 1])]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    @pytest.mark.parametrize("direction", [("left"), ("right")])
+    def test_dont_push_y_gates(self, direction):
+        """Test that Y-basis gates next to controlled-Y-type gates on controls don't get pushed."""
+
+        def qfunc():
+            qml.CRY(-0.2, wires=["a", 2])
+            qml.PauliY(wires="a")
+            qml.CNOT(wires=[1, 2])
+            qml.CY(wires=["a", 1])
+            qml.RY(0.3, wires="a")
+
+        transformed_qfunc = commute_controlled()(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["CRY", "PauliY", "CNOT", "CY", "RY"]
+        wires_expected = [Wires(["a", 2]), Wires("a"), Wires([1, 2]), Wires(["a", 1]), Wires("a")]
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_z_gates_right(self):
+        """Test that Z-basis gates before controlled-Z-type gates on controls *and* targets get pushed ahead."""
+
+        def qfunc():
+            qml.PauliZ(wires=2)
+            qml.S(wires=0)
+            qml.CZ(wires=[0, 2])
+
+            qml.CNOT(wires=[0, 1])
+
+            qml.PhaseShift(0.2, wires=2)
+            qml.T(wires=0)
+            qml.PauliZ(wires=0)
+            qml.CRZ(0.5, wires=[0, 1])
+
+        transformed_qfunc = commute_controlled()(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["CZ", "PauliZ", "CNOT", "PhaseShift", "CRZ", "S", "T", "PauliZ"]
+        wires_expected = (
+            [Wires([0, 2]), Wires(2), Wires([0, 1]), Wires(2)] + [Wires([0, 1])] + [Wires(0)] * 3
+        )
+
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    def test_push_z_gates_left(self):
+        """Test that Z-basis after before controlled-Z-type gates on controls *and*
+        targets get pushed behind."""
+
+        def qfunc():
+            qml.CZ(wires=[0, 2])
+            qml.PauliZ(wires=2)
+            qml.S(wires=0)
+
+            qml.CNOT(wires=[0, 1])
+
+            qml.CRZ(0.5, wires=[0, 1])
+            qml.RZ(0.2, wires=2)
+            qml.T(wires=0)
+            qml.PauliZ(wires=0)
+
+        transformed_qfunc = commute_controlled(direction="left")(qfunc)
+
+        ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        names_expected = ["PauliZ", "S", "RZ", "T", "PauliZ", "CZ", "CNOT", "CRZ"]
+        wires_expected = [Wires(2), Wires(0), Wires(2), Wires(0), Wires(0), Wires([0, 2])] + [
+            Wires([0, 1])
+        ] * 2
+
+        compare_operation_lists(ops, names_expected, wires_expected)
+
+    @pytest.mark.parametrize("direction", [("left"), ("right")])
+    def test_push_mixed_with_matrix(self, direction):
+        """Test that arbitrary gates after controlled gates on controls *and*
+        targets get properly pushed."""
+
+        def qfunc():
+            qml.PauliX(wires=1)
+            qml.S(wires=0)
+            qml.CZ(wires=[0, 1])
+            qml.CNOT(wires=[1, 0])
+            qml.PauliY(wires=1)
+            qml.CRY(0.5, wires=[1, 0])
+            qml.PhaseShift(0.2, wires=0)
+            qml.PauliY(wires=1)
+            qml.T(wires=0)
+            qml.CRZ(-0.3, wires=[0, 1])
+            qml.RZ(0.2, wires=0)
+            qml.PauliZ(wires=0)
+            qml.PauliX(wires=1)
+            qml.CRY(0.2, wires=[1, 0])
+
+        transformed_qfunc = commute_controlled()(qfunc)
+
+        original_ops = qml.transforms.make_tape(qfunc)().operations
+        transformed_ops = qml.transforms.make_tape(transformed_qfunc)().operations
+
+        assert len(original_ops) == len(transformed_ops)
+
+        # Compare matrices
+        matrix_expected = compute_matrix_from_ops_two_qubit(original_ops, wire_order=[0, 1])
+        matrix_obtained = compute_matrix_from_ops_two_qubit(transformed_ops, wire_order=[0, 1])
+        assert check_matrix_equivalence(matrix_expected, matrix_obtained)
+
+
+# Example QNode and device for interface testing
+dev = qml.device("default.qubit", wires=3)
+
+
+def qfunc(theta):
+    qml.PauliX(wires=2)
+    qml.S(wires=0)
+    qml.CNOT(wires=[0, 1])
+    qml.PauliY(wires=1)
+    qml.CRY(theta[0], wires=[2, 1])
+    qml.PhaseShift(theta[1], wires=0)
+    qml.T(wires=0)
+    qml.Toffoli(wires=[0, 1, 2])
+    return qml.expval(qml.PauliZ(0))
+
+
+transformed_qfunc = commute_controlled()(qfunc)
+
+expected_op_list = ["PauliX", "CNOT", "CRY", "PauliY", "Toffoli", "S", "PhaseShift", "T"]
+expected_wires_list = [
+    Wires(2),
+    Wires([0, 1]),
+    Wires([2, 1]),
+    Wires(1),
+    Wires([0, 1, 2]),
+    Wires(0),
+    Wires(0),
+    Wires(0),
+]
+
+
+class TestCommuteControlledInterfaces:
+    """Test that single-qubit gates can be pushed through controlled gates in all interfaces."""
+
+    def test_commute_controlled_autograd(self):
+        """Test QNode and gradient in autograd interface."""
+
+        original_qnode = qml.QNode(qfunc, dev, interface="autograd")
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, interface="autograd")
+
+        input = np.array([0.1, 0.2], requires_grad=True)
+
+        # Check that the numerical output is the same
+        assert qml.math.allclose(original_qnode(input), transformed_qnode(input))
+
+        # Check that the gradient is the same
+        assert qml.math.allclose(
+            qml.grad(original_qnode)(input), qml.grad(transformed_qnode)(input)
+        )
+
+        # Check operation list
+        ops = transformed_qnode.qtape.operations
+        compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    def test_commute_controlled_torch(self):
+        """Test QNode and gradient in torch interface."""
+        torch = pytest.importorskip("torch", minversion="1.8")
+
+        original_qnode = qml.QNode(qfunc, dev, interface="torch")
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, interface="torch")
+
+        original_input = torch.tensor([1.2, -0.35], requires_grad=True)
+        transformed_input = torch.tensor([1.2, -0.35], requires_grad=True)
+
+        original_result = original_qnode(original_input)
+        transformed_result = transformed_qnode(transformed_input)
+
+        # Check that the numerical output is the same
+        assert qml.math.allclose(original_result, transformed_result)
+
+        # Check that the gradient is the same
+        original_result.backward()
+        transformed_result.backward()
+
+        assert qml.math.allclose(original_input.grad, transformed_input.grad)
+
+        # Check operation list
+        ops = transformed_qnode.qtape.operations
+        compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    def test_commute_controlled_tf(self):
+        """Test QNode and gradient in tensorflow interface."""
+        tf = pytest.importorskip("tensorflow")
+
+        original_qnode = qml.QNode(qfunc, dev, interface="tf")
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, interface="tf")
+
+        original_input = tf.Variable([0.8, -0.6])
+        transformed_input = tf.Variable([0.8, -0.6])
+
+        original_result = original_qnode(original_input)
+        transformed_result = transformed_qnode(transformed_input)
+
+        # Check that the numerical output is the same
+        assert qml.math.allclose(original_result, transformed_result)
+
+        # Check that the gradient is the same
+        with tf.GradientTape() as tape:
+            loss = original_qnode(original_input)
+        original_grad = tape.gradient(loss, original_input)
+
+        with tf.GradientTape() as tape:
+            loss = transformed_qnode(transformed_input)
+        transformed_grad = tape.gradient(loss, transformed_input)
+
+        assert qml.math.allclose(original_grad, transformed_grad)
+
+        # Check operation list
+        ops = transformed_qnode.qtape.operations
+        compare_operation_lists(ops, expected_op_list, expected_wires_list)
+
+    def test_commute_controlled_jax(self):
+        """Test QNode and gradient in JAX interface."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        from jax.config import config
+
+        remember = config.read("jax_enable_x64")
+        config.update("jax_enable_x64", True)
+
+        original_qnode = qml.QNode(qfunc, dev, interface="jax")
+        transformed_qnode = qml.QNode(transformed_qfunc, dev, interface="jax")
+
+        input = jnp.array([0.3, 0.4], dtype=jnp.float64)
+
+        # Check that the numerical output is the same
+        assert qml.math.allclose(original_qnode(input), transformed_qnode(input))
+
+        # Check that the gradient is the same
+        assert qml.math.allclose(
+            jax.grad(original_qnode)(input), jax.grad(transformed_qnode)(input)
+        )
+
+        # Check operation list
+        ops = transformed_qnode.qtape.operations
+        compare_operation_lists(ops, expected_op_list, expected_wires_list)
diff --git a/tests/transforms/test_optimization/utils.py b/tests/transforms/test_optimization/utils.py
index 2bf198a05..b6dc9f8cf 100644
--- a/tests/transforms/test_optimization/utils.py
+++ b/tests/transforms/test_optimization/utils.py
@@ -40,8 +40,10 @@ def compute_matrix_from_ops_two_qubit(ops, wire_order):
             else:
                 mat = qml.math.dot(qml.math.kron(I, op.matrix), mat)
         else:
-            mat = qml.math.dot(op.matrix, mat)
-
+            if op.wires == qml.wires.Wires(wire_order):
+                mat = qml.math.dot(op.matrix, mat)
+            else:
+                mat = qml.math.dot(qml.SWAP(wires=[0, 1]).matrix, qml.math.dot(op.matrix, mat))
     return mat