From da80676cbe365a2aa7972f9a480341b44aef415f Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 24 Apr 2024 11:53:44 -0700
Subject: [PATCH 1/9] costing framework

---
 qualtran/_infra/bloq.py                     |  16 +-
 qualtran/bloqs/for_testing/costing.py       |  56 ++++++
 qualtran/bloqs/for_testing/costing_test.py  |  32 ++++
 qualtran/resource_counting/__init__.py      |   4 +-
 qualtran/resource_counting/_call_graph.py   |  11 +-
 qualtran/resource_counting/_costing.py      | 197 ++++++++++++++++++++
 qualtran/resource_counting/_costing_test.py | 112 +++++++++++
 7 files changed, 422 insertions(+), 6 deletions(-)
 create mode 100644 qualtran/bloqs/for_testing/costing.py
 create mode 100644 qualtran/bloqs/for_testing/costing_test.py
 create mode 100644 qualtran/resource_counting/_costing.py
 create mode 100644 qualtran/resource_counting/_costing_test.py

diff --git a/qualtran/_infra/bloq.py b/qualtran/_infra/bloq.py
index 8ef231b82..350b79777 100644
--- a/qualtran/_infra/bloq.py
+++ b/qualtran/_infra/bloq.py
@@ -37,7 +37,7 @@
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, GeneralizerT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountT, CostKey, GeneralizerT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -294,6 +294,20 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         """
         return self.decompose_bloq().build_call_graph(ssa)
 
+    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+        """Override this method to provide static costs.
+
+        The system will query a particular cost by asking for a `cost_key`. This method
+        can optionally provide a value, which will be preferred over a computed cost.
+
+        Static costs can be provided if the particular cost cannot be easily computed or
+        as a performance optimization.
+
+        This method must return `NotImplemented` if a value cannot be provided for the specified
+        CostKey.
+        """
+        return NotImplemented
+
     def call_graph(
         self,
         generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
new file mode 100644
index 000000000..5d793d4e2
--- /dev/null
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -0,0 +1,56 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import Any, Sequence, Set, Tuple, Union
+
+from attrs import field, frozen
+
+from qualtran import Bloq, Signature
+from qualtran.resource_counting import BloqCountT, CostKey, SympySymbolAllocator
+
+
+@frozen
+class CostingBloq(Bloq):
+    """A bloq that lets you set the costs via attributes."""
+
+    name: str
+    num_qubits: int
+    callees: Sequence[BloqCountT] = field(converter=tuple, factory=tuple)
+    static_costs: Sequence[Tuple[CostKey, Any]] = field(converter=tuple, factory=tuple)
+
+    @property
+    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 my_static_costs(self, k: 'CostKey') -> Union[Any, NotImplemented]:
+        return dict(self.static_costs).get(k, NotImplemented)
+
+    def pretty_name(self):
+        return self.name
+
+    def __str__(self):
+        return self.name
+
+
+def make_example_costing_bloqs():
+    from qualtran.bloqs.basic_gates import Hadamard, TGate, Toffoli
+
+    func1 = CostingBloq(
+        'Func1', num_qubits=10, callees=[(TGate(), 10), (TGate().adjoint(), 10), (Hadamard(), 10)]
+    )
+    func2 = CostingBloq('Func2', num_qubits=3, callees=[(Toffoli(), 100)])
+    algo = CostingBloq('Algo', num_qubits=100, callees=[(func1, 1), (func2, 1)])
+    return algo
diff --git a/qualtran/bloqs/for_testing/costing_test.py b/qualtran/bloqs/for_testing/costing_test.py
new file mode 100644
index 000000000..b6b43598b
--- /dev/null
+++ b/qualtran/bloqs/for_testing/costing_test.py
@@ -0,0 +1,32 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from qualtran.bloqs.for_testing.costing import make_example_costing_bloqs
+from qualtran.resource_counting import format_call_graph_debug_text
+
+
+def test_costing_bloqs():
+    algo = make_example_costing_bloqs()
+    g, _ = algo.call_graph()
+    assert (
+        format_call_graph_debug_text(g)
+        == """\
+Algo -- 1 -> Func1
+Algo -- 1 -> Func2
+Func1 -- 10 -> Hadamard()
+Func1 -- 10 -> TGate()
+Func1 -- 10 -> TGate(is_adjoint=True)
+Func2 -- 100 -> Toffoli()
+Toffoli() -- 4 -> TGate()"""
+    )
diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py
index dc77fd876..0f9d1885f 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -25,8 +25,10 @@
     SympySymbolAllocator,
     get_bloq_callee_counts,
     get_bloq_call_graph,
-    print_counts_graph,
     build_cbloq_call_graph,
+    format_call_graph_debug_text,
 )
 
+from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
+
 from . import generalizers
diff --git a/qualtran/resource_counting/_call_graph.py b/qualtran/resource_counting/_call_graph.py
index 79c180ca1..864310c41 100644
--- a/qualtran/resource_counting/_call_graph.py
+++ b/qualtran/resource_counting/_call_graph.py
@@ -240,8 +240,11 @@ def get_bloq_call_graph(
     return g, sigma
 
 
-def print_counts_graph(g: nx.DiGraph):
+def format_call_graph_debug_text(g: nx.DiGraph) -> str:
     """Print the graph returned from `get_bloq_counts_graph`."""
-    for b in nx.topological_sort(g):
-        for succ in g.succ[b]:
-            print(b, '--', g.edges[b, succ]['n'], '->', succ)
+    lines = []
+    for gen in nx.topological_generations(g):
+        for b in sorted(gen, key=str):
+            for succ in sorted(g.succ[b], key=str):
+                lines.append(f"{b} -- {g.edges[b, succ]['n']} -> {succ}")
+    return '\n'.join(lines)
diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
new file mode 100644
index 000000000..19372f7ba
--- /dev/null
+++ b/qualtran/resource_counting/_costing.py
@@ -0,0 +1,197 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import abc
+import logging
+import time
+from collections import defaultdict
+from typing import (
+    Callable,
+    Dict,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    TypeVar,
+    Union,
+)
+
+from ._generalization import _make_composite_generalizer, GeneralizerT
+
+if TYPE_CHECKING:
+    from qualtran import Bloq
+
+logger = logging.getLogger(__name__)
+
+CostValT = TypeVar('CostValT')
+
+
+class CostKey(Generic[CostValT], metaclass=abc.ABCMeta):
+    """Abstract base class for different types of costs.
+
+    One important aspect of a bloq is the resources required to execute it on an error
+    corrected quantum computer. Since we're usually trying to minimize these resource requirements
+    we will generally use the catch-all term "costs".
+
+    There are a variety of different types or flavors of costs. Each is represented by an
+    instance of a sublcass of `CostKey`. For example, gate counts (including T-gate counts),
+    qubit requirements, and circuit depth are all cost metrics that may be of interest.
+
+    Each `CostKey` primarily encodes the behavior required to compute a cost value from a
+    bloq. Often, these costs are defined recursively: a bloq's costs is some combination
+    of the costs of the bloqs in its decomposition (i.e. the bloq 'callees'). Implementors
+    must override the `compute` method to define the cost computation.
+
+    Each cost key has an associated CostValT. For example, the CostValT of a "t count"
+    CostKey could be an integer. For a more complicated gateset, the value could be a mapping
+    from gate to count. This abstract base class is generic w.r.t. `CostValT`. Subclasses
+    should have a concrete value type. The `validate_val` method can optionally be overridden
+    to raise an exception if a bad value type is encountered. The `zero` method must return
+    the zero (additive identity) cost value of the correct type.
+    """
+
+    @abc.abstractmethod
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], CostValT]) -> CostValT:
+        """Compute this type of cost.
+
+        When implementing a new CostKey, this method must be overridden.
+        Users should not call this method directly. Instead: use the `qualtran.resource_counting`
+        functions like `get_cost_value`, `get_cost_cache`, or `query_costs`. These provide
+        caching, logging, generalizers, and support for static costs.
+
+        For recursive computations, use the provided callable to recurse.
+
+        Args:
+            bloq: The bloq to compute the cost of.
+            get_callee_cost: A qualtran-provided function for computing costs for "callees"
+                of the bloq; i.e. bloqs in the decomposition. Use this function to accurately
+                cache intermediate cost values and respect bloqs' static costs.
+
+        Returns:
+            A value of the generic type `CostValT`. Subclasses should define their value type.
+        """
+
+    @abc.abstractmethod
+    def zero(self) -> CostValT:
+        """The value corresponding to zero cost."""
+
+    def validate_val(self, val: CostValT):
+        """Assert that `val` is a valid `CostValT`.
+
+        This method can be optionally overridden to raise an error if an invalid value
+        is encountered. By default, no validation is performed.
+        """
+
+
+def _get_cost_value(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    *,
+    costs_cache: Dict['Bloq', CostValT],
+    generalizer: 'GeneralizerT',
+) -> CostValT:
+    """Helper function for `query_costs`.
+
+    This function tries the following strategies
+     1. Use the value found in `costs_cache`, if it exists.
+     2. Use the value returned by `Bloq.my_static_costs` if one is returned.
+     3. Use `cost_key.compute()` and cache the result in `costs_cache`.
+
+    Args:
+        bloq: The bloq.
+        cost_key: The cost key to get the value for.
+        costs_cache: A dictionary to use as a cache for computed bloq costs. This cache
+            will be mutated by this function.
+        generalizer: The generalizer to operate on each bloq before computing its cost.
+    """
+    bloq = generalizer(bloq)
+    if bloq is None:
+        return cost_key.zero()
+
+    # Strategy 1: Use cached value
+    if bloq in costs_cache:
+        logger.debug("Using cached %s for %s", cost_key, bloq)
+        return costs_cache[bloq]
+
+    # Strategy 2: Static costs
+    static_cost = bloq.my_static_costs(cost_key)
+    if static_cost is not NotImplemented:
+        cost_key.validate_val(static_cost)
+        logger.info("Using static %s for %s", cost_key, bloq)
+        costs_cache[bloq] = static_cost
+        return static_cost
+
+    # Strategy 3: Compute
+    # part a. set up caching of computed costs. Using the callable will use the cache if possible
+    #         and only recurse if the bloq has not been seen before. The result of a computation
+    #         will be cached.
+    def _get_cost_val_internal(callee: 'Bloq'):
+        return _get_cost_value(callee, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+
+    tstart = time.perf_counter()
+    computed_cost = cost_key.compute(bloq, _get_cost_val_internal)
+    tdur = time.perf_counter() - tstart
+    logger.info("Computed %s for %s in %g s", cost_key, bloq, tdur)
+    costs_cache[bloq] = computed_cost
+    return computed_cost
+
+
+def get_cost_value(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    costs_cache: Optional[Dict['Bloq', CostValT]] = None,
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> CostValT:
+    if costs_cache is None:
+        costs_cache = {}
+    if generalizer is None:
+        generalizer = lambda b: b
+    if isinstance(generalizer, (list, tuple)):
+        generalizer = _make_composite_generalizer(*generalizer)
+
+    cost_val = _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+    return cost_val
+
+
+def get_cost_cache(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    costs_cache: Optional[Dict['Bloq', CostValT]] = None,
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> Dict['Bloq', CostValT]:
+    if costs_cache is None:
+        costs_cache = {}
+    if generalizer is None:
+        generalizer = lambda b: b
+    if isinstance(generalizer, (list, tuple)):
+        generalizer = _make_composite_generalizer(*generalizer)
+
+    _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+    return costs_cache
+
+
+def query_costs(
+    bloq: 'Bloq',
+    cost_keys: Iterable[CostKey],
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> Dict['Bloq', Dict[CostKey, CostValT]]:
+
+    costs = defaultdict(dict)
+    for cost_key in cost_keys:
+        cost_for_bloqs = get_cost_cache(bloq, cost_key, generalizer=generalizer)
+        for bloq, val in cost_for_bloqs.items():
+            costs[bloq][cost_key] = val
+    return dict(costs)
diff --git a/qualtran/resource_counting/_costing_test.py b/qualtran/resource_counting/_costing_test.py
new file mode 100644
index 000000000..0fc9f7ac0
--- /dev/null
+++ b/qualtran/resource_counting/_costing_test.py
@@ -0,0 +1,112 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import Callable, List
+
+import attrs
+
+from qualtran import Bloq
+from qualtran.bloqs.basic_gates import Hadamard, TGate
+from qualtran.bloqs.for_testing.costing import CostingBloq, make_example_costing_bloqs
+from qualtran.resource_counting import (
+    CostKey,
+    CostValT,
+    get_bloq_callee_counts,
+    get_cost_cache,
+    get_cost_value,
+)
+from qualtran.resource_counting.generalizers import generalize_rotation_angle
+
+
+class TestCostKey(CostKey[int]):
+    def __init__(self):
+        # For testing, keep a log of all the bloqs for which we called 'compute' on.
+        self._log: List[Bloq] = []
+
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], int]) -> int:
+        self._log.append(bloq)
+
+        total = 1
+        for callee, n_times_called in get_bloq_callee_counts(bloq):
+            total += n_times_called * get_callee_cost(callee)
+
+        return total
+
+    def zero(self) -> CostValT:
+        return 0
+
+    def __hash__(self):
+        return hash(self.__class__)
+
+    def __eq__(self, other):
+        return isinstance(other, self.__class__)
+
+
+def test_get_cost_value_caching():
+    cost = TestCostKey()
+    algo = make_example_costing_bloqs()
+    assert isinstance(algo, CostingBloq)
+    _ = get_cost_value(algo, cost)
+    n_times_compute_called_on_t = sum(b == TGate() for b in cost._log)
+    assert n_times_compute_called_on_t == 1, 'should use cached value'
+
+
+def test_get_cost_value_static():
+    algo = make_example_costing_bloqs()
+
+    # Modify "func1" to have static costs
+    func1 = algo.callees[0][0]
+    func1_mod = attrs.evolve(func1, static_costs=[(TestCostKey(), 123)])
+    algo_mod = attrs.evolve(algo, callees=[(func1_mod, 1), algo.callees[1]])
+    assert get_cost_value(func1_mod, TestCostKey()) == 123
+
+    # Should not call "compute" for Func1, since it has static costs
+    # Should not have to recurse into H, T^dag; which is only used by Func1
+    cost = TestCostKey()
+    _ = get_cost_value(algo_mod, cost)
+    assert len(cost._log) == 4
+    assert 'Func2' in [str(b) for b in cost._log]
+    assert 'Func1' not in [str(b) for b in cost._log]
+    assert TGate().adjoint() not in cost._log
+    assert Hadamard() not in cost._log
+
+
+def test_get_cost_value_static_user_provided():
+    cost = TestCostKey()
+    algo = make_example_costing_bloqs()
+
+    # Provide cached costs up front for func1
+    func1 = algo.callees[0][0]
+
+    # Should not call "compute" for Func1, since we supplied an existing cache
+    # Should not have to recurse into H, T^dag; which is only used by Func1
+    _ = get_cost_value(algo, cost, costs_cache={func1: 0})
+    assert len(cost._log) == 4
+    assert 'Func2' in [str(b) for b in cost._log]
+    assert 'Func1' not in [str(b) for b in cost._log]
+    assert TGate().adjoint() not in cost._log
+    assert Hadamard() not in cost._log
+
+
+def test_costs_generalizer():
+    assert generalize_rotation_angle(TGate().adjoint()) == TGate()
+
+    algo = CostingBloq(name='algo', num_qubits=1, callees=[(TGate(), 1), (TGate().adjoint(), 1)])
+    cost_cache = get_cost_cache(algo, TestCostKey())
+    assert cost_cache[algo] == 3
+
+    cost_cache_gen = get_cost_cache(algo, TestCostKey(), generalizer=generalize_rotation_angle)
+    assert TGate() in cost_cache_gen
+    assert TGate().adjoint() not in cost_cache_gen
+    assert cost_cache_gen[algo] == 3

From 79b85b1b387d77772d01c51543d5fd8f2671a4f1 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Tue, 30 Apr 2024 15:03:12 -0700
Subject: [PATCH 2/9] docstrings and linting

---
 qualtran/bloqs/for_testing/costing.py  |  4 +-
 qualtran/resource_counting/_costing.py | 56 +++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index 5d793d4e2..0350feed2 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -35,8 +35,8 @@ def signature(self) -> 'Signature':
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return set(self.callees)
 
-    def my_static_costs(self, k: 'CostKey') -> Union[Any, NotImplemented]:
-        return dict(self.static_costs).get(k, NotImplemented)
+    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+        return dict(self.static_costs).get(cost_key, NotImplemented)
 
     def pretty_name(self):
         return self.name
diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
index 19372f7ba..28a1dfe0f 100644
--- a/qualtran/resource_counting/_costing.py
+++ b/qualtran/resource_counting/_costing.py
@@ -23,7 +23,6 @@
     Iterable,
     Optional,
     Sequence,
-    Tuple,
     TYPE_CHECKING,
     TypeVar,
     Union,
@@ -103,7 +102,7 @@ def _get_cost_value(
     costs_cache: Dict['Bloq', CostValT],
     generalizer: 'GeneralizerT',
 ) -> CostValT:
-    """Helper function for `query_costs`.
+    """Helper function for getting costs.
 
     This function tries the following strategies
      1. Use the value found in `costs_cache`, if it exists.
@@ -155,6 +154,22 @@ def get_cost_value(
     costs_cache: Optional[Dict['Bloq', CostValT]] = None,
     generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
 ) -> CostValT:
+    """Compute the specified cost of the provided bloq.
+
+    Args:
+        bloq: The bloq to compute the cost of.
+        cost_key: A CostKey that specifies which cost to compute.
+        costs_cache: If provided, use this dictionary of cached cost values. Values in this
+            dictionary will be preferred over computed values (even if they disagree). This
+            dictionary will be mutated by the function.
+        generalizer: If provided, run this function on each bloq in the call graph to dynamically
+            modify attributes. If the function returns `None`, the bloq is ignored in the
+            cost computation. If a sequence of generalizers is provided, each generalizer
+            will be run in order.
+
+    Returns:
+        The cost value. Its type depends on the provided `cost_key`.
+    """
     if costs_cache is None:
         costs_cache = {}
     if generalizer is None:
@@ -172,6 +187,26 @@ def get_cost_cache(
     costs_cache: Optional[Dict['Bloq', CostValT]] = None,
     generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
 ) -> Dict['Bloq', CostValT]:
+    """Build a cache of cost values for the bloq and its callees.
+
+    This can be useful to inspect how callees' costs flow upwards in a given cost computation.
+
+    Args:
+        bloq: The bloq to seed the cost computation.
+        cost_key: A CostKey that specifies which cost to compute.
+        costs_cache: If provided, use this dictionary for initial cached cost values. Values in this
+            dictionary will be preferred over computed values (even if they disagree). This
+            dictionary will be mutated by the function. This dictionary will be returned by the
+            function.
+        generalizer: If provided, run this function on each bloq in the call graph to dynamically
+            modify attributes. If the function returns `None`, the bloq is ignored in the
+            cost computation. If a sequence of generalizers is provided, each generalizer
+            will be run in order.
+
+    Returns:
+        A dictionary mapping bloqs to cost values. The value type depends on the `cost_key`.
+        The bloqs in the mapping depend on the recursive nature of the cost key.
+    """
     if costs_cache is None:
         costs_cache = {}
     if generalizer is None:
@@ -188,7 +223,24 @@ def query_costs(
     cost_keys: Iterable[CostKey],
     generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
 ) -> Dict['Bloq', Dict[CostKey, CostValT]]:
+    """Compute a selection of costs for a bloq and its callees.
+
+    This function can be used to annotate a call graph diagram with multiple costs
+    for each bloq. Specifically, the return value of this function can be used as the
+    `bloq_data` argument to `GraphvizCallGraph`.
 
+    Args:
+        bloq: The bloq to seed the cost computation.
+        cost_key: A sequence of CostKey that specifies which costs to compute.
+        generalizer: If provided, run this function on each bloq in the call graph to dynamically
+            modify attributes. If the function returns `None`, the bloq is ignored in the
+            cost computation. If a sequence of generalizers is provided, each generalizer
+            will be run in order.
+
+    Returns:
+        A dictionary of dictionaries forming a table of multiple costs for multiple bloqs.
+        This is indexed by bloq, then cost key.
+    """
     costs = defaultdict(dict)
     for cost_key in cost_keys:
         cost_for_bloqs = get_cost_cache(bloq, cost_key, generalizer=generalizer)

From a3a2f094a9224b1cd0153e24f7c1a9546bdf3e01 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Tue, 30 Apr 2024 16:19:16 -0700
Subject: [PATCH 3/9] lint

---
 qualtran/resource_counting/_costing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
index 28a1dfe0f..1909bcaf0 100644
--- a/qualtran/resource_counting/_costing.py
+++ b/qualtran/resource_counting/_costing.py
@@ -231,7 +231,7 @@ def query_costs(
 
     Args:
         bloq: The bloq to seed the cost computation.
-        cost_key: A sequence of CostKey that specifies which costs to compute.
+        cost_keys: A sequence of CostKey that specifies which costs to compute.
         generalizer: If provided, run this function on each bloq in the call graph to dynamically
             modify attributes. If the function returns `None`, the bloq is ignored in the
             cost computation. If a sequence of generalizers is provided, each generalizer

From cce03f3bf335a48ac5133b1c4b7b6144f8ddc18d Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 1 May 2024 15:01:30 -0700
Subject: [PATCH 4/9] python 3.10 type issues

---
 qualtran/_infra/bloq.py                                    | 2 +-
 qualtran/bloqs/for_testing/costing.py                      | 2 +-
 qualtran/bloqs/phase_estimation/lp_resource_state.py       | 2 +-
 qualtran/bloqs/rotations/quantum_variable_rotation_test.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/qualtran/_infra/bloq.py b/qualtran/_infra/bloq.py
index 350b79777..da74b99b9 100644
--- a/qualtran/_infra/bloq.py
+++ b/qualtran/_infra/bloq.py
@@ -294,7 +294,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         """
         return self.decompose_bloq().build_call_graph(ssa)
 
-    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+    def my_static_costs(self, cost_key: 'CostKey'):
         """Override this method to provide static costs.
 
         The system will query a particular cost by asking for a `cost_key`. This method
diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index 0350feed2..a2bd19401 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -35,7 +35,7 @@ def signature(self) -> 'Signature':
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return set(self.callees)
 
-    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+    def my_static_costs(self, cost_key: 'CostKey'):
         return dict(self.static_costs).get(cost_key, NotImplemented)
 
     def pretty_name(self):
diff --git a/qualtran/bloqs/phase_estimation/lp_resource_state.py b/qualtran/bloqs/phase_estimation/lp_resource_state.py
index 7cc129c70..5343d7d8f 100644
--- a/qualtran/bloqs/phase_estimation/lp_resource_state.py
+++ b/qualtran/bloqs/phase_estimation/lp_resource_state.py
@@ -19,7 +19,7 @@
 import attrs
 import cirq
 import numpy as np
-from numpy._typing import NDArray
+from numpy.typing import NDArray
 
 from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, QUInt, Register, Side, Signature
 from qualtran.bloqs.basic_gates import CZPowGate, GlobalPhase, Hadamard, OnEach, Ry, Rz, XGate
diff --git a/qualtran/bloqs/rotations/quantum_variable_rotation_test.py b/qualtran/bloqs/rotations/quantum_variable_rotation_test.py
index 6630877f6..ddd1f5254 100644
--- a/qualtran/bloqs/rotations/quantum_variable_rotation_test.py
+++ b/qualtran/bloqs/rotations/quantum_variable_rotation_test.py
@@ -16,7 +16,7 @@
 import numpy as np
 import pytest
 from fxpmath import Fxp
-from numpy._typing import NDArray
+from numpy.typing import NDArray
 
 from qualtran import GateWithRegisters, QFxp, Register, Signature
 from qualtran.bloqs.rotations.phase_gradient import PhaseGradientUnitary

From 5b9ad8b601fdeadebeb8822d02152f6882541459 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Tue, 7 May 2024 16:18:33 -0700
Subject: [PATCH 5/9] lint

---
 qualtran/bloqs/for_testing/costing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index a2bd19401..fa2fc073f 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.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 Any, Sequence, Set, Tuple, Union
+from typing import Any, Sequence, Set, Tuple
 
 from attrs import field, frozen
 

From 88ea840359ac0700a177b464374ff77e553f2496 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Thu, 9 May 2024 15:01:22 -0700
Subject: [PATCH 6/9] part b

---
 qualtran/resource_counting/_costing.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
index 1909bcaf0..f51c3e2f3 100644
--- a/qualtran/resource_counting/_costing.py
+++ b/qualtran/resource_counting/_costing.py
@@ -134,12 +134,11 @@ def _get_cost_value(
         return static_cost
 
     # Strategy 3: Compute
-    # part a. set up caching of computed costs. Using the callable will use the cache if possible
-    #         and only recurse if the bloq has not been seen before. The result of a computation
-    #         will be cached.
+    # part a. set up caching of computed costs by currying the costs_cache.
     def _get_cost_val_internal(callee: 'Bloq'):
         return _get_cost_value(callee, cost_key, costs_cache=costs_cache, generalizer=generalizer)
 
+    # part b. call the compute method and cache the result.
     tstart = time.perf_counter()
     computed_cost = cost_key.compute(bloq, _get_cost_val_internal)
     tdur = time.perf_counter() - tstart

From d3c54610aba95f4ec0ee24c63bfec39efcd477a2 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Thu, 9 May 2024 15:22:11 -0700
Subject: [PATCH 7/9] copyright

---
 qualtran/bloqs/for_testing/costing.py       | 2 +-
 qualtran/bloqs/for_testing/costing_test.py  | 2 +-
 qualtran/resource_counting/_costing.py      | 2 +-
 qualtran/resource_counting/_costing_test.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index fa2fc073f..dc70f3d9f 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
diff --git a/qualtran/bloqs/for_testing/costing_test.py b/qualtran/bloqs/for_testing/costing_test.py
index b6b43598b..fb8340b74 100644
--- a/qualtran/bloqs/for_testing/costing_test.py
+++ b/qualtran/bloqs/for_testing/costing_test.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
index f51c3e2f3..8703676c6 100644
--- a/qualtran/resource_counting/_costing.py
+++ b/qualtran/resource_counting/_costing.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
diff --git a/qualtran/resource_counting/_costing_test.py b/qualtran/resource_counting/_costing_test.py
index 0fc9f7ac0..cb5436300 100644
--- a/qualtran/resource_counting/_costing_test.py
+++ b/qualtran/resource_counting/_costing_test.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.

From 2c81d448113756e19cc11e242a320a8c5f3135db Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Thu, 9 May 2024 15:30:00 -0700
Subject: [PATCH 8/9] mypy

---
 qualtran/bloqs/for_testing/costing.py       | 7 ++++++-
 qualtran/resource_counting/_call_graph.py   | 4 ++--
 qualtran/resource_counting/_costing.py      | 7 ++++---
 qualtran/resource_counting/_costing_test.py | 2 +-
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index dc70f3d9f..cc0e24114 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -19,13 +19,18 @@
 from qualtran.resource_counting import BloqCountT, CostKey, SympySymbolAllocator
 
 
+def _convert_callees(callees: Sequence[BloqCountT]) -> Tuple[BloqCountT, ...]:
+    # Convert to tuples in a type-checked way.
+    return tuple(callees)
+
+
 @frozen
 class CostingBloq(Bloq):
     """A bloq that lets you set the costs via attributes."""
 
     name: str
     num_qubits: int
-    callees: Sequence[BloqCountT] = field(converter=tuple, factory=tuple)
+    callees: Sequence[BloqCountT] = field(converter=_convert_callees, factory=tuple)
     static_costs: Sequence[Tuple[CostKey, Any]] = field(converter=tuple, factory=tuple)
 
     @property
diff --git a/qualtran/resource_counting/_call_graph.py b/qualtran/resource_counting/_call_graph.py
index 09c6d56b2..0e6017a97 100644
--- a/qualtran/resource_counting/_call_graph.py
+++ b/qualtran/resource_counting/_call_graph.py
@@ -14,7 +14,7 @@
 
 """Functionality for the `Bloq.call_graph()` protocol."""
 
-import collections.abc as abc
+import collections.abc
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
@@ -231,7 +231,7 @@ def get_bloq_call_graph(
         keep = lambda b: False
     if generalizer is None:
         generalizer = lambda b: b
-    if isinstance(generalizer, abc.Sequence):
+    if isinstance(generalizer, collections.abc.Sequence):
         generalizer = _make_composite_generalizer(*generalizer)
 
     g = nx.DiGraph()
diff --git a/qualtran/resource_counting/_costing.py b/qualtran/resource_counting/_costing.py
index 8703676c6..da866d024 100644
--- a/qualtran/resource_counting/_costing.py
+++ b/qualtran/resource_counting/_costing.py
@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 import abc
+import collections
 import logging
 import time
 from collections import defaultdict
@@ -173,7 +174,7 @@ def get_cost_value(
         costs_cache = {}
     if generalizer is None:
         generalizer = lambda b: b
-    if isinstance(generalizer, (list, tuple)):
+    if isinstance(generalizer, collections.abc.Sequence):
         generalizer = _make_composite_generalizer(*generalizer)
 
     cost_val = _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
@@ -210,7 +211,7 @@ def get_cost_cache(
         costs_cache = {}
     if generalizer is None:
         generalizer = lambda b: b
-    if isinstance(generalizer, (list, tuple)):
+    if isinstance(generalizer, collections.abc.Sequence):
         generalizer = _make_composite_generalizer(*generalizer)
 
     _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
@@ -240,7 +241,7 @@ def query_costs(
         A dictionary of dictionaries forming a table of multiple costs for multiple bloqs.
         This is indexed by bloq, then cost key.
     """
-    costs = defaultdict(dict)
+    costs: Dict['Bloq', Dict[CostKey, CostValT]] = defaultdict(dict)
     for cost_key in cost_keys:
         cost_for_bloqs = get_cost_cache(bloq, cost_key, generalizer=generalizer)
         for bloq, val in cost_for_bloqs.items():
diff --git a/qualtran/resource_counting/_costing_test.py b/qualtran/resource_counting/_costing_test.py
index cb5436300..4925e7cc7 100644
--- a/qualtran/resource_counting/_costing_test.py
+++ b/qualtran/resource_counting/_costing_test.py
@@ -43,7 +43,7 @@ def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], int]) -> int
 
         return total
 
-    def zero(self) -> CostValT:
+    def zero(self) -> int:
         return 0
 
     def __hash__(self):

From 32a30a73ec08e5c69f412ae7a08e23039e45d24e Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 15 May 2024 10:11:46 -0700
Subject: [PATCH 9/9] lint

---
 qualtran/resource_counting/_costing_test.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/qualtran/resource_counting/_costing_test.py b/qualtran/resource_counting/_costing_test.py
index 4925e7cc7..30cfbd2ba 100644
--- a/qualtran/resource_counting/_costing_test.py
+++ b/qualtran/resource_counting/_costing_test.py
@@ -21,7 +21,6 @@
 from qualtran.bloqs.for_testing.costing import CostingBloq, make_example_costing_bloqs
 from qualtran.resource_counting import (
     CostKey,
-    CostValT,
     get_bloq_callee_counts,
     get_cost_cache,
     get_cost_value,