Strategy for rendering bloq_counts from a CirqGateAsBloq object?

[fdmalone]

Currently rendering something like QROM as a nice bloq_counts picture requires a bit of hackery (see below).

What I have done below also seems wrong, our bloqs are frozen but I'm overwriting their bloq counts on the fly (which is somewhat related to the question of how one would modify the decomposition of an existing bloq but with one primitive substituted out.) I could have defined a bloq_counts on QROM alternatively but this seems to defeat the purpose of automatically determining the counts from the decomposition.

import numpy as np
from cirq_ft.algos import QROM
from qualtran.cirq_interop import CirqGateAsBloq

from qualtran.resource_counting import get_bloq_counts_graph, GraphvizCounts
from qualtran.bloqs.util_bloqs import Split, Join, Allocate, Free
from cirq_ft.algos.and_gate import And
import sympy

from typing import Any, Dict, Optional, Set, Tuple

from qualtran import Bloq, Register, Side, Signature, Soquet, SoquetT
from qualtran.bloqs.basic_gates import TGate
from qualtran.bloqs.util_bloqs import ArbitraryClifford
from qualtran.resource_counting import big_O, SympySymbolAllocator

def bloq_counts(self, ssa: Optional['SympySymbolAllocator'] = None) -> Set[Tuple[int, Bloq]]:
    cv1, cv2 = self.gate.cv
    if isinstance(cv1, sympy.Expr) or isinstance(cv2, sympy.Expr):
        pre_post_cliffords = big_O(1)
    else:
        pre_post_cliffords = 2 - cv1 - cv2
    if self.gate.adjoint:
        return {(4 + 2 * pre_post_cliffords, ArbitraryClifford(n=2))}

    return {(9 + 2 * pre_post_cliffords, ArbitraryClifford(n=2)), (4, TGate())}
# 
ignore = ['H, XPow', 'Y', 'X', 'Free', 'CNOT', 'T', 'H', 'S', 'Measurement', 'reset']
cirq_ignore = [f'cirq.{g}' for g in ignore]

def generalize(bloq):
    if 'cirq.And' in bloq.pretty_name():
        bloq.__class__.bloq_counts = bloq_counts
    if any([g in bloq.pretty_name() for g in cirq_ignore]):
        return None
    if isinstance(bloq, Allocate):
        return None
    if isinstance(bloq, Free):
        return None
    return bloq

cirq_qrom = QROM.build(*[np.random.randint(0, 10, 100)])
def custom_repr(self):
    selection_repr = repr(self.selection_bitsizes)
    target_repr = repr(self.target_bitsizes)
    return (f"cirq_ft.QROM(selection_bitsizes={selection_repr}, "
        f"target_bitsizes={target_repr}, num_controls={self.num_controls})"
    )
cirq_qrom.__class__.__repr__ = custom_repr
qrom = CirqGateAsBloq(cirq_qrom)
graph, sigma = get_bloq_counts_graph(qrom, generalizer=generalize)
GraphvizCounts(graph).get_svg()

[fdmalone]

[fdmalone]

Modifying the bloq directly with the generalizer is definitely not the way to go.

[fdmalone]

A maybe better way to do it is

from qualtran.bloqs.and_bloq import And
and_cv0 = ssa.new_symbol('cv0')
and_cv1 = ssa.new_symbol('cv1')
def generalize(bloq):
    if ('cirq.And' in bloq.pretty_name()) and (len(bloq.gate.cv) == 2):
        return And(cv1=and_cv0, cv2=and_cv1, adjoint=bloq.gate.adjoint)

[fdmalone]

@mpharrigan is this last way something sensible, or is there a better way?

[mpharrigan]

yes

strictly speaking you may want a more defensible condition than if "cirq.And" in bloq.pretty_name() but that's details

[mpharrigan]

there's almost no CirqGateAsBloq in the call graphs in the standard library. Otherwise, CirqGateAsBloq should use its __str__ method to control how it is displayed