Support symbolic decomposition for state prep via alias sampling

qualtran/bloqs/state_preparation/state_preparation_alias_sampling.ipynb

@@ -62,23 +62,6 @@
     "selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;\n",
     "otherwise returning `alt[l]`.\n",
     "\n",
-    "#### Parameters\n",
-    " - `selection_registers`: The input/output registers to prepare the state on (see Signature).\n",
-    " - `keep`: The discretized `keep` probabilities for alias sampling.\n",
-    " - `alt`: The alternate/alias values to swap.\n",
-    " - `mu`: The number of bits to approximate the `keep` probabilities.\n",
-    " - `sum_of_unnormalized_probabilities`: The total of the input unnormalized probabilities, i.e., $\\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property. \n",
-    "\n",
-    "Signature:\n",
-    "    selection: The input/output register $|\\ell\\rangle$ of size lg(L) where the desired\n",
-    "        coefficient state is prepared.\n",
-    "    temp: Work space comprised of sub signature:\n",
-    "        - sigma: A mu-sized register containing uniform probabilities for comparison against\n",
-    "            `keep`.\n",
-    "        - alt: A lg(L)-sized register of alternate indices\n",
-    "        - keep: a mu-sized register of probabilities of keeping the initially sampled index.\n",
-    "        - one bit for the result of the comparison.\n",
-    "\n",
     "This gate corresponds to the following operations:\n",
     " - UNIFORM_L on the selection register\n",
     " - H^mu on the sigma register\n",
@@ -90,6 +73,20 @@
     "Total space will be (2 * log(L) + 2 mu + 1) work qubits + log(L) ancillas for QROM.\n",
     "The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.\n",
     "\n",
+    "#### Registers\n",
+    " - `selection`: The input/output register $|\\mathrm{ind}_l\\rangle$ of size lg(L) where the desired coefficient state is prepared. Default name is 'selection' if the builder methods on the class are used. Or else, users can specify custom named registers\n",
+    " - `sigma_mu`: A mu-sized register containing uniform probabilities for comparison against `keep`.\n",
+    " - `alt`: A lg(L)-sized register of alternate indices\n",
+    " - `keep`: a mu-sized register of probabilities of keeping the initially sampled index.\n",
+    " - `less_than_equal`: one bit for the result of the comparison. \n",
+    "\n",
+    "#### Parameters\n",
+    " - `selection_registers`: The input/output registers to prepare the state on (see Registers section).\n",
+    " - `keep`: The discretized `keep` probabilities for alias sampling.\n",
+    " - `alt`: The alternate/alias values to swap.\n",
+    " - `mu`: The number of bits to approximate the `keep` probabilities.\n",
+    " - `sum_of_unnormalized_probabilities`: The total of the input unnormalized probabilities, i.e., $\\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property. \n",
+    "\n",
     "#### References\n",
     " - [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity](https://arxiv.org/abs/1805.03662). Babbush et. al. (2018). Section III.D. and Figure 11.\n"
    ]

qualtran/bloqs/state_preparation/state_preparation_alias_sampling.py

@@ -20,10 +20,9 @@
 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
 import numpy as np
 from numpy.typing import NDArray
 
@@ -48,7 +47,7 @@
 from qualtran.symbolics import bit_length, is_symbolic, Shaped, slen, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
-    from qualtran import BloqBuilder, SoquetT
+    from qualtran import BloqBuilder, Soquet, SoquetT
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 
 
@@ -81,24 +80,6 @@ class StatePreparationAliasSampling(PrepareOracle):
     selecting `l` uniformly at random and then returning it with probability `keep[l] / 2**mu`;
     otherwise returning `alt[l]`.
 
-    Args:
-        selection_registers: The input/output registers to prepare the state on (see Signature).
-        keep: The discretized `keep` probabilities for alias sampling.
-        alt: The alternate/alias values to swap.
-        mu: The number of bits to approximate the `keep` probabilities.
-        sum_of_unnormalized_probabilities: The total of the input unnormalized probabilities,
-            i.e., $\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property.
-
-    Signature:
-        selection: The input/output register $|\ell\rangle$ of size lg(L) where the desired
-            coefficient state is prepared.
-        temp: Work space comprised of sub signature:
-            - sigma: A mu-sized register containing uniform probabilities for comparison against
-                `keep`.
-            - alt: A lg(L)-sized register of alternate indices
-            - keep: a mu-sized register of probabilities of keeping the initially sampled index.
-            - one bit for the result of the comparison.
-
     This gate corresponds to the following operations:
      - UNIFORM_L on the selection register
      - H^mu on the sigma register
@@ -110,6 +91,23 @@ class StatePreparationAliasSampling(PrepareOracle):
     Total space will be (2 * log(L) + 2 mu + 1) work qubits + log(L) ancillas for QROM.
     The 1 ancilla in work qubits is for the `LessThanEqualGate` followed by coherent swap.
 
+    Registers:
+        selection: The input/output register $|\mathrm{ind}_l\rangle$ of size lg(L) where the desired
+            coefficient state is prepared. Default name is 'selection' if the builder methods on
+            the class are used. Or else, users can specify custom named registers
+        sigma_mu: A mu-sized register containing uniform probabilities for comparison against `keep`.
+        alt: A lg(L)-sized register of alternate indices
+        keep: a mu-sized register of probabilities of keeping the initially sampled index.
+        less_than_equal: one bit for the result of the comparison.
+
+    Args:
+        selection_registers: The input/output registers to prepare the state on (see Registers section).
+        keep: The discretized `keep` probabilities for alias sampling.
+        alt: The alternate/alias values to swap.
+        mu: The number of bits to approximate the `keep` probabilities.
+        sum_of_unnormalized_probabilities: The total of the input unnormalized probabilities,
+            i.e., $\lambda$. This is used as the `PrepareOracle.l1_norm_of_coeffs` property.
+
     References:
         [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity](https://arxiv.org/abs/1805.03662).
         Babbush et. al. (2018). Section III.D. and Figure 11.
@@ -123,6 +121,14 @@ class StatePreparationAliasSampling(PrepareOracle):
     mu: SymbolicInt
     sum_of_unnormalized_probabilities: SymbolicFloat
 
+    def __attrs_post_init__(self):
+        if not is_symbolic(self.mu) and self.mu <= 0:
+            raise ValueError(f"{self.mu=} must be greater than 0.")
+        if len(self.selection_registers) != 1:
+            raise ValueError(
+                f"{type(self)} only supports 1D state preparation. Found multiple {self.selection_registers=}."
+            )
+
     @classmethod
     def from_probabilities(
         cls, unnormalized_probabilities: Sequence[float], *, precision: float = 1.0e-5
@@ -235,21 +241,35 @@ def qrom_bloq(self) -> QROM:
             (self.alternates_bitsize, self.keep_bitsize),
         )
 
-    def decompose_from_registers(
+    def build_composite_bloq(
         self,
-        *,
-        context: cirq.DecompositionContext,
-        **quregs: NDArray[cirq.Qid],  # type:ignore[type-var]
-    ) -> Iterator[cirq.OP_TREE]:
-        yield PrepareUniformSuperposition(self.n_coeff).on(*quregs['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)
+        bb: 'BloqBuilder',
+        sigma_mu: 'SoquetT',
+        alt: 'SoquetT',
+        keep: 'SoquetT',
+        less_than_equal: 'Soquet',
+        **soqs: 'SoquetT',
+    ):
+        selection = soqs.pop(self.selection_registers[0].name)
+        assert not soqs
+        selection = bb.add(PrepareUniformSuperposition(self.n_coeff), target=selection)
+        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 {
+            self.selection_registers[0].name: 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

@@ -35,6 +35,7 @@
 
 def test_state_prep_alias_sampling_autotest(bloq_autotester):
     bloq_autotester(_state_prep_alias)
+    bloq_autotester(_state_prep_alias_symb)
 
 
 def test_sparse_state_prep_alias_sampling_autotest(bloq_autotester):
@@ -89,6 +90,8 @@ def test_state_prep_alias_sampling_symb():
         )
     )
     np.testing.assert_allclose(concrete_t_counts, symb_t_counts, rtol=5e-4)
+    # Ensure the symbolic bloq can decomposes into a composite bloq
+    assert_valid_bloq_decomposition(bloq)
 
 
 def assert_state_preparation_valid_for_coefficient(
@@ -166,11 +169,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)───
                                  │        │           │

qualtran/conftest.py

@@ -103,6 +103,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
         'apply_lth_bloq',
         'linear_combination_block_encoding',
         'phase_block_encoding',
+        'state_prep_alias_symb',  # cannot serialize Shaped
         'sparse_matrix_block_encoding',
         'sparse_matrix_symb_block_encoding',
         'sparse_state_prep_alias_symb',  # cannot serialize Shaped