Split util_bloqs.py into individual files (#1006)

* Split

* Join

* Bookkeeping

* Partition

* Allocate

* Free

* Arbitrary Clifford

* Cast

* Power

* Complete the split

* update resolver dict

Author: tanujkhattar

qualtran/bloqs/arithmetic/addition.py

@@ -56,7 +56,7 @@
     SoquetT,
 )
 from qualtran.bloqs.basic_gates import CNOT, XGate
-from qualtran.bloqs.bookkeeping import util_bloqs
+from qualtran.bloqs.bookkeeping import ArbitraryClifford
 from qualtran.bloqs.mcmt.and_bloq import And
 from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlX
 from qualtran.cirq_interop import decompose_from_cirq_style_method
@@ -352,7 +352,7 @@ def _t_complexity_(self) -> TComplexity:
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return {
             (And(uncompute=self.is_adjoint), self.bitsize),
-            (util_bloqs.ArbitraryClifford(n=2), 5 * self.bitsize),
+            (ArbitraryClifford(n=2), 5 * self.bitsize),
         }
 
     def __pow__(self, power: int):

qualtran/bloqs/arithmetic/subtraction.py

@@ -33,7 +33,7 @@
 )
 from qualtran.bloqs.arithmetic.addition import Add, AddK
 from qualtran.bloqs.basic_gates import XGate
-from qualtran.bloqs.bookkeeping import util_bloqs
+from qualtran.bloqs.bookkeeping import Join, Split
 from qualtran.drawing import Text
 
 if TYPE_CHECKING:
@@ -129,8 +129,8 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
             (XGate(), self.b_dtype.bitsize),
             (AddK(self.b_dtype.bitsize, k=1), 1),
             (Add(a_dtype, b_dtype), 1),
-            (util_bloqs.Split(self.b_dtype), 1),
-            (util_bloqs.Join(self.b_dtype), 1),
+            (Split(self.b_dtype), 1),
+            (Join(self.b_dtype), 1),
         }
 
     def build_composite_bloq(self, bb: 'BloqBuilder', a: Soquet, b: Soquet) -> Dict[str, 'SoquetT']:

qualtran/bloqs/bookkeeping/__init__.py

@@ -11,13 +11,13 @@
 #  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 qualtran.bloqs.bookkeeping.util_bloqs import (
-    Allocate,
-    ArbitraryClifford,
-    Cast,
-    Free,
-    Join,
-    Partition,
-    Power,
-    Split,
-)
+"""Bloqs for virtual operations and register reshaping."""
+
+from qualtran.bloqs.bookkeeping.allocate import Allocate
+from qualtran.bloqs.bookkeeping.arbitrary_clifford import ArbitraryClifford
+from qualtran.bloqs.bookkeeping.cast import Cast
+from qualtran.bloqs.bookkeeping.free import Free
+from qualtran.bloqs.bookkeeping.join import Join
+from qualtran.bloqs.bookkeeping.partition import Partition
+from qualtran.bloqs.bookkeeping.power import Power
+from qualtran.bloqs.bookkeeping.split import Split

qualtran/bloqs/bookkeeping/_bookkeeping_bloq.py

@@ -0,0 +1,47 @@
+#  Copyright 2023 Google LLC
+#
+#  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
+from typing import Dict, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING
+
+from qualtran import Bloq, BloqBuilder, SoquetT
+from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+
+if TYPE_CHECKING:
+
+    from qualtran import AddControlledT, CtrlSpec
+
+
+class _BookkeepingBloq(Bloq, metaclass=abc.ABCMeta):
+    """Base class for utility bloqs used for bookkeeping.
+
+    This bloq:
+    - has trivial controlled versions, which pass through the control register.
+    - does not affect T complexity.
+    """
+
+    def get_ctrl_system(
+        self, ctrl_spec: Optional['CtrlSpec'] = None
+    ) -> Tuple['Bloq', 'AddControlledT']:
+        def add_controlled(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            # ignore `ctrl_soq` and pass it through for bookkeeping operation.
+            out_soqs = bb.add_t(self, **in_soqs)
+            return ctrl_soqs, out_soqs
+
+        return self, add_controlled
+
+    def _t_complexity_(self) -> 'TComplexity':
+        return TComplexity()

qualtran/bloqs/bookkeeping/allocate.py

@@ -0,0 +1,68 @@
+#  Copyright 2023 Google LLC
+#
+#  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.
+from functools import cached_property
+from typing import Any, Dict, Tuple, TYPE_CHECKING
+
+import numpy as np
+from attrs import frozen
+
+from qualtran import Bloq, QDType, Register, Side, Signature, SoquetT
+from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
+from qualtran.drawing import directional_text_box, Text, WireSymbol
+
+if TYPE_CHECKING:
+    import quimb.tensor as qtn
+
+
+@frozen
+class Allocate(_BookkeepingBloq):
+    """Allocate an `n` bit register.
+
+    Attributes:
+          dtype: the quantum data type of the allocated register.
+    """
+
+    dtype: QDType
+
+    @cached_property
+    def signature(self) -> Signature:
+        return Signature([Register('reg', self.dtype, side=Side.RIGHT)])
+
+    def adjoint(self) -> 'Bloq':
+        from qualtran.bloqs.bookkeeping.free import Free
+
+        return Free(self.dtype)
+
+    def on_classical_vals(self) -> Dict[str, int]:
+        return {'reg': 0}
+
+    def add_my_tensors(
+        self,
+        tn: 'qtn.TensorNetwork',
+        tag: Any,
+        *,
+        incoming: Dict[str, 'SoquetT'],
+        outgoing: Dict[str, 'SoquetT'],
+    ):
+        import quimb.tensor as qtn
+
+        data = np.zeros(1 << self.dtype.num_qubits)
+        data[0] = 1
+        tn.add(qtn.Tensor(data=data, inds=(outgoing['reg'],), tags=['Allocate', tag]))
+
+    def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text('')
+        assert reg.name == 'reg'
+        return directional_text_box('alloc', Side.RIGHT)

qualtran/bloqs/bookkeeping/allocate_test.py

@@ -0,0 +1,41 @@
+#  Copyright 2023 Google LLC
+#
+#  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 subprocess
+
+import numpy as np
+import pytest
+
+from qualtran import BloqBuilder, QAny
+from qualtran.bloqs.bookkeeping import Allocate, Free
+from qualtran.testing import execute_notebook
+
+
+def test_free_nonzero_state_vector_leads_to_unnormalized_state():
+    from qualtran.bloqs.basic_gates.hadamard import Hadamard
+    from qualtran.bloqs.basic_gates.on_each import OnEach
+
+    bb = BloqBuilder()
+    qs1 = bb.add(Allocate(QAny(10)))
+    qs2 = bb.add(OnEach(10, Hadamard()), q=qs1)
+    no_return = bb.add(Free(QAny(10)), reg=qs2)
+    assert np.allclose(bb.finalize().tensor_contract(), np.sqrt(1 / 2**10))
+
+
+@pytest.mark.notebook
+def test_notebook():
+    execute_notebook('bookkeeping_bloqs')
+
+
+def test_no_circular_import():
+    subprocess.check_call(['python', '-c', 'from qualtran.bloqs.bookkeeping import allocate'])

qualtran/bloqs/bookkeeping/arbitrary_clifford.py

@@ -0,0 +1,44 @@
+#  Copyright 2023 Google LLC
+#
+#  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.
+from functools import cached_property
+from typing import Union
+
+from attrs import frozen
+from sympy import Expr
+
+from qualtran import Bloq, QAny, Register, Signature
+from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+
+
+@frozen
+class ArbitraryClifford(Bloq):
+    """A bloq representing an arbitrary `n`-qubit clifford operation.
+
+    In the surface code architecture, clifford operations are generally considered
+    cheaper than non-clifford gates. Each clifford also has roughly the same cost independent
+    of what particular operation it is doing.
+
+    You can use this to bloq to represent an arbitrary clifford operation e.g. in bloq_counts
+    resource estimates where the details are unimportant for the resource estimation task
+    at hand.
+    """
+
+    n: Union[int, Expr]
+
+    @cached_property
+    def signature(self) -> Signature:
+        return Signature([Register('x', QAny(bitsize=self.n))])
+
+    def _t_complexity_(self) -> 'TComplexity':
+        return TComplexity(clifford=1)

qualtran/bloqs/bookkeeping/arbitrary_clifford_test.py

@@ -0,0 +1,20 @@
+#  Copyright 2023 Google LLC
+#
+#  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 subprocess
+
+
+def test_no_circular_import():
+    subprocess.check_call(
+        ['python', '-c', 'from qualtran.bloqs.bookkeeping import arbitrary_clifford']
+    )

…ltran/bloqs/bookkeeping/util_bloqs.ipynb …loqs/bookkeeping/bookkeeping_bloqs.ipynbqualtran/bloqs/bookkeeping/util_bloqs.ipynb renamed to qualtran/bloqs/bookkeeping/bookkeeping_bloqs.ipynb qualtran/bloqs/bookkeeping/util_bloqs.ipynb renamed to qualtran/bloqs/bookkeeping/bookkeeping_bloqs.ipynb

@@ -0,0 +1,92 @@
+#  Copyright 2023 Google LLC
+#
+#  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.
+from functools import cached_property
+from typing import Any, Dict, Tuple, TYPE_CHECKING
+
+import attrs
+import numpy as np
+from attrs import frozen
+
+from qualtran import Bloq, QDType, Register, Side, Signature, SoquetT
+from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
+
+if TYPE_CHECKING:
+    import quimb.tensor as qtn
+
+    from qualtran.cirq_interop import CirqQuregT
+    from qualtran.simulation.classical_sim import ClassicalValT
+
+
+@frozen
+class Cast(_BookkeepingBloq):
+    """Cast a register from one n-bit QDType to another QDType.
+
+
+    Args:
+        inp_dtype: Input QDType to cast from.
+        out_dtype: Output QDType to cast to.
+        shape: shape of the register to cast.
+
+    Registers:
+        in: input register to cast from.
+        out: input register to cast to.
+    """
+
+    inp_dtype: QDType
+    out_dtype: QDType
+    shape: Tuple[int, ...] = attrs.field(
+        default=tuple(), converter=lambda v: (v,) if isinstance(v, int) else tuple(v)
+    )
+
+    def __attrs_post_init__(self):
+        if isinstance(self.inp_dtype.num_qubits, int):
+            if self.inp_dtype.num_qubits != self.out_dtype.num_qubits:
+                raise ValueError("Casting only permitted between same sized registers.")
+
+    def adjoint(self) -> 'Bloq':
+        return Cast(inp_dtype=self.out_dtype, out_dtype=self.inp_dtype)
+
+    @cached_property
+    def signature(self) -> Signature:
+        return Signature(
+            [
+                Register('reg', dtype=self.inp_dtype, shape=self.shape, side=Side.LEFT),
+                Register('reg', dtype=self.out_dtype, shape=self.shape, side=Side.RIGHT),
+            ]
+        )
+
+    def add_my_tensors(
+        self,
+        tn: 'qtn.TensorNetwork',
+        tag: Any,
+        *,
+        incoming: Dict[str, 'SoquetT'],
+        outgoing: Dict[str, 'SoquetT'],
+    ):
+        import quimb.tensor as qtn
+
+        tn.add(
+            qtn.Tensor(
+                data=np.eye(2**self.inp_dtype.num_qubits, 2**self.inp_dtype.num_qubits),
+                inds=[outgoing['reg']] + [incoming['reg']],
+                tags=['Cast', tag],
+            )
+        )
+
+    def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
+        # TODO: Actually cast the values https://github.com/quantumlib/Qualtran/issues/734
+        return {'reg': reg}
+
+    def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]:
+        return None, {'reg': reg}