Support symbolic decomposition for state prep via alias sampling

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 Iterator, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
@@ -30,7 +30,7 @@
 from qualtran import bloq_example, BloqDocSpec, BoundedQUInt, Register, Signature
 from qualtran._infra.gate_with_registers import total_bits
 from qualtran.bloqs.arithmetic import LessThanEqual
-from qualtran.bloqs.basic_gates import CSwap, Hadamard
+from qualtran.bloqs.basic_gates import CSwap, Hadamard, OnEach
 from qualtran.bloqs.block_encoding.lcu_select_and_prepare import PrepareOracle
 from qualtran.bloqs.data_loading.qrom import QROM
 from qualtran.bloqs.state_preparation.prepare_uniform_superposition import (
@@ -209,21 +209,31 @@ def qrom_bloq(self) -> QROM:
             (self.alternates_bitsize, self.keep_bitsize),
         )
 
-    def decompose_from_registers(
-        self,
-        *,
-        context: cirq.DecompositionContext,
-        **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
-    ) -> Iterator[cirq.OP_TREE]:
-        yield PrepareUniformSuperposition(self.n_coeff).on(*quregs['selection'])
+    def build_composite_bloq(self, bb, **soqs):
+        soqs['selection'] = bb.add(
+            PrepareUniformSuperposition(self.n_coeff), target=soqs['selection']
+        )
         if self.mu == 0:
-            return
-        selection, less_than_equal = quregs['selection'], quregs['less_than_equal']
-        sigma_mu, alt, keep = quregs['sigma_mu'], quregs['alt'], quregs['keep']
-        yield cirq.H.on_each(*sigma_mu)
-        yield self.qrom_bloq.on_registers(selection=selection, target0_=alt, target1_=keep)
-        yield LessThanEqual(self.mu, self.mu).on(*keep, *sigma_mu, *less_than_equal)
-        yield CSwap.make_on(ctrl=less_than_equal, x=alt, y=selection)
+            return soqs
+        selection, less_than_equal = soqs['selection'], soqs['less_than_equal']
+        sigma_mu, alt, keep = soqs['sigma_mu'], soqs['alt'], soqs['keep']
+        sigma_mu = bb.add(OnEach(self.mu, Hadamard()), q=sigma_mu)
+        selection, alt, keep = bb.add(
+            self.qrom_bloq, selection=selection, target0_=alt, target1_=keep
+        )
+        keep, sigma_mu, less_than_equal = bb.add(
+            LessThanEqual(self.mu, self.mu), x=keep, y=sigma_mu, target=less_than_equal
+        )
+        less_than_equal, alt, selection = bb.add(
+            CSwap(self.selection_bitsize), ctrl=less_than_equal, x=alt, y=selection
+        )
+        return {
+            'selection': selection,
+            'less_than_equal': less_than_equal,
+            'sigma_mu': sigma_mu,
+            'alt': alt,
+            'keep': keep,
+        }
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return {

qualtran/bloqs/state_preparation/state_preparation_alias_sampling_test.py

@@ -53,6 +53,8 @@ def test_state_prep_alias_sampling_symb():
     # Symbolic T-counts
     symb_t_counts = int(expected_t_count_expr.subs({L: N, sympy.Symbol(r"\epsilon"): epsilon}))
     np.testing.assert_allclose(concrete_t_counts, symb_t_counts, rtol=1e-4)
+    # Ensure the symbolic bloq can decomposes into a composite bloq
+    _ = bloq.decompose_bloq()
 
 
 def assert_state_preparation_valid_for_coefficient(lcu_coefficients: np.ndarray, epsilon: float):
@@ -119,11 +121,11 @@ def test_state_preparation_via_coherent_alias_sampling_diagram():
                     │            │                    │
 selection1: ────────target───────In───────────────────×(y)───
                                  │                    │
-sigma_mu0: ─────────H────────────┼────────In(y)───────┼──────
-                                 │        │           │
-sigma_mu1: ─────────H────────────┼────────In(y)───────┼──────
-                                 │        │           │
-sigma_mu2: ─────────H────────────┼────────In(y)───────┼──────
+sigma_mu0: ─────────H⨂3──────────┼────────In(y)───────┼──────
+                    │            │        │           │
+sigma_mu1: ─────────H⨂3──────────┼────────In(y)───────┼──────
+                    │            │        │           │
+sigma_mu2: ─────────H⨂3──────────┼────────In(y)───────┼──────
                                  │        │           │
 alt0: ───────────────────────────QROM_a───┼───────────×(x)───
                                  │        │           │