Repository: quantumlib/Cirq

Issue:
Implement ControlledOperation._apply_unitary_
- Copy over the implementation from ControlledGate and tweak it to work in the new context
- Then rewrite ControlledGate's apply unitary to delegate to ControlledOperation's


Discussion:
Happy to pick this up.
(Comes after https://github.com/quantumlib/Cirq/pull/1320)
I understand the first bit :) 
The second bit I have a problem with - not sure how to delegate to ControlledOperation's method as `ApplyUnitaryArgs` only knows the indices of the qubits - hence there are no actual qubits passed to  `_apply_unitary_` - how do I create an `ControlledOperation` instance without qubits from the gate? Maybe I'm missing something and qubits can be recovered from `ApplyUnitaryArgs`? 

Resolving patch:
diff --git a/cirq/ops/controlled_gate.py b/cirq/ops/controlled_gate.py
--- a/cirq/ops/controlled_gate.py
+++ b/cirq/ops/controlled_gate.py
@@ -16,6 +16,7 @@
 
 import numpy as np
 
+import cirq
 from cirq import linalg, protocols, value
 from cirq.ops import raw_types, controlled_operation as cop
 from cirq.type_workarounds import NotImplementedType
@@ -54,35 +55,11 @@ def _value_equality_values_(self):
         return self.sub_gate
 
     def _apply_unitary_(self, args: protocols.ApplyUnitaryArgs) -> np.ndarray:
-        control = args.axes[0]
-        rest = args.axes[1:]
-        active = linalg.slice_for_qubits_equal_to([control], 1)
-        sub_axes = [r - int(r > control) for r in rest]
-        target_view = args.target_tensor[active]
-        buffer_view = args.available_buffer[active]
-        result = protocols.apply_unitary(
-            self.sub_gate,
-            protocols.ApplyUnitaryArgs(
-                target_view,
-                buffer_view,
-                sub_axes),
-            default=NotImplemented)
+        qubits = cirq.LineQubit.range(1 + self.sub_gate.num_qubits())
+        c_op = cop.ControlledOperation(qubits[0],
+                                       self.sub_gate.on(*qubits[1:]))
 
-        if result is NotImplemented:
-            return NotImplemented
-
-        if result is target_view:
-            return args.target_tensor
-
-        if result is buffer_view:
-            inactive = linalg.slice_for_qubits_equal_to([control], 0)
-            args.available_buffer[inactive] = args.target_tensor[inactive]
-            return args.available_buffer
-
-        # HACK: assume they didn't somehow escape the slice view and edit the
-        # rest of target_tensor.
-        args.target_tensor[active] = result
-        return args.target_tensor
+        return protocols.apply_unitary(c_op, args, default=NotImplemented)
 
     def _has_unitary_(self) -> bool:
         return protocols.has_unitary(self.sub_gate)
@@ -138,4 +115,3 @@ def __str__(self):
 
     def __repr__(self):
         return 'cirq.ControlledGate(sub_gate={!r})'.format(self.sub_gate)
-
diff --git a/cirq/ops/controlled_operation.py b/cirq/ops/controlled_operation.py
--- a/cirq/ops/controlled_operation.py
+++ b/cirq/ops/controlled_operation.py
@@ -47,6 +47,37 @@ def _decompose_(self):
     def _value_equality_values_(self):
         return tuple([self.control, self.sub_operation])
 
+    def _apply_unitary_(self, args: protocols.ApplyUnitaryArgs) -> np.ndarray:
+        control = args.axes[0]
+        rest = args.axes[1:]
+        active = linalg.slice_for_qubits_equal_to([control], 1)
+        sub_axes = [r - int(r > control) for r in rest]
+        target_view = args.target_tensor[active]
+        buffer_view = args.available_buffer[active]
+        result = protocols.apply_unitary(
+            self.sub_operation,
+            protocols.ApplyUnitaryArgs(
+                target_view,
+                buffer_view,
+                sub_axes),
+            default=NotImplemented)
+
+        if result is NotImplemented:
+            return NotImplemented
+
+        if result is target_view:
+            return args.target_tensor
+
+        if result is buffer_view:
+            inactive = linalg.slice_for_qubits_equal_to([control], 0)
+            args.available_buffer[inactive] = args.target_tensor[inactive]
+            return args.available_buffer
+
+        # HACK: assume they didn't somehow escape the slice view and edit the
+        # rest of target_tensor.
+        args.target_tensor[active] = result
+        return args.target_tensor
+
     def _has_unitary_(self) -> bool:
         return protocols.has_unitary(self.sub_operation)
 


Test patch:
diff --git a/cirq/ops/controlled_gate_test.py b/cirq/ops/controlled_gate_test.py
--- a/cirq/ops/controlled_gate_test.py
+++ b/cirq/ops/controlled_gate_test.py
@@ -11,6 +11,7 @@
 # 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.
+
 from typing import Union, Tuple, cast
 
 import numpy as np
@@ -20,6 +21,48 @@
 from cirq.type_workarounds import NotImplementedType
 
 
+class GateUsingWorkspaceForApplyUnitary(cirq.SingleQubitGate):
+    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
+                        ) -> Union[np.ndarray, NotImplementedType]:
+        args.available_buffer[...] = args.target_tensor
+        args.target_tensor[...] = 0
+        return args.available_buffer
+
+    def _unitary_(self):
+        return np.eye(2)
+
+    def __eq__(self, other):
+        return isinstance(other, type(self))
+
+    def __repr__(self):
+        return ('cirq.ops.controlled_gate_test.'
+                'GateUsingWorkspaceForApplyUnitary()')
+
+
+class GateAllocatingNewSpaceForResult(cirq.SingleQubitGate):
+    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
+                        ) -> Union[np.ndarray, NotImplementedType]:
+        assert len(args.axes) == 1
+        a = args.axes[0]
+        seed = cast(Tuple[Union[int, slice, 'ellipsis'], ...],
+                    (slice(None),))
+        zero = seed*a + (0, Ellipsis)
+        one = seed*a + (1, Ellipsis)
+        result = np.zeros(args.target_tensor.shape, args.target_tensor.dtype)
+        result[zero] = args.target_tensor[zero]*2 + args.target_tensor[one]*3
+        result[one] = args.target_tensor[zero]*5 + args.target_tensor[one]*7
+        return result
+
+    def _unitary_(self):
+        return np.array([[2, 3], [5, 7]])
+
+    def __eq__(self, other):
+        return isinstance(other, type(self))
+
+    def __repr__(self):
+        return ('cirq.ops.controlled_gate_test.'
+                'GateAllocatingNewSpaceForResult()')
+
 class RestrictedGate(cirq.SingleQubitGate):
     pass
 
@@ -101,54 +144,6 @@ def test_unitary():
         atol=1e-8)
 
 
-class GateUsingWorkspaceForApplyUnitary(cirq.SingleQubitGate):
-    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
-                        ) -> Union[np.ndarray, NotImplementedType]:
-        args.available_buffer[...] = args.target_tensor
-        args.target_tensor[...] = 0
-        return args.available_buffer
-
-    def _unitary_(self):
-        return np.eye(2)
-
-    def __pow__(self, exponent):
-        return self
-
-    def __eq__(self, other):
-        return isinstance(other, type(self))
-
-    def __repr__(self):
-        return ('cirq.ops.controlled_gate_test.'
-                'GateUsingWorkspaceForApplyUnitary()')
-
-
-class GateAllocatingNewSpaceForResult(cirq.SingleQubitGate):
-    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
-                        ) -> Union[np.ndarray, NotImplementedType]:
-        assert len(args.axes) == 1
-        a = args.axes[0]
-        seed = cast(Tuple[Union[int, slice, 'ellipsis'], ...],
-                    (slice(None),))
-        zero = seed*a + (0, Ellipsis)
-        one = seed*a + (1, Ellipsis)
-        result = np.zeros(args.target_tensor.shape, args.target_tensor.dtype)
-        result[zero] = args.target_tensor[zero]*2 + args.target_tensor[one]*3
-        result[one] = args.target_tensor[zero]*5 + args.target_tensor[one]*7
-        return result
-
-    def _unitary_(self):
-        return np.array([[2, 3], [5, 7]])
-
-    def __pow__(self, factor):
-        return self
-
-    def __eq__(self, other):
-        return isinstance(other, type(self))
-
-    def __repr__(self):
-        return ('cirq.ops.controlled_gate_test.'
-                'GateAllocatingNewSpaceForResult()')
-
 
 @pytest.mark.parametrize('gate', [
     cirq.X,
diff --git a/cirq/ops/controlled_operation_test.py b/cirq/ops/controlled_operation_test.py
--- a/cirq/ops/controlled_operation_test.py
+++ b/cirq/ops/controlled_operation_test.py
@@ -11,11 +11,58 @@
 # 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 typing import Union, Tuple, cast
 
 import numpy as np
+import pytest
+
 import cirq
 from cirq import protocols
+from cirq.type_workarounds import NotImplementedType
+
+
+class GateUsingWorkspaceForApplyUnitary(cirq.SingleQubitGate):
+    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
+                        ) -> Union[np.ndarray, NotImplementedType]:
+        args.available_buffer[...] = args.target_tensor
+        args.target_tensor[...] = 0
+        return args.available_buffer
+
+    def _unitary_(self):
+        return np.eye(2)
+
+
+    def __eq__(self, other):
+        return isinstance(other, type(self))
+
+    def __repr__(self):
+        return ('cirq.ops.controlled_operation_test.'
+                'GateUsingWorkspaceForApplyUnitary()')
+
+
+class GateAllocatingNewSpaceForResult(cirq.SingleQubitGate):
+    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs
+                        ) -> Union[np.ndarray, NotImplementedType]:
+        assert len(args.axes) == 1
+        a = args.axes[0]
+        seed = cast(Tuple[Union[int, slice, 'ellipsis'], ...],
+                    (slice(None),))
+        zero = seed*a + (0, Ellipsis)
+        one = seed*a + (1, Ellipsis)
+        result = np.zeros(args.target_tensor.shape, args.target_tensor.dtype)
+        result[zero] = args.target_tensor[zero]*2 + args.target_tensor[one]*3
+        result[one] = args.target_tensor[zero]*5 + args.target_tensor[one]*7
+        return result
+
+    def _unitary_(self):
+        return np.array([[2, 3], [5, 7]])
+
+    def __eq__(self, other):
+        return isinstance(other, type(self))
+
+    def __repr__(self):
+        return ('cirq.ops.controlled_operation_test.'
+                'GateAllocatingNewSpaceForResult()')
 
 
 def test_controlled_operation_init():
@@ -141,7 +188,10 @@ class UndiagrammableGate(cirq.SingleQubitGate):
     cirq.SWAP(cirq.NamedQubit('q1'), cirq.NamedQubit('q2')),
     cirq.CCZ(cirq.NamedQubit('q1'), cirq.NamedQubit('q2'),
              cirq.NamedQubit('q3')),
-    cirq.ControlledGate(cirq.ControlledGate(cirq.CCZ))(*cirq.LineQubit.range(5))
+    cirq.ControlledGate(cirq.ControlledGate(cirq.CCZ))(
+        *cirq.LineQubit.range(5)),
+    GateUsingWorkspaceForApplyUnitary()(cirq.NamedQubit('q1')),
+    GateAllocatingNewSpaceForResult()(cirq.NamedQubit('q1')),
 ])
 def test_controlled_operation_is_consistent(gate: cirq.GateOperation):
     cb = cirq.NamedQubit('ctr')