cirq-core 1.1.0.dev20221219200817__py3-none-any.whl → 1.2.0__py3-none-any.whl

cirq/ops/qubit_manager.py

@@ -0,0 +1,97 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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
+#
+#     https://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 abc
+import dataclasses
+from typing import Iterable, List, TYPE_CHECKING
+from cirq.ops import raw_types
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class QubitManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def qalloc(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        """Allocate `n` clean qubits, i.e. qubits guaranteed to be in state |0>."""
+
+    @abc.abstractmethod
+    def qborrow(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        """Allocate `n` dirty qubits, i.e. the returned qubits can be in any state."""
+
+    @abc.abstractmethod
+    def qfree(self, qubits: Iterable['cirq.Qid']) -> None:
+        """Free pre-allocated clean or dirty qubits managed by this qubit manager."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _BaseAncillaQid(raw_types.Qid):
+    id: int
+    dim: int = 2
+    prefix: str = ''
+
+    def _comparison_key(self) -> int:
+        return self.id
+
+    @property
+    def dimension(self) -> int:
+        return self.dim
+
+    def __repr__(self) -> str:
+        dim_str = f', dim={self.dim}' if self.dim != 2 else ''
+        prefix_str = f', prefix={self.prefix}' if self.prefix != '' else ''
+        return f"cirq.ops.{type(self).__name__}({self.id}{dim_str}{prefix_str})"
+
+
+class CleanQubit(_BaseAncillaQid):
+    """An internal qid type that represents a clean ancilla allocation."""
+
+    def __str__(self) -> str:
+        dim_str = f' (d={self.dimension})' if self.dim != 2 else ''
+        return f"{self.prefix}_c({self.id}){dim_str}"
+
+
+class BorrowableQubit(_BaseAncillaQid):
+    """An internal qid type that represents a dirty ancilla allocation."""
+
+    def __str__(self) -> str:
+        dim_str = f' (d={self.dimension})' if self.dim != 2 else ''
+        return f"{self.prefix}_b({self.id}){dim_str}"
+
+
+class SimpleQubitManager(QubitManager):
+    """Allocates a new `CleanQubit`/`BorrowableQubit` for every `qalloc`/`qborrow` request."""
+
+    def __init__(self, prefix: str = ''):
+        self._clean_id = 0
+        self._borrow_id = 0
+        self._prefix = prefix
+
+    def qalloc(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        self._clean_id += n
+        return [CleanQubit(i, dim, self._prefix) for i in range(self._clean_id - n, self._clean_id)]
+
+    def qborrow(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        self._borrow_id = self._borrow_id + n
+        return [
+            BorrowableQubit(i, dim, self._prefix)
+            for i in range(self._borrow_id - n, self._borrow_id)
+        ]
+
+    def qfree(self, qubits: Iterable['cirq.Qid']) -> None:
+        for q in qubits:
+            good = isinstance(q, CleanQubit) and q.id < self._clean_id
+            good |= isinstance(q, BorrowableQubit) and q.id < self._borrow_id
+            if not good:
+                raise ValueError(f"{q} was not allocated by {self}")

cirq/ops/qubit_manager_test.py

@@ -0,0 +1,66 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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
+#
+#     https://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 cirq
+from cirq.ops import qubit_manager as cqi
+import pytest
+
+
+def test_clean_qubits():
+    q = cqi.CleanQubit(1)
+    assert q.id == 1
+    assert q.dimension == 2
+    assert str(q) == '_c(1)'
+    assert repr(q) == 'cirq.ops.CleanQubit(1)'
+
+    q = cqi.CleanQubit(2, dim=3)
+    assert q.id == 2
+    assert q.dimension == 3
+    assert str(q) == '_c(2) (d=3)'
+    assert repr(q) == 'cirq.ops.CleanQubit(2, dim=3)'
+
+    assert cqi.CleanQubit(1) < cqi.CleanQubit(2)
+
+
+def test_borrow_qubits():
+    q = cqi.BorrowableQubit(10)
+    assert q.id == 10
+    assert q.dimension == 2
+    assert str(q) == '_b(10)'
+    assert repr(q) == 'cirq.ops.BorrowableQubit(10)'
+
+    q = cqi.BorrowableQubit(20, dim=4)
+    assert q.id == 20
+    assert q.dimension == 4
+    assert str(q) == '_b(20) (d=4)'
+    assert repr(q) == 'cirq.ops.BorrowableQubit(20, dim=4)'
+
+    assert cqi.BorrowableQubit(1) < cqi.BorrowableQubit(2)
+
+
+@pytest.mark.parametrize('_', range(2))
+def test_simple_qubit_manager(_):
+    qm = cirq.ops.SimpleQubitManager()
+    assert qm.qalloc(1) == [cqi.CleanQubit(0)]
+    assert qm.qalloc(2) == [cqi.CleanQubit(1), cqi.CleanQubit(2)]
+    assert qm.qalloc(1, dim=3) == [cqi.CleanQubit(3, dim=3)]
+    assert qm.qborrow(1) == [cqi.BorrowableQubit(0)]
+    assert qm.qborrow(2) == [cqi.BorrowableQubit(1), cqi.BorrowableQubit(2)]
+    assert qm.qborrow(1, dim=3) == [cqi.BorrowableQubit(3, dim=3)]
+    qm.qfree([cqi.CleanQubit(i) for i in range(3)] + [cqi.CleanQubit(3, dim=3)])
+    qm.qfree([cqi.BorrowableQubit(i) for i in range(3)] + [cqi.BorrowableQubit(3, dim=3)])
+    with pytest.raises(ValueError, match="not allocated"):
+        qm.qfree([cqi.CleanQubit(10)])
+    with pytest.raises(ValueError, match="not allocated"):
+        qm.qfree([cqi.BorrowableQubit(10)])

cirq/ops/raw_types.py

@@ -17,6 +17,7 @@
 import abc
 import functools
 from typing import (
+    cast,
     AbstractSet,
     Any,
     Callable,
@@ -30,16 +31,17 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
-    TypeVar,
     TYPE_CHECKING,
     Union,
 )
+from typing_extensions import Self
 
 import numpy as np
 import sympy
 
 from cirq import protocols, value
 from cirq._import import LazyLoader
+from cirq._compat import __cirq_debug__, cached_method
 from cirq.type_workarounds import NotImplementedType
 from cirq.ops import control_values as cv
 
@@ -109,7 +111,8 @@ class Qid(metaclass=abc.ABCMeta):
     def _cmp_tuple(self):
         return (type(self).__name__, repr(type(self)), self._comparison_key(), self.dimension)
 
-    def __hash__(self):
+    @cached_method
+    def __hash__(self) -> int:
         return hash((Qid, self._comparison_key()))
 
     def __eq__(self, other):
@@ -215,7 +218,8 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
         Raises:
             ValueError: The gate can't be applied to the qubits.
         """
-        _validate_qid_shape(self, qubits)
+        if __cirq_debug__.get():
+            _validate_qid_shape(self, qubits)
 
     def on(self, *qubits: Qid) -> 'Operation':
         """Returns an application of this gate to the given qubits.
@@ -254,19 +258,33 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
                     raise TypeError(f'{targets[0]} object is not iterable.')
                 t0 = list(targets[0])
                 iterator = [t0] if t0 and isinstance(t0[0], Qid) else t0
-            for target in iterator:
-                if not isinstance(target, Sequence):
-                    raise ValueError(
-                        f'Inputs to multi-qubit gates must be Sequence[Qid].'
-                        f' Type: {type(target)}'
-                    )
-                if not all(isinstance(x, Qid) for x in target):
-                    raise ValueError(f'All values in sequence should be Qids, but got {target}')
-                if len(target) != self._num_qubits_():
-                    raise ValueError(f'Expected {self._num_qubits_()} qubits, got {target}')
-                operations.append(self.on(*target))
+            if __cirq_debug__.get():
+                for target in iterator:
+                    if not isinstance(target, Sequence):
+                        raise ValueError(
+                            f'Inputs to multi-qubit gates must be Sequence[Qid].'
+                            f' Type: {type(target)}'
+                        )
+                    if not all(isinstance(x, Qid) for x in target):
+                        raise ValueError(f'All values in sequence should be Qids, but got {target}')
+                    if len(target) != self._num_qubits_():
+                        raise ValueError(f'Expected {self._num_qubits_()} qubits, got {target}')
+                    operations.append(self.on(*target))
+            else:
+                operations = [self.on(*target) for target in iterator]
             return operations
 
+        if not __cirq_debug__.get():
+            return [
+                op
+                for q in targets
+                for op in (
+                    self.on_each(*q)
+                    if isinstance(q, Iterable) and not isinstance(q, str)
+                    else [self.on(cast('cirq.Qid', q))]
+                )
+            ]
+
         for target in targets:
             if isinstance(target, Qid):
                 operations.append(self.on(target))
@@ -342,7 +360,7 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
         return self.on(*qubits)
 
     def with_probability(self, probability: 'cirq.TParamVal') -> 'cirq.Gate':
-        """Creates a probabalistic channel with this gate.
+        """Creates a probabilistic channel with this gate.
 
         Args:
             probability: floating point value between 0 and 1, giving the
@@ -358,7 +376,7 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
 
     def controlled(
         self,
-        num_controls: int = None,
+        num_controls: Optional[int] = None,
         control_values: Optional[
             Union[cv.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
         ] = None,
@@ -406,6 +424,7 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
     @value.alternative(requires='_num_qubits_', implementation=_backwards_compatibility_num_qubits)
     def num_qubits(self) -> int:
         """The number of qubits this gate acts on."""
+        raise NotImplementedError
 
     def _num_qubits_from_shape(self) -> int:
         shape = self._qid_shape_()
@@ -420,6 +439,7 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
     @value.alternative(requires='_qid_shape_', implementation=_num_qubits_from_shape)
     def _num_qubits_(self) -> int:
         """The number of qubits this gate acts on."""
+        raise NotImplementedError
 
     def _default_shape_from_num_qubits(self) -> Tuple[int, ...]:
         num_qubits = self._num_qubits_()
@@ -433,6 +453,7 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
         the gate acts on.  E.g. (2, 2, 2) for the three-qubit CCZ gate and
         (3, 3) for a 2-qutrit ternary gate.
         """
+        raise NotImplementedError
 
     def _commutes_on_qids_(
         self, qids: 'Sequence[cirq.Qid]', other: Any, *, atol: float = 1e-8
@@ -462,9 +483,6 @@ class Gate(metaclass=value.ABCMetaImplementAnyOneOf):
         return protocols.obj_to_dict_helper(self, attribute_names=[])
 
 
-TSelf = TypeVar('TSelf', bound='Operation')
-
-
 class Operation(metaclass=abc.ABCMeta):
     """An effect applied to a collection of qubits.
 
@@ -488,11 +506,12 @@ class Operation(metaclass=abc.ABCMeta):
         """
         return len(self.qubits)
 
+    @cached_method
     def _qid_shape_(self) -> Tuple[int, ...]:
         return protocols.qid_shape(self.qubits)
 
     @abc.abstractmethod
-    def with_qubits(self: TSelf, *new_qubits: 'cirq.Qid') -> TSelf:
+    def with_qubits(self, *new_qubits: 'cirq.Qid') -> Self:
         """Returns the same operation, but applied to different qubits.
 
         Args:
@@ -534,9 +553,8 @@ class Operation(metaclass=abc.ABCMeta):
         return TaggedOperation(self, *new_tags)
 
     def transform_qubits(
-        self: TSelf,
-        qubit_map: Union[Dict['cirq.Qid', 'cirq.Qid'], Callable[['cirq.Qid'], 'cirq.Qid']],
-    ) -> TSelf:
+        self, qubit_map: Union[Dict['cirq.Qid', 'cirq.Qid'], Callable[['cirq.Qid'], 'cirq.Qid']]
+    ) -> Self:
         """Returns the same operation, but with different qubits.
 
         Args:
@@ -583,7 +601,7 @@ class Operation(metaclass=abc.ABCMeta):
         return ops.controlled_operation.ControlledOperation(control_qubits, self, control_values)
 
     def with_probability(self, probability: 'cirq.TParamVal') -> 'cirq.Operation':
-        """Creates a probabalistic channel with this operation.
+        """Creates a probabilistic channel with this operation.
 
         Args:
             probability: floating point value between 0 and 1, giving the
@@ -617,7 +635,8 @@ class Operation(metaclass=abc.ABCMeta):
         Raises:
             ValueError: The operation had qids that don't match it's qid shape.
         """
-        _validate_qid_shape(self, qubits)
+        if __cirq_debug__.get():
+            _validate_qid_shape(self, qubits)
 
     def _commutes_(
         self, other: Any, *, atol: float = 1e-8
@@ -811,7 +830,14 @@ class TaggedOperation(Operation):
         return protocols.obj_to_dict_helper(self, ['sub_operation', 'tags'])
 
     def _decompose_(self) -> 'cirq.OP_TREE':
-        return protocols.decompose_once(self.sub_operation, default=None)
+        return self._decompose_with_context_()
+
+    def _decompose_with_context_(
+        self, context: Optional['cirq.DecompositionContext'] = None
+    ) -> 'cirq.OP_TREE':
+        return protocols.decompose_once(
+            self.sub_operation, default=None, flatten=False, context=context
+        )
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:
         return protocols.pauli_expansion(self.sub_operation)
@@ -821,6 +847,7 @@ class TaggedOperation(Operation):
     ) -> Union[np.ndarray, None, NotImplementedType]:
         return protocols.apply_unitary(self.sub_operation, args, default=None)
 
+    @cached_method
     def _has_unitary_(self) -> bool:
         return protocols.has_unitary(self.sub_operation)
 
@@ -832,30 +859,36 @@ class TaggedOperation(Operation):
     ) -> Union[bool, NotImplementedType, None]:
         return protocols.commutes(self.sub_operation, other, atol=atol)
 
+    @cached_method
     def _has_mixture_(self) -> bool:
         return protocols.has_mixture(self.sub_operation)
 
     def _mixture_(self) -> Sequence[Tuple[float, Any]]:
         return protocols.mixture(self.sub_operation, NotImplemented)
 
+    @cached_method
     def _has_kraus_(self) -> bool:
         return protocols.has_kraus(self.sub_operation)
 
     def _kraus_(self) -> Union[Tuple[np.ndarray], NotImplementedType]:
         return protocols.kraus(self.sub_operation, NotImplemented)
 
+    @cached_method
     def _measurement_key_names_(self) -> FrozenSet[str]:
         return protocols.measurement_key_names(self.sub_operation)
 
+    @cached_method
     def _measurement_key_objs_(self) -> FrozenSet['cirq.MeasurementKey']:
         return protocols.measurement_key_objs(self.sub_operation)
 
+    @cached_method
     def _is_measurement_(self) -> bool:
         sub = getattr(self.sub_operation, "_is_measurement_", None)
         if sub is not None:
             return sub()
         return NotImplemented
 
+    @cached_method
     def _is_parameterized_(self) -> bool:
         return protocols.is_parameterized(self.sub_operation) or any(
             protocols.is_parameterized(tag) for tag in self.tags
@@ -867,6 +900,7 @@ class TaggedOperation(Operation):
             return sub(sim_state)
         return NotImplemented
 
+    @cached_method
     def _parameter_names_(self) -> AbstractSet[str]:
         tag_params = {name for tag in self.tags for name in protocols.parameter_names(tag)}
         return protocols.parameter_names(self.sub_operation) | tag_params
@@ -891,6 +925,7 @@ class TaggedOperation(Operation):
             ) + sub_op_info.wire_symbols[1:]
         return sub_op_info
 
+    @cached_method
     def _trace_distance_bound_(self) -> float:
         return protocols.trace_distance_bound(self.sub_operation)
 
@@ -951,7 +986,14 @@ class _InverseCompositeGate(Gate):
         return NotImplemented
 
     def _decompose_(self, qubits):
-        return protocols.inverse(protocols.decompose_once_with_qubits(self._original, qubits))
+        return self._decompose_with_context_(qubits)
+
+    def _decompose_with_context_(
+        self, qubits: Sequence['cirq.Qid'], context: Optional['cirq.DecompositionContext'] = None
+    ) -> 'cirq.OP_TREE':
+        return protocols.inverse(
+            protocols.decompose_once_with_qubits(self._original, qubits, context=context)
+        )
 
     def _has_unitary_(self):
         from cirq import protocols, devices
@@ -962,9 +1004,11 @@ class _InverseCompositeGate(Gate):
             for op in protocols.decompose_once_with_qubits(self._original, qubits)
         )
 
+    @cached_method
     def _is_parameterized_(self) -> bool:
         return protocols.is_parameterized(self._original)
 
+    @cached_method
     def _parameter_names_(self) -> AbstractSet[str]:
         return protocols.parameter_names(self._original)
 
@@ -998,15 +1042,13 @@ def _validate_qid_shape(val: Any, qubits: Sequence['cirq.Qid']) -> None:
     qid_shape = protocols.qid_shape(val)
     if len(qubits) != len(qid_shape):
         raise ValueError(
-            'Wrong number of qubits for <{!r}>. '
-            'Expected {} qubits but got <{!r}>.'.format(val, len(qid_shape), qubits)
+            f'Wrong number of qubits for <{val!r}>. '
+            f'Expected {len(qid_shape)} qubits but got <{qubits!r}>.'
         )
     if any(qid.dimension != dimension for qid, dimension in zip(qubits, qid_shape)):
         raise ValueError(
-            'Wrong shape of qids for <{!r}>. '
-            'Expected {} but got {} <{!r}>.'.format(
-                val, qid_shape, tuple(qid.dimension for qid in qubits), qubits
-            )
+            f'Wrong shape of qids for <{val!r}>. '
+            f'Expected {qid_shape} but got {tuple(qid.dimension for qid in qubits)} <{qubits!r}>.'
         )
     if len(set(qubits)) != len(qubits):
         raise ValueError(

cirq/ops/raw_types_test.py

@@ -151,6 +151,29 @@ def test_op_validate():
         op2.validate_args([cirq.LineQid(1, 2), cirq.LineQid(1, 2)])
 
 
+def test_disable_op_validation():
+    q0, q1 = cirq.LineQubit.range(2)
+    h_op = cirq.H(q0)
+
+    # Fails normally.
+    with pytest.raises(ValueError, match='Wrong number'):
+        _ = cirq.H(q0, q1)
+    with pytest.raises(ValueError, match='Wrong number'):
+        h_op.validate_args([q0, q1])
+
+    # Passes, skipping validation.
+    with cirq.with_debug(False):
+        op = cirq.H(q0, q1)
+        assert op.qubits == (q0, q1)
+        h_op.validate_args([q0, q1])
+
+    # Fails again when validation is re-enabled.
+    with pytest.raises(ValueError, match='Wrong number'):
+        _ = cirq.H(q0, q1)
+    with pytest.raises(ValueError, match='Wrong number'):
+        h_op.validate_args([q0, q1])
+
+
 def test_default_validation_and_inverse():
     class TestGate(cirq.Gate):
         def _num_qubits_(self):
@@ -178,6 +201,8 @@ def test_default_validation_and_inverse():
     assert i**-1 == t
     assert t**-1 == i
     assert cirq.decompose(i) == [cirq.X(a), cirq.S(b) ** -1, cirq.Z(a)]
+    assert [*i._decompose_()] == [cirq.X(a), cirq.S(b) ** -1, cirq.Z(a)]
+    assert [*i.gate._decompose_([a, b])] == [cirq.X(a), cirq.S(b) ** -1, cirq.Z(a)]
     cirq.testing.assert_allclose_up_to_global_phase(
         cirq.unitary(i), cirq.unitary(t).conj().T, atol=1e-8
     )
@@ -595,6 +620,7 @@ def test_tagged_operation_forwards_protocols():
     np.testing.assert_equal(cirq.unitary(tagged_h), cirq.unitary(h))
     assert cirq.has_unitary(tagged_h)
     assert cirq.decompose(tagged_h) == cirq.decompose(h)
+    assert [*tagged_h._decompose_()] == cirq.decompose(h)
     assert cirq.pauli_expansion(tagged_h) == cirq.pauli_expansion(h)
     assert cirq.equal_up_to_global_phase(h, tagged_h)
     assert np.isclose(cirq.kraus(h), cirq.kraus(tagged_h)).all()
@@ -787,6 +813,10 @@ def test_single_qubit_gate_validates_on_each():
     test_non_qubits = [str(i) for i in range(3)]
     with pytest.raises(ValueError):
         _ = g.on_each(*test_non_qubits)
+
+    with cirq.with_debug(False):
+        assert g.on_each(*test_non_qubits)[0].qubits == ('0',)
+
     with pytest.raises(ValueError):
         _ = g.on_each(*test_non_qubits)
 
@@ -853,6 +883,10 @@ def test_on_each_two_qubits():
         g.on_each([(a,)])
     with pytest.raises(ValueError, match='Expected 2 qubits'):
         g.on_each([(a, b, a)])
+
+    with cirq.with_debug(False):
+        assert g.on_each([(a, b, a)])[0].qubits == (a, b, a)
+
     with pytest.raises(ValueError, match='Expected 2 qubits'):
         g.on_each(zip([a, a]))
     with pytest.raises(ValueError, match='Expected 2 qubits'):

cirq/ops/three_qubit_gates.py

@@ -176,8 +176,9 @@ class CCZPowGate(gate_features.InterchangeableQubitsGate, eigen_gate.EigenGate):
             if self._exponent == 1:
                 return 'cirq.CCZ'
             return f'(cirq.CCZ**{proper_repr(self._exponent)})'
-        return 'cirq.CCZPowGate(exponent={}, global_shift={!r})'.format(
-            proper_repr(self._exponent), self._global_shift
+        return (
+            f'cirq.CCZPowGate(exponent={proper_repr(self._exponent)}, '
+            f'global_shift={self._global_shift!r})'
         )
 
     def __str__(self) -> str:
@@ -190,7 +191,7 @@ class CCZPowGate(gate_features.InterchangeableQubitsGate, eigen_gate.EigenGate):
 
     def controlled(
         self,
-        num_controls: int = None,
+        num_controls: Optional[int] = None,
         control_values: Optional[
             Union[cv.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
         ] = None,
@@ -383,9 +384,8 @@ class ThreeQubitDiagonalGate(raw_types.Gate):
         return protocols.obj_to_dict_helper(self, attribute_names=["diag_angles_radians"])
 
     def __repr__(self) -> str:
-        return 'cirq.ThreeQubitDiagonalGate([{}])'.format(
-            ','.join(proper_repr(angle) for angle in self._diag_angles_radians)
-        )
+        angles = ','.join(proper_repr(angle) for angle in self._diag_angles_radians)
+        return f'cirq.ThreeQubitDiagonalGate([{angles}])'
 
     def _num_qubits_(self) -> int:
         return 3
@@ -488,8 +488,9 @@ class CCXPowGate(gate_features.InterchangeableQubitsGate, eigen_gate.EigenGate):
             if self._exponent == 1:
                 return 'cirq.TOFFOLI'
             return f'(cirq.TOFFOLI**{proper_repr(self._exponent)})'
-        return 'cirq.CCXPowGate(exponent={}, global_shift={!r})'.format(
-            proper_repr(self._exponent), self._global_shift
+        return (
+            f'cirq.CCXPowGate(exponent={proper_repr(self._exponent)}, '
+            f'global_shift={self._global_shift!r})'
         )
 
     def __str__(self) -> str:
@@ -502,7 +503,7 @@ class CCXPowGate(gate_features.InterchangeableQubitsGate, eigen_gate.EigenGate):
 
     def controlled(
         self,
-        num_controls: int = None,
+        num_controls: Optional[int] = None,
         control_values: Optional[
             Union[cv.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
         ] = None,
@@ -661,6 +662,11 @@ class CSwapGate(gate_features.InterchangeableQubitsGate, raw_types.Gate):
     def _value_equality_values_(self):
         return ()
 
+    def __pow__(self, power):
+        if power == 1 or power == -1:
+            return self
+        return NotImplemented
+
     def __str__(self) -> str:
         return 'FREDKIN'
 
@@ -672,7 +678,7 @@ class CSwapGate(gate_features.InterchangeableQubitsGate, raw_types.Gate):
 
     def controlled(
         self,
-        num_controls: int = None,
+        num_controls: Optional[int] = None,
         control_values: Optional[
             Union[cv.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
         ] = None,

cirq/ops/three_qubit_gates_test.py

@@ -94,8 +94,9 @@ def test_unitary():
     )
 
     assert cirq.has_unitary(cirq.CSWAP)
+    u = cirq.unitary(cirq.CSWAP)
     np.testing.assert_allclose(
-        cirq.unitary(cirq.CSWAP),
+        u,
         np.array(
             [
                 [1, 0, 0, 0, 0, 0, 0, 0],
@@ -110,6 +111,7 @@ def test_unitary():
         ),
         atol=1e-8,
     )
+    np.testing.assert_allclose(u @ u, np.eye(8))
 
     diagonal_angles = [2, 3, 5, 7, 11, 13, 17, 19]
     assert cirq.has_unitary(cirq.ThreeQubitDiagonalGate(diagonal_angles))
@@ -156,7 +158,7 @@ def test_eq():
     eq.add_equality_group(cirq.TOFFOLI(a, b, c), cirq.CCX(a, b, c))
     eq.add_equality_group(cirq.TOFFOLI(a, c, b), cirq.TOFFOLI(c, a, b))
     eq.add_equality_group(cirq.TOFFOLI(a, b, d))
-    eq.add_equality_group(cirq.CSWAP(a, b, c), cirq.FREDKIN(a, b, c))
+    eq.add_equality_group(cirq.CSWAP(a, b, c), cirq.FREDKIN(a, b, c), cirq.FREDKIN(a, b, c) ** -1)
     eq.add_equality_group(cirq.CSWAP(b, a, c), cirq.CSWAP(b, c, a))
 
 

cirq/ops/two_qubit_diagonal_gate.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.
+
 """Creates the gate instance for a two qubit diagonal gate.
 
 The gate is used to create a 4x4 matrix with the diagonal elements
@@ -134,9 +135,8 @@ class TwoQubitDiagonalGate(raw_types.Gate):
         return tuple(self._diag_angles_radians)
 
     def __repr__(self) -> str:
-        return 'cirq.TwoQubitDiagonalGate([{}])'.format(
-            ','.join(proper_repr(angle) for angle in self._diag_angles_radians)
-        )
+        angles = ','.join(proper_repr(angle) for angle in self._diag_angles_radians)
+        return f'cirq.TwoQubitDiagonalGate([{angles}])'
 
     def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(self, attribute_names=["diag_angles_radians"])

cirq/ops/wait_gate.py

@@ -34,7 +34,7 @@ class WaitGate(raw_types.Gate):
         self,
         duration: 'cirq.DURATION_LIKE',
         num_qubits: Optional[int] = None,
-        qid_shape: Tuple[int, ...] = None,
+        qid_shape: Optional[Tuple[int, ...]] = None,
     ) -> None:
         """Initialize a wait gate with the given duration.
 

cirq/protocols/__init__.py

@@ -50,6 +50,7 @@ from cirq.protocols.decompose_protocol import (
     decompose,
     decompose_once,
     decompose_once_with_qubits,
+    DecompositionContext,
     SupportsDecompose,
     SupportsDecomposeWithQubits,
 )