ArithmeticGate implementation (quantumlib#4702)

Migrate ArithmeticOperation to ArithmeticGate.

As implemented, ArithmeticGate.on(qubits) returns a GateOperation. In other words, ArithmeticGate is completely unrelated to ArithmeticOperation. They can be considered two different ways to do the same thing. This was done in order to enable ArithmeticOperation to be deprecated.

Additionally this PR implements QuirkArithmeticGate, deprecating QuirkArithmeticOperation, and rewrites ModularExp as a gate without a deprecation cycle since it's in /examples.

The code for ArithmeticGate (and subclasses) is *basically* identical with ArithmeticOperation, except instead of `Sequence[Qid]]`, a quantum register is a `Sequence[int]`, where the int represents the dimension. It implements the _qid_shape_ protocol from this data, and the GateOperation constructor already has logic to ensure the appropriate number/dimension of qubits are applied.

Tests are added to that effect.

Closes quantumlib#4683

Author: daxfohl

cirq/__init__.py

@@ -185,6 +185,7 @@
     AmplitudeDampingChannel,
     AnyIntegerPowerGateFamily,
     AnyUnitaryGateFamily,
+    ArithmeticGate,
     ArithmeticOperation,
     asymmetric_depolarize,
     AsymmetricDepolarizingChannel,

cirq/interop/quirk/__init__.py

@@ -20,6 +20,7 @@
 
 # Imports from cells are only to ensure operation reprs work correctly.
 from cirq.interop.quirk.cells import (
+    QuirkArithmeticGate,
     QuirkArithmeticOperation,
     QuirkInputRotationOperation,
     QuirkQubitPermutationGate,

cirq/interop/quirk/cells/__init__.py

@@ -21,7 +21,7 @@
 
 from cirq.interop.quirk.cells.qubit_permutation_cells import QuirkQubitPermutationGate
 
-from cirq.interop.quirk.cells.arithmetic_cells import QuirkArithmeticOperation
+from cirq.interop.quirk.cells.arithmetic_cells import QuirkArithmeticGate, QuirkArithmeticOperation
 
 from cirq.interop.quirk.cells.input_rotation_cells import QuirkInputRotationOperation
 

cirq/interop/quirk/cells/arithmetic_cells.py

@@ -28,12 +28,14 @@
 )
 
 from cirq import ops, value
+from cirq._compat import deprecated_class
 from cirq.interop.quirk.cells.cell import Cell, CellMaker, CELL_SIZES
 
 if TYPE_CHECKING:
     import cirq
 
 
+@deprecated_class(deadline='v0.15', fix='Use cirq.QuirkArithmeticGate')
 @value.value_equality
 class QuirkArithmeticOperation(ops.ArithmeticOperation):
     """Applies arithmetic to a target and some inputs.
@@ -148,6 +150,110 @@ def __repr__(self) -> str:
         )
 
 
+@value.value_equality
+class QuirkArithmeticGate(ops.ArithmeticGate):
+    """Applies arithmetic to a target and some inputs.
+
+    Implements Quirk-specific implicit effects like assuming that the presence
+    of an 'r' input implies modular arithmetic.
+
+    In Quirk, modular operations have no effect on values larger than the
+    modulus. This convention is used because unitarity forces *some* convention
+    on out-of-range values (they cannot simply disappear or raise exceptions),
+    and the simplest is to do nothing. This call handles ensuring that happens,
+    and ensuring the new target register value is normalized modulo the modulus.
+    """
+
+    def __init__(
+        self, identifier: str, target: Sequence[int], inputs: Sequence[Union[Sequence[int], int]]
+    ):
+        """Inits QuirkArithmeticGate.
+
+        Args:
+            identifier: The quirk identifier string for this operation.
+            target: The target qubit register.
+            inputs: Qubit registers, which correspond to the qid shape of the
+                qubits from which the input will be read, or classical
+                constants, that determine what happens to the target.
+
+        Raises:
+            ValueError: If the target is too small for a modular operation with
+                too small modulus.
+        """
+        self.identifier = identifier
+        self.target: Tuple[int, ...] = tuple(target)
+        self.inputs: Tuple[Union[Sequence[int], int], ...] = tuple(
+            e if isinstance(e, int) else tuple(e) for e in inputs
+        )
+
+        if self.operation.is_modular:
+            r = inputs[-1]
+            if isinstance(r, int):
+                over = r > 1 << len(target)
+            else:
+                over = len(cast(Sequence, r)) > len(target)
+            if over:
+                raise ValueError(f'Target too small for modulus.\nTarget: {target}\nModulus: {r}')
+
+    @property
+    def operation(self) -> '_QuirkArithmeticCallable':
+        return ARITHMETIC_OP_TABLE[self.identifier]
+
+    def _value_equality_values_(self) -> Any:
+        return self.identifier, self.target, self.inputs
+
+    def registers(self) -> Sequence[Union[int, Sequence[int]]]:
+        return [self.target, *self.inputs]
+
+    def with_registers(self, *new_registers: Union[int, Sequence[int]]) -> 'QuirkArithmeticGate':
+        if len(new_registers) != len(self.inputs) + 1:
+            raise ValueError(
+                'Wrong number of registers.\n'
+                f'New registers: {repr(new_registers)}\n'
+                f'Operation: {repr(self)}'
+            )
+
+        if isinstance(new_registers[0], int):
+            raise ValueError(
+                'The first register is the mutable target. '
+                'It must be a list of qubits, not the constant '
+                f'{new_registers[0]}.'
+            )
+
+        return QuirkArithmeticGate(self.identifier, new_registers[0], new_registers[1:])
+
+    def apply(self, *registers: int) -> Union[int, Iterable[int]]:
+        return self.operation(*registers)
+
+    def _circuit_diagram_info_(self, args: 'cirq.CircuitDiagramInfoArgs') -> List[str]:
+        lettered_args = list(zip(self.operation.letters, self.inputs))
+
+        result: List[str] = []
+
+        # Target register labels.
+        consts = ''.join(
+            f',{letter}={reg}' for letter, reg in lettered_args if isinstance(reg, int)
+        )
+        result.append(f'Quirk({self.identifier}{consts})')
+        result.extend(f'#{i}' for i in range(2, len(self.target) + 1))
+
+        # Input register labels.
+        for letter, reg in lettered_args:
+            if not isinstance(reg, int):
+                result.extend(f'{letter.upper()}{i}' for i in range(len(cast(Sequence, reg))))
+
+        return result
+
+    def __repr__(self) -> str:
+        return (
+            'cirq.interop.quirk.QuirkArithmeticGate(\n'
+            f'    {repr(self.identifier)},\n'
+            f'    target={repr(self.target)},\n'
+            f'    inputs={_indented_list_lines_repr(self.inputs)},\n'
+            ')'
+        )
+
+
 _IntsToIntCallable = Union[
     Callable[[int], int],
     Callable[[int, int], int],
@@ -244,11 +350,13 @@ def operations(self) -> 'cirq.OP_TREE':
         if missing_inputs:
             raise ValueError(f'Missing input: {sorted(missing_inputs)}')
 
-        return QuirkArithmeticOperation(
+        inputs = cast(Sequence[Union[Sequence['cirq.Qid'], int]], self.inputs)
+        qubits = self.target + tuple(q for i in self.inputs if isinstance(i, Sequence) for q in i)
+        return QuirkArithmeticGate(
             self.identifier,
-            self.target,
-            cast(Sequence[Union[Sequence['cirq.Qid'], int]], self.inputs),
-        )
+            [q.dimension for q in self.target],
+            [i if isinstance(i, int) else [q.dimension for q in i] for i in inputs],
+        ).on(*qubits)
 
 
 def _indented_list_lines_repr(items: Sequence[Any]) -> str:

cirq/interop/quirk/cells/arithmetic_cells_test.py

@@ -344,7 +344,7 @@ def test_with_registers():
         '["+=AB3",1,1,"inputB2"]'
         ']}'
     )
-    op = cast(cirq.ArithmeticOperation, circuit[0].operations[0])
+    op = cast(cirq.ArithmeticGate, circuit[0].operations[0].gate)
 
     with pytest.raises(ValueError, match='number of registers'):
         _ = op.with_registers()
@@ -353,11 +353,11 @@ def test_with_registers():
         _ = op.with_registers(1, 2, 3)
 
     op2 = op.with_registers([], 5, 5)
-    np.testing.assert_allclose(cirq.unitary(cirq.Circuit(op2)), np.array([[1]]), atol=1e-8)
+    np.testing.assert_allclose(cirq.unitary(cirq.Circuit(op2())), np.array([[1]]), atol=1e-8)
 
-    op2 = op.with_registers([*cirq.LineQubit.range(3)], 5, 5)
+    op2 = op.with_registers([2, 2, 2], 5, 5)
     np.testing.assert_allclose(
-        cirq.final_state_vector(cirq.Circuit(op2), initial_state=0),
+        cirq.final_state_vector(cirq.Circuit(op2(*cirq.LineQubit.range(3))), initial_state=0),
         cirq.one_hot(index=25 % 8, shape=8, dtype=np.complex64),
         atol=1e-8,
     )

cirq/interop/quirk/cells/composite_cell_test.py

@@ -92,13 +92,17 @@ def test_custom_circuit_gate():
     # With internal input.
     assert_url_to_circuit_returns(
         '{"cols":[["~a5ls"]],"gates":[{"id":"~a5ls","circuit":{"cols":[["inputA1","+=A1"]]}}]}',
-        cirq.Circuit(cirq.interop.quirk.QuirkArithmeticOperation('+=A1', target=[b], inputs=[[a]])),
+        cirq.Circuit(
+            cirq.interop.quirk.QuirkArithmeticGate('+=A1', target=[2], inputs=[[2]]).on(b, a)
+        ),
     )
 
     # With external input.
     assert_url_to_circuit_returns(
         '{"cols":[["inputA1","~r79k"]],"gates":[{"id":"~r79k","circuit":{"cols":[["+=A1"]]}}]}',
-        cirq.Circuit(cirq.interop.quirk.QuirkArithmeticOperation('+=A1', target=[b], inputs=[[a]])),
+        cirq.Circuit(
+            cirq.interop.quirk.QuirkArithmeticGate('+=A1', target=[2], inputs=[[2]]).on(b, a)
+        ),
     )
 
     # With external control.
@@ -127,9 +131,15 @@ def test_custom_circuit_gate():
         '{"cols":[["~q1fh",1,1,"inputA2"]],"gates":[{"id":"~q1fh",'
         '"circuit":{"cols":[["+=A2"],[1,"+=A2"],[1,"+=A2"]]}}]}',
         cirq.Circuit(
-            cirq.interop.quirk.QuirkArithmeticOperation('+=A2', target=[a, b], inputs=[[d, e]]),
-            cirq.interop.quirk.QuirkArithmeticOperation('+=A2', target=[b, c], inputs=[[d, e]]),
-            cirq.interop.quirk.QuirkArithmeticOperation('+=A2', target=[b, c], inputs=[[d, e]]),
+            cirq.interop.quirk.QuirkArithmeticGate('+=A2', target=[2, 2], inputs=[[2, 2]]).on(
+                a, b, d, e
+            ),
+            cirq.interop.quirk.QuirkArithmeticGate('+=A2', target=[2, 2], inputs=[[2, 2]]).on(
+                b, c, d, e
+            ),
+            cirq.interop.quirk.QuirkArithmeticGate('+=A2', target=[2, 2], inputs=[[2, 2]]).on(
+                b, c, d, e
+            ),
         ),
     )
 

cirq/interop/quirk/cells/input_cells_test.py

@@ -35,7 +35,7 @@ def test_input_cell():
     )
 
     # Overlaps with effect.
-    with pytest.raises(ValueError, match='Overlapping registers'):
+    with pytest.raises(ValueError, match='Duplicate qids'):
         _ = quirk_url_to_circuit(
             'https://algassert.com/quirk#circuit={"cols":[["+=A3","inputA3"]]}'
         )
@@ -53,7 +53,7 @@ def test_reversed_input_cell():
     )
 
     # Overlaps with effect.
-    with pytest.raises(ValueError, match='Overlapping registers'):
+    with pytest.raises(ValueError, match='Duplicate qids'):
         _ = quirk_url_to_circuit(
             'https://algassert.com/quirk#circuit={"cols":[["+=A3","revinputA3"]]}'
         )

cirq/ops/__init__.py

@@ -14,7 +14,7 @@
 """Gates (unitary and non-unitary), operations, base types, and gate sets.
 """
 
-from cirq.ops.arithmetic_operation import ArithmeticOperation
+from cirq.ops.arithmetic_operation import ArithmeticGate, ArithmeticOperation
 
 from cirq.ops.clifford_gate import CliffordGate, PauliTransform, SingleQubitCliffordGate