diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py index d0b774142..f55ebb282 100644 --- a/qualtran/_infra/composite_bloq.py +++ b/qualtran/_infra/composite_bloq.py @@ -51,7 +51,7 @@ from qualtran.bloqs.bookkeeping.auto_partition import Unused from qualtran.cirq_interop._cirq_to_bloq import CirqQuregInT, CirqQuregT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT # NDArrays must be bound to np.generic @@ -237,7 +237,7 @@ def decompose_bloq(self) -> 'CompositeBloq': "Consider using the composite bloq directly or using `.flatten()`." ) - def build_call_graph(self, ssa: Optional['SympySymbolAllocator']) -> Set['BloqCountT']: + def build_call_graph(self, ssa: Optional['SympySymbolAllocator']) -> 'BloqCountDictT': """Return the bloq counts by counting up all the subbloqs.""" from qualtran.resource_counting import build_cbloq_call_graph diff --git a/qualtran/bloqs/arithmetic/_shims.py b/qualtran/bloqs/arithmetic/_shims.py index 586fface5..0d40daba7 100644 --- a/qualtran/bloqs/arithmetic/_shims.py +++ b/qualtran/bloqs/arithmetic/_shims.py @@ -19,13 +19,12 @@ will be fleshed out and moved to their final organizational location soon (written: 2024-05-06). """ from functools import cached_property -from typing import Set from attrs import frozen from qualtran import Bloq, QBit, QUInt, Register, Signature from qualtran.bloqs.basic_gates import Toffoli -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -36,8 +35,8 @@ class MultiCToffoli(Bloq): def signature(self) -> 'Signature': return Signature([Register('ctrl', QBit(), shape=(self.n,)), Register('target', QBit())]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Toffoli(), self.n - 2)} + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: + return {Toffoli(): self.n - 2} @frozen @@ -51,9 +50,9 @@ def signature(self) -> 'Signature': [Register('x', QUInt(self.n)), Register('y', QUInt(self.n)), Register('out', QBit())] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: # litinski - return {(Toffoli(), self.n)} + return {Toffoli(): self.n} @frozen diff --git a/qualtran/bloqs/arithmetic/addition.py b/qualtran/bloqs/arithmetic/addition.py index 37ac7630a..17ca331a3 100644 --- a/qualtran/bloqs/arithmetic/addition.py +++ b/qualtran/bloqs/arithmetic/addition.py @@ -62,7 +62,12 @@ if TYPE_CHECKING: from qualtran.drawing import WireSymbol - from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import ( + BloqCountDictT, + BloqCountT, + MutableBloqCountDictT, + SympySymbolAllocator, + ) from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import SymbolicInt @@ -209,10 +214,10 @@ def decompose_from_registers( yield CNOT().on(input_bits[0], output_bits[0]) context.qubit_manager.qfree(ancillas) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n = self.b_dtype.bitsize n_cnot = (n - 2) * 6 + 3 - return {(And(), n - 1), (And().adjoint(), n - 1), (CNOT(), n_cnot)} + return {And(): n - 1, And().adjoint(): n - 1, CNOT(): n_cnot} @bloq_example(generalizer=ignore_split_join) @@ -327,8 +332,8 @@ def decompose_from_registers( ] return cirq.inverse(optree) if self.is_adjoint else optree - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(And(uncompute=self.is_adjoint), self.bitsize), (CNOT(), 5 * self.bitsize)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {And(uncompute=self.is_adjoint): self.bitsize, CNOT(): 5 * self.bitsize} def __pow__(self, power: int): if power == 1: @@ -505,16 +510,16 @@ def build_composite_bloq( def build_call_graph( self, ssa: 'SympySymbolAllocator' ) -> Union['BloqCountDictT', Set['BloqCountT']]: - loading_cost: Tuple[Bloq, SymbolicInt] + loading_cost: MutableBloqCountDictT if len(self.cvs) == 0: - loading_cost = (XGate(), self.bitsize) # upper bound; depends on the data. + loading_cost = {XGate(): self.bitsize} # upper bound; depends on the data. elif len(self.cvs) == 1: - loading_cost = (CNOT(), self.bitsize) # upper bound; depends on the data. + loading_cost = {CNOT(): self.bitsize} # upper bound; depends on the data. else: # Otherwise, use the decomposition return super().build_call_graph(ssa=ssa) - - return {loading_cost, (Add(QUInt(self.bitsize)), 1)} + loading_cost[Add(QUInt(self.bitsize))] = 1 + return loading_cost def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']: if self.cvs: diff --git a/qualtran/bloqs/arithmetic/bitwise.py b/qualtran/bloqs/arithmetic/bitwise.py index fb5f8264c..87dc2d84b 100644 --- a/qualtran/bloqs/arithmetic/bitwise.py +++ b/qualtran/bloqs/arithmetic/bitwise.py @@ -39,7 +39,7 @@ from qualtran.symbolics import is_symbolic, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -90,9 +90,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'Soq return {'x': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_flips = self.bitsize if self.is_symbolic() else sum(self._bits_k) - return {(XGate(), num_flips)} + return {XGate(): num_flips} def on_classical_vals(self, x: 'ClassicalValT') -> dict[str, 'ClassicalValT']: if isinstance(self.k, sympy.Expr): @@ -156,8 +156,8 @@ def build_composite_bloq(self, bb: BloqBuilder, x: Soquet, y: Soquet) -> dict[st return {'x': bb.join(xs, dtype=self.dtype), 'y': bb.join(ys, dtype=self.dtype)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: - return {(CNOT(), self.dtype.num_qubits)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {CNOT(): self.dtype.num_qubits} def on_classical_vals( self, x: 'ClassicalValT', y: 'ClassicalValT' diff --git a/qualtran/bloqs/arithmetic/comparison.py b/qualtran/bloqs/arithmetic/comparison.py index 9da8ed052..ff4986262 100644 --- a/qualtran/bloqs/arithmetic/comparison.py +++ b/qualtran/bloqs/arithmetic/comparison.py @@ -14,18 +14,7 @@ from collections import defaultdict from functools import cached_property -from typing import ( - Dict, - Iterable, - Iterator, - List, - Optional, - Sequence, - Set, - Tuple, - TYPE_CHECKING, - Union, -) +from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -65,7 +54,11 @@ if TYPE_CHECKING: from qualtran import BloqBuilder - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import ( + BloqCountDictT, + MutableBloqCountDictT, + SympySymbolAllocator, + ) from qualtran.simulation.classical_sim import ClassicalValT @@ -183,22 +176,22 @@ def decompose_from_registers( def _has_unitary_(self): return True - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if ( not is_symbolic(self.less_than_val, self.bitsize) and self.less_than_val >= 2**self.bitsize ): - return {(XGate(), 1)} + return {XGate(): 1} num_set_bits = ( int(self.less_than_val).bit_count() if not is_symbolic(self.less_than_val) else self.bitsize ) return { - (And(), self.bitsize), - (And().adjoint(), self.bitsize), - (CNOT(), num_set_bits + 2 * self.bitsize), - (XGate(), 2 * (1 + num_set_bits)), + And(): self.bitsize, + And().adjoint(): self.bitsize, + CNOT(): num_set_bits + 2 * self.bitsize, + XGate(): 2 * (1 + num_set_bits), } @@ -307,8 +300,8 @@ def __pow__(self, power: int) -> 'BiQubitsMixer': return self.adjoint() return NotImplemented # pragma: no cover - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(XGate(), 1), (CNOT(), 9), (And(uncompute=self.is_adjoint), 2)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {XGate(): 1, CNOT(): 9, And(uncompute=self.is_adjoint): 2} def _has_unitary_(self): return not self.is_adjoint @@ -380,8 +373,8 @@ def __pow__(self, power: int) -> Union['SingleQubitCompare', cirq.Gate]: return self.adjoint() return self - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(XGate(), 1), (CNOT(), 4), (And(uncompute=self.is_adjoint), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {XGate(): 1, CNOT(): 4, And(uncompute=self.is_adjoint): 1} @bloq_example @@ -575,13 +568,13 @@ def decompose_from_registers( all_ancilla = set([q for op in adjoint for q in op.qubits if q not in input_qubits]) context.qubit_manager.qfree(all_ancilla) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if is_symbolic(self.x_bitsize, self.y_bitsize): return { - (BiQubitsMixer(), self.x_bitsize), - (BiQubitsMixer().adjoint(), self.x_bitsize), - (SingleQubitCompare(), 1), - (SingleQubitCompare().adjoint(), 1), + BiQubitsMixer(): self.x_bitsize, + BiQubitsMixer().adjoint(): self.x_bitsize, + SingleQubitCompare(): 1, + SingleQubitCompare().adjoint(): 1, } n = min(self.x_bitsize, self.y_bitsize) @@ -613,7 +606,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: ret[And(1, 0).adjoint()] += 1 ret[CNOT()] += 1 - return set(ret.items()) + return ret def _has_unitary_(self): return True @@ -691,8 +684,8 @@ def build_composite_bloq( target = bb.add(XGate(), q=target) return {'a': a, 'b': b, 'target': target} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(LessThanEqual(self.a_bitsize, self.b_bitsize), 1), (XGate(), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {LessThanEqual(self.a_bitsize, self.b_bitsize): 1, XGate(): 1} @bloq_example @@ -885,23 +878,23 @@ def wire_symbol( return TextBox('t⨁(a>b)') raise ValueError(f'Unknown register name {reg.name}') - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.bitsize == 1: - return {(MultiControlX(cvs=(1, 0)), 1)} + return {MultiControlX(cvs=(1, 0)): 1} if self.signed: return { - (CNOT(), 6 * self.bitsize - 7), - (XGate(), 2 * self.bitsize + 2), - (And(), self.bitsize - 1), - (And(uncompute=True), self.bitsize - 1), + CNOT(): 6 * self.bitsize - 7, + XGate(): 2 * self.bitsize + 2, + And(): self.bitsize - 1, + And(uncompute=True): self.bitsize - 1, } return { - (CNOT(), 6 * self.bitsize - 1), - (XGate(), 2 * self.bitsize + 4), - (And(), self.bitsize), - (And(uncompute=True), self.bitsize), + CNOT(): 6 * self.bitsize - 1, + XGate(): 2 * self.bitsize + 4, + And(): self.bitsize, + And(uncompute=True): self.bitsize, } @@ -941,8 +934,8 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - return TextBox(f"⨁(x > {self.val})") raise ValueError(f'Unknown register symbol {reg.name}') - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(LessThanConstant(self.bitsize, less_than_val=self.val), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {LessThanConstant(self.bitsize, less_than_val=self.val): 1} @bloq_example @@ -1007,8 +1000,8 @@ def build_composite_bloq( x = bb.join(xs) return {'x': x, 'target': target} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(MultiControlX(self.bits_k), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {MultiControlX(self.bits_k): 1} def _make_equals_a_constant(): @@ -1134,21 +1127,22 @@ def on_classical_vals( return {'ctrl': ctrl, 'a': a, 'b': b, 'target': target ^ (a > b)} return {'ctrl': ctrl, 'a': a, 'b': b, 'target': target} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - signed_ops = [] + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + signed_ops: 'MutableBloqCountDictT' = {} if isinstance(self.dtype, QInt): - signed_ops = [ - (SignExtend(self.dtype, QInt(self.dtype.bitsize + 1)), 2), - (SignExtend(self.dtype, QInt(self.dtype.bitsize + 1)).adjoint(), 2), - ] + signed_ops = { + SignExtend(self.dtype, QInt(self.dtype.bitsize + 1)): 2, + SignExtend(self.dtype, QInt(self.dtype.bitsize + 1)).adjoint(): 2, + } dtype = attrs.evolve(self.dtype, bitsize=self.dtype.bitsize + 1) return { - (BitwiseNot(dtype), 2), - (BitwiseNot(QUInt(dtype.bitsize + 1)), 2), - (OutOfPlaceAdder(self.dtype.bitsize + 1).adjoint(), 1), - (OutOfPlaceAdder(self.dtype.bitsize + 1), 1), - (MultiControlX((self.cv, 1)), 1), - }.union(signed_ops) + BitwiseNot(dtype): 2, + BitwiseNot(QUInt(dtype.bitsize + 1)): 2, + OutOfPlaceAdder(self.dtype.bitsize + 1).adjoint(): 1, + OutOfPlaceAdder(self.dtype.bitsize + 1): 1, + MultiControlX((self.cv, 1)): 1, + **signed_ops, + } @bloq_example(generalizer=ignore_split_join) diff --git a/qualtran/bloqs/arithmetic/controlled_addition.py b/qualtran/bloqs/arithmetic/controlled_addition.py index aa3f744a6..cfe4bcc5d 100644 --- a/qualtran/bloqs/arithmetic/controlled_addition.py +++ b/qualtran/bloqs/arithmetic/controlled_addition.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Set, TYPE_CHECKING, Union +from typing import Dict, TYPE_CHECKING, Union import numpy as np import sympy @@ -42,7 +42,7 @@ import quimb.tensor as qtn from qualtran.drawing import WireSymbol - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -155,11 +155,11 @@ def build_composite_bloq( ctrl = bb.join(np.array([ctrl_q])) return {'ctrl': ctrl, 'a': a, 'b': b} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (And(self.cv, 1), self.a_dtype.bitsize), - (Add(self.a_dtype, self.b_dtype), 1), - (And(self.cv, 1).adjoint(), self.a_dtype.bitsize), + And(self.cv, 1): self.a_dtype.bitsize, + Add(self.a_dtype, self.b_dtype): 1, + And(self.cv, 1).adjoint(): self.a_dtype.bitsize, } diff --git a/qualtran/bloqs/arithmetic/conversions/contiguous_index.py b/qualtran/bloqs/arithmetic/conversions/contiguous_index.py index be2ca4d7b..5444e62d3 100644 --- a/qualtran/bloqs/arithmetic/conversions/contiguous_index.py +++ b/qualtran/bloqs/arithmetic/conversions/contiguous_index.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING from attrs import frozen @@ -23,7 +23,7 @@ from qualtran.drawing.musical_score import Text, TextBox if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -85,9 +85,9 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - text = r'⊕ν(ν-1)/2+μ' return TextBox(text) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_toffoli = self.bitsize**2 + self.bitsize - 1 - return {(Toffoli(), num_toffoli)} + return {Toffoli(): num_toffoli} @bloq_example diff --git a/qualtran/bloqs/arithmetic/conversions/ones_complement_to_twos_complement.py b/qualtran/bloqs/arithmetic/conversions/ones_complement_to_twos_complement.py index 1fdda0419..19a9ebd1c 100644 --- a/qualtran/bloqs/arithmetic/conversions/ones_complement_to_twos_complement.py +++ b/qualtran/bloqs/arithmetic/conversions/ones_complement_to_twos_complement.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -51,10 +51,10 @@ def signature(self) -> Signature: ] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Take the sign qubit as a control and cnot the remaining qubits, then # add it to the remaining n-1 bits. - return {(Toffoli(), (self.bitsize - 2))} + return {Toffoli(): (self.bitsize - 2)} @bloq_example diff --git a/qualtran/bloqs/arithmetic/conversions/sign_extension.py b/qualtran/bloqs/arithmetic/conversions/sign_extension.py index b9a45e0cd..16e8f2152 100644 --- a/qualtran/bloqs/arithmetic/conversions/sign_extension.py +++ b/qualtran/bloqs/arithmetic/conversions/sign_extension.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, Soquet, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -91,8 +91,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'Soq return {'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: - return {(MultiTargetCNOT(self.extend_bitsize), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {MultiTargetCNOT(self.extend_bitsize): 1} def on_classical_vals(self, x: 'ClassicalValT') -> dict[str, 'ClassicalValT']: return {'y': x} @@ -173,8 +173,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'Soq return {'y': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: - return {(MultiTargetCNOT(self.truncate_bitsize), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {MultiTargetCNOT(self.truncate_bitsize): 1} def on_classical_vals(self, x: 'ClassicalValT') -> dict[str, 'ClassicalValT']: bits = self.inp_dtype.to_bits(int(x)) diff --git a/qualtran/bloqs/arithmetic/hamming_weight.py b/qualtran/bloqs/arithmetic/hamming_weight.py index 0b99ba5f6..2322fd840 100644 --- a/qualtran/bloqs/arithmetic/hamming_weight.py +++ b/qualtran/bloqs/arithmetic/hamming_weight.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Iterator, List, Set, TYPE_CHECKING +from typing import Iterator, List, TYPE_CHECKING import cirq from attrs import frozen @@ -25,7 +25,7 @@ from qualtran.symbolics import bit_length, is_symbolic, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -114,7 +114,7 @@ def decompose_from_registers( out: List[cirq.Qid] = [*quregs['out'][::-1]] yield self._decompose_using_three_to_two_adders(x, junk, out) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_and = self.junk_bitsize num_cnot = num_and * 5 + self.bit_count_of_bitsize - return {(And(), num_and), (CNOT(), num_cnot)} + return {And(): num_and, CNOT(): num_cnot} diff --git a/qualtran/bloqs/arithmetic/multiplication.py b/qualtran/bloqs/arithmetic/multiplication.py index aff7af746..5408c3d9c 100644 --- a/qualtran/bloqs/arithmetic/multiplication.py +++ b/qualtran/bloqs/arithmetic/multiplication.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Iterable, List, Sequence, Set, TYPE_CHECKING, Union +from typing import Dict, Iterable, List, Sequence, TYPE_CHECKING, Union import cirq import numpy as np @@ -39,7 +39,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -131,9 +131,9 @@ def my_tensors( return _my_tensors_from_gate(self, self.signature, incoming=incoming, outgoing=outgoing) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # TODO: The T-complexity here is approximate. - return {(TGate(), 8 * smax(self.a_bitsize, self.b_bitsize) ** 2)} + return {TGate(): 8 * smax(self.a_bitsize, self.b_bitsize) ** 2} @bloq_example @@ -189,12 +189,12 @@ def on_classical_vals(self, **vals: int) -> Dict[str, 'ClassicalValT']: def pretty_name(self) -> str: return "a^2" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # TODO Determine precise clifford count and/or ignore. # See: https://github.com/quantumlib/Qualtran/issues/219 # See: https://github.com/quantumlib/Qualtran/issues/217 num_toff = self.bitsize * (self.bitsize - 1) - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} def my_tensors( self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] @@ -275,11 +275,11 @@ def signature(self): def pretty_name(self) -> str: return "SOS" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_toff = self.k * self.bitsize**2 - self.bitsize if self.k % 3 == 0: num_toff -= 1 - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} @bloq_example @@ -329,12 +329,12 @@ def signature(self): def pretty_name(self) -> str: return "a*b" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # TODO Determine precise clifford count and/or ignore. # See: https://github.com/quantumlib/Qualtran/issues/219 # See: https://github.com/quantumlib/Qualtran/issues/217 num_toff = 2 * self.a_bitsize * self.b_bitsize - max(self.a_bitsize, self.b_bitsize) - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} @bloq_example @@ -390,12 +390,12 @@ def signature(self): def pretty_name(self) -> str: return "r*i" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Eq. D8, we are assuming dA(r_bitsize) and dB(i_bitsize) are inputs and # the user has ensured these are large enough for their desired # precision. num_toff = self.r_bitsize * (2 * self.i_bitsize - 1) - self.i_bitsize**2 - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} @bloq_example @@ -447,10 +447,10 @@ def signature(self): def pretty_name(self) -> str: return "a*b" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Eq. D13, there it is suggested keeping both registers the same size is optimal. num_toff = self.bitsize**2 - self.bitsize - 1 - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} @bloq_example @@ -506,10 +506,10 @@ def signature(self): def pretty_name(self) -> str: return "a^2" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Bottom of page 74 num_toff = self.bitsize**2 // 2 - 4 - return {(Toffoli(), num_toff)} + return {Toffoli(): num_toff} @bloq_example @@ -562,7 +562,7 @@ def signature(self): def pretty_name(self) -> str: return "1/a" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # initial approximation: Figure 4 num_int = self.bitsize - self.num_frac # Newton-Raphson: Eq. (1) @@ -571,14 +571,14 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # TODO: When decomposing we will potentially need to use larger registers. # Related issue: https://github.com/quantumlib/Qualtran/issues/655 return { - (Toffoli(), num_int - 1), - (CNOT(), 2 + num_int - 1), - (MultiControlX(cvs=HasLength(num_int)), 1), - (XGate(), 1), - (SquareRealNumber(self.bitsize), num_iters), # x^2 - (MultiplyTwoReals(self.bitsize), num_iters), # a * x^2 - (ScaleIntByReal(self.bitsize, 2), num_iters), # 2 * x - (Subtract(QUInt(self.bitsize)), num_iters), # 2 * x - a * x^2 + Toffoli(): num_int - 1, + CNOT(): 2 + num_int - 1, + MultiControlX(cvs=HasLength(num_int)): 1, + XGate(): 1, + SquareRealNumber(self.bitsize): num_iters, # x^2 + MultiplyTwoReals(self.bitsize): num_iters, # a * x^2 + ScaleIntByReal(self.bitsize, 2): num_iters, # 2 * x + Subtract(QUInt(self.bitsize)): num_iters, # 2 * x - a * x^2 } diff --git a/qualtran/bloqs/arithmetic/permutation.py b/qualtran/bloqs/arithmetic/permutation.py index 093c8a297..f3ad8f21f 100644 --- a/qualtran/bloqs/arithmetic/permutation.py +++ b/qualtran/bloqs/arithmetic/permutation.py @@ -129,8 +129,8 @@ def build_call_graph( x = ssa.new_symbol('x') cycle_len = slen(self.cycle) return { - (EqualsAConstant(self.bitsize, x), cycle_len + 1), - (XorK(QUInt(self.bitsize), x).controlled(), cycle_len), + EqualsAConstant(self.bitsize, x): cycle_len + 1, + XorK(QUInt(self.bitsize), x).controlled(): cycle_len, } return super().build_call_graph(ssa) @@ -275,7 +275,7 @@ def build_call_graph( if is_symbolic(self.cycles): # worst case cost: single cycle of length N cycle = Shaped((self.N,)) - return {(PermutationCycle(self.N, cycle), 1)} + return {PermutationCycle(self.N, cycle): 1} return super().build_call_graph(ssa) diff --git a/qualtran/bloqs/arithmetic/sorting.py b/qualtran/bloqs/arithmetic/sorting.py index baa99413b..d18db924e 100644 --- a/qualtran/bloqs/arithmetic/sorting.py +++ b/qualtran/bloqs/arithmetic/sorting.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, Set +from typing import Dict import numpy as np import sympy @@ -34,7 +34,7 @@ ) from qualtran.bloqs.arithmetic import GreaterThan from qualtran.bloqs.basic_gates import CSwap -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import bit_length, is_symbolic, SymbolicInt @@ -174,8 +174,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', xs: 'SoquetT') -> Dict[str, 'S return {'xs': xs, 'junk': np.array(junk)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Comparator(self.bitsize), self.num_comparisons)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> BloqCountDictT: + return {Comparator(self.bitsize): self.num_comparisons} @bloq_example @@ -278,8 +278,8 @@ def build_composite_bloq( return {'result': result, 'junk': np.array(junk)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Comparator(self.bitsize), self.num_comparisons)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> BloqCountDictT: + return {Comparator(self.bitsize): self.num_comparisons} @bloq_example @@ -379,8 +379,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', xs: 'SoquetT') -> dict[str, 'S return {'xs': xs, 'junk': junk} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: - return {(Comparator(self.bitsize), self.num_comparisons)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> BloqCountDictT: + return {Comparator(self.bitsize): self.num_comparisons} @bloq_example diff --git a/qualtran/bloqs/arithmetic/subtraction.py b/qualtran/bloqs/arithmetic/subtraction.py index c61429621..f9bb91437 100644 --- a/qualtran/bloqs/arithmetic/subtraction.py +++ b/qualtran/bloqs/arithmetic/subtraction.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Optional, Tuple, TYPE_CHECKING, Union import numpy as np import sympy @@ -43,7 +43,7 @@ if TYPE_CHECKING: from qualtran.drawing import WireSymbol - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -150,18 +150,18 @@ def wire_symbol( else: raise ValueError() - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': delta = self.b_dtype.bitsize - self.a_dtype.bitsize - return ( - { - (OnEach(self.b_dtype.bitsize, XGate()), 3), - (Add(QUInt(self.b_dtype.bitsize), QUInt(self.b_dtype.bitsize)), 1), - } - .union( - [(MultiTargetCNOT(delta), 2)] if delta and isinstance(self.a_dtype, QInt) else [] - ) - .union([(Allocate(QAny(delta)), 1), (Free(QAny(delta)), 1)] if delta else []) - ) + costs = { + OnEach(self.b_dtype.bitsize, XGate()): 3, + Add(QUInt(self.b_dtype.bitsize), QUInt(self.b_dtype.bitsize)): 1, + } + if delta: + costs[Allocate(QAny(delta))] = 1 + costs[Free(QAny(delta))] = 1 + if isinstance(self.a_dtype, QInt): + costs[MultiTargetCNOT(delta)] = 2 + return costs def build_composite_bloq(self, bb: 'BloqBuilder', a: Soquet, b: Soquet) -> Dict[str, 'SoquetT']: delta = self.b_dtype.bitsize - self.a_dtype.bitsize @@ -288,8 +288,8 @@ def wire_symbol( else: raise ValueError() - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(BitwiseNot(self.dtype), 2), (Add(self.dtype, self.dtype), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {BitwiseNot(self.dtype): 2, Add(self.dtype, self.dtype): 1} def build_composite_bloq(self, bb: 'BloqBuilder', a: Soquet, b: Soquet) -> Dict[str, 'SoquetT']: b = bb.add(BitwiseNot(self.dtype), x=b) # a, -1 - b diff --git a/qualtran/bloqs/arithmetic/trigonometric/arcsin.py b/qualtran/bloqs/arithmetic/trigonometric/arcsin.py index 232359bc5..2405e65ef 100644 --- a/qualtran/bloqs/arithmetic/trigonometric/arcsin.py +++ b/qualtran/bloqs/arithmetic/trigonometric/arcsin.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Set +from typing import Dict import numpy as np from attrs import frozen @@ -20,7 +20,7 @@ from qualtran import Bloq, bloq_example, BloqDocSpec, QFxp, Register, Signature from qualtran.bloqs.basic_gates import Toffoli from qualtran.bloqs.rotations.phase_gradient import _fxp -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import is_symbolic, SymbolicInt @@ -81,7 +81,7 @@ def on_classical_vals( result ^= int(np.arcsin(x_fxp) * 2**self.bitsize) return {'x': x, 'result': result} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: n = self.bitsize p = self.bitsize - self.num_frac d = self.degree @@ -97,7 +97,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - 9 * p**2 + 2 ) - return {(Toffoli(), toffolis)} + return {Toffoli(): toffolis} @bloq_example diff --git a/qualtran/bloqs/arithmetic/trigonometric/arctan.py b/qualtran/bloqs/arithmetic/trigonometric/arctan.py index 07a0aba7c..746411d9f 100644 --- a/qualtran/bloqs/arithmetic/trigonometric/arctan.py +++ b/qualtran/bloqs/arithmetic/trigonometric/arctan.py @@ -20,7 +20,7 @@ from qualtran import GateWithRegisters, QFxp, Signature from qualtran.bloqs.arithmetic import PlusEqualProduct -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import is_symbolic, SymbolicInt @@ -76,12 +76,12 @@ def apply(self, *register_values: int) -> Union[int, Iterable[int]]: ) return input_val, target_sign ^ output_sign, target_val ^ output_bin - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: # Hack to propagate bigO(...). Here `c` is some constant. c = ssa.new_symbol("c") return { - (PlusEqualProduct(self.target_bitsize, self.target_bitsize, 2 * self.target_bitsize), c) + PlusEqualProduct(self.target_bitsize, self.target_bitsize, 2 * self.target_bitsize): c } def adjoint(self) -> 'ArcTan': diff --git a/qualtran/bloqs/basic_gates/on_each.py b/qualtran/bloqs/basic_gates/on_each.py index b6938326a..0d8a19070 100644 --- a/qualtran/bloqs/basic_gates/on_each.py +++ b/qualtran/bloqs/basic_gates/on_each.py @@ -14,7 +14,7 @@ """Classes to apply single qubit bloq to multiple qubits.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple +from typing import Dict, Optional, Tuple import attrs import sympy @@ -32,7 +32,7 @@ ) from qualtran.drawing import Text, WireSymbol from qualtran.drawing.musical_score import TextBox -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import SymbolicInt @@ -78,8 +78,8 @@ def build_composite_bloq(self, bb: BloqBuilder, *, q: Soquet) -> Dict[str, Soque qs[i] = bb.add(self.gate, q=qs[i]) return {'q': bb.join(qs, self.target_dtype)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(self.gate, self.n)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {self.gate: self.n} def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> WireSymbol: one_reg = self.gate.wire_symbol(reg=reg, idx=idx) diff --git a/qualtran/bloqs/basic_gates/power.py b/qualtran/bloqs/basic_gates/power.py index d7cad38c5..5331c3afd 100644 --- a/qualtran/bloqs/basic_gates/power.py +++ b/qualtran/bloqs/basic_gates/power.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING +from typing import Dict, TYPE_CHECKING import numpy as np from attrs import frozen @@ -23,7 +23,7 @@ if TYPE_CHECKING: import cirq - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -66,8 +66,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str soqs = bb.add_d(self.bloq, **soqs) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(self.bloq, self.power)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {self.bloq: self.power} def __pow__(self, power) -> 'Power': bloq = self.bloq.adjoint() if power < 0 else self.bloq diff --git a/qualtran/bloqs/basic_gates/swap.py b/qualtran/bloqs/basic_gates/swap.py index fece5642d..3242a1abb 100644 --- a/qualtran/bloqs/basic_gates/swap.py +++ b/qualtran/bloqs/basic_gates/swap.py @@ -13,18 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import ( - Dict, - Iterable, - Iterator, - List, - Optional, - Sequence, - Set, - Tuple, - TYPE_CHECKING, - Union, -) +from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union import cirq import numpy as np @@ -59,7 +48,7 @@ import quimb.tensor as qtn from qualtran import AddControlledT, CompositeBloq - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -203,8 +192,8 @@ def on_classical_vals( def _t_complexity_(self) -> 'TComplexity': return TComplexity(t=7, clifford=10) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(TGate(), 7), (CNOT(), 8), (Hadamard(), 2)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {TGate(): 7, CNOT(): 8, Hadamard(): 2} def adjoint(self) -> 'Bloq': return self @@ -250,8 +239,8 @@ def build_composite_bloq( return {'x': bb.join(xs), 'y': bb.join(ys)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(TwoBitSwap(), self.bitsize)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {TwoBitSwap(): self.bitsize} def on_classical_vals( self, x: 'ClassicalValT', y: 'ClassicalValT' @@ -329,8 +318,8 @@ def build_composite_bloq( return {'ctrl': ctrl, 'x': bb.join(xs), 'y': bb.join(ys)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(TwoBitCSwap(), self.bitsize)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {TwoBitCSwap(): self.bitsize} def on_classical_vals( self, ctrl: 'ClassicalValT', x: 'ClassicalValT', y: 'ClassicalValT' diff --git a/qualtran/bloqs/basic_gates/z_basis.py b/qualtran/bloqs/basic_gates/z_basis.py index deffb4a6e..e7b017eeb 100644 --- a/qualtran/bloqs/basic_gates/z_basis.py +++ b/qualtran/bloqs/basic_gates/z_basis.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import cast, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union import attrs import numpy as np @@ -49,7 +49,7 @@ import quimb.tensor as qtn from qualtran.cirq_interop import CirqQuregT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator _ZERO = np.array([1, 0], dtype=np.complex128) _ONE = np.array([0, 1], dtype=np.complex128) @@ -431,8 +431,8 @@ def on_classical_vals(self, *, val: Optional[int] = None) -> Dict[str, Union[int assert val == self.val, val return {} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(ArbitraryClifford(self.bitsize), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {ArbitraryClifford(self.bitsize): 1} def __str__(self) -> str: s = f'{self.val}' diff --git a/qualtran/bloqs/block_encoding/chebyshev_polynomial.py b/qualtran/bloqs/block_encoding/chebyshev_polynomial.py index d05610558..0ac54b717 100644 --- a/qualtran/bloqs/block_encoding/chebyshev_polynomial.py +++ b/qualtran/bloqs/block_encoding/chebyshev_polynomial.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Tuple, TYPE_CHECKING, Union import attrs import numpy as np @@ -36,7 +36,11 @@ from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import ( + BloqCountDictT, + MutableBloqCountDictT, + SympySymbolAllocator, + ) @attrs.frozen @@ -138,14 +142,14 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So soqs |= bb.add_d(self.block_encoding, **soqs) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n = self.order - s: Set['BloqCountT'] = { - (self.block_encoding, n // 2 + n % 2), - (self.block_encoding.adjoint(), n // 2), + s: 'MutableBloqCountDictT' = { + self.block_encoding: n // 2 + n % 2, + self.block_encoding.adjoint(): n // 2, } if is_symbolic(self.ancilla_bitsize) or self.ancilla_bitsize > 0: - s.add((self.reflection_bloq, n - n % 2)) + s[self.reflection_bloq] = n - n % 2 return s diff --git a/qualtran/bloqs/block_encoding/phase.py b/qualtran/bloqs/block_encoding/phase.py index 33786a1fd..3d55b06c9 100644 --- a/qualtran/bloqs/block_encoding/phase.py +++ b/qualtran/bloqs/block_encoding/phase.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set +from typing import Dict from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import GlobalPhase from qualtran.bloqs.block_encoding import BlockEncoding from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import SymbolicFloat, SymbolicInt @@ -79,8 +79,8 @@ def signature(self) -> Signature: def signal_state(self) -> BlackBoxPrepare: return self.block_encoding.signal_state - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: - return {(self.block_encoding, 1), (GlobalPhase(exponent=self.phi, eps=self.eps), 1)} + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: + return {self.block_encoding: 1, GlobalPhase(exponent=self.phi, eps=self.eps): 1} def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, SoquetT]: bb.add(GlobalPhase(exponent=self.phi, eps=self.eps)) diff --git a/qualtran/bloqs/block_encoding/product.py b/qualtran/bloqs/block_encoding/product.py index 44156197b..b4169cc10 100644 --- a/qualtran/bloqs/block_encoding/product.py +++ b/qualtran/bloqs/block_encoding/product.py @@ -14,7 +14,7 @@ from collections import Counter from functools import cached_property -from typing import cast, Dict, List, Sequence, Set, Tuple, Union +from typing import cast, Dict, List, Sequence, Tuple, Union from attrs import field, frozen, validators from numpy.typing import NDArray @@ -40,7 +40,7 @@ from qualtran.bloqs.mcmt import MultiControlX from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.resource_counting.generalizers import ignore_split_join from qualtran.symbolics import HasLength, is_symbolic, prod, smax, ssum, SymbolicFloat, SymbolicInt from qualtran.symbolics.math_funcs import is_zero @@ -171,7 +171,7 @@ def constituents(self) -> Sequence[Bloq]: ret.append(AutoPartition(u, partition, left_only=False)) return ret - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: counts = Counter[Bloq]() for bloq in self.constituents: counts[bloq] += 1 @@ -180,7 +180,7 @@ def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: if not is_zero(u.ancilla_bitsize) and n - 1 > 0 and i != n - 1: counts[MultiControlX(HasLength(u.ancilla_bitsize))] += 1 counts[XGate()] += 1 - return set(counts.items()) + return counts def build_composite_bloq( self, bb: BloqBuilder, system: SoquetT, **soqs: SoquetT diff --git a/qualtran/bloqs/block_encoding/sparse_matrix.py b/qualtran/bloqs/block_encoding/sparse_matrix.py index 8cfe3cff5..60f8a39cf 100644 --- a/qualtran/bloqs/block_encoding/sparse_matrix.py +++ b/qualtran/bloqs/block_encoding/sparse_matrix.py @@ -14,7 +14,7 @@ import abc from functools import cached_property -from typing import Dict, Iterable, Set, Tuple +from typing import Dict, Iterable, Tuple import numpy as np import sympy @@ -46,7 +46,7 @@ from qualtran.bloqs.state_preparation.prepare_uniform_superposition import ( PrepareUniformSuperposition, ) -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt from qualtran.symbolics.math_funcs import bit_length @@ -220,14 +220,14 @@ def diffusion(self): ], ) - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: return { - (self.diffusion, 1), - (self.col_oracle, 1), - (self.entry_oracle, 1), - (Swap(self.system_bitsize), 1), - (self.row_oracle.adjoint(), 1), - (self.diffusion.adjoint(), 1), + self.diffusion: 1, + self.col_oracle: 1, + self.entry_oracle: 1, + Swap(self.system_bitsize): 1, + self.row_oracle.adjoint(): 1, + self.diffusion.adjoint(): 1, } def build_composite_bloq( @@ -331,10 +331,10 @@ def call_classically(self, l: ClassicalValT, i: ClassicalValT) -> Tuple[Classica raise IndexError("l out of bounds") return ((l + i - self.bandsize) % (2**self.system_bitsize), i) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set[BloqCountT]: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> BloqCountDictT: return { - (Add(QUInt(self.system_bitsize), QUInt(self.system_bitsize)), 1), - (AddK(self.system_bitsize, -self.bandsize, signed=True), 1), + Add(QUInt(self.system_bitsize), QUInt(self.system_bitsize)): 1, + AddK(self.system_bitsize, -self.bandsize, signed=True): 1, } def build_composite_bloq(self, bb: BloqBuilder, l: SoquetT, i: SoquetT) -> Dict[str, SoquetT]: diff --git a/qualtran/bloqs/block_encoding/tensor_product.py b/qualtran/bloqs/block_encoding/tensor_product.py index 24b28e078..9150c4516 100644 --- a/qualtran/bloqs/block_encoding/tensor_product.py +++ b/qualtran/bloqs/block_encoding/tensor_product.py @@ -14,7 +14,7 @@ from collections import Counter from functools import cached_property -from typing import Dict, Set, Tuple +from typing import Dict, Tuple from attrs import evolve, field, frozen, validators from typing_extensions import Self @@ -32,7 +32,7 @@ from qualtran.bloqs.bookkeeping import Partition from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import is_symbolic, prod, ssum, SymbolicFloat, SymbolicInt @@ -106,8 +106,8 @@ def signal_state(self) -> BlackBoxPrepare: # TODO: implement by taking tensor product of component signal states raise NotImplementedError - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: - return set(Counter(self.block_encodings).items()) + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: + return Counter(self.block_encodings) def build_composite_bloq( self, bb: BloqBuilder, system: SoquetT, **soqs: SoquetT diff --git a/qualtran/bloqs/block_encoding/unitary.py b/qualtran/bloqs/block_encoding/unitary.py index 60821402b..96da03307 100644 --- a/qualtran/bloqs/block_encoding/unitary.py +++ b/qualtran/bloqs/block_encoding/unitary.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set +from typing import Dict from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.block_encoding import BlockEncoding from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import SymbolicFloat, SymbolicInt @@ -74,8 +74,8 @@ def signature(self) -> Signature: def signal_state(self) -> BlackBoxPrepare: return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize])) - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set[BloqCountT]: - return {(self.U, 1)} + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: + return {self.U: 1} def build_composite_bloq( self, bb: BloqBuilder, system: SoquetT, **soqs: SoquetT diff --git a/qualtran/bloqs/chemistry/black_boxes.py b/qualtran/bloqs/chemistry/black_boxes.py index 7829fd644..f9670a2eb 100644 --- a/qualtran/bloqs/chemistry/black_boxes.py +++ b/qualtran/bloqs/chemistry/black_boxes.py @@ -16,7 +16,7 @@ These are for temporary convenience to lock-in the quoted literature costs. """ from functools import cached_property -from typing import Optional, Set, Tuple, TYPE_CHECKING +from typing import Optional, Tuple, TYPE_CHECKING import attrs import numpy as np @@ -28,7 +28,7 @@ from qualtran.drawing import Circle, Text, TextBox, WireSymbol if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def qroam_cost(x, data_size: int, bitsize: int, adjoint: bool = False): @@ -108,14 +108,14 @@ class QROAM(Bloq): def signature(self) -> Signature: return Signature.build(sel=(self.data_size - 1).bit_length(), trg=self.target_bitsize) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cost = get_qroam_cost_clean_ancilla( self.data_size, self.target_bitsize, adjoint=self.is_adjoint, qroam_block_size=self.qroam_block_size, ) - return {(Toffoli(), cost)} + return {Toffoli(): cost} def adjoint(self) -> 'Bloq': return attrs.evolve(self, is_adjoint=not self.is_adjoint) @@ -156,14 +156,14 @@ class QROAMTwoRegs(Bloq): def signature(self) -> Signature: return Signature.build(sel=(self.data_a_size - 1).bit_length(), trg=self.target_bitsize) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cost = int(np.ceil(self.data_a_size / self.data_a_block_size)) cost *= int(np.ceil(self.data_b_size / self.data_b_block_size)) if self.is_adjoint: cost += self.data_a_block_size * self.data_b_block_size else: cost += self.target_bitsize * (self.data_a_block_size * self.data_b_block_size - 1) - return {(Toffoli(), cost)} + return {Toffoli(): cost} def adjoint(self) -> 'Bloq': return attrs.evolve(self, is_adjoint=not self.is_adjoint) @@ -218,6 +218,6 @@ def build_composite_bloq(self, bb: 'BloqBuilder', ctrls: SoquetT, system: Soquet system = bb.join(split_sys) return {'ctrls': ctrls, 'system': system} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # remove this method once https://github.com/quantumlib/Qualtran/issues/528 is resolved. - return {(Toffoli(), len(self.cvs) - 1)} + return {Toffoli(): len(self.cvs) - 1} diff --git a/qualtran/bloqs/chemistry/df/double_factorization.py b/qualtran/bloqs/chemistry/df/double_factorization.py index d895662a6..4cc9b47b7 100644 --- a/qualtran/bloqs/chemistry/df/double_factorization.py +++ b/qualtran/bloqs/chemistry/df/double_factorization.py @@ -31,7 +31,7 @@ where $\Xi^{(l)} $ is the rank of second factorization. """ from functools import cached_property -from typing import Dict, Iterable, Set, TYPE_CHECKING +from typing import Dict, Iterable, TYPE_CHECKING import numpy as np from attrs import frozen @@ -63,7 +63,7 @@ from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -258,7 +258,7 @@ def build_composite_bloq( 'sys': sys, } - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': in_prep = InnerPrepareDoubleFactorization( num_aux=self.num_aux, num_spin_orb=self.num_spin_orb, @@ -276,16 +276,14 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: rot_dag = rot.adjoint() # 2*In-prep_l, addition, Rotations, 2*H, 2*SWAPS, subtraction return { - (in_prep, 1), # in-prep_l listing 3 page 52/53 - (in_prep_dag, 1), # in_prep_l^dag - (rot, 1), # rotate into system basis listing 4 pg 54 - ( - Toffoli(), - 1, - ), # apply CCZ first then CCCZ, the cost is 1 + 2 Toffolis (step 4e, and 7) - (rot_dag, 1), # Undo rotations - (CSwap(self.num_spin_orb // 2), 2), # Swaps for spins - (ArbitraryClifford(n=1), 1), # 2 Hadamards for spin superposition + in_prep: 1, # in-prep_l listing 3 page 52/53 + in_prep_dag: 1, # in_prep_l^dag + rot: 1, # rotate into system basis listing 4 pg 54 + # apply CCZ first then CCCZ, the cost is 1 + 2 Toffolis (step 4e, and 7) + Toffoli(): 1, + rot_dag: 1, # Undo rotations + CSwap(self.num_spin_orb // 2): 2, # Swaps for spins + ArbitraryClifford(n=1): 1, # 2 Hadamards for spin superposition } def __str__(self) -> str: diff --git a/qualtran/bloqs/chemistry/df/prepare.py b/qualtran/bloqs/chemistry/df/prepare.py index 2c0ee2323..49ea17ef2 100644 --- a/qualtran/bloqs/chemistry/df/prepare.py +++ b/qualtran/bloqs/chemistry/df/prepare.py @@ -14,7 +14,7 @@ r"""Bloqs for the state preparation of the DF Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -27,7 +27,7 @@ ) if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -68,12 +68,12 @@ def signature(self) -> Signature: nxi = (self.num_spin_orb // 2 - 1).bit_length() return Signature.build(xi=nxi, offset=nlxi, rot=self.num_bits_rot_aa, succ_p=1, p=nxi) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Step 3. num_bits_xi = (self.num_spin_orb // 2 - 1).bit_length() # uniform superposition state, requires controlled hadamards # https://github.com/quantumlib/Qualtran/issues/237 - cost_a = (Toffoli(), 7 * num_bits_xi + 2 * self.num_bits_rot_aa - 6) + cost_a = 7 * num_bits_xi + 2 * self.num_bits_rot_aa - 6 # add offset to get correct bit of QROM from [l + offset^l, l+offset^l+Xi^l] num_bits_lxi = (self.num_eig + self.num_spin_orb // 2 - 1).bit_length() cost_b = (Add(QUInt(num_bits_lxi)), 1) @@ -81,8 +81,9 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: bp = num_bits_xi + self.num_bits_state_prep + 2 # C31 cost_c = (QROAM(self.num_eig + self.num_spin_orb // 2, bp), 1) # inequality tests + CSWAP - cost_d = (Toffoli(), self.num_bits_state_prep + num_bits_xi) - return {cost_a, cost_b, cost_c, cost_d} + cost_d = self.num_bits_state_prep + num_bits_xi + cost_a_and_d = (Toffoli(), cost_a + cost_d) + return dict([cost_a_and_d, cost_b, cost_c]) def __str__(self) -> str: return "In-Prep" @@ -126,7 +127,7 @@ class OuterPrepareDoubleFactorization(Bloq): def signature(self) -> Signature: return Signature.build(l=self.num_aux.bit_length(), succ_l=1) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Listing 1 steps a-d num_bits_l = self.num_aux.bit_length() cost_a = (PrepareUniformSuperposition(self.num_aux + 1), 1) @@ -134,10 +135,10 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: output_size = num_bits_l + self.num_bits_state_prep cost_b = (QROAM(self.num_aux + 1, output_size), 1) # inequality test - cost_c = (Toffoli(), self.num_bits_state_prep) + cost_c = self.num_bits_state_prep # controlled swaps - cost_d = (Toffoli(), num_bits_l) - return {cost_a, cost_b, cost_c, cost_d} + cost_d = num_bits_l + return dict([cost_a, cost_b, (Toffoli(), cost_c + cost_d)]) def __str__(self) -> str: return "OuterPrep" @@ -189,12 +190,12 @@ def signature(self) -> Signature: offset=nlxi, ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # listing 2 C29/30 page 52 num_bits_lxi = (self.num_eig + self.num_spin_orb // 2 - 1).bit_length() num_bits_xi = (self.num_spin_orb // 2 - 1).bit_length() bo = num_bits_xi + num_bits_lxi + self.num_bits_rot_aa + 1 - return {(QROAM(self.num_aux + 1, bo), 1)} + return {QROAM(self.num_aux + 1, bo): 1} def __str__(self) -> str: return "In_l-data_l" diff --git a/qualtran/bloqs/chemistry/df/select_bloq.py b/qualtran/bloqs/chemistry/df/select_bloq.py index 2d620edf9..9630186cd 100644 --- a/qualtran/bloqs/chemistry/df/select_bloq.py +++ b/qualtran/bloqs/chemistry/df/select_bloq.py @@ -14,7 +14,7 @@ r"""Bloqs for the applying number operators to system for the DF Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -23,7 +23,7 @@ from qualtran.bloqs.chemistry.black_boxes import QROAM if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -70,7 +70,7 @@ def signature(self) -> Signature: ] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Step 4 in the reference. nlxi = (self.num_eig + self.num_spin_orb // 2 - 1).bit_length() nxi = (self.num_spin_orb // 2 - 1).bit_length() @@ -78,6 +78,6 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: data_size = self.num_eig + self.num_spin_orb // 2 cost_c = self.num_spin_orb * (self.num_bits_rot - 2) # apply rotations return { - (Toffoli(), (cost_a + cost_c)), - (QROAM(data_size, self.num_spin_orb * self.num_bits_rot // 2), 1), + Toffoli(): (cost_a + cost_c), + QROAM(data_size, self.num_spin_orb * self.num_bits_rot // 2): 1, } diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare.py b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare.py index 9b1bad254..d85665de0 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare.py @@ -13,7 +13,7 @@ # limitations under the License. r"""PREPARE for the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -50,7 +50,7 @@ def signature(self) -> Signature: n_eta = (self.eta - 1).bit_length() return Signature.build(i=n_eta, j=n_eta) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n_eta = (self.eta - 1).bit_length() # Half of Eq. 62 which is the cost for prep and prep^\dagger - return {(Toffoli(), (7 * n_eta + 4 * self.num_bits_rot_aa - 18))} + return {Toffoli(): (7 * n_eta + 4 * self.num_bits_rot_aa - 18)} diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_nu.py b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_nu.py index 8c160f39c..74998758e 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_nu.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_nu.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for preparation of the U and V parts of the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING from attrs import evolve, frozen @@ -26,7 +26,7 @@ from qualtran.drawing import Text, WireSymbol if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -58,8 +58,8 @@ def signature(self) -> Signature: [Register("mu", QAny(self.num_bits_p)), Register("flag", QBit(), side=Side.RIGHT)] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Toffoli(), (self.num_bits_p - 1))} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {Toffoli(): (self.num_bits_p - 1)} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -108,9 +108,9 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # controlled hadamards which cannot be inverted at zero Toffoli cost. - return {(Toffoli(), (3 * (self.num_bits_p - 1)))} + return {Toffoli(): (3 * (self.num_bits_p - 1))} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -152,14 +152,14 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: # This can be inverted with cliffords. - return set() + return {} else: # Controlled Toffoli each having n_p + 1 controls and 2 Toffolis to # check the result of the Toffolis. - return {(Toffoli(), (3 * self.num_bits_p + 2))} + return {Toffoli(): (3 * self.num_bits_p + 2)} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -202,13 +202,13 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: # This can be inverted with cliffords. - return {(Toffoli(), 0)} + return {Toffoli(): 0} else: # n_p controlled Toffolis with four controls. - return {(Toffoli(), 3 * self.num_bits_p)} + return {Toffoli(): 3 * self.num_bits_p} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -277,9 +277,9 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: - return {(Toffoli(), 0)} + return {Toffoli(): 0} else: # 1. Compute $\nu_x^2 + \nu_y^2 + \nu_z^2$ cost_1 = (SumOfSquares(self.num_bits_p, k=3), 1) @@ -290,7 +290,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: cost_3 = (GreaterThan(self.num_bits_m, 2 * self.num_bits_p + 2), 1) # 4. 3 Toffoli for overall success cost_4 = (Toffoli(), 3) - return {cost_1, cost_2, cost_3, cost_4} + return dict([cost_1, cost_2, cost_3, cost_4]) def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -382,7 +382,7 @@ def build_composite_bloq( ) return {'mu': mu, 'nu': nu, 'm': m, 'flag_nu': flag_nu} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # 1. Prepare unary encoded superposition state (Eq 77) cost_1 = (PrepareMuUnaryEncodedOneHot(self.num_bits_p), 1) n_m = (self.m_param - 1).bit_length() @@ -395,7 +395,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # 5. Prepare superposition over $m$ which is a power of two so only clifford. # 6. Test that $(2^{\mu-2})^2\mathcal{M} > m (\nu_x^2 + \nu_y^2 + \nu_z^2)$ cost_6 = (TestNuInequality(self.num_bits_p, n_m), 1) - return {cost_1, cost_2, cost_3, cost_4, cost_6} + return dict([cost_1, cost_2, cost_3, cost_4, cost_6]) def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_t.py b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_t.py index bee32c8de..26e9f24f1 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_t.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_t.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for PREPARE T for the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING from attrs import frozen @@ -25,7 +25,7 @@ from qualtran.drawing import Text, WireSymbol if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -56,8 +56,8 @@ class PreparePowerTwoState(Bloq): def signature(self) -> Signature: return Signature.build(r=self.bitsize) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Toffoli(), (self.bitsize - 2))} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {Toffoli(): (self.bitsize - 2)} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -114,14 +114,13 @@ def build_composite_bloq( s = bb.add(PreparePowerTwoState(self.num_bits_p), r=s) return {'w': w, 'r': r, 's': s} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # there is a cost for the uniform state preparation for the $w$ # register. Adding a bloq is sort of overkill, should just tag the # correct cost on UniformSuperPosition bloq - # 13 is from assuming 8 bits for the rotation, and n = 2. - uni_prep_w = (Toffoli(), 13) - # Factor of two for r and s registers. - return {uni_prep_w, (PreparePowerTwoState(bitsize=self.num_bits_p), 2)} + # 13 Toffolis is from assuming 8 bits for the rotation, and n = 2. + # Factor of two for PreparePowerTwoState for r and s registers. + return {Toffoli(): 13, PreparePowerTwoState(bitsize=self.num_bits_p): 2} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_uv.py b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_uv.py index 2f11804de..de5273458 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_uv.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_uv.py @@ -13,7 +13,7 @@ # limitations under the License. r"""PREPARE the potential energy terms of the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING from attrs import frozen @@ -33,7 +33,7 @@ from qualtran.drawing import Text, WireSymbol if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -91,12 +91,12 @@ def build_composite_bloq( l = bb.add(PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos), l=l) return {'mu': mu, 'nu': nu, 'm': m, 'l': l, 'flag_nu': flag_nu} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # 1. Prepare the nu state # 2. Prepare the zeta_l state return { - (PrepareNuState(self.num_bits_p, self.m_param), 1), - (PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos), 1), + PrepareNuState(self.num_bits_p, self.m_param): 1, + PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos): 1, } def wire_symbol( diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta.py b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta.py index 69e8db895..48df42a2f 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta.py @@ -14,7 +14,7 @@ r"""PREPARE the superposition over nuclear weights for the first quantized chemistry Hamiltonian. """ from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING import numpy as np from attrs import evolve, frozen @@ -23,7 +23,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -56,11 +56,11 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: # Really Er(x), eq 91. In practice we will replace this with the # appropriate qrom call down the line. - return {(Toffoli(), int(np.ceil(self.lambda_zeta**0.5)))} + return {Toffoli(): int(np.ceil(self.lambda_zeta**0.5))} else: - return {(Toffoli(), self.lambda_zeta)} + return {Toffoli(): self.lambda_zeta} diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_nu.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_nu.py index ff2f9fb56..4d06eab32 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_nu.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_nu.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for preparing the $\nu$ state for the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING +from typing import Dict, TYPE_CHECKING from attrs import evolve, frozen @@ -30,7 +30,7 @@ ) if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -77,11 +77,11 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: - return {(Toffoli(), (self.bitsize_n - 1) + 1)} + return {Toffoli(): (self.bitsize_n - 1) + 1} else: - return {(Toffoli(), (self.bitsize_n - 1) + (self.bitsize_n - self.bitsize_p - 1) + 1)} + return {Toffoli(): (self.bitsize_n - 1) + (self.bitsize_n - self.bitsize_p - 1) + 1} @frozen @@ -154,7 +154,7 @@ def build_composite_bloq( ) return {'mu': mu, 'nu': nu, 'm': m, 'flag_nu': flag_nu} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # 1. Prepare unary encoded superposition state (Eq 77) cost_1 = (PrepareMuUnaryEncodedOneHotWithProj(self.num_bits_n, self.num_bits_p), 1) n_m = (self.m_param - 1).bit_length() @@ -167,7 +167,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # 5. Prepare superposition over $m$ which is a power of two so only clifford. # 6. Test that $(2^{\mu-2})^2\mathcal{M} > m (\nu_x^2 + \nu_y^2 + \nu_z^2)$ cost_6 = (TestNuInequality(self.num_bits_n, n_m), 1) - return {cost_1, cost_2, cost_3, cost_4, cost_6} + return dict([cost_1, cost_2, cost_3, cost_4, cost_6]) @bloq_example diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_t.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_t.py index 4caf03922..2603b51c7 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_t.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_t.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for PREPARE T for the first quantized chemistry Hamiltonian with a quantum projectile.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import evolve, frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -68,14 +68,14 @@ def adjoint(self) -> 'Bloq': def signature(self) -> Signature: return Signature.build(r=self.bitsize_n) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: - return {(Toffoli(), (self.bitsize_n - 2))} + return {Toffoli(): (self.bitsize_n - 2)} else: # The doubly controlled hadamard can be converted to and And and a # controlled Hadamard, with the And gate being inverted at zero # Toffoli cost. - return {(Toffoli(), (self.bitsize_n - 2) + self.bitsize_n - self.bitsize_p)} + return {Toffoli(): (self.bitsize_n - 2) + self.bitsize_n - self.bitsize_p} @frozen @@ -128,32 +128,29 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return evolve(self, is_adjoint=not self.is_adjoint) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # there is a cost for the uniform state preparation for the $w$ # register. Adding a bloq is sort of overkill, should just tag the # correct cost on UniformSuperPosition bloq # 13 is from assuming 8 bits for the rotation, and n = 2. - uni_prep_w = (Toffoli(), 13) + uni_prep_w = 13 # Factor of two for r and s registers. - ctrl_mom = ( - PreparePowerTwoStateWithProj( - bitsize_n=self.num_bits_n, bitsize_p=self.num_bits_p, is_adjoint=self.is_adjoint - ), - 2, + ctrl_mom_gate = PreparePowerTwoStateWithProj( + bitsize_n=self.num_bits_n, bitsize_p=self.num_bits_p, is_adjoint=self.is_adjoint ) # Inequality test can be inverted at zero cost if self.is_adjoint: # pg 31 (Appendix A. Sec 2 c) - k_k_proj = (Toffoli(), 0) + k_k_proj = 0 else: # Cost for preparing a state for selecting the components of k_p^w k_proj^w # Prepare a uniform superposition over 8 states and do 2 inequality # tests to select between x, y and z. # built on w_proj above - k_k_proj = (Toffoli(), 16) + k_k_proj = 16 # pg 31 (Appendix A. Sec 2 c) - ctrl_swap = (Toffoli(), 2) - return {uni_prep_w, ctrl_mom, k_k_proj, ctrl_swap} + ctrl_swap = 2 + return {Toffoli(): uni_prep_w + k_k_proj + ctrl_swap, ctrl_mom_gate: 2} @bloq_example diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_uv.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_uv.py index 54ce6a111..c9c5f573c 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_uv.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_uv.py @@ -14,7 +14,7 @@ r"""PREPARE the potential energy terms of the first quantized chemistry Hamiltonian with projectile. """ from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING +from typing import Dict, TYPE_CHECKING from attrs import frozen @@ -25,7 +25,7 @@ ) if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -91,12 +91,12 @@ def build_composite_bloq( l = bb.add(PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos), l=l) return {'mu': mu, 'nu': nu, 'm': m, 'l': l, 'flag_nu': flag_nu} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # 1. Prepare the nu state # 2. Prepare the zeta_l state return { - (PrepareNuStateWithProj(self.num_bits_p, self.num_bits_n, self.m_param), 1), - (PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos), 1), + PrepareNuStateWithProj(self.num_bits_p, self.num_bits_n, self.m_param): 1, + PrepareZetaState(self.num_atoms, self.lambda_zeta, self.num_bits_nuc_pos): 1, } diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_and_prepare.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_and_prepare.py index 93dbf744f..1d135ec84 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_and_prepare.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_and_prepare.py @@ -13,7 +13,7 @@ # limitations under the License. r"""SELECT and PREPARE for the first quantized chemistry Hamiltonian with a quantum projectile.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import evolve, field, frozen @@ -55,7 +55,7 @@ from qualtran.drawing import Circle, Text, TextBox, WireSymbol if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -106,12 +106,12 @@ def adjoint(self) -> 'Bloq': def pretty_name(self) -> str: return 'PREP TUV' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.is_adjoint: # inverting inequality tests at zero Toffoli. - return set() + return {} else: - return {(Toffoli(), 6 * self.num_bits_t + 2)} + return {Toffoli(): 6 * self.num_bits_t + 2} @frozen diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_t.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_t.py index 7f27abda4..478e4c90c 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_t.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_t.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for SELECT T for the first quantized chemistry Hamiltonian with a quantum projectile.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -72,11 +72,11 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return r'SEL T' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Modification of the SEL T costs from the first quantized bloq with n_p replace with n_n. # The + 1 is from an additional Toffoli for the selection between the # square and the product of the momentum offset of the projectile. - return {(Toffoli(), (5 * (self.num_bits_n - 1) + 2 + 1))} + return {Toffoli(): (5 * (self.num_bits_n - 1) + 2 + 1)} @bloq_example diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_uv.py b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_uv.py index 000e325f9..fccda7af0 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_uv.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_uv.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. r"""Bloqs for SELECT for the U and V parts of the first quantized chemistry Hamiltonian.""" +from collections import Counter from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -23,7 +24,7 @@ from qualtran.bloqs.chemistry.pbc.first_quantization.select_uv import ApplyNuclearPhase if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator @frozen @@ -77,32 +78,26 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return r'SEL UV' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + cost = Counter['Bloq']() + # tc_p and tc_n # C8 and C9, one of the registers of size num_bits_n so need to account for this. - cost_tc_p = (SignedIntegerToTwosComplement(self.num_bits_p), 3) - cost_tc_n = (SignedIntegerToTwosComplement(self.num_bits_n), 3) + cost[SignedIntegerToTwosComplement(self.num_bits_p)] += 3 + cost[SignedIntegerToTwosComplement(self.num_bits_n)] += 3 # Adding nu into p / q. Nu is one bit larger than p. - cost_add_p = (Add(QInt(self.num_bits_p + 1)), 3) - cost_add_n = (Add(QInt(self.num_bits_n + 1)), 3) - cost_ctrl_add_p = (Toffoli(), 3 * (self.num_bits_p + 1)) - cost_ctrl_add_n = (Toffoli(), 3 * (self.num_bits_n + 1)) + cost[Add(QInt(self.num_bits_p + 1))] += 3 + cost[Add(QInt(self.num_bits_n + 1))] += 3 + # ctrl_add_p and ctrl_add_n + cost[Toffoli()] += 3 * (self.num_bits_p + 1) + cost[Toffoli()] += 3 * (self.num_bits_n + 1) + # inv_tc_p and inv_tc_n # + 2 as these numbers are larger from addition of $\nu$ - cost_inv_tc_p = (SignedIntegerToTwosComplement(self.num_bits_p + 2), 3) - cost_inv_tc_n = (SignedIntegerToTwosComplement(self.num_bits_n + 2), 3) + cost[SignedIntegerToTwosComplement(self.num_bits_p + 2)] += 3 + cost[SignedIntegerToTwosComplement(self.num_bits_n + 2)] += 3 + # cost for phase: # 2. Phase by $e^{ik\cdot R}$ in the case of $U$ only. - cost_phase = (ApplyNuclearPhase(self.num_bits_n, self.num_bits_nuc_pos), 1) - return { - cost_tc_p, - cost_tc_n, - cost_tc_n, - cost_add_p, - cost_add_n, - cost_ctrl_add_n, - cost_ctrl_add_p, - cost_inv_tc_p, - cost_inv_tc_n, - cost_phase, - } + cost[ApplyNuclearPhase(self.num_bits_n, self.num_bits_nuc_pos)] += 1 + return cost @bloq_example diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/select_and_prepare.py b/qualtran/bloqs/chemistry/pbc/first_quantization/select_and_prepare.py index 8303440ea..7bb407889 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/select_and_prepare.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/select_and_prepare.py @@ -13,7 +13,7 @@ # limitations under the License. r"""SELECT and PREPARE for the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import frozen @@ -45,7 +45,7 @@ if TYPE_CHECKING: from qualtran import Soquet - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -76,11 +76,11 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return 'PREP TUV' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n_eta_zeta = (self.eta + 2 * self.lambda_zeta - 1).bit_length() # The cost arises from rotating a qubit, and uniform state preparation # over eta + 2 lambda_zeta numbers along. - return {(Toffoli(), self.num_bits_t + 4 * n_eta_zeta + 2 * self.num_bits_rot_aa - 12)} + return {Toffoli(): self.num_bits_t + 4 * n_eta_zeta + 2 * self.num_bits_rot_aa - 12} @frozen @@ -108,10 +108,10 @@ def signature(self) -> Signature: n_eta = (self.eta - 1).bit_length() return Signature.build(i=n_eta, j=n_eta) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n_eta = (self.eta - 1).bit_length() # Half of Eq. 62 which is the cost for prep and prep^\dagger - return {(Toffoli(), (7 * n_eta + 4 * self.num_bits_rot_aa - 18))} + return {Toffoli(): (7 * n_eta + 4 * self.num_bits_rot_aa - 18)} @frozen diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/select_t.py b/qualtran/bloqs/chemistry/pbc/first_quantization/select_t.py index df2b218d8..af09a1620 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/select_t.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/select_t.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for SELECT T for the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -60,15 +60,15 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return r'SEL T' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Cost is $5(n_{p} - 1) + 2$ which comes from copying each $w$ component of $p$ # into an ancilla register ($3(n_{p}-1)$), copying the $r$ and $s$ bit of into an # ancilla ($2(n_{p}-1)$), controlling on both those bit perform phase flip on an # ancilla $|+\rangle$ state. This requires $1$ Toffoli, Then erase which costs # only Cliffords. There is an additional control bit controlling the application # of $T$ thus we come to our total. - # Eq 73. page - return {(Toffoli(), (5 * (self.num_bits_p - 1) + 2))} + # Eq 73, page 20. + return {Toffoli(): (5 * (self.num_bits_p - 1) + 2)} @bloq_example diff --git a/qualtran/bloqs/chemistry/pbc/first_quantization/select_uv.py b/qualtran/bloqs/chemistry/pbc/first_quantization/select_uv.py index 0915d0a36..672bc2187 100644 --- a/qualtran/bloqs/chemistry/pbc/first_quantization/select_uv.py +++ b/qualtran/bloqs/chemistry/pbc/first_quantization/select_uv.py @@ -13,7 +13,7 @@ # limitations under the License. r"""Bloqs for SELECT for the U and V parts of the first quantized chemistry Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -22,7 +22,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -59,7 +59,7 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return r'-e^(-k_ν⋅R_l)' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n_p = self.num_bits_p n_n = self.num_bits_nuc # This is some complicated application of phase gradient gates. @@ -68,7 +68,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: cost = 3 * (2 * n_p * n_n - n_p * (n_p + 1) - 1) else: cost = 3 * n_n * (n_n - 1) - return {(Toffoli(), cost)} + return {Toffoli(): cost} @frozen @@ -114,7 +114,7 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return r'SEL UV' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cost_tc = (SignedIntegerToTwosComplement(self.num_bits_p), 6) cost_add = (Add(QInt(self.num_bits_p + 1)), 6) # + 2? cost_ctrl_add = (Toffoli(), 6 * (self.num_bits_p + 1)) @@ -122,7 +122,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: cost_inv_tc = (SignedIntegerToTwosComplement(self.num_bits_p + 2), 6) # 2. Phase by $e^{ik\cdot R}$ in the case of $U$ only. cost_phase = (ApplyNuclearPhase(self.num_bits_p, self.num_bits_nuc_pos), 1) - return {cost_tc, cost_add, cost_ctrl_add, cost_inv_tc, cost_phase} + return dict([cost_tc, cost_add, cost_ctrl_add, cost_inv_tc, cost_phase]) @bloq_example diff --git a/qualtran/bloqs/chemistry/sf/prepare.py b/qualtran/bloqs/chemistry/sf/prepare.py index 57e38973f..fd7542fa9 100644 --- a/qualtran/bloqs/chemistry/sf/prepare.py +++ b/qualtran/bloqs/chemistry/sf/prepare.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -27,7 +27,7 @@ ) if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -77,7 +77,7 @@ def signature(self) -> Signature: n = (self.num_spin_orb // 2 - 1).bit_length() return Signature.build(l=self.num_aux.bit_length(), p=n, q=n, succ_pq=1) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n = (self.num_spin_orb // 2 - 1).bit_length() # 2. a prep uniform upper triangular. cost_up_tri = (Toffoli(), 6 * n + 2 * self.num_bits_rot_aa - 7) @@ -100,7 +100,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # inequality test cost_ineq = (LessThanEqual(self.num_bits_state_prep, self.num_bits_state_prep), 1) cost_swap = (CSwap(2 * n), 1) - return {cost_up_tri, cost_contg_indx, cost_qroam, cost_ineq, cost_swap} + return dict([cost_up_tri, cost_contg_indx, cost_qroam, cost_ineq, cost_swap]) @frozen @@ -140,7 +140,7 @@ def pretty_name(self) -> str: def signature(self) -> Signature: return Signature.build(l=self.num_aux.bit_length(), succ_l=1, l_ne_zero=1, rot_aa=1) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # 1.a cost_uni = (PrepareUniformSuperposition(self.num_aux + 1), 1) num_bits_l = (self.num_aux + 1).bit_length() @@ -151,7 +151,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: cost_ineq = (LessThanEqual(self.num_bits_state_prep, self.num_bits_state_prep), 1) # 1.d swap alt/keep values cost_swap = (CSwap(num_bits_l + 1), 1) - return {cost_uni, cost_qroam, cost_ineq, cost_swap} + return dict([cost_uni, cost_qroam, cost_ineq, cost_swap]) @bloq_example diff --git a/qualtran/bloqs/chemistry/sf/select_bloq.py b/qualtran/bloqs/chemistry/sf/select_bloq.py index dbca7a2ad..030a498a7 100644 --- a/qualtran/bloqs/chemistry/sf/select_bloq.py +++ b/qualtran/bloqs/chemistry/sf/select_bloq.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -21,7 +21,7 @@ from qualtran.bloqs.basic_gates import Toffoli if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -54,8 +54,8 @@ def signature(self) -> Signature: n = (self.num_spin_orb // 2 - 1).bit_length() return Signature.build(p=n, q=n, spin=1, succ_pq=1, succ_l=1) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Toffoli(), 2 * (self.num_spin_orb - 2))} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {Toffoli(): 2 * (self.num_spin_orb - 2)} @bloq_example diff --git a/qualtran/bloqs/chemistry/sf/single_factorization.py b/qualtran/bloqs/chemistry/sf/single_factorization.py index 173156fab..5e1104df8 100644 --- a/qualtran/bloqs/chemistry/sf/single_factorization.py +++ b/qualtran/bloqs/chemistry/sf/single_factorization.py @@ -23,7 +23,7 @@ """ from functools import cached_property -from typing import Dict, Iterable, Set, TYPE_CHECKING +from typing import Dict, Iterable, TYPE_CHECKING import numpy as np from attrs import evolve, frozen @@ -55,7 +55,7 @@ from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -250,7 +250,7 @@ def build_composite_bloq( 'sys': sys, } - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': iprep = InnerPrepareSingleFactorization( self.num_aux, self.num_spin_orb, @@ -269,11 +269,11 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: ).adjoint() n = (self.num_spin_orb // 2 - 1).bit_length() return { - (iprep, 1), - (iprep_dag, 1), - (CSwap(n), 2), - (SelectSingleFactorization(num_spin_orb=self.num_spin_orb), 1), - (Hadamard(), 4), + iprep: 1, + iprep_dag: 1, + CSwap(n): 2, + SelectSingleFactorization(num_spin_orb=self.num_spin_orb): 1, + Hadamard(): 4, } diff --git a/qualtran/bloqs/chemistry/sparse/prepare.py b/qualtran/bloqs/chemistry/sparse/prepare.py index c5e08be72..eae8a51b1 100644 --- a/qualtran/bloqs/chemistry/sparse/prepare.py +++ b/qualtran/bloqs/chemistry/sparse/prepare.py @@ -15,7 +15,7 @@ import itertools from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import attrs import numpy as np @@ -51,7 +51,7 @@ if TYPE_CHECKING: from qualtran import Bloq - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def get_sparse_inputs_from_integrals( @@ -422,15 +422,15 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str ) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (PrepareUniformSuperposition(self.num_non_zero), 1), - (self.build_qrom_bloq(), 1), - (OnEach(self.num_bits_state_prep, Hadamard()), 1), - (Hadamard(), 3), - (CSwap(1), 1), - (CSwap(self.num_bits_spat_orb), 4 + 4), - (LessThanEqual(self.num_bits_state_prep, self.num_bits_state_prep), 1), + PrepareUniformSuperposition(self.num_non_zero): 1, + self.build_qrom_bloq(): 1, + OnEach(self.num_bits_state_prep, Hadamard()): 1, + Hadamard(): 3, + CSwap(1): 1, + CSwap(self.num_bits_spat_orb): 4 + 4, + LessThanEqual(self.num_bits_state_prep, self.num_bits_state_prep): 1, } diff --git a/qualtran/bloqs/chemistry/sparse/select_bloq.py b/qualtran/bloqs/chemistry/sparse/select_bloq.py index 8b632ac6f..06d9c42c1 100644 --- a/qualtran/bloqs/chemistry/sparse/select_bloq.py +++ b/qualtran/bloqs/chemistry/sparse/select_bloq.py @@ -14,7 +14,7 @@ """SELECT for the sparse chemistry Hamiltonian in second quantization.""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import cirq from attrs import frozen @@ -26,7 +26,7 @@ from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -143,7 +143,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str out_soqs['control'] = soqs['control'] return out_soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Pg 30, enumeration 1: 2 applications of SELECT in Fig. 13, one of # which is not controlled (for the two body part of the Ham). The figure # is a bit misleading as applying that circuit twice would square the @@ -155,7 +155,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: maj_y = SelectedMajoranaFermion(sel_pa, target_gate=cirq.Y, control_regs=()) c_maj_x = SelectedMajoranaFermion(sel_pa, target_gate=cirq.X) c_maj_y = SelectedMajoranaFermion(sel_pa, target_gate=cirq.Y) - return {(SGate(), 1), (maj_x, 1), (c_maj_x, 1), (maj_y, 1), (c_maj_y, 1)} + return {SGate(): 1, maj_x: 1, c_maj_x: 1, maj_y: 1, c_maj_y: 1} @bloq_example diff --git a/qualtran/bloqs/chemistry/thc/prepare.py b/qualtran/bloqs/chemistry/thc/prepare.py index dfc38fee9..1af85c05a 100644 --- a/qualtran/bloqs/chemistry/thc/prepare.py +++ b/qualtran/bloqs/chemistry/thc/prepare.py @@ -13,7 +13,7 @@ # limitations under the License. """PREPARE for the molecular tensor hypercontraction (THC) hamiltonian""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import cirq import numpy as np @@ -53,7 +53,7 @@ from qualtran.symbolics import SymbolicFloat if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -462,7 +462,7 @@ def build_composite_bloq( } return out_regs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cost_1 = (UniformSuperpositionTHC(self.num_mu, self.num_spin_orb), 1) nmu = self.num_mu.bit_length() data_size = self.num_spin_orb // 2 + self.num_mu * (self.num_mu + 1) // 2 @@ -478,7 +478,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: cost_5 = (LessThanEqual(self.keep_bitsize, self.keep_bitsize), 2) cost_6 = (CSwap(nmu), 3) cost_7 = (Toffoli(), 1) - return {cost_1, cost_2, cost_3, cost_4, cost_5, cost_6, cost_7} + return dict([cost_1, cost_2, cost_3, cost_4, cost_5, cost_6, cost_7]) @bloq_example diff --git a/qualtran/bloqs/chemistry/thc/select_bloq.py b/qualtran/bloqs/chemistry/thc/select_bloq.py index d172a3d72..ed039e676 100644 --- a/qualtran/bloqs/chemistry/thc/select_bloq.py +++ b/qualtran/bloqs/chemistry/thc/select_bloq.py @@ -14,7 +14,7 @@ """SELECT for the molecular tensor hypercontraction (THC) hamiltonian""" from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import evolve, frozen @@ -37,7 +37,7 @@ from qualtran.bloqs.multiplexers.select_base import SelectOracle if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -95,7 +95,7 @@ def pretty_name(self) -> str: dag = '†' if self.is_adjoint else '' return f"In_mu-R{dag}" - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # from listings on page 17 of Ref. [1] num_data_sets = self.num_mu + self.num_spin_orb // 2 if self.is_adjoint: @@ -116,7 +116,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # xref https://github.com/quantumlib/Qualtran/issues/370, the cost below # assume a phase gradient. rot_cost = self.num_spin_orb * (self.num_bits_theta - 2) - return {(Toffoli(), (rot_cost + toff_cost_qrom))} + return {Toffoli(): (rot_cost + toff_cost_qrom)} @frozen diff --git a/qualtran/bloqs/chemistry/trotter/grid_ham/inverse_sqrt.py b/qualtran/bloqs/chemistry/trotter/grid_ham/inverse_sqrt.py index c4d6280ad..f953f298a 100644 --- a/qualtran/bloqs/chemistry/trotter/grid_ham/inverse_sqrt.py +++ b/qualtran/bloqs/chemistry/trotter/grid_ham/inverse_sqrt.py @@ -13,7 +13,7 @@ # limitations under the License. """Bloqs for computing the inverse Square root of a fixed point number.""" from functools import cached_property -from typing import Set, Tuple, TYPE_CHECKING +from typing import Tuple, TYPE_CHECKING import numpy as np from attrs import frozen @@ -23,7 +23,7 @@ from qualtran.bloqs.arithmetic import Add, MultiplyTwoReals, ScaleIntByReal, SquareRealNumber if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def get_inverse_square_root_poly_coeffs() -> Tuple[NDArray, NDArray]: @@ -167,7 +167,7 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return 'y = x^{-1/2}' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # y * ((2 + b^2 + delta) + y^2 x) # 1. square y # 2. scale y^2 by x @@ -175,12 +175,12 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # 4. multiply y^2 x by y # 5. add 3. and 4. return { - (SquareRealNumber(self.poly_bitsize), 1), + SquareRealNumber(self.poly_bitsize): 1, # TODO: When decomposing we will potentially need to cast into a larger register. # See: https://github.com/quantumlib/Qualtran/issues/655 - (ScaleIntByReal(self.poly_bitsize, self.x_sq_bitsize), 1), - (MultiplyTwoReals(self.target_bitsize), 2), - (Add(QInt(self.target_bitsize)), 1), + ScaleIntByReal(self.poly_bitsize, self.x_sq_bitsize): 1, + MultiplyTwoReals(self.target_bitsize): 2, + Add(QInt(self.target_bitsize)): 1, } @@ -218,10 +218,10 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return 'y ~ x^{-1/2}' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # This should probably be scale int by float rather than 3 real # multiplications as x in Eq. 49 of the reference is an integer. - return {(MultiplyTwoReals(self.poly_bitsize), 3), (Add(QInt(self.poly_bitsize)), 3)} + return {MultiplyTwoReals(self.poly_bitsize): 3, Add(QInt(self.poly_bitsize)): 3} @bloq_example diff --git a/qualtran/bloqs/chemistry/trotter/grid_ham/qvr.py b/qualtran/bloqs/chemistry/trotter/grid_ham/qvr.py index e9ab587e6..e10f3d9a6 100644 --- a/qualtran/bloqs/chemistry/trotter/grid_ham/qvr.py +++ b/qualtran/bloqs/chemistry/trotter/grid_ham/qvr.py @@ -14,7 +14,7 @@ """Quantum Variable Rotation.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -22,7 +22,7 @@ from qualtran.bloqs.basic_gates import Rz if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -55,10 +55,10 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return 'e^{i*phi}' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': theta = ssa.new_symbol('theta') # need to update rotation bloq. - return {(Rz(theta), self.phi_bitsize)} + return {Rz(theta): self.phi_bitsize} @bloq_example diff --git a/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py b/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py index 8db597c8d..3caf30d03 100644 --- a/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py +++ b/qualtran/bloqs/chemistry/trotter/hubbard/hopping.py @@ -13,7 +13,7 @@ # limitations under the License. """Bloqs implementing unitary evolution under the one-body hopping Hamiltonian in 2D.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -24,7 +24,7 @@ from qualtran.symbolics import SymbolicFloat, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -74,12 +74,12 @@ class HoppingPlaquette(Bloq): def signature(self) -> Signature: return Signature([Register('qubits', QBit(), shape=(4,))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # The TwoBitFFFT in the reference is F(k=0, n=arbitrary) # page 14, discussion after E13 # There are 4 flanking f-gates and a e^{iXX}e^{iYY} rotation, which can # be rotated to single rotation + cliffords. - return {(TwoBitFFFT(0, 1, eps=self.eps), 4), (Rz(self.kappa, eps=self.eps), 2)} + return {TwoBitFFFT(0, 1, eps=self.eps): 4, Rz(self.kappa, eps=self.eps): 2} @frozen @@ -128,11 +128,9 @@ def pretty_name(self) -> str: def signature(self) -> Signature: return Signature([Register('system', QAny(self.length), shape=(2,))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Page 5, text after Eq. 22. There are L^2 / 4 plaquettes of a given colour and x2 for spin. - return { - (HoppingPlaquette(kappa=self.tau * self.angle, eps=self.eps), self.length**2 // 2) - } + return {HoppingPlaquette(kappa=self.tau * self.angle, eps=self.eps): self.length**2 // 2} @frozen @@ -181,19 +179,14 @@ def short_name(self) -> str: l = 'p' if self.pink else 'g' return f'H_h^{l}(HWP)' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Page 5, text after Eq. 22. There are L^2 / 4 plaquettes of a given colour and x2 for spin. # Each plaquette contributes 4 TwoBitFFFT gates and two arbitrary rotations. # We use Hamming weight phasing to apply all 2 * L^2/4 (two for spin # here) for both of these rotations. return { - (TwoBitFFFT(0, 1, self.eps), 4 * self.length**2 // 2), - ( - HammingWeightPhasing( - 2 * self.length**2 // 4, self.tau * self.angle, eps=self.eps - ), - 2, - ), + TwoBitFFFT(0, 1, self.eps): 4 * self.length**2 // 2, + HammingWeightPhasing(2 * self.length**2 // 4, self.tau * self.angle, eps=self.eps): 2, } diff --git a/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py b/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py index a910cedcd..30b515299 100644 --- a/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py +++ b/qualtran/bloqs/chemistry/trotter/hubbard/interaction.py @@ -14,7 +14,7 @@ r"""Bloqs implementing unitary evolution under the interacting part of the Hubbard Hamiltonian.""" from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -24,7 +24,7 @@ from qualtran.symbolics import SymbolicFloat, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -62,9 +62,9 @@ class Interaction(Bloq): def signature(self) -> Signature: return Signature([Register('system', QAny(self.length), shape=(2,))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Page 13 paragraph 1. - return {(Rz(angle=self.angle * self.hubb_u, eps=self.eps), self.length**2)} + return {Rz(angle=self.angle * self.hubb_u, eps=self.eps): self.length**2} @frozen @@ -105,9 +105,9 @@ class InteractionHWP(Bloq): def signature(self) -> Signature: return Signature([Register('system', QAny(self.length), shape=(2,))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (HammingWeightPhasing(self.length**2 // 2, self.angle * self.hubb_u, eps=self.eps), 2) + HammingWeightPhasing(self.length**2 // 2, self.angle * self.hubb_u, eps=self.eps): 2 } diff --git a/qualtran/bloqs/data_loading/qroam_clean.py b/qualtran/bloqs/data_loading/qroam_clean.py index 69e5a21bf..30e2ec25d 100644 --- a/qualtran/bloqs/data_loading/qroam_clean.py +++ b/qualtran/bloqs/data_loading/qroam_clean.py @@ -14,7 +14,7 @@ import numbers from collections import defaultdict from functools import cached_property -from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, Union +from typing import cast, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union import attrs import numpy as np @@ -30,7 +30,7 @@ if TYPE_CHECKING: from qualtran import Bloq, BloqBuilder, SoquetT, QDType from qualtran.simulation.classical_sim import ClassicalValT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.bloqs.data_loading.select_swap_qrom import _alloc_anc_for_reg, SelectSwapQROM @@ -174,11 +174,11 @@ def with_log_block_sizes( assert all(1 <= 2**bs <= ilen for bs, ilen in zip(log_block_sizes, self.data_shape)) return attrs.evolve(self, log_block_sizes=log_block_sizes) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': block_sizes = prod([2**k for k in self.log_block_sizes]) data_size = prod(self.data_shape) n_toffoli = ceil(data_size / block_sizes) + block_sizes - return {(Toffoli(), n_toffoli)} + return {Toffoli(): n_toffoli} @cached_property def signature(self) -> Signature: @@ -260,11 +260,11 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str soqs = bb.add_d(self.qroam_clean_adjoint_bloq, **soqs) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': block_sizes = prod([2**k for k in self.log_block_sizes]) data_size = prod(self.qroam_clean.data_shape) n_toffoli = ceil(data_size / block_sizes) + block_sizes - return {(Toffoli(), n_toffoli)} + return {Toffoli(): n_toffoli} def adjoint(self) -> 'QROAMClean': return self.qroam_clean @@ -437,13 +437,13 @@ def junk_registers(self) -> Tuple[Register, ...]: junk_regs += [attrs.evolve(reg, name='junk_' + reg.name, shape=(block_size - 1,))] return tuple(junk_regs) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': ret: Dict[Bloq, SymbolicInt] = defaultdict(lambda: 0) ret[self.qrom_bloq] += 1 for swz in self.swap_with_zero_bloqs: if any(is_symbolic(s) or s > 0 for s in swz.selection_bitsizes): ret[swz] += 1 - return set(ret.items()) + return ret def _build_composite_bloq_with_swz_clean( self, diff --git a/qualtran/bloqs/data_loading/qrom.py b/qualtran/bloqs/data_loading/qrom.py index 7f3197582..61ac5a7b3 100644 --- a/qualtran/bloqs/data_loading/qrom.py +++ b/qualtran/bloqs/data_loading/qrom.py @@ -217,7 +217,7 @@ def build_call_graph( n_cnot = prod( bitsize * prod(sh) for bitsize, sh in zip(self.target_bitsizes, self.target_shapes) ) * prod(self.data_shape) - return {(And(), n_and), (And().adjoint(), n_and), (CNOT(), n_cnot)} + return {And(): n_and, And().adjoint(): n_and, CNOT(): n_cnot} def adjoint(self) -> 'QROM': return self diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.py b/qualtran/bloqs/data_loading/select_swap_qrom.py index 590fccaf0..52952719e 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom.py @@ -14,7 +14,7 @@ import numbers from collections import defaultdict from functools import cached_property -from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, TypeVar, Union +from typing import cast, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, TypeVar, Union import attrs import cirq @@ -33,7 +33,7 @@ if TYPE_CHECKING: from qualtran import Bloq, BloqBuilder, QDType, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator SelSwapQROM_T = TypeVar('SelSwapQROM_T', bound='SelectSwapQROM') @@ -255,7 +255,7 @@ def swap_with_zero_bloqs(self) -> List[SwapWithZero]: for target_bitsize in self.target_bitsizes ] - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': ret: Dict[Bloq, SymbolicInt] = defaultdict(lambda: 0) toggle_overhead = 2 if self.use_dirty_ancilla else 1 ret[self.qrom_bloq] += 1 @@ -266,7 +266,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: if any(is_symbolic(s) or s > 0 for s in swz.selection_bitsizes): ret[swz] += toggle_overhead ret[swz.adjoint()] += toggle_overhead - return set(ret.items()) + return ret def _add_qrom_bloq( self, diff --git a/qualtran/bloqs/factoring/ecc/ec_add.py b/qualtran/bloqs/factoring/ecc/ec_add.py index 102ca3619..e0ff64e6a 100644 --- a/qualtran/bloqs/factoring/ecc/ec_add.py +++ b/qualtran/bloqs/factoring/ecc/ec_add.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Set import sympy from attrs import frozen @@ -21,7 +20,7 @@ from qualtran.bloqs.arithmetic._shims import MultiCToffoli from qualtran.bloqs.mod_arithmetic import CModAdd, CModNeg, CModSub, ModAdd, ModNeg, ModSub from qualtran.bloqs.mod_arithmetic._shims import ModDbl, ModInv, ModMul -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -64,19 +63,19 @@ def signature(self) -> 'Signature': ] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # litinksi return { - (MultiCToffoli(n=self.n), 18), - (ModAdd(bitsize=self.n, mod=self.mod), 3), - (CModAdd(QUInt(self.n), mod=self.mod), 2), - (ModSub(QUInt(self.n), mod=self.mod), 2), - (CModSub(QUInt(self.n), mod=self.mod), 4), - (ModNeg(QUInt(self.n), mod=self.mod), 2), - (CModNeg(QUInt(self.n), mod=self.mod), 1), - (ModDbl(QUInt(self.n), mod=self.mod), 2), - (ModMul(n=self.n, mod=self.mod), 10), - (ModInv(n=self.n, mod=self.mod), 4), + MultiCToffoli(n=self.n): 18, + ModAdd(bitsize=self.n, mod=self.mod): 3, + CModAdd(QUInt(self.n), mod=self.mod): 2, + ModSub(QUInt(self.n), mod=self.mod): 2, + CModSub(QUInt(self.n), mod=self.mod): 4, + ModNeg(QUInt(self.n), mod=self.mod): 2, + CModNeg(QUInt(self.n), mod=self.mod): 1, + ModDbl(QUInt(self.n), mod=self.mod): 2, + ModMul(n=self.n, mod=self.mod): 10, + ModInv(n=self.n, mod=self.mod): 4, } diff --git a/qualtran/bloqs/factoring/ecc/ec_phase_estimate_r.py b/qualtran/bloqs/factoring/ecc/ec_phase_estimate_r.py index c5b09558f..f4ddb9d19 100644 --- a/qualtran/bloqs/factoring/ecc/ec_phase_estimate_r.py +++ b/qualtran/bloqs/factoring/ecc/ec_phase_estimate_r.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set +from typing import Dict import sympy from attrs import frozen @@ -31,7 +31,7 @@ SoquetT, ) from qualtran.bloqs.basic_gates import PlusState -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from ._ecc_shims import MeasureQFT from .ec_add_r import ECAddR @@ -67,8 +67,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet, y: Soquet) -> Dict[ bb.add(MeasureQFT(n=self.n), x=ctrl) return {'x': x, 'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(ECAddR(n=self.n, R=self.point), self.n), (MeasureQFT(n=self.n), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {ECAddR(n=self.n, R=self.point): self.n, MeasureQFT(n=self.n): 1} def __str__(self) -> str: return f'PE${self.point}$' diff --git a/qualtran/bloqs/factoring/ecc/find_ecc_private_key.py b/qualtran/bloqs/factoring/ecc/find_ecc_private_key.py index 61de8a3e4..efaa42ed3 100644 --- a/qualtran/bloqs/factoring/ecc/find_ecc_private_key.py +++ b/qualtran/bloqs/factoring/ecc/find_ecc_private_key.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set +from typing import Dict import sympy from attrs import frozen @@ -21,7 +21,7 @@ from qualtran import Bloq, bloq_example, BloqBuilder, BloqDocSpec, QUInt, Signature, SoquetT from qualtran.bloqs.basic_gates import IntState from qualtran.bloqs.bookkeeping import Free -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import SymbolicInt from .ec_phase_estimate_r import ECPhaseEstimateR @@ -103,12 +103,12 @@ def build_composite_bloq(self, bb: 'BloqBuilder') -> Dict[str, 'SoquetT']: bb.add(Free(QUInt(self.n)), reg=y) return {} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': Rx = ssa.new_symbol('Rx') Ry = ssa.new_symbol('Ry') generic_point = ECPoint(Rx, Ry, mod=self.mod, curve_a=self.curve_a) - return {(ECPhaseEstimateR(n=self.n, point=generic_point), 2)} + return {ECPhaseEstimateR(n=self.n, point=generic_point): 2} def cost_attrs(self): return [('n', self.n)] diff --git a/qualtran/bloqs/factoring/mod_exp.py b/qualtran/bloqs/factoring/mod_exp.py index a4db2a0c5..cabc6e0bd 100644 --- a/qualtran/bloqs/factoring/mod_exp.py +++ b/qualtran/bloqs/factoring/mod_exp.py @@ -13,7 +13,7 @@ # limitations under the License. import math from functools import cached_property -from typing import Dict, Optional, Set, Tuple, Union +from typing import Dict, Optional, Tuple, Union import attrs import numpy as np @@ -36,7 +36,7 @@ from qualtran.bloqs.basic_gates import IntState from qualtran.bloqs.mod_arithmetic import CModMulK from qualtran.drawing import Text, WireSymbol -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.resource_counting.generalizers import ignore_split_join @@ -118,12 +118,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', exponent: 'Soquet') -> Dict[st return {'exponent': bb.join(exponent, dtype=QUInt(self.exp_bitsize)), 'x': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': k = ssa.new_symbol('k') - return { - (IntState(val=1, bitsize=self.x_bitsize), 1), - (self._CtrlModMul(k=k), self.exp_bitsize), - } + return {IntState(val=1, bitsize=self.x_bitsize): 1, self._CtrlModMul(k=k): self.exp_bitsize} def on_classical_vals(self, exponent: int): return {'exponent': exponent, 'x': (self.base**exponent) % self.mod} diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py index 20cce0e95..bef382696 100644 --- a/qualtran/bloqs/for_testing/costing.py +++ b/qualtran/bloqs/for_testing/costing.py @@ -11,12 +11,12 @@ # 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 Any, Sequence, Set, Tuple +from typing import Any, Sequence, Tuple from attrs import field, frozen from qualtran import Bloq, Signature -from qualtran.resource_counting import BloqCountT, CostKey, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, BloqCountT, CostKey, SympySymbolAllocator def _convert_callees(callees: Sequence[BloqCountT]) -> Tuple[BloqCountT, ...]: @@ -46,8 +46,8 @@ class CostingBloq(Bloq): def signature(self) -> 'Signature': return Signature.build(register=self.num_qubits) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return set(self.callees) + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return dict(self.callees) def my_static_costs(self, cost_key: 'CostKey'): return dict(self.static_costs).get(cost_key, NotImplemented) diff --git a/qualtran/bloqs/for_testing/with_call_graph.py b/qualtran/bloqs/for_testing/with_call_graph.py index c459451b1..10659af0a 100644 --- a/qualtran/bloqs/for_testing/with_call_graph.py +++ b/qualtran/bloqs/for_testing/with_call_graph.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Set, TYPE_CHECKING +from typing import TYPE_CHECKING from attrs import frozen @@ -22,7 +22,7 @@ from qualtran.bloqs.for_testing.with_decomposition import TestParallelCombo, TestSerialCombo if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -33,6 +33,6 @@ class TestBloqWithCallGraph(Bloq): def signature(self) -> Signature: return Signature.build() - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n = ssa.new_symbol('n') - return {(TestParallelCombo(), 1), (TestSerialCombo(), 1), (TestAtom(), n)} + return {TestParallelCombo(): 1, TestSerialCombo(): 1, TestAtom(): n} diff --git a/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py b/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py index 0f8db837c..d9134e4bc 100644 --- a/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py +++ b/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py @@ -13,7 +13,7 @@ # limitations under the License. from collections import Counter from functools import cached_property -from typing import cast, Dict, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Tuple, TYPE_CHECKING, Union import numpy as np from attrs import field, frozen @@ -32,7 +32,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -182,7 +182,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': counts = Counter[Bloq]() d = self.degree @@ -192,7 +192,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: counts[self.walk_operator.adjoint().controlled()] += d counts[SU2RotationGate.arbitrary(ssa)] += 2 * d + 1 - return set(counts.items()) + return counts @bloq_example diff --git a/qualtran/bloqs/mcmt/and_bloq.py b/qualtran/bloqs/mcmt/and_bloq.py index 71738a707..29f8d4188 100644 --- a/qualtran/bloqs/mcmt/and_bloq.py +++ b/qualtran/bloqs/mcmt/and_bloq.py @@ -23,7 +23,7 @@ """ import itertools from functools import cached_property -from typing import cast, Dict, Iterable, Iterator, List, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Iterable, Iterator, List, Optional, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -50,7 +50,12 @@ from qualtran.cirq_interop import decompose_from_cirq_style_method from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.drawing import Circle, directional_text_box, Text, WireSymbol -from qualtran.resource_counting import big_O, BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import ( + big_O, + BloqCountDictT, + MutableBloqCountDictT, + SympySymbolAllocator, +) from qualtran.resource_counting.generalizers import ( cirq_to_bloqs, generalize_cvs, @@ -108,7 +113,7 @@ def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic.") return decompose_from_cirq_style_method(self) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if FLAG_AND_AS_LEAF: raise DecomposeTypeError(f"{self} is atomic.") @@ -117,9 +122,9 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: else: pre_post_cliffords = 2 - self.cv1 - self.cv2 if self.uncompute: - return {(ArbitraryClifford(n=2), 4 + 2 * pre_post_cliffords)} + return {ArbitraryClifford(n=2): 4 + 2 * pre_post_cliffords} - return {(ArbitraryClifford(n=2), 9 + 2 * pre_post_cliffords), (TGate(), 4)} + return {ArbitraryClifford(n=2): 9 + 2 * pre_post_cliffords, TGate(): 4} def _t_complexity_(self) -> 'TComplexity': if not FLAG_AND_AS_LEAF: @@ -399,16 +404,16 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - def __str__(self): return f'MultiAnd(n={self.n_ctrls})' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - pre_post_cliffords = set() + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + cost: 'MutableBloqCountDictT' = {And(): self.n_ctrls - 1} if not ( is_symbolic(self.cvs) or is_symbolic(*self.concrete_cvs) or (self.n_ctrls == sum(self.concrete_cvs)) ): - pre_post_cliffords = {(XGate(), 2 * (self.n_ctrls - sum(self.concrete_cvs)))} + cost[XGate()] = 2 * (self.n_ctrls - sum(self.concrete_cvs)) - return {(And(), self.n_ctrls - 1)} | pre_post_cliffords + return cost @bloq_example(generalizer=(ignore_cliffords, generalize_cvs)) diff --git a/qualtran/bloqs/mcmt/controlled_via_and.py b/qualtran/bloqs/mcmt/controlled_via_and.py index 8fd751336..a6778483d 100644 --- a/qualtran/bloqs/mcmt/controlled_via_and.py +++ b/qualtran/bloqs/mcmt/controlled_via_and.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -112,7 +112,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': counts: Counter[Bloq] = Counter() counts[self.subbloq.controlled()] += 1 @@ -124,7 +124,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: counts[ctrl] += 1 counts[ctrl.adjoint()] += 1 - return set(counts.items()) + return counts @bloq_example diff --git a/qualtran/bloqs/mcmt/ctrl_spec_and.py b/qualtran/bloqs/mcmt/ctrl_spec_and.py index 587670594..d22b47f4e 100644 --- a/qualtran/bloqs/mcmt/ctrl_spec_and.py +++ b/qualtran/bloqs/mcmt/ctrl_spec_and.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -171,13 +171,13 @@ def wire_symbol(self, reg: Optional[Register], idx: tuple[int, ...] = tuple()) - pretty_text = reg.name return directional_text_box(text=pretty_text, side=reg.side) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if not is_symbolic(self.n_ctrl_qubits) and self.n_ctrl_qubits == 2: assert isinstance(self._flat_cvs, tuple) cv1, cv2 = self._flat_cvs - return {(And(cv1, cv2), 1)} + return {And(cv1, cv2): 1} - return {(MultiAnd(self._flat_cvs), 1)} + return {MultiAnd(self._flat_cvs): 1} @bloq_example(generalizer=ignore_split_join) diff --git a/qualtran/bloqs/mcmt/multi_control_pauli.py b/qualtran/bloqs/mcmt/multi_control_pauli.py index a64bb2cf2..f95dfdbbd 100644 --- a/qualtran/bloqs/mcmt/multi_control_pauli.py +++ b/qualtran/bloqs/mcmt/multi_control_pauli.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Tuple, TYPE_CHECKING, Union import cirq import numpy as np @@ -39,7 +39,7 @@ from qualtran.symbolics import HasLength, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -129,9 +129,9 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT return {'controls': controls, 'target': target} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': if self.n_ctrls == 0: - return {(self.target_bloq, 1)} + return {self.target_bloq: 1} if is_symbolic(self.cvs): # TODO CtrlSpec does not support symbolic cvs yet. @@ -140,19 +140,15 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd if self.n_ctrls == 1: - return {(self.target_bloq.controlled(), 1)} + return {self.target_bloq.controlled(): 1} elif self.n_ctrls == 2: and_bloq = And(ssa.new_symbol('cv1'), ssa.new_symbol('cv2')) - return {(self.target_bloq.controlled(), 1), (and_bloq, 1), (and_bloq.adjoint(), 1)} + return {self.target_bloq.controlled(): 1, and_bloq: 1, and_bloq.adjoint(): 1} else: m_and_bloq = MultiAnd(self.cvs) - return { - (self.target_bloq.controlled(), 1), - (m_and_bloq, 1), - (m_and_bloq.adjoint(), 1), - } + return {self.target_bloq.controlled(): 1, m_and_bloq: 1, m_and_bloq.adjoint(): 1} - return {(self._multi_ctrl_bloq, 1)} + return {self._multi_ctrl_bloq: 1} def _apply_unitary_(self, args: 'cirq.ApplyUnitaryArgs') -> np.ndarray: cpauli = ( diff --git a/qualtran/bloqs/mod_arithmetic/_shims.py b/qualtran/bloqs/mod_arithmetic/_shims.py index bab116711..c6e630235 100644 --- a/qualtran/bloqs/mod_arithmetic/_shims.py +++ b/qualtran/bloqs/mod_arithmetic/_shims.py @@ -22,7 +22,7 @@ from collections import defaultdict from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, TYPE_CHECKING from attrs import frozen @@ -35,7 +35,7 @@ from qualtran.symbolics import ceil, log2 if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -47,7 +47,7 @@ class _ModInvInner(Bloq): def signature(self) -> 'Signature': return Signature([Register('x', QUInt(self.n)), Register('out', QUInt(self.n))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # This listing is based off of Haner 2023, fig 15. The order of operations # matches the order in the figure listing = [ @@ -71,7 +71,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: summer: Dict[Bloq, int] = defaultdict(lambda: 0) for bloq, n in listing: summer[bloq] += n - return set(summer.items()) + return summer def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() @@ -94,14 +94,14 @@ class ModInv(Bloq): def signature(self) -> 'Signature': return Signature([Register('x', QUInt(self.n)), Register('out', QUInt(self.n))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Roetteler # return {(Toffoli(), 32 * self.n**2 * log2(self.n))} return { - (_ModInvInner(n=self.n, mod=self.mod), 2 * self.n), - (Negate(QUInt(self.n)), 1), - (AddK(self.n, k=self.mod), 1), - (Swap(self.n), 1), + _ModInvInner(n=self.n, mod=self.mod): 2 * self.n, + Negate(QUInt(self.n)): 1, + AddK(self.n, k=self.mod): 1, + Swap(self.n): 1, } def wire_symbol( @@ -131,9 +131,9 @@ def signature(self) -> 'Signature': ] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Roetteler montgomery - return {(Toffoli(), ceil(16 * self.n**2 * log2(self.n) - 26.3 * self.n**2))} + return {Toffoli(): ceil(16 * self.n**2 * log2(self.n) - 26.3 * self.n**2)} def wire_symbol( self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple() diff --git a/qualtran/bloqs/mod_arithmetic/mod_addition.py b/qualtran/bloqs/mod_arithmetic/mod_addition.py index a1ce1c021..31f66a847 100644 --- a/qualtran/bloqs/mod_arithmetic/mod_addition.py +++ b/qualtran/bloqs/mod_arithmetic/mod_addition.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Optional, Tuple, TYPE_CHECKING, Union import numpy as np import sympy @@ -38,7 +38,7 @@ from qualtran.bloqs.basic_gates import XGate from qualtran.bloqs.bookkeeping import Cast from qualtran.drawing import Circle, Text, TextBox, WireSymbol -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.resource_counting.generalizers import ignore_split_join from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import is_symbolic @@ -143,13 +143,13 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet, y: Soquet) -> Dict[ # Return the output registers. return {'x': x, 'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (Add(QUInt(self.bitsize + 1)), 1), - (AddK(self.bitsize + 1, k=-self.mod), 1), - (AddK(self.bitsize, k=self.mod).controlled(), 1), - (LinearDepthGreaterThan(self.bitsize), 1), - (XGate(), 1), + Add(QUInt(self.bitsize + 1)): 1, + AddK(self.bitsize + 1, k=-self.mod): 1, + AddK(self.bitsize, k=self.mod).controlled(): 1, + LinearDepthGreaterThan(self.bitsize): 1, + XGate(): 1, } def short_name(self) -> str: @@ -229,8 +229,8 @@ def on_classical_vals( def __pow__(self, power: int) -> 'ModAddK': return ModAddK(self.bitsize, self.mod, add_val=self.add_val * power, cvs=self.cvs) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(Add(QUInt(self.bitsize), QUInt(self.bitsize)), 5)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {Add(QUInt(self.bitsize), QUInt(self.bitsize)): 5} @bloq_example @@ -279,9 +279,9 @@ class CModAddK(Bloq): def signature(self) -> 'Signature': return Signature([Register('ctrl', QBit()), Register('x', QUInt(self.bitsize))]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': k = ssa.new_symbol('k') - return {(AddK(k=k, bitsize=self.bitsize).controlled(), 5)} + return {AddK(k=k, bitsize=self.bitsize).controlled(): 5} def short_name(self) -> str: return f'x += {self.k} % {self.mod}' @@ -316,9 +316,9 @@ def signature(self) -> 'Signature': ] ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': k = ssa.new_symbol('k') - return {(CModAddK(k=k, bitsize=self.bitsize, mod=self.mod), self.bitsize)} + return {CModAddK(k=k, bitsize=self.bitsize, mod=self.mod): self.bitsize} def on_classical_vals( self, ctrl: 'ClassicalValT', x: 'ClassicalValT', y: 'ClassicalValT' @@ -424,13 +424,13 @@ def build_composite_bloq( return {'ctrl': ctrl, 'x': x, 'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (CAdd(QUInt(self.dtype.bitsize), QUInt(self.dtype.bitsize + 1), cv=self.cv), 1), - (AddK(self.dtype.bitsize + 1, -self.mod, signed=False), 1), - (AddK(self.dtype.bitsize, self.mod, cvs=(1,), signed=False), 1), - (CLinearDepthGreaterThan(QUInt(self.dtype.bitsize), cv=self.cv), 1), - (XGate(), 1), + CAdd(QUInt(self.dtype.bitsize), QUInt(self.dtype.bitsize + 1), cv=self.cv): 1, + AddK(self.dtype.bitsize + 1, -self.mod, signed=False): 1, + AddK(self.dtype.bitsize, self.mod, cvs=(1,), signed=False): 1, + CLinearDepthGreaterThan(QUInt(self.dtype.bitsize), cv=self.cv): 1, + XGate(): 1, } diff --git a/qualtran/bloqs/mod_arithmetic/mod_multiplication.py b/qualtran/bloqs/mod_arithmetic/mod_multiplication.py index 71aedcd38..cdf389a52 100644 --- a/qualtran/bloqs/mod_arithmetic/mod_multiplication.py +++ b/qualtran/bloqs/mod_arithmetic/mod_multiplication.py @@ -14,7 +14,7 @@ import numbers from functools import cached_property -from typing import Dict, Optional, Set, Tuple, Union +from typing import Dict, Optional, Tuple, Union import attrs import numpy as np @@ -38,7 +38,7 @@ from qualtran.bloqs.basic_gates import CNOT, CSwap, XGate from qualtran.bloqs.mod_arithmetic.mod_addition import CtrlScaleModAdd from qualtran.drawing import Circle, directional_text_box, Text, WireSymbol -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import is_symbolic @@ -134,12 +134,12 @@ def wire_symbol( return Text(f'x = 2 * x mod {self.mod}') return super().wire_symbol(reg, idx) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: return { - (AddK(self.dtype.bitsize + 2, -self.mod, signed=False), 1), - (AddK(self.dtype.bitsize + 1, self.mod, cvs=(1,), signed=False), 1), - (CNOT(), 1), - (XGate(), 2), + AddK(self.dtype.bitsize + 2, -self.mod, signed=False): 1, + AddK(self.dtype.bitsize + 1, self.mod, cvs=(1,), signed=False): 1, + CNOT(): 1, + XGate(): 2, } @@ -221,9 +221,9 @@ def build_composite_bloq( bb.free(y) return {'ctrl': ctrl, 'x': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: k = ssa.new_symbol('k') - return {(self._Add(k=k), 2), (CSwap(self.dtype.bitsize), 1)} + return {self._Add(k=k): 2, CSwap(self.dtype.bitsize): 1} def on_classical_vals(self, ctrl, x) -> Dict[str, ClassicalValT]: if ctrl and x < self.mod: diff --git a/qualtran/bloqs/mod_arithmetic/mod_subtraction.py b/qualtran/bloqs/mod_arithmetic/mod_subtraction.py index 122a0872e..01863dc04 100644 --- a/qualtran/bloqs/mod_arithmetic/mod_subtraction.py +++ b/qualtran/bloqs/mod_arithmetic/mod_subtraction.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Optional, Tuple, TYPE_CHECKING, Union import sympy from attrs import frozen @@ -39,7 +39,7 @@ from qualtran.symbolics import HasLength if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import SymbolicInt @@ -100,18 +100,18 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet) -> Dict[str, 'Soque bb.free(ancilla) return {'x': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cvs: Union[list[int], HasLength] if isinstance(self.dtype.bitsize, int): cvs = [0] * self.dtype.bitsize else: cvs = HasLength(self.dtype.bitsize) return { - (MultiControlX(cvs), 1), - (MultiControlX(cvs).adjoint(), 1), - (MultiTargetCNOT(self.dtype.bitsize), 1), - (AddK(self.dtype.bitsize, k=self.mod + 1, cvs=(1), signed=False), 1), - (XGate(), 2), + MultiControlX(cvs): 1, + MultiControlX(cvs).adjoint(): 1, + MultiTargetCNOT(self.dtype.bitsize): 1, + AddK(self.dtype.bitsize, k=self.mod + 1, cvs=(1), signed=False): 1, + XGate(): 2, } def wire_symbol( @@ -194,20 +194,20 @@ def build_composite_bloq( bb.free(ancilla) return {'ctrl': ctrl, 'x': x} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': cvs: Union[list[int], HasLength] if isinstance(self.dtype.bitsize, int): cvs = [0] * self.dtype.bitsize else: cvs = HasLength(self.dtype.bitsize) return { - (MultiControlX(cvs), 1), - (MultiControlX(cvs).adjoint(), 1), - (And(self.cv, 1), 1), - (And(self.cv, 1).adjoint(), 1), - (MultiTargetCNOT(self.dtype.bitsize), 1), - (AddK(self.dtype.bitsize, k=self.mod + 1, cvs=(1,), signed=False), 1), - (XGate(), 2), + MultiControlX(cvs): 1, + MultiControlX(cvs).adjoint(): 1, + And(self.cv, 1): 1, + And(self.cv, 1).adjoint(): 1, + MultiTargetCNOT(self.dtype.bitsize): 1, + AddK(self.dtype.bitsize, k=self.mod + 1, cvs=(1,), signed=False): 1, + XGate(): 2, } def wire_symbol( @@ -285,12 +285,12 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet, y: Soquet) -> Dict[ x = bb.add(BitwiseNot(self.dtype), x=x) return {'x': x, 'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (BitwiseNot(self.dtype), 2), - (AddK(self.dtype.bitsize, self.mod + 1, signed=False), 1), - (AddK(self.dtype.bitsize, self.mod + 1, signed=False).adjoint(), 1), - (ModAdd(self.dtype.bitsize, self.mod), 1), + BitwiseNot(self.dtype): 2, + AddK(self.dtype.bitsize, self.mod + 1, signed=False): 1, + AddK(self.dtype.bitsize, self.mod + 1, signed=False).adjoint(): 1, + ModAdd(self.dtype.bitsize, self.mod): 1, } def wire_symbol( @@ -365,12 +365,12 @@ def build_composite_bloq( x = bb.add(BitwiseNot(self.dtype), x=x) return {'ctrl': ctrl, 'x': x, 'y': y} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (BitwiseNot(self.dtype), 2), - (AddK(self.dtype.bitsize, self.mod + 1, signed=False), 1), - (AddK(self.dtype.bitsize, self.mod + 1, signed=False).adjoint(), 1), - (CModAdd(self.dtype, self.mod, self.cv), 1), + BitwiseNot(self.dtype): 2, + AddK(self.dtype.bitsize, self.mod + 1, signed=False): 1, + AddK(self.dtype.bitsize, self.mod + 1, signed=False).adjoint(): 1, + CModAdd(self.dtype, self.mod, self.cv): 1, } def wire_symbol( diff --git a/qualtran/bloqs/multiplexers/unary_iteration_bloq.py b/qualtran/bloqs/multiplexers/unary_iteration_bloq.py index dceddb99c..f5b5308b6 100644 --- a/qualtran/bloqs/multiplexers/unary_iteration_bloq.py +++ b/qualtran/bloqs/multiplexers/unary_iteration_bloq.py @@ -625,6 +625,6 @@ def unary_iteration_loops( try: unary_iteration_loops(0, {}, total_bits(self.control_registers)) - return {(bloq, count) for bloq, count in bloq_counts.items()} + return bloq_counts except NotImplementedError: return super().build_call_graph(ssa) diff --git a/qualtran/bloqs/phase_estimation/kaiser_window_state.py b/qualtran/bloqs/phase_estimation/kaiser_window_state.py index fe313e5af..b82a2493f 100644 --- a/qualtran/bloqs/phase_estimation/kaiser_window_state.py +++ b/qualtran/bloqs/phase_estimation/kaiser_window_state.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, List, Set, TYPE_CHECKING +from typing import Dict, List, TYPE_CHECKING import attrs import numpy as np @@ -35,7 +35,7 @@ import quimb.tensor as qtn from qualtran import ConnectionT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @attrs.frozen @@ -118,7 +118,7 @@ def my_tensors( ) ] - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': from qualtran.bloqs.state_preparation.state_preparation_via_rotation import ( StatePreparationViaRotations, ) @@ -127,7 +127,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: state_prep_coeff = HasLength(2**self.bitsize) else: state_prep_coeff = self.kaiser_state_coeff.tolist() - return {(StatePreparationViaRotations(state_prep_coeff, self.bitsize), 1)} + return {StatePreparationViaRotations(state_prep_coeff, self.bitsize): 1} @bloq_example diff --git a/qualtran/bloqs/phase_estimation/lp_resource_state.py b/qualtran/bloqs/phase_estimation/lp_resource_state.py index aa1047c1f..b9dcd387e 100644 --- a/qualtran/bloqs/phase_estimation/lp_resource_state.py +++ b/qualtran/bloqs/phase_estimation/lp_resource_state.py @@ -15,7 +15,7 @@ """Resource states proposed by A. Luis and J. Peřina (1996) for optimal phase measurements""" from collections import Counter from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING +from typing import Dict, TYPE_CHECKING import attrs import numpy as np @@ -29,7 +29,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, Soquet, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @attrs.frozen @@ -73,7 +73,7 @@ def build_composite_bloq( anc = bb.add(Hadamard(), q=anc) return {'m': bb.join(q[::-1]), 'anc': anc} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': rz_angle = -2 * pi(self.bitsize) / (2**self.bitsize + 1) ret: Counter['Bloq'] = Counter() ret[Rz(angle=rz_angle)] += 1 @@ -84,7 +84,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: else: for i in range(self.bitsize): ret[Rz(angle=rz_angle * (2**i)).controlled()] += 1 - return set(ret.items()) + return ret @attrs.frozen @@ -172,19 +172,19 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str bb.free(anc) return {'qpe_reg': qpe_reg} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': flag_angle = acos(1 / (1 + 2**self.bitsize)) reflection_bloq: 'Bloq' = ReflectionUsingPrepare.reflection_around_zero( [1, 1, self.bitsize], global_phase=1j ) return { - (LPRSInterimPrep(self.bitsize), 2), - (LPRSInterimPrep(self.bitsize).adjoint(), 1), - (Ry(angle=flag_angle), 2), - (Ry(angle=-1 * flag_angle), 1), - (reflection_bloq, 1), - (XGate(), 2), - (CZ(), 1), + LPRSInterimPrep(self.bitsize): 2, + LPRSInterimPrep(self.bitsize).adjoint(): 1, + Ry(angle=flag_angle): 2, + Ry(angle=-1 * flag_angle): 1, + reflection_bloq: 1, + XGate(): 2, + CZ(): 1, } diff --git a/qualtran/bloqs/phase_estimation/qubitization_qpe.py b/qualtran/bloqs/phase_estimation/qubitization_qpe.py index fcbbeabc8..9cc1b3a8b 100644 --- a/qualtran/bloqs/phase_estimation/qubitization_qpe.py +++ b/qualtran/bloqs/phase_estimation/qubitization_qpe.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Iterator, Set, Tuple, TYPE_CHECKING +from typing import Iterator, Tuple, TYPE_CHECKING import attrs import cirq @@ -33,7 +33,7 @@ from qualtran.symbolics import is_symbolic, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @attrs.frozen @@ -128,15 +128,15 @@ def decompose_from_registers( yield reflect_controlled.on_registers(control=qpre_reg[i], **reflect_regs) yield self.qft_inv.on(*qpre_reg) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Assumes self.unitary is not fast forwardable. M = 2**self.m_bits return { - (self.ctrl_state_prep, 1), - (self.walk.controlled(), 1), - (self.walk.reflect.controlled(ctrl_spec=CtrlSpec(cvs=0)), 2 * (self.m_bits - 1)), - (self.walk, M - 2), - (self.qft_inv, 1), + self.ctrl_state_prep: 1, + self.walk.controlled(): 1, + self.walk.reflect.controlled(ctrl_spec=CtrlSpec(cvs=0)): 2 * (self.m_bits - 1), + self.walk: M - 2, + self.qft_inv: 1, } def __str__(self) -> str: diff --git a/qualtran/bloqs/phase_estimation/text_book_qpe.py b/qualtran/bloqs/phase_estimation/text_book_qpe.py index 6ab432f79..9349ca329 100644 --- a/qualtran/bloqs/phase_estimation/text_book_qpe.py +++ b/qualtran/bloqs/phase_estimation/text_book_qpe.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Iterator, Set, Tuple, TYPE_CHECKING +from typing import Iterator, Tuple, TYPE_CHECKING import attrs import cirq @@ -23,7 +23,7 @@ from qualtran.symbolics import is_symbolic, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @attrs.frozen @@ -180,12 +180,12 @@ def decompose_from_registers( yield cirq.pow(unitary_op.controlled_by(qbit), 2**i) yield self.qft_inv.on(*phase_qubits) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': # Assumes self.unitary is not fast forwardable. return { - (self.ctrl_state_prep, 1), - (self.unitary.controlled(), (2**self.m_bits) - 1), - (self.qft_inv, 1), + self.ctrl_state_prep: 1, + self.unitary.controlled(): (2**self.m_bits) - 1, + self.qft_inv: 1, } diff --git a/qualtran/bloqs/qft/approximate_qft.py b/qualtran/bloqs/qft/approximate_qft.py index 6abb82694..e6dfcee86 100644 --- a/qualtran/bloqs/qft/approximate_qft.py +++ b/qualtran/bloqs/qft/approximate_qft.py @@ -13,7 +13,7 @@ # limitations under the License. from collections import defaultdict from functools import cached_property -from typing import Dict, Iterator, Set, TYPE_CHECKING +from typing import Iterator, TYPE_CHECKING import attrs import cirq @@ -28,7 +28,11 @@ from qualtran.symbolics import ceil, is_symbolic, log2, SymbolicFloat, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import ( + BloqCountDictT, + MutableBloqCountDictT, + SympySymbolAllocator, + ) @attrs.frozen @@ -130,8 +134,8 @@ def decompose_from_registers( for i in range(self.bitsize // 2): yield cirq.SWAP(q[i], q[-i - 1]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - phase_dict: Dict[AddIntoPhaseGrad, SymbolicInt] = defaultdict(int) + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + phase_dict: 'MutableBloqCountDictT' = defaultdict(int) if is_symbolic(self.bitsize, self.phase_bitsize): phase_dict[ AddIntoPhaseGrad( @@ -142,10 +146,10 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: for i in range(1, int(self.bitsize)): b = min(i, self.phase_bitsize - 1) phase_dict[AddIntoPhaseGrad(b, b + 1, right_shift=1, controlled_by=1)] += 1 - ret = {(Hadamard(), self.bitsize), *phase_dict.items()} + phase_dict[Hadamard()] = self.bitsize if self.with_reverse: - ret |= {(TwoBitSwap(), self.bitsize // 2)} - return ret + phase_dict[TwoBitSwap()] = self.bitsize // 2 + return phase_dict @bloq_example diff --git a/qualtran/bloqs/qft/qft_phase_gradient.py b/qualtran/bloqs/qft/qft_phase_gradient.py index cf89f0145..6597eba42 100644 --- a/qualtran/bloqs/qft/qft_phase_gradient.py +++ b/qualtran/bloqs/qft/qft_phase_gradient.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Iterator, Set +from typing import Iterator import attrs import cirq @@ -23,7 +23,7 @@ from qualtran.bloqs.arithmetic.multiplication import PlusEqualProduct from qualtran.bloqs.basic_gates import Hadamard from qualtran.bloqs.basic_gates.swap import Swap -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, MutableBloqCountDictT, SympySymbolAllocator from qualtran.symbolics.types import is_symbolic @@ -96,37 +96,27 @@ def decompose_from_registers( for i in range(self.bitsize // 2): yield cirq.SWAP(q[i], q[-i - 1]) - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> 'BloqCountDictT': if is_symbolic(self.bitsize): # TODO: The T-gate cost here is an upper bound constructed off of the recurrence # relation for the QFT as used in decompose_from_registers above. - return set( - [ - ( - PlusEqualProduct( - self.bitsize // 2, self.bitsize - (self.bitsize // 2), self.bitsize - ), - log(self.bitsize, 2), - ) - ] - ) + return { + PlusEqualProduct( + self.bitsize // 2, self.bitsize - (self.bitsize // 2), self.bitsize + ): log(self.bitsize, 2) + } if self.bitsize == 1: - return set([(Hadamard(), 1)]) - ret: Set['BloqCountT'] = set( - [ - (QFTPhaseGradient(self.bitsize // 2), 1), - (QFTPhaseGradient(self.bitsize - (self.bitsize // 2)), 1), - ( - PlusEqualProduct( - self.bitsize // 2, self.bitsize - (self.bitsize // 2), self.bitsize - ), - 1, - ), - ] - ) + return {Hadamard(): 1} + ret: 'MutableBloqCountDictT' = { + QFTPhaseGradient(self.bitsize // 2): 1, + QFTPhaseGradient(self.bitsize - (self.bitsize // 2)): 1, + PlusEqualProduct( + self.bitsize // 2, self.bitsize - (self.bitsize // 2), self.bitsize + ): 1, + } if self.with_reverse: - ret |= set([(Swap(1), self.bitsize // 2)]) + ret[Swap(1)] = self.bitsize // 2 return ret diff --git a/qualtran/bloqs/qft/qft_text_book.py b/qualtran/bloqs/qft/qft_text_book.py index 96c021b1c..c560deab2 100644 --- a/qualtran/bloqs/qft/qft_text_book.py +++ b/qualtran/bloqs/qft/qft_text_book.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Iterator, Set +from typing import Iterator import attrs import cirq @@ -23,7 +23,7 @@ from qualtran.bloqs.basic_gates.hadamard import Hadamard from qualtran.bloqs.basic_gates.swap import TwoBitSwap from qualtran.bloqs.rotations.phase_gradient import PhaseGradientUnitary -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, MutableBloqCountDictT, SympySymbolAllocator from qualtran.symbolics import SymbolicInt from qualtran.symbolics.types import is_symbolic @@ -84,20 +84,17 @@ def decompose_from_registers( for i in range(self.bitsize // 2): yield cirq.SWAP(q[i], q[-i - 1]) - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set['BloqCountT']: - ret: Set['BloqCountT'] = {(Hadamard(), self.bitsize)} + def build_call_graph(self, ssa: SympySymbolAllocator) -> 'BloqCountDictT': + ret: 'MutableBloqCountDictT' = {Hadamard(): self.bitsize} if is_symbolic(self.bitsize): - ret |= { - ( - PhaseGradientUnitary(self.bitsize - 1, exponent=0.5, is_controlled=True), - self.bitsize // 2, - ) - } + ret[PhaseGradientUnitary(self.bitsize - 1, exponent=0.5, is_controlled=True)] = ( + self.bitsize // 2 + ) else: for i in range(1, self.bitsize): - ret |= {(PhaseGradientUnitary(i, exponent=0.5, is_controlled=True), 1)} + ret[PhaseGradientUnitary(i, exponent=0.5, is_controlled=True)] = 1 if self.with_reverse: - ret |= {(TwoBitSwap(), self.bitsize // 2)} + ret[TwoBitSwap()] = self.bitsize // 2 return ret diff --git a/qualtran/bloqs/qft/two_bit_ffft.py b/qualtran/bloqs/qft/two_bit_ffft.py index 96bccc8bb..e50c1c6dd 100644 --- a/qualtran/bloqs/qft/two_bit_ffft.py +++ b/qualtran/bloqs/qft/two_bit_ffft.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, List, Set, TYPE_CHECKING +from typing import Dict, List, TYPE_CHECKING import numpy as np from attrs import frozen @@ -41,7 +41,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def _fkn_matrix(k: int, n: int) -> NDArray[np.complex128]: @@ -105,13 +105,13 @@ def my_tensors( qtn.Tensor(data=matrix.reshape((2,) * 4), inds=out_inds + in_inds, tags=[str(self)]) ] - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (Rz(2 * np.pi * self.k / self.n, eps=self.eps), 1), - (SGate(), 3), - (Hadamard(), 6), - (TGate(), 2), - (CNOT(), 3), + Rz(2 * np.pi * self.k / self.n, eps=self.eps): 1, + SGate(): 3, + Hadamard(): 6, + TGate(): 2, + CNOT(): 3, } def build_composite_bloq( diff --git a/qualtran/bloqs/qsp/generalized_qsp.py b/qualtran/bloqs/qsp/generalized_qsp.py index e8b71e5b7..186147ee9 100644 --- a/qualtran/bloqs/qsp/generalized_qsp.py +++ b/qualtran/bloqs/qsp/generalized_qsp.py @@ -13,7 +13,7 @@ # limitations under the License. from collections import Counter from functools import cached_property -from typing import Iterable, Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Iterable, Iterator, Sequence, Tuple, TYPE_CHECKING, Union import numpy as np from attrs import field, frozen @@ -47,7 +47,7 @@ if TYPE_CHECKING: import cirq - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def qsp_complementary_polynomial( @@ -369,7 +369,7 @@ def decompose_from_registers( def is_symbolic(self) -> bool: return is_symbolic(self.P, self.Q, self.negative_power) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': counts = Counter[Bloq]() degree = slen(self.P) - 1 @@ -380,7 +380,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: counts[self.U.adjoint()] += smax(0, self.negative_power - degree) counts[self.U.adjoint().controlled()] += smin(degree, self.negative_power) - return set((bloq, count) for bloq, count in counts.items() if not is_zero(count)) + return {bloq: count for bloq, count in counts.items() if not is_zero(count)} @bloq_example diff --git a/qualtran/bloqs/reflections/reflection_using_prepare.py b/qualtran/bloqs/reflections/reflection_using_prepare.py index bdb37bd39..ac2ec1337 100644 --- a/qualtran/bloqs/reflections/reflection_using_prepare.py +++ b/qualtran/bloqs/reflections/reflection_using_prepare.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Iterator, Optional, Sequence, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -33,7 +33,11 @@ if TYPE_CHECKING: from qualtran.bloqs.block_encoding.lcu_block_encoding import BlackBoxPrepare from qualtran.bloqs.state_preparation.prepare_base import PrepareOracle - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import ( + BloqCountDictT, + MutableBloqCountDictT, + SympySymbolAllocator, + ) @attrs.frozen(cache_hash=True) @@ -162,21 +166,21 @@ def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.Circ wire_symbols += ['R_L'] * total_bits(self.selection_registers) return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n_phase_control = sum(reg.total_bits() for reg in self.selection_registers) cvs = HasLength(n_phase_control) if is_symbolic(n_phase_control) else [0] * n_phase_control - costs: Set['BloqCountT'] = { - (self.prepare_gate, 1), - (self.prepare_gate.adjoint(), 1), - (MultiControlZ(cvs), 1), + costs: 'MutableBloqCountDictT' = { + self.prepare_gate: 1, + self.prepare_gate.adjoint(): 1, + MultiControlZ(cvs): 1, } if self.control_val is None: - costs.add((XGate(), 2)) + costs[XGate()] = 2 if self.global_phase != 1: phase_op: Bloq = GlobalPhase.from_coefficient(self.global_phase, eps=self.eps) if self.control_val is not None: phase_op = phase_op.controlled(ctrl_spec=CtrlSpec(cvs=self.control_val)) - costs.add((phase_op, 1)) + costs[phase_op] = 1 return costs def adjoint(self) -> 'ReflectionUsingPrepare': diff --git a/qualtran/bloqs/rotations/hamming_weight_phasing.py b/qualtran/bloqs/rotations/hamming_weight_phasing.py index cfdfa3b4e..39feeacb1 100644 --- a/qualtran/bloqs/rotations/hamming_weight_phasing.py +++ b/qualtran/bloqs/rotations/hamming_weight_phasing.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Set, TYPE_CHECKING +from typing import Dict, TYPE_CHECKING import attrs import numpy as np @@ -26,7 +26,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @attrs.frozen @@ -96,14 +96,13 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str def pretty_name(self) -> str: return f'HWP_{self.bitsize}(Z^{self.exponent})' - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (HammingWeightCompute(self.bitsize), 1), - (HammingWeightCompute(self.bitsize).adjoint(), 1), - ( - ZPowGate(exponent=self.exponent, eps=self.eps / self.bitsize.bit_length()), - self.bitsize.bit_length(), - ), + HammingWeightCompute(self.bitsize): 1, + HammingWeightCompute(self.bitsize).adjoint(): 1, + ZPowGate( + exponent=self.exponent, eps=self.eps / self.bitsize.bit_length() + ): self.bitsize.bit_length(), } diff --git a/qualtran/bloqs/rotations/phase_gradient.py b/qualtran/bloqs/rotations/phase_gradient.py index 2546e9322..ab981c17a 100644 --- a/qualtran/bloqs/rotations/phase_gradient.py +++ b/qualtran/bloqs/rotations/phase_gradient.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, Iterator, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -42,7 +42,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import SymbolicFloat, SymbolicInt @@ -128,20 +128,19 @@ def __pow__(self, power): return self return attrs.evolve(self, exponent=self.exponent * power) - def build_call_graph(self, ssa: SympySymbolAllocator) -> Set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> 'BloqCountDictT': gate = CZPowGate if self.is_controlled else ZPowGate if is_symbolic(self.bitsize): return { - ( - gate(exponent=self.exponent / 2 ** (self.bitsize), eps=self.eps / self.bitsize), - self.bitsize, - ) + gate( + exponent=self.exponent / 2 ** (self.bitsize), eps=self.eps / self.bitsize + ): self.bitsize } - ret: Set['BloqCountT'] = set() - for i in range(self.bitsize): - ret.add((gate(exponent=self.exponent / 2**i, eps=self.eps / self.bitsize), 1)) - return ret + return { + gate(exponent=self.exponent / 2**i, eps=self.eps / self.bitsize): 1 + for i in range(self.bitsize) + } @bloq_example @@ -324,12 +323,12 @@ def on_classical_vals(self, **kwargs) -> Dict[str, 'ClassicalValT']: phase_grad_out = (phase_grad + self.sign * self.scaled_val(x)) % (2**self.phase_bitsize) return {'x': x, 'phase_grad': phase_grad_out} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_toffoli = self.phase_bitsize - 2 if self.controlled_by is not None: - return {(Toffoli(), 2 * num_toffoli)} + return {Toffoli(): 2 * num_toffoli} - return {(Toffoli(), num_toffoli)} + return {Toffoli(): num_toffoli} def adjoint(self) -> 'AddIntoPhaseGrad': return attrs.evolve(self, sign=-self.sign) @@ -514,7 +513,7 @@ def on_classical_vals(self, x: int, phase_grad: int) -> Dict[str, 'ClassicalValT phase_grad_out = (phase_grad + self.scaled_val(x)) % 2**self.phase_bitsize return {'x': x, 'phase_grad': phase_grad_out} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_additions = (self.gamma_dtype.bitsize + 2) // 2 if not isinstance(self.gamma, sympy.Basic): num_additions_naive = 0 @@ -525,7 +524,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: if -(self.phase_bitsize + self.x_dtype.num_int) < shift < self.x_dtype.num_frac: num_additions_naive += 1 num_additions = min(num_additions_naive, num_additions) - return {(AddIntoPhaseGrad(self.x_dtype.bitsize, self.phase_bitsize), num_additions)} + return {AddIntoPhaseGrad(self.x_dtype.bitsize, self.phase_bitsize): num_additions} @bloq_example diff --git a/qualtran/bloqs/rotations/programmable_ancilla_rotation.py b/qualtran/bloqs/rotations/programmable_ancilla_rotation.py index 9d9285453..e65779cb1 100644 --- a/qualtran/bloqs/rotations/programmable_ancilla_rotation.py +++ b/qualtran/bloqs/rotations/programmable_ancilla_rotation.py @@ -21,7 +21,7 @@ from qualtran import Bloq, bloq_example, BloqBuilder, QBit, Register, Side, Signature, SoquetT from qualtran.bloqs.basic_gates import CNOT, Hadamard, XGate, ZPowGate from qualtran.bloqs.basic_gates._shims import Measure -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.symbolics import ceil, is_symbolic, log2, SymbolicFloat, SymbolicInt @@ -131,7 +131,7 @@ def from_failure_probability( n_rounds = ceil(log2(1 / max_fail_probability)) return cls(exponent=exponent, eps=eps, n_rounds=n_rounds) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: resources: Counter[Bloq] = Counter({CNOT(): self.n_rounds, XGate(): self.n_rounds}) n_rz = self.n_rounds + (1 if self.apply_final_correction else 0) @@ -149,7 +149,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: resources[Measure()] += self.n_rounds - return set(resources.items()) + return resources @bloq_example diff --git a/qualtran/bloqs/rotations/quantum_variable_rotation.py b/qualtran/bloqs/rotations/quantum_variable_rotation.py index 4dccd8c7e..869882f96 100644 --- a/qualtran/bloqs/rotations/quantum_variable_rotation.py +++ b/qualtran/bloqs/rotations/quantum_variable_rotation.py @@ -56,7 +56,7 @@ import abc from functools import cached_property -from typing import cast, Dict, Sequence, Set, TYPE_CHECKING +from typing import cast, Dict, Sequence, TYPE_CHECKING import attrs import numpy as np @@ -90,7 +90,7 @@ if TYPE_CHECKING: from qualtran import SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator class QvrInterface(GateWithRegisters, metaclass=abc.ABCMeta): @@ -199,11 +199,11 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str out[-(i + offset)] = bb.add(ZPowGate(exponent=exp, eps=eps), q=out[-(i + offset)]) return {self.cost_reg.name: bb.join(out, self.cost_reg.dtype)} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': zpow = ZPowGate( exponent=self.gamma / (2**self.num_frac_rotations), eps=self.eps / self.num_rotations ) - return {(zpow, self.num_rotations)} + return {zpow: self.num_rotations} @bloq_example @@ -481,14 +481,9 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str ) return {self.cost_reg.name: out, 'phase_grad': phase_grad} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - ( - AddScaledValIntoPhaseReg( - self.cost_dtype, self.b_grad, self.gamma, self.gamma_dtype - ), - 1, - ) + AddScaledValIntoPhaseReg(self.cost_dtype, self.b_grad, self.gamma, self.gamma_dtype): 1 } diff --git a/qualtran/bloqs/rotations/rz_via_phase_gradient.py b/qualtran/bloqs/rotations/rz_via_phase_gradient.py index abe7c707d..00e71e2ee 100644 --- a/qualtran/bloqs/rotations/rz_via_phase_gradient.py +++ b/qualtran/bloqs/rotations/rz_via_phase_gradient.py @@ -26,7 +26,7 @@ ) from qualtran.bloqs.arithmetic.controlled_add_or_subtract import ControlledAddOrSubtract from qualtran.bloqs.bookkeeping import Cast -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -99,8 +99,8 @@ def build_composite_bloq( return {'q': q, 'angle': angle, 'phase_grad': phase_grad} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: - return {(ControlledAddOrSubtract(self._angle_int_dtype, self._phasegrad_int_dtype), 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {ControlledAddOrSubtract(self._angle_int_dtype, self._phasegrad_int_dtype): 1} @bloq_example diff --git a/qualtran/bloqs/rotations/zpow_via_phase_gradient.py b/qualtran/bloqs/rotations/zpow_via_phase_gradient.py index 309f08d39..5e4e9c412 100644 --- a/qualtran/bloqs/rotations/zpow_via_phase_gradient.py +++ b/qualtran/bloqs/rotations/zpow_via_phase_gradient.py @@ -30,7 +30,7 @@ ) from qualtran.bloqs.arithmetic import XorK from qualtran.bloqs.rotations.phase_gradient import AddIntoPhaseGrad -from qualtran.resource_counting import BloqCountT, SympySymbolAllocator +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.resource_counting.generalizers import ignore_alloc_free from qualtran.symbolics import ceil, is_symbolic, log2, pi, SymbolicFloat, SymbolicInt @@ -123,10 +123,10 @@ def build_composite_bloq( return {'q': q, 'phase_grad': phase_grad} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> set['BloqCountT']: + def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT: return { - (self._load_bloq.controlled(), 2), - (AddIntoPhaseGrad(self.phase_grad_bitsize, self.phase_grad_bitsize), 1), + self._load_bloq.controlled(): 2, + AddIntoPhaseGrad(self.phase_grad_bitsize, self.phase_grad_bitsize): 1, } def __str__(self) -> str: diff --git a/qualtran/bloqs/state_preparation/prepare_uniform_superposition.py b/qualtran/bloqs/state_preparation/prepare_uniform_superposition.py index a44ff6297..c3a7002bc 100644 --- a/qualtran/bloqs/state_preparation/prepare_uniform_superposition.py +++ b/qualtran/bloqs/state_preparation/prepare_uniform_superposition.py @@ -157,12 +157,12 @@ def build_call_graph( _, l, logL = self.k_l_logL() theta = acos(1 - (2 ** floor(log2(l))) / l) return { - (Hadamard(), 3 * logL), - (LessThanConstant(logL, l), 1), - (LessThanConstant(logL, l).adjoint(), 1), - (Rz(theta), 2), - (MultiAnd(HasLength(logL)), 1), - (MultiAnd(HasLength(logL)).adjoint(), 1), + Hadamard(): 3 * logL, + LessThanConstant(logL, l): 1, + LessThanConstant(logL, l).adjoint(): 1, + Rz(theta): 2, + MultiAnd(HasLength(logL)): 1, + MultiAnd(HasLength(logL)).adjoint(): 1, } diff --git a/qualtran/bloqs/state_preparation/sparse_state_preparation_via_rotations.py b/qualtran/bloqs/state_preparation/sparse_state_preparation_via_rotations.py index cd0890a1c..37939becf 100644 --- a/qualtran/bloqs/state_preparation/sparse_state_preparation_via_rotations.py +++ b/qualtran/bloqs/state_preparation/sparse_state_preparation_via_rotations.py @@ -11,7 +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 Sequence, Set, TYPE_CHECKING, Union +from typing import Sequence, TYPE_CHECKING, Union import numpy as np import sympy @@ -31,7 +31,7 @@ import scipy from qualtran import BloqBuilder, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator @frozen @@ -193,8 +193,8 @@ def build_composite_bloq( return {'target_state': target_state, 'phase_gradient': phase_gradient} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(self._dense_stateprep_bloq, 1), (self._basis_permutation_bloq, 1)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {self._dense_stateprep_bloq: 1, self._basis_permutation_bloq: 1} @bloq_example diff --git a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py index 22e344c2b..99bc134f0 100644 --- a/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py +++ b/qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py @@ -20,7 +20,7 @@ largest absolute error that one can tolerate in the prepared amplitudes. """ from functools import cached_property -from typing import Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Sequence, Tuple, TYPE_CHECKING, Union import attrs import numpy as np @@ -48,7 +48,7 @@ if TYPE_CHECKING: from qualtran import BloqBuilder, Soquet, SoquetT - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def _data_or_shape_to_tuple(data_or_shape: Union[NDArray, Shaped]) -> Tuple: @@ -271,13 +271,13 @@ def build_composite_bloq( 'keep': keep, } - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': return { - (PrepareUniformSuperposition(self.n_coeff), 1), - (self.qrom_bloq, 1), - (LessThanEqual(self.mu, self.mu), 1), - (CSwap(self.selection_bitsize), 1), - (Hadamard(), self.mu), + PrepareUniformSuperposition(self.n_coeff): 1, + self.qrom_bloq: 1, + LessThanEqual(self.mu, self.mu): 1, + CSwap(self.selection_bitsize): 1, + Hadamard(): self.mu, } diff --git a/qualtran/bloqs/state_preparation/state_preparation_via_rotation.py b/qualtran/bloqs/state_preparation/state_preparation_via_rotation.py index 5fdf8a8e4..2466a11c3 100644 --- a/qualtran/bloqs/state_preparation/state_preparation_via_rotation.py +++ b/qualtran/bloqs/state_preparation/state_preparation_via_rotation.py @@ -74,7 +74,7 @@ """ from collections import Counter -from typing import cast, Dict, Iterable, List, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Iterable, List, Tuple, TYPE_CHECKING, Union import attrs import numpy as np @@ -99,7 +99,7 @@ from qualtran.symbolics import bit_length, HasLength, is_symbolic, Shaped, slen, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator def _to_tuple_or_has_length( @@ -229,7 +229,7 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So soqs = self._prepare_phases(bb, **soqs) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': ret: 'Counter[Bloq]' = Counter() ret[Rx(angle=-np.pi / 2)] += self.state_bitsize ret[Rx(angle=np.pi / 2)] += self.state_bitsize @@ -237,7 +237,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: ret[self.prga_prepare_phases] += 1 for bloq in self.prga_prepare_amplitude: ret[bloq] += 1 - return set(ret.items()) + return ret def _prepare_amplitudes(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, SoquetT]: r"""Parameters into soqs: @@ -456,12 +456,12 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So bb.free(cast(Soquet, soqs.pop("target0_"))) return soqs - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': ret: 'Counter[Bloq]' = Counter() ret[self.qrom_bloq] += 1 ret[self.qrom_bloq.adjoint()] += 1 ret[self.add_into_phase_grad] += 1 - return set(ret.items()) + return ret def __repr__(self): return ( diff --git a/qualtran/bloqs/swap_network/cswap_approx.py b/qualtran/bloqs/swap_network/cswap_approx.py index 40ecb9557..b2f72a504 100644 --- a/qualtran/bloqs/swap_network/cswap_approx.py +++ b/qualtran/bloqs/swap_network/cswap_approx.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, Iterator, Set, TYPE_CHECKING +from typing import Dict, Iterator, TYPE_CHECKING import cirq from attrs import frozen @@ -32,7 +32,7 @@ from qualtran.symbolics import SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -105,16 +105,12 @@ def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.Circ ("@(approx)",) + ("×(x)",) * self.bitsize + ("×(y)",) * self.bitsize ) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': n = self.bitsize # 4 * n: G gates, each wth 1 T and 4 single qubit cliffords # 4 * n: CNOTs # 2 * n - 1: CNOTs from 1 MultiTargetCNOT - return { - (TGate(), 4 * n), - (ArbitraryClifford(n=1), 16 * n), - (ArbitraryClifford(n=2), 6 * n - 1), - } + return {TGate(): 4 * n, ArbitraryClifford(n=1): 16 * n, ArbitraryClifford(n=2): 6 * n - 1} @bloq_example diff --git a/qualtran/bloqs/swap_network/swap_with_zero.py b/qualtran/bloqs/swap_network/swap_with_zero.py index c4d7c4358..1f14a973d 100644 --- a/qualtran/bloqs/swap_network/swap_with_zero.py +++ b/qualtran/bloqs/swap_network/swap_with_zero.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import cast, Dict, Iterable, Iterator, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Iterable, Iterator, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -38,7 +38,7 @@ from qualtran.symbolics import is_symbolic, prod, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -180,9 +180,9 @@ def build_composite_bloq( } return sel | {'targets': targets} - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': num_swaps = prod(x for x in self.n_target_registers) - 1 - return {(self.cswap_n, num_swaps)} + return {self.cswap_n: num_swaps} def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo: from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py index 131c9058b..6823999c1 100644 --- a/qualtran/resource_counting/__init__.py +++ b/qualtran/resource_counting/__init__.py @@ -23,6 +23,7 @@ BloqCountDictT, BloqCountT, big_O, + MutableBloqCountDictT, SympySymbolAllocator, get_bloq_callee_counts, get_bloq_call_graph, diff --git a/qualtran/resource_counting/_call_graph.py b/qualtran/resource_counting/_call_graph.py index d398a95e0..cf3e1923a 100644 --- a/qualtran/resource_counting/_call_graph.py +++ b/qualtran/resource_counting/_call_graph.py @@ -23,6 +23,7 @@ Iterable, List, Mapping, + MutableMapping, Optional, Sequence, Set, @@ -37,6 +38,7 @@ BloqCountT = Tuple[Bloq, Union[int, sympy.Expr]] BloqCountDictT = Mapping[Bloq, Union[int, sympy.Expr]] +MutableBloqCountDictT = MutableMapping[Bloq, Union[int, sympy.Expr]] from ._generalization import _make_composite_generalizer, GeneralizerT @@ -66,7 +68,7 @@ def new_symbol(self, prefix: str) -> sympy.Symbol: return s -def build_cbloq_call_graph(cbloq: CompositeBloq) -> Set[BloqCountT]: +def build_cbloq_call_graph(cbloq: CompositeBloq) -> BloqCountDictT: """Count all the subbloqs in a composite bloq. This is the function underpinning `CompositeBloq.build_call_graph`. @@ -78,7 +80,7 @@ def build_cbloq_call_graph(cbloq: CompositeBloq) -> Set[BloqCountT]: for binst in cbloq.bloq_instances: counts[binst.bloq] += 1 - return {(bloq, n) for bloq, n in counts.items()} + return counts def _generalize_callees( diff --git a/qualtran/resource_counting/_call_graph_test.py b/qualtran/resource_counting/_call_graph_test.py index c0dc001d6..af540e50c 100644 --- a/qualtran/resource_counting/_call_graph_test.py +++ b/qualtran/resource_counting/_call_graph_test.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections import defaultdict from functools import cached_property -from typing import Dict, Iterable, Optional, Sequence, Set, Tuple +from typing import Dict, Iterable, Optional, Sequence, Tuple import attrs import networkx as nx @@ -26,9 +27,11 @@ from qualtran.bloqs.basic_gates import TGate from qualtran.bloqs.bookkeeping import ArbitraryClifford, Join, Split from qualtran.resource_counting import ( + BloqCountDictT, BloqCountT, get_bloq_call_graph, get_bloq_callee_counts, + MutableBloqCountDictT, SympySymbolAllocator, ) from qualtran.resource_counting.generalizers import generalize_rotation_angle @@ -43,8 +46,8 @@ class BigBloq(Bloq): def signature(self) -> 'Signature': return Signature.build(x=self.bitsize) - def build_call_graph(self, ssa: Optional['SympySymbolAllocator']) -> Set['BloqCountT']: - return {(SubBloq(unrelated_param=0.5), sympy.log(self.bitsize))} + def build_call_graph(self, ssa: Optional['SympySymbolAllocator']) -> 'BloqCountDictT': + return {SubBloq(unrelated_param=0.5): sympy.log(self.bitsize)} @frozen @@ -71,8 +74,8 @@ class SubBloq(Bloq): def signature(self) -> 'Signature': return Signature.build(q=1) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return {(TGate(), 3)} + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {TGate(): 3} def get_big_bloq_counts_graph_1(bloq: Bloq) -> Tuple[nx.DiGraph, Dict[Bloq, SymbolicInt]]: @@ -148,8 +151,11 @@ class OnlyCallGraphBloqShim(Bloq): def signature(self) -> 'Signature': return Signature([]) - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return set(self.callees) + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + counts: 'MutableBloqCountDictT' = defaultdict(int) + for bloq, count in self.callees: + counts[bloq] += count + return counts def __str__(self): return self.name diff --git a/qualtran/resource_counting/classify_bloqs_test.py b/qualtran/resource_counting/classify_bloqs_test.py index dc0ae3db1..cd51f6ca6 100644 --- a/qualtran/resource_counting/classify_bloqs_test.py +++ b/qualtran/resource_counting/classify_bloqs_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Set, Tuple, TYPE_CHECKING +from typing import Tuple import attrs import numpy as np @@ -30,7 +30,13 @@ from qualtran.bloqs.state_preparation.prepare_uniform_superposition import ( PrepareUniformSuperposition, ) -from qualtran.resource_counting import BloqCountT, get_cost_value, QECGatesCost +from qualtran.resource_counting import ( + BloqCountDictT, + BloqCountT, + get_cost_value, + QECGatesCost, + SympySymbolAllocator, +) from qualtran.resource_counting.classify_bloqs import ( _get_basic_bloq_classification, bloq_is_rotation, @@ -39,9 +45,6 @@ classify_t_count_by_bloq_type, ) -if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator - @attrs.frozen class TestBundleOfBloqs(Bloq): @@ -53,8 +56,8 @@ class TestBundleOfBloqs(Bloq): def signature(self) -> 'Signature': return Signature.build() - def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: - return set(self.bloqs) + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return dict(self.bloqs) @pytest.mark.parametrize(