SelectSwapQROM revamp and upgrades

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -75,6 +75,7 @@
 import qualtran.bloqs.chemistry.trotter.ising.unitaries
 import qualtran.bloqs.chemistry.trotter.trotterized_unitary
 import qualtran.bloqs.data_loading.qrom
+import qualtran.bloqs.data_loading.select_swap_qrom
 import qualtran.bloqs.factoring.ecc
 import qualtran.bloqs.factoring.mod_exp
 import qualtran.bloqs.hamiltonian_simulation.hamiltonian_simulation_by_gqsp
@@ -498,6 +499,11 @@
         module=qualtran.bloqs.data_loading.qrom,
         bloq_specs=[qualtran.bloqs.data_loading.qrom._QROM_DOC],
     ),
+    NotebookSpecV2(
+        title='SelectSwapQROM',
+        module=qualtran.bloqs.data_loading.select_swap_qrom,
+        bloq_specs=[qualtran.bloqs.data_loading.select_swap_qrom._SELECT_SWAP_QROM_DOC],
+    ),
     NotebookSpecV2(
         title='Block Encoding',
         module=qualtran.bloqs.block_encoding,

docs/bloqs/index.rst

@@ -95,6 +95,7 @@ Bloqs Library
     hubbard_model.ipynb
     multiplexers/apply_gate_to_lth_target.ipynb
     data_loading/qrom.ipynb
+    data_loading/select_swap_qrom.ipynb
     block_encoding.ipynb
     reflection.ipynb
     mcmt/multi_control_multi_target_pauli.ipynb

qualtran/bloqs/chemistry/sparse/prepare.py

@@ -43,6 +43,7 @@
     PrepareUniformSuperposition,
 )
 from qualtran.linalg.lcu_util import preprocess_lcu_coefficients_for_reversible_sampling
+from qualtran.symbolics.math_funcs import ceil, log2
 
 if TYPE_CHECKING:
     from qualtran import Bloq
@@ -313,10 +314,10 @@ def build_qrom_bloq(self) -> 'Bloq':
             (n_n,) * 4 + (1,) * 2 + (n_n,) * 4 + (1,) * 2 + (self.num_bits_state_prep,)
         )
         if self.qroam_block_size is None:
-            block_size = 2 ** find_optimal_log_block_size(self.num_non_zero, sum(target_bitsizes))
+            log_block_sizes = find_optimal_log_block_size(self.num_non_zero, sum(target_bitsizes))
         else:
-            block_size = self.qroam_block_size
-        qrom = SelectSwapQROM(
+            log_block_sizes = ceil(log2(self.qroam_block_size))
+        qrom = SelectSwapQROM.build_from_data(
             self.ind_pqrs[0],
             self.ind_pqrs[1],
             self.ind_pqrs[2],
@@ -331,7 +332,7 @@ def build_qrom_bloq(self) -> 'Bloq':
             self.alt_one_body,
             self.keep,
             target_bitsizes=target_bitsizes,
-            block_size=block_size,
+            log_block_sizes=log_block_sizes,
         )
         return qrom
 

qualtran/bloqs/chemistry/thc/notebook_utils.py

@@ -48,8 +48,8 @@
 
 
 def custom_qroam_repr(self) -> str:
-    target_repr = repr(self._target_bitsizes)
-    return f"SelectSwapQROM(target_bitsizes={target_repr}, block_size={self.block_size})"
+    target_repr = repr(self.target_bitsizes)
+    return f"SelectSwapQROM(target_bitsizes={target_repr}, block_sizes={self.block_sizes})"
 
 
 # TODO: better way of customizing label

qualtran/bloqs/chemistry/thc/prepare.py

@@ -382,7 +382,7 @@ def build_composite_bloq(
         # 2. Make contiguous register from mu and nu and store in register `s`.
         mu, nu, s = bb.add(ToContiguousIndex(log_mu, log_d), mu=mu, nu=nu, s=s)
         # 3. Load alt / keep values
-        qroam = SelectSwapQROM(
+        qroam = SelectSwapQROM.build_from_data(
             *(self.theta, self.alt_theta, self.alt_mu, self.alt_nu, self.keep),
             target_bitsizes=(1, 1, log_mu, log_mu, self.keep_bitsize),
         )
@@ -444,7 +444,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         data_size = self.num_spin_orb // 2 + self.num_mu * (self.num_mu + 1) // 2
         nd = (data_size - 1).bit_length()
         cost_2 = (ToContiguousIndex(nmu, nd), 1)
-        qroam = SelectSwapQROM(
+        qroam = SelectSwapQROM.build_from_data(
             *(self.theta, self.alt_theta, self.alt_mu, self.alt_nu, self.keep),
             target_bitsizes=(1, 1, nmu, nmu, self.keep_bitsize),
         )

qualtran/bloqs/chemistry/writing_algorithms.ipynb

@@ -375,15 +375,15 @@
     "        target_bitsizes = (n_n,) * 4 + (self.num_bits_state_prep,)\n",
     "        ns = self.num_spin_orb // 2\n",
     "        data_size = ns ** 2 + ns**4\n",
-    "        block_size = 2 ** find_optimal_log_block_size(data_size, sum(target_bitsizes))\n",
-    "        qroam = SelectSwapQROM(\n",
+    "        log_block_size = find_optimal_log_block_size(data_size, sum(target_bitsizes))\n",
+    "        qroam = SelectSwapQROM.build_from_data(\n",
     "            self.alt_pqrs[0],\n",
     "            self.alt_pqrs[1],\n",
     "            self.alt_pqrs[2],\n",
     "            self.alt_pqrs[3],\n",
     "            self.keep,\n",
     "            target_bitsizes=target_bitsizes,\n",
-    "            block_size=block_size,\n",
+    "            log_block_sizes=[log_block_size],\n",
     "        )\n",
     "        (l, alt_pqrs[0], alt_pqrs[1], alt_pqrs[2], alt_pqrs[3], keep) = bb.add(\n",
     "            qroam,\n",

qualtran/bloqs/data_loading/qrom.py

@@ -13,10 +13,10 @@
 #  limitations under the License.
 
 """Quantum read-only memory."""
-from functools import cached_property
+import numbers
 from typing import (
     Callable,
-    Dict,
+    cast,
     Iterable,
     Iterator,
     Optional,
@@ -33,18 +33,17 @@
 import sympy
 from numpy.typing import ArrayLike, NDArray
 
-from qualtran import bloq_example, BloqDocSpec, BoundedQUInt, QAny, Register
+from qualtran import bloq_example, BloqDocSpec, Register
 from qualtran._infra.gate_with_registers import merge_qubits
 from qualtran.bloqs.basic_gates import CNOT
+from qualtran.bloqs.data_loading.qrom_base import QROMBase
 from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd
 from qualtran.bloqs.multiplexers.unary_iteration_bloq import UnaryIterationGate
 from qualtran.drawing import Circle, Text, TextBox, WireSymbol
-from qualtran.resource_counting import BloqCountT
-from qualtran.simulation.classical_sim import ClassicalValT
-from qualtran.symbolics import bit_length, is_symbolic, prod, shape, Shaped, SymbolicInt
+from qualtran.symbolics import prod, SymbolicInt
 
 if TYPE_CHECKING:
-    from qualtran.resource_counting import SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 
 
 def _to_tuple(x: Iterable[NDArray]) -> Sequence[NDArray]:
@@ -54,7 +53,7 @@ def _to_tuple(x: Iterable[NDArray]) -> Sequence[NDArray]:
 
 @cirq.value_equality()
 @attrs.frozen
-class QROM(UnaryIterationGate):
+class QROM(QROMBase, UnaryIterationGate):  # type: ignore[misc]
     r"""Bloq to load `data[l]` in the target register when the selection stores an index `l`.
 
     ## Overview
@@ -146,55 +145,15 @@ class QROM(UnaryIterationGate):
             Babbush et. al. (2020). Figure 3.
     """
 
-    data_or_shape: Union[NDArray, Shaped] = attrs.field(
-        converter=lambda x: np.array(x) if isinstance(x, (list, tuple)) else x
-    )
-    selection_bitsizes: Tuple[SymbolicInt, ...] = attrs.field(
-        converter=lambda x: tuple(x.tolist() if isinstance(x, np.ndarray) else x)
-    )
-    target_bitsizes: Tuple[SymbolicInt, ...] = attrs.field(
-        converter=lambda x: tuple(x.tolist() if isinstance(x, np.ndarray) else x)
-    )
-    num_controls: SymbolicInt = 0
-
-    def has_data(self) -> bool:
-        return not isinstance(self.data_or_shape, Shaped)
-
-    @property
-    def data_shape(self) -> Tuple[SymbolicInt, ...]:
-        return shape(self.data_or_shape)[1:]
-
-    @property
-    def data(self) -> np.ndarray:
-        if not self.has_data():
-            raise ValueError(f"Data not available for symbolic QROM {self}")
-        assert isinstance(self.data_or_shape, np.ndarray)
-        return self.data_or_shape
-
-    def __attrs_post_init__(self):
-        assert all([is_symbolic(s) or isinstance(s, int) for s in self.selection_bitsizes])
-        assert all([is_symbolic(t) or isinstance(t, int) for t in self.target_bitsizes])
-        assert len(self.target_bitsizes) == self.data_or_shape.shape[0], (
-            f"len(self.target_bitsizes)={len(self.target_bitsizes)} should be same as "
-            f"len(self.data)={self.data_or_shape.shape[0]}"
-        )
-        if isinstance(self.data_or_shape, np.ndarray) and not is_symbolic(*self.target_bitsizes):
-            assert all(
-                t >= int(np.max(d)).bit_length() for t, d in zip(self.target_bitsizes, self.data)
-            )
-        assert isinstance(self.selection_bitsizes, tuple)
-        assert isinstance(self.target_bitsizes, tuple)
-
     @classmethod
-    def build_from_data(cls, *data: ArrayLike, num_controls: SymbolicInt = 0) -> 'QROM':
-        _data = np.array([np.array(d, dtype=int) for d in data])
-        selection_bitsizes = tuple((s - 1).bit_length() for s in _data.shape[1:])
-        target_bitsizes = tuple(max(int(np.max(d)).bit_length(), 1) for d in data)
-        return QROM(
-            data_or_shape=_data,
-            selection_bitsizes=selection_bitsizes,
-            target_bitsizes=target_bitsizes,
-            num_controls=num_controls,
+    def build_from_data(
+        cls,
+        *data: ArrayLike,
+        target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None,
+        num_controls: SymbolicInt = 0,
+    ) -> 'QROM':
+        return cls._build_from_data(
+            *data, target_bitsizes=target_bitsizes, num_controls=num_controls
         )
 
     @classmethod
@@ -203,59 +162,33 @@ def build_from_bitsize(
         data_len_or_shape: Union[SymbolicInt, Tuple[SymbolicInt, ...]],
         target_bitsizes: Union[SymbolicInt, Tuple[SymbolicInt, ...]],
         *,
+        target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = (),
         selection_bitsizes: Tuple[SymbolicInt, ...] = (),
         num_controls: SymbolicInt = 0,
     ) -> 'QROM':
-        data_shape = (
-            (data_len_or_shape,) if isinstance(data_len_or_shape, int) else data_len_or_shape
-        )
-        if not isinstance(target_bitsizes, tuple):
-            target_bitsizes = (target_bitsizes,)
-        _data = Shaped((len(target_bitsizes),) + data_shape)
-        if selection_bitsizes is ():
-            selection_bitsizes = tuple(bit_length(s - 1) for s in _data.shape[1:])
-        assert len(selection_bitsizes) == len(_data.shape) - 1
-        return QROM(
-            data_or_shape=_data,
+        return cls._build_from_bitsize(
+            data_len_or_shape,
+            target_bitsizes,
+            target_shapes=target_shapes,
             selection_bitsizes=selection_bitsizes,
-            target_bitsizes=target_bitsizes,
             num_controls=num_controls,
         )
 
-    @cached_property
-    def control_registers(self) -> Tuple[Register, ...]:
-        return () if not self.num_controls else (Register('control', QAny(self.num_controls)),)
-
-    @cached_property
-    def selection_registers(self) -> Tuple[Register, ...]:
-        types = [
-            BoundedQUInt(sb, l)
-            for l, sb in zip(self.data_or_shape.shape[1:], self.selection_bitsizes)
-            if is_symbolic(sb) or sb > 0
-        ]
-        if len(types) == 1:
-            return (Register('selection', types[0]),)
-        return tuple(Register(f'selection{i}', qdtype) for i, qdtype in enumerate(types))
-
-    @cached_property
-    def target_registers(self) -> Tuple[Register, ...]:
-        return tuple(
-            Register(f'target{i}_', QAny(l))
-            for i, l in enumerate(self.target_bitsizes)
-            if is_symbolic(l) or l
-        )
-
     def _load_nth_data(
         self,
         selection_idx: Tuple[int, ...],
         gate: Callable[[cirq.Qid], cirq.Operation],
         **target_regs: NDArray[cirq.Qid],  # type: ignore[type-var]
     ) -> Iterator[cirq.OP_TREE]:
         for i, d in enumerate(self.data):
-            target = target_regs.get(f'target{i}_', ())
-            for q, bit in zip(target, f'{int(d[selection_idx]):0{len(target)}b}'):
-                if int(bit):
-                    yield gate(q)
+            target = target_regs.get(f'target{i}_', np.array([]))
+            target_bitsize, target_shape = self.target_bitsizes[i], self.target_shapes[i]
+            assert all(isinstance(x, (int, numbers.Integral)) for x in target_shape)
+            for idx in np.ndindex(cast(Tuple[int, ...], target_shape)):
+                data_to_load = int(d[selection_idx + idx])
+                for q, bit in zip(target[idx], f'{data_to_load:0{target_bitsize}b}'):
+                    if int(bit):
+                        yield gate(q)
 
     def decompose_zero_selection(
         self, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
@@ -301,32 +234,6 @@ def nth_operation(
         target_regs = {reg.name: kwargs[reg.name] for reg in self.target_registers}
         yield self._load_nth_data(selection_idx, lambda q: CNOT().on(control, q), **target_regs)
 
-    def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
-        if not self.has_data():
-            raise NotImplementedError(f'Symbolic {self} does not support classical simulation')
-
-        if self.num_controls > 0:
-            control = vals['control']
-            if control != 2**self.num_controls - 1:
-                return vals
-            controls = {'control': control}
-        else:
-            controls = {}
-
-        n_dim = len(self.selection_bitsizes)
-        if n_dim == 1:
-            idx = vals['selection']
-            selections = {'selection': idx}
-        else:
-            # Multidimensional
-            idx = tuple(vals[f'selection{i}'] for i in range(n_dim))  # type: ignore[assignment]
-            selections = {f'selection{i}': idx[i] for i in range(n_dim)}  # type: ignore[index]
-
-        # Retrieve the data; bitwise add them in to the input target values
-        targets = {f'target{d_i}_': d[idx] for d_i, d in enumerate(self.data)}
-        targets = {k: v ^ vals[k] for k, v in targets.items()}
-        return controls | selections | targets
-
     def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo:
         from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info
 
@@ -352,17 +259,6 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return Circle()
         raise ValueError(f'Unrecognized register name {name}')
 
-    def __pow__(self, power: int):
-        if power in [1, -1]:
-            return self
-        return NotImplemented  # pragma: no cover
-
-    def _value_equality_values_(self):
-        data_tuple = (
-            tuple(tuple(d.flatten()) for d in self.data) if self.has_data() else self.data_or_shape
-        )
-        return (self.selection_registers, self.target_registers, self.control_registers, data_tuple)
-
     def nth_operation_callgraph(self, **kwargs: int) -> Set['BloqCountT']:
         selection_idx = tuple(kwargs[reg.name] for reg in self.selection_registers)
         return {(CNOT(), sum(int(d[selection_idx]).bit_count() for d in self.data))}