From cdf09f2a2405b35c0f71644505acb4585a3d2650 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Tue, 28 May 2024 13:00:06 -0700 Subject: [PATCH] `SelectSwapQROM` revamp and upgrades (#986) * SwapWithZero learns how to swap multidimensional selection index. * Clear output * Fix mypy and use wire_symbol instead of cirq diagram. Fix a bug in _wire_symbol_to_cirq_diagram_info * SelectSwapQROM revamp and upgrades * Fix pylint and typos * Fix mypy errors and address comments * Add type ignores, fix notebook and nits * Update docstrings and regenerate ipynb files --- dev_tools/autogenerate-bloqs-notebooks-v2.py | 15 +- docs/bloqs/index.rst | 1 + qualtran/bloqs/chemistry/sparse/prepare.py | 9 +- .../bloqs/chemistry/thc/notebook_utils.py | 4 +- qualtran/bloqs/chemistry/thc/prepare.py | 4 +- .../bloqs/chemistry/writing_algorithms.ipynb | 6 +- qualtran/bloqs/data_loading/qrom.ipynb | 214 +++++---- qualtran/bloqs/data_loading/qrom.py | 259 ++-------- qualtran/bloqs/data_loading/qrom_base.py | 326 +++++++++++++ .../bloqs/data_loading/select_swap_qrom.ipynb | 368 +++++++++++++++ .../bloqs/data_loading/select_swap_qrom.py | 443 ++++++++++++------ .../data_loading/select_swap_qrom_test.py | 52 +- 12 files changed, 1218 insertions(+), 483 deletions(-) create mode 100644 qualtran/bloqs/data_loading/qrom_base.py create mode 100644 qualtran/bloqs/data_loading/select_swap_qrom.ipynb diff --git a/dev_tools/autogenerate-bloqs-notebooks-v2.py b/dev_tools/autogenerate-bloqs-notebooks-v2.py index 5bae86bf7..be64f2ee9 100644 --- a/dev_tools/autogenerate-bloqs-notebooks-v2.py +++ b/dev_tools/autogenerate-bloqs-notebooks-v2.py @@ -75,6 +75,8 @@ 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.qrom_base +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 @@ -496,7 +498,18 @@ NotebookSpecV2( title='QROM', module=qualtran.bloqs.data_loading.qrom, - bloq_specs=[qualtran.bloqs.data_loading.qrom._QROM_DOC], + bloq_specs=[ + qualtran.bloqs.data_loading.qrom_base._QROM_BASE_DOC, + qualtran.bloqs.data_loading.qrom._QROM_DOC, + ], + ), + NotebookSpecV2( + title='SelectSwapQROM', + module=qualtran.bloqs.data_loading.select_swap_qrom, + bloq_specs=[ + qualtran.bloqs.data_loading.qrom_base._QROM_BASE_DOC, + qualtran.bloqs.data_loading.select_swap_qrom._SELECT_SWAP_QROM_DOC, + ], ), NotebookSpecV2( title='Block Encoding', diff --git a/docs/bloqs/index.rst b/docs/bloqs/index.rst index a05519bc2..360da844f 100644 --- a/docs/bloqs/index.rst +++ b/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 diff --git a/qualtran/bloqs/chemistry/sparse/prepare.py b/qualtran/bloqs/chemistry/sparse/prepare.py index 86668df7a..79f6929a3 100644 --- a/qualtran/bloqs/chemistry/sparse/prepare.py +++ b/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 diff --git a/qualtran/bloqs/chemistry/thc/notebook_utils.py b/qualtran/bloqs/chemistry/thc/notebook_utils.py index a0cfdd691..c5d70ca19 100644 --- a/qualtran/bloqs/chemistry/thc/notebook_utils.py +++ b/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 diff --git a/qualtran/bloqs/chemistry/thc/prepare.py b/qualtran/bloqs/chemistry/thc/prepare.py index 87eb01f46..b86671140 100644 --- a/qualtran/bloqs/chemistry/thc/prepare.py +++ b/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), ) diff --git a/qualtran/bloqs/chemistry/writing_algorithms.ipynb b/qualtran/bloqs/chemistry/writing_algorithms.ipynb index 2807bf554..b0803c7f6 100644 --- a/qualtran/bloqs/chemistry/writing_algorithms.ipynb +++ b/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", diff --git a/qualtran/bloqs/data_loading/qrom.ipynb b/qualtran/bloqs/data_loading/qrom.ipynb index 389736633..0841947de 100644 --- a/qualtran/bloqs/data_loading/qrom.ipynb +++ b/qualtran/bloqs/data_loading/qrom.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "ffe91e97", + "id": "6a30c548", "metadata": { "cq.autogen": "title_cell" }, @@ -15,7 +15,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2737c79f", + "id": "20c4cf8d", "metadata": { "cq.autogen": "top_imports" }, @@ -32,13 +32,13 @@ }, { "cell_type": "markdown", - "id": "4ea4344b", + "id": "ee1b29d3", "metadata": { - "cq.autogen": "QROM.bloq_doc.md" + "cq.autogen": "QROMBase.bloq_doc.md" }, "source": [ - "## `QROM`\n", - "Bloq to load `data[l]` in the target register when the selection stores an index `l`.\n", + "## `QROMBase`\n", + "Interface for Bloqs to load `data[l]` when the selection register stores index `l`.\n", "\n", "## Overview\n", "The action of a QROM can be described as\n", @@ -53,65 +53,127 @@ " |d_L[s_1, s_2, \\dots, s_k]\\rangle\n", "$$\n", "\n", - "There two high level parameters that control the behavior of a QROM are -\n", + "A behavior of a QROM can be understood in terms of its classical analogue, where a for-loop\n", + "over one or more (selection) indices can be used to load one or more classical datasets, where\n", + "each of the classical dataset can be multidimensional.\n", + "\n", + "```\n", + ">>> # N, M, P, Q, R, S, T are pre-initialized integer parameters.\n", + ">>> output = [np.zeros((P, Q)), np.zeros((R, S, T))]\n", + ">>> # Load two different classical datasets; each of different shape.\n", + ">>> data = [np.random.rand(N, M, P, Q), np.random.rand(N, M, R, S, T)]\n", + ">>> for i in range(N): # For loop over two selection indices i and j.\n", + ">>> for j in range(M):\n", + ">>> # Load two multidimensional classical datasets data[0] and data[1] s.t.\n", + ">>> # |i, j⟩|0⟩ -> |i, j⟩|data[0][i, j, :]⟩|data[1][i, j, :]⟩\n", + ">>> output[0] = data[0][i, j, :]\n", + ">>> output[1] = data[1][i, j, :]\n", + "```\n", + "\n", + "The parameters that control the behavior and costs of a QROM are -\n", + "\n", + "1. Number of selection registers (eg: $i$, $j$) and their iteration lengths (eg: $N$, $M$).\n", + "2. Number of target registers, their quantum datatype and shape.\n", + " - Number of target registers: One for each classical dataset to load (eg: $\\text{data}[0]$\n", + " and $\\text{data}[1]$)\n", + " - QDType of target registers: Depends on `dtype` of the $i$'th classical dataset\n", + " - Shape of target registers: Depends on shape of classical data (eg: $(P, Q)$ and\n", + " $(R, S, T)$ above)\n", + "\n", + "### Specification of classical data via `data_or_shape`\n", + "Users can specify the classical data to load via QROM by passing in an appropriate value\n", + "for `data_or_shape` attribute. This is a list of numpy arrays or `Shaped` objects, where\n", + "each item of the list corresponds to a classical dataset to load.\n", + "\n", + "Each classical dataset to load can be specified as a numpy array (or a `Shaped` object for\n", + "symbolic bloqs). The shape of the dataset is a union of the selection shape and target shape,\n", + "s.t.\n", + "$$\n", + " \\text{data[i].shape} = \\text{selection\\_shape} + \\text{target\\_shape[i]}\n", + "$$\n", + "\n", + "Note that the $\\text{selection\\_shape}$ should be same across all classical datasets to be\n", + "loaded and correspond to a tuple of iteration lengths of selection indices (i.e. $(N, M)$\n", + "in the example above).\n", + "\n", + "The target shape of each classical dataset can be different and parameterizes the size of\n", + "the desired output that should be loaded in a target register.\n", + "\n", + "### Number of selection registers and their iteration lengths\n", + "As describe in the previous section, the number of selection registers and their iteration\n", + "lengths can be inferred from the shape of the classical dataset. All classical datasets\n", + "to be loaded must have the same $\\text{selection\\_shape}$, which is a tuple of iteration\n", + "lengths over each dimension of the dataset (i.e. the range for each nested for-loop).\n", "\n", - "1. Shape of the classical dataset to be loaded ($\\text{data.shape} = (S_1, S_2, ..., S_K)$).\n", - "2. Number of distinct datasets to be loaded ($\\text{data.bitsizes} = (b_1, b_2, ..., b_L)$).\n", + "In order to load a data set with $\\text{selection\\_shape} == (P, Q, R, S)$ the QROM bloq\n", + "needs four selection registers with bitsizes $(p, q, r, s)$ where each of\n", + "$p,q,r,s \\geq \\log_2{P}, \\log_2{Q}, \\log_2{R}, \\log_2{S}$.\n", "\n", - "Each of these have an effect on the cost of the QROM. The `data_or_shape` parameter stores\n", - "either\n", - "1. A numpy array of shape $(L, S_1, S_2, ..., S_K)$ when $L$ classical datasets, each of\n", - " shape $(S_1, S_2, ..., S_K)$ and bitsizes $(b_1, b_2, ..., b_L)$ are to be loaded and\n", - " the classical data is available to instantiate the QROM bloq. In this case, the helper\n", - " builder `QROM.build_from_data(data_1, data_2, ..., data_L)` can be used to build the QROM.\n", + "In general, to load $K$ dimensional data, we use $K$ named selection registers\n", + "$(\\text{selection}_0, \\text{selection}_1, ..., \\text{selection}_k)$ to index and\n", + "load the data. For the $i$'th selection register, its size is configured using\n", + "attribute $\\text{selection\\_bitsizes[i]}$ and the iteration range is configued\n", + "using $\\text{data\\_or\\_shape[0].shape[i]}$.\n", "\n", - "2. A `Shaped` object that stores a (potentially symbolic) tuple $(L, S_1, S_2, ..., S_K)$\n", - " that represents the number of classical datasets `L=data_or_shape.shape[0]` and\n", - " their shape `data_shape=data_or_shape.shape[1:]` to be loaded by this QROM. This is used\n", - " to instantiate QROM bloqs for symbolic cost analysis where the exact data to be loaded\n", - " is not known. In this case, the helper builder `QROM.build_from_bitsize` can be used\n", - " to build the QROM.\n", + "### Number of target registers, their quantum datatype and shape\n", + "QROM bloq uses one target register for each entry corresponding to classical dataset in the\n", + "tuple `data_or_shape`. Thus, to load $L$ classical datsets, we use $L$ names target registers\n", + "$(\\text{target}_0, \\text{target}_1, ..., \\text{target}_L)$\n", "\n", - "### Shape of the classical dataset to be loaded.\n", - "QROM bloq supports loading multidimensional classical datasets. In order to load a data\n", - "set of shape $\\mathrm{data.shape} == (P, Q, R, S)$ the QROM bloq needs four selection\n", - "registers with bitsizes $(p, q, r, s)$ where\n", - "$p,q,r,s=\\log_2{P}, \\log_2{Q}, \\log_2{R}, \\log_2{S}$.\n", + "Each named target register has a bitsize $b_{i}=\\text{target\\_bitsizes[i]}$ that represents\n", + "the size of the register and depends upon the maximum value of individual elements in the\n", + "$i$'th classical dataset.\n", + "\n", + "Each named target register has a shape that can be configured using attribute\n", + "$\\text{target\\_shape[i]}$ that represents the number of target registers if the output to load\n", + "is multidimensional.\n", + "\n", + "#### Parameters\n", + " - `data_or_shape`: List of numpy ndarrays specifying the data to load. If the length of this list ($L$) is greater than one then we use the same selection indices to load each dataset. The shape of a classical dataset is a concatenation of selection_shape and target_shape[i]; i.e. `data_or_shape[i].shape = selection_shape + target_shape[i]`. Thus, each data set is required to have the same selection shape $(S_1, S_2, ..., S_K)$ and can have a different target shape given by `target_shapes[i]`. For symbolic QROMs, pass a list of `Shaped` objects instead with shape $(S_1, S_2, ..., S_K) + target_shape[i]$.\n", + " - `selection_bitsizes`: The number of bits used to represent each selection register corresponding to the size of each dimension of the selection_shape $(S_1, S_2, ..., S_K)$. Should be the same length as the selection shape of each of the datasets and $2**\\text{selection\\_bitsizes[i]} >= S_i$\n", + " - `target_shapes`: Shape of target registers for each classical dataset to be loaded. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_shapes)` and `data_or_shape[-len(target_shapes[i]):] == target_shapes[i]`.\n", + " - `target_bitsizes`: Bitsize (or qdtype) of the target registers for each classical dataset to be loaded. This can be deduced from the maximum element of each of the datasets. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_bitsizes)` and `target_bitsizes[i] >= max(data[i]).bitsize`.\n", + " - `num_controls`: The number of controls to instanstiate a controlled version of this bloq.\n" + ] + }, + { + "cell_type": "markdown", + "id": "9d6450ad", + "metadata": { + "cq.autogen": "QROM.bloq_doc.md" + }, + "source": [ + "## `QROM`\n", + "Bloq to load `data[l]` in the target register when the selection stores an index `l`.\n", "\n", - "In general, to load K dimensional data, we use K named selection registers `(selection0,\n", - "selection1, ..., selection{k})` to index and load the data.\n", + "See docstrings of `QROMBase` for an overview of the QROM primitive and the various attributes.\n", "\n", - "The T/Toffoli cost of the QROM scales linearly with the number of elements in the dataset\n", - "(i.e. $\\mathcal{O}(\\mathrm{np.prod(data.shape)}$).\n", + "This bloq is an implementation of the `QROMBase` interface that uses the unary iteration based\n", + "approach described in Ref [1].\n", "\n", - "### Number of distinct datasets to be loaded, and their corresponding target bitsize.\n", - "To load a classical dataset into a target register of bitsize $b$, the clifford cost of a QROM\n", - "scales as $\\mathcal{O}(b \\mathrm{np.prod}(\\mathrm{data.shape}))$. This is because we need\n", - "$\\mathcal{O}(b)$ CNOT gates to load the ith data element in the target register when the\n", - "selection register stores index $i$.\n", + "## Cost of this (unary iteration based) QROM\n", "\n", - "If you have multiple classical datasets `(data_1, data_2, data_3, ..., data_L)` to be loaded\n", - "and each of them has the same shape `(data_1.shape == data_2.shape == ... == data_L.shape)`\n", - "and different target bitsizes `(b_1, b_2, ..., b_L)`, then one construct a single classical\n", - "dataset `data = merge(data_1, data_2, ..., data_L)` where\n", + "### T / Toffoli cost\n", + "The T/Toffoli cost of this QROM scales linearly with the product of iteration lengths over\n", + "all dimensions (i.e. $\\mathcal{O}(\\mathrm{np.prod(\\text{selection\\_shape})}$).\n", "\n", - "- `data.shape == data_1.shape == data_2.shape == ... == data_L` and\n", - "- `data[idx] = f'{data_1[idx]!0{b_1}b}' + f'{data_2[idx]!0{b_2}b}' + ... + f'{data_L[idx]!0{b_L}b}'`\n", + "### Clifford Cost\n", + "To load a classical dataset into a target register of bitsize $b$ and shape\n", + "$\\text{target\\_shape}$, the clifford cost of this QROM scales as\n", + "$\\mathcal{O}(b \\cdot \\text{np.prod(selection\\_shape+target\\_shape)})\n", + "=\\mathcal{O}(b \\cdot \\text{np.prod(data.shape)})$. This is because we need $\\mathcal{O}(b)$\n", + "CNOT gates to load 1 classical data element in the target register and for each of the\n", + "$\\text{np.prod(selection\\_shape)}$ selection indices, we have $\\text{np.prod(target\\_shape)}$\n", + "such data elements to load.\n", "\n", - "Thus, the target bitsize of the merged dataset is $b = b_1 + b_2 + \\dots + b_L$ and clifford\n", - "cost of loading merged dataset scales as\n", - "$\\mathcal{O}((b_1 + b_2 + \\dots + b_L) \\mathrm{np.prod}(\\mathrm{data.shape}))$.\n", + "### Ancilla cost\n", + "The number of clean ancilla required by this QROM scales linearly with the size of the\n", + "selection registers + number of controls.\n", "\n", "## Variable spaced QROM\n", "When the input classical data contains consecutive entries of identical data elements to\n", "load, the QROM also implements the \"variable-spaced\" QROM optimization described in Ref [2].\n", "\n", - "#### Parameters\n", - " - `data_or_shape`: List of numpy ndarrays specifying the data to load. If the length of this list ($L$) is greater than one then we use the same selection indices to load each dataset. Each data set is required to have the same shape $(S_1, S_2, ..., S_K)$ and to be of integer type. For symbolic QROMs, pass a `Shaped` object instead with shape $(L, S_1, S_2, ..., S_K)$.\n", - " - `selection_bitsizes`: The number of bits used to represent each selection register corresponding to the size of each dimension of the array $(S_1, S_2, ..., S_K)$. Should be the same length as the shape of each of the datasets.\n", - " - `target_bitsizes`: The number of bits used to represent the data signature. This can be deduced from the maximum element of each of the datasets. Should be a tuple $(b_1, b_2, ..., b_L)$ of length `L = len(data)`, i.e. the number of datasets to be loaded.\n", - " - `num_controls`: The number of controls. \n", - "\n", "#### References\n", " - [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity](https://arxiv.org/abs/1805.03662). Babbush et. al. (2018). Figure 1.\n", " - [Compilation of Fault-Tolerant Quantum Heuristics for Combinatorial Optimization](https://arxiv.org/abs/2007.07391). Babbush et. al. (2020). Figure 3.\n" @@ -120,7 +182,7 @@ { "cell_type": "code", "execution_count": null, - "id": "af317b30", + "id": "2c3d471e", "metadata": { "cq.autogen": "QROM.bloq_doc.py" }, @@ -131,7 +193,7 @@ }, { "cell_type": "markdown", - "id": "d5d3a19d", + "id": "0b256e78", "metadata": { "cq.autogen": "QROM.example_instances.md" }, @@ -142,7 +204,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5f02d641", + "id": "c457f9f9", "metadata": { "cq.autogen": "QROM.qrom_small" }, @@ -155,7 +217,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c2f8b350", + "id": "2b97c379", "metadata": { "cq.autogen": "QROM.qrom_multi_data" }, @@ -169,7 +231,7 @@ { "cell_type": "code", "execution_count": null, - "id": "036cf220", + "id": "fc65e4e0", "metadata": { "cq.autogen": "QROM.qrom_multi_dim" }, @@ -183,18 +245,19 @@ { "cell_type": "code", "execution_count": null, - "id": "a084ca8c-0c89-4439-86d9-51cf91e972c4", - "metadata": {}, + "id": "01c89177", + "metadata": { + "cq.autogen": "QROM.qrom_symb" + }, "outputs": [], "source": [ "N, M, b1, b2, c = sympy.symbols('N M b1 b2 c')\n", - "qrom_symb = QROM.build_from_bitsize((N, M), (b1, b2), num_controls=c)\n", - "qrom_symb" + "qrom_symb = QROM.build_from_bitsize((N, M), (b1, b2), num_controls=c)" ] }, { "cell_type": "markdown", - "id": "b92d1c8e", + "id": "d063ce85", "metadata": { "cq.autogen": "QROM.graphical_signature.md" }, @@ -205,7 +268,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9681cfed", + "id": "8a3ea28a", "metadata": { "cq.autogen": "QROM.graphical_signature.py" }, @@ -218,7 +281,7 @@ }, { "cell_type": "markdown", - "id": "8eb02cb9", + "id": "b63db87f", "metadata": { "cq.autogen": "QROM.call_graph.md" }, @@ -229,7 +292,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7f29f104", + "id": "65639e51", "metadata": { "cq.autogen": "QROM.call_graph.py" }, @@ -240,31 +303,6 @@ "show_call_graph(qrom_small_g)\n", "show_counts_sigma(qrom_small_sigma)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5392752e-276e-434c-9d60-fadcf6478077", - "metadata": {}, - "outputs": [], - "source": [ - "qrom_symb_g, qrom_symb_sigma = qrom_symb.call_graph(generalizer=ignore_split_join)\n", - "show_call_graph(qrom_symb_g)\n", - "show_counts_sigma(qrom_symb_sigma)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "52a98d72", - "metadata": { - "cq.autogen": "QROM.qrom_symb" - }, - "outputs": [], - "source": [ - "N, M, b1, b2, c = sympy.symbols('N M b1 b2 c')\n", - "qrom_symb = QROM.build_from_bitsize((N, M), (b1, b2), num_controls=c)" - ] } ], "metadata": { diff --git a/qualtran/bloqs/data_loading/qrom.py b/qualtran/bloqs/data_loading/qrom.py index 4a7388431..f00dbab78 100644 --- a/qualtran/bloqs/data_loading/qrom.py +++ b/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,90 +53,37 @@ 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 - The action of a QROM can be described as - $$ - \text{QROM}_{s_1, s_2, \dots, s_K}^{d_1, d_2, \dots, d_L} - |s_1\rangle |s_2\rangle \dots |s_K\rangle - |0\rangle^{\otimes b_1} |0\rangle^{\otimes b_2} \dots |0\rangle^{\otimes b_L} - \rightarrow - |s_1\rangle |s_2\rangle \dots |s_K\rangle - |d_1[s_1, s_2, \dots, s_k]\rangle - |d_2[s_1, s_2, \dots, s_k]\rangle \dots - |d_L[s_1, s_2, \dots, s_k]\rangle - $$ - - There two high level parameters that control the behavior of a QROM are - - - 1. Shape of the classical dataset to be loaded ($\text{data.shape} = (S_1, S_2, ..., S_K)$). - 2. Number of distinct datasets to be loaded ($\text{data.bitsizes} = (b_1, b_2, ..., b_L)$). - - Each of these have an effect on the cost of the QROM. The `data_or_shape` parameter stores - either - 1. A numpy array of shape $(L, S_1, S_2, ..., S_K)$ when $L$ classical datasets, each of - shape $(S_1, S_2, ..., S_K)$ and bitsizes $(b_1, b_2, ..., b_L)$ are to be loaded and - the classical data is available to instantiate the QROM bloq. In this case, the helper - builder `QROM.build_from_data(data_1, data_2, ..., data_L)` can be used to build the QROM. - - 2. A `Shaped` object that stores a (potentially symbolic) tuple $(L, S_1, S_2, ..., S_K)$ - that represents the number of classical datasets `L=data_or_shape.shape[0]` and - their shape `data_shape=data_or_shape.shape[1:]` to be loaded by this QROM. This is used - to instantiate QROM bloqs for symbolic cost analysis where the exact data to be loaded - is not known. In this case, the helper builder `QROM.build_from_bitsize` can be used - to build the QROM. - - ### Shape of the classical dataset to be loaded. - QROM bloq supports loading multidimensional classical datasets. In order to load a data - set of shape $\mathrm{data.shape} == (P, Q, R, S)$ the QROM bloq needs four selection - registers with bitsizes $(p, q, r, s)$ where - $p,q,r,s=\log_2{P}, \log_2{Q}, \log_2{R}, \log_2{S}$. - - In general, to load K dimensional data, we use K named selection registers `(selection0, - selection1, ..., selection{k})` to index and load the data. - - The T/Toffoli cost of the QROM scales linearly with the number of elements in the dataset - (i.e. $\mathcal{O}(\mathrm{np.prod(data.shape)}$). - - ### Number of distinct datasets to be loaded, and their corresponding target bitsize. - To load a classical dataset into a target register of bitsize $b$, the clifford cost of a QROM - scales as $\mathcal{O}(b \mathrm{np.prod}(\mathrm{data.shape}))$. This is because we need - $\mathcal{O}(b)$ CNOT gates to load the ith data element in the target register when the - selection register stores index $i$. - - If you have multiple classical datasets `(data_1, data_2, data_3, ..., data_L)` to be loaded - and each of them has the same shape `(data_1.shape == data_2.shape == ... == data_L.shape)` - and different target bitsizes `(b_1, b_2, ..., b_L)`, then one construct a single classical - dataset `data = merge(data_1, data_2, ..., data_L)` where - - - `data.shape == data_1.shape == data_2.shape == ... == data_L` and - - `data[idx] = f'{data_1[idx]!0{b_1}b}' + f'{data_2[idx]!0{b_2}b}' + ... + f'{data_L[idx]!0{b_L}b}'` - - Thus, the target bitsize of the merged dataset is $b = b_1 + b_2 + \dots + b_L$ and clifford - cost of loading merged dataset scales as - $\mathcal{O}((b_1 + b_2 + \dots + b_L) \mathrm{np.prod}(\mathrm{data.shape}))$. + See docstrings of `QROMBase` for an overview of the QROM primitive and the various attributes. + + This bloq is an implementation of the `QROMBase` interface that uses the unary iteration based + approach described in Ref [1]. + + ## Cost of this (unary iteration based) QROM + + ### T / Toffoli cost + The T/Toffoli cost of this QROM scales linearly with the product of iteration lengths over + all dimensions (i.e. $\mathcal{O}(\mathrm{np.prod(\text{selection\_shape})}$). + + ### Clifford Cost + To load a classical dataset into a target register of bitsize $b$ and shape + $\text{target\_shape}$, the clifford cost of this QROM scales as + $\mathcal{O}(b \cdot \text{np.prod(selection\_shape+target\_shape)}) + =\mathcal{O}(b \cdot \text{np.prod(data.shape)})$. This is because we need $\mathcal{O}(b)$ + CNOT gates to load 1 classical data element in the target register and for each of the + $\text{np.prod(selection\_shape)}$ selection indices, we have $\text{np.prod(target\_shape)}$ + such data elements to load. + + ### Ancilla cost + The number of clean ancilla required by this QROM scales linearly with the size of the + selection registers + number of controls. ## Variable spaced QROM When the input classical data contains consecutive entries of identical data elements to load, the QROM also implements the "variable-spaced" QROM optimization described in Ref [2]. - Args: - data_or_shape: List of numpy ndarrays specifying the data to load. If the length - of this list ($L$) is greater than one then we use the same selection indices - to load each dataset. Each data set is required to have the same - shape $(S_1, S_2, ..., S_K)$ and to be of integer type. For symbolic QROMs, - pass a `Shaped` object instead with shape $(L, S_1, S_2, ..., S_K)$. - selection_bitsizes: The number of bits used to represent each selection register - corresponding to the size of each dimension of the array $(S_1, S_2, ..., S_K)$. - Should be the same length as the shape of each of the datasets. - target_bitsizes: The number of bits used to represent the data signature. This can be - deduced from the maximum element of each of the datasets. Should be a tuple - $(b_1, b_2, ..., b_L)$ of length `L = len(data)`, i.e. the number of datasets to - be loaded. - num_controls: The number of controls. - References: [Encoding Electronic Spectra in Quantum Circuits with Linear T Complexity](https://arxiv.org/abs/1805.03662). Babbush et. al. (2018). Figure 1. @@ -146,55 +92,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,48 +109,18 @@ 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, ...], @@ -252,10 +128,14 @@ def _load_nth_data( **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 +181,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 +206,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))} diff --git a/qualtran/bloqs/data_loading/qrom_base.py b/qualtran/bloqs/data_loading/qrom_base.py new file mode 100644 index 000000000..eb48a27a8 --- /dev/null +++ b/qualtran/bloqs/data_loading/qrom_base.py @@ -0,0 +1,326 @@ +# 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. +# 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. + +"""Base class for Bloqs implementing a QROM (Quantum read-only memory) circuit.""" +import abc +import numbers +from functools import cached_property +from typing import cast, Dict, Optional, Tuple, Type, TypeVar, Union + +import attrs +import cirq +import numpy as np +import sympy +from numpy.typing import ArrayLike, NDArray + +from qualtran import BloqDocSpec, BoundedQUInt, QAny, Register +from qualtran.simulation.classical_sim import ClassicalValT +from qualtran.symbolics import bit_length, is_symbolic, shape, Shaped, SymbolicInt + +QROM_T = TypeVar('QROM_T', bound='QROMBase') + + +@cirq.value_equality(distinct_child_types=True) +@attrs.frozen +class QROMBase(metaclass=abc.ABCMeta): + r"""Interface for Bloqs to load `data[l]` when the selection register stores index `l`. + + ## Overview + The action of a QROM can be described as + $$ + \text{QROM}_{s_1, s_2, \dots, s_K}^{d_1, d_2, \dots, d_L} + |s_1\rangle |s_2\rangle \dots |s_K\rangle + |0\rangle^{\otimes b_1} |0\rangle^{\otimes b_2} \dots |0\rangle^{\otimes b_L} + \rightarrow + |s_1\rangle |s_2\rangle \dots |s_K\rangle + |d_1[s_1, s_2, \dots, s_k]\rangle + |d_2[s_1, s_2, \dots, s_k]\rangle \dots + |d_L[s_1, s_2, \dots, s_k]\rangle + $$ + + A behavior of a QROM can be understood in terms of its classical analogue, where a for-loop + over one or more (selection) indices can be used to load one or more classical datasets, where + each of the classical dataset can be multidimensional. + + ``` + >>> # N, M, P, Q, R, S, T are pre-initialized integer parameters. + >>> output = [np.zeros((P, Q)), np.zeros((R, S, T))] + >>> # Load two different classical datasets; each of different shape. + >>> data = [np.random.rand(N, M, P, Q), np.random.rand(N, M, R, S, T)] + >>> for i in range(N): # For loop over two selection indices i and j. + >>> for j in range(M): + >>> # Load two multidimensional classical datasets data[0] and data[1] s.t. + >>> # |i, j⟩|0⟩ -> |i, j⟩|data[0][i, j, :]⟩|data[1][i, j, :]⟩ + >>> output[0] = data[0][i, j, :] + >>> output[1] = data[1][i, j, :] + ``` + + The parameters that control the behavior and costs of a QROM are - + + 1. Number of selection registers (eg: $i$, $j$) and their iteration lengths (eg: $N$, $M$). + 2. Number of target registers, their quantum datatype and shape. + - Number of target registers: One for each classical dataset to load (eg: $\text{data}[0]$ + and $\text{data}[1]$) + - QDType of target registers: Depends on `dtype` of the $i$'th classical dataset + - Shape of target registers: Depends on shape of classical data (eg: $(P, Q)$ and + $(R, S, T)$ above) + + ### Specification of classical data via `data_or_shape` + Users can specify the classical data to load via QROM by passing in an appropriate value + for `data_or_shape` attribute. This is a list of numpy arrays or `Shaped` objects, where + each item of the list corresponds to a classical dataset to load. + + Each classical dataset to load can be specified as a numpy array (or a `Shaped` object for + symbolic bloqs). The shape of the dataset is a union of the selection shape and target shape, + s.t. + $$ + \text{data[i].shape} = \text{selection\_shape} + \text{target\_shape[i]} + $$ + + Note that the $\text{selection\_shape}$ should be same across all classical datasets to be + loaded and correspond to a tuple of iteration lengths of selection indices (i.e. $(N, M)$ + in the example above). + + The target shape of each classical dataset can be different and parameterizes the size of + the desired output that should be loaded in a target register. + + ### Number of selection registers and their iteration lengths + As describe in the previous section, the number of selection registers and their iteration + lengths can be inferred from the shape of the classical dataset. All classical datasets + to be loaded must have the same $\text{selection\_shape}$, which is a tuple of iteration + lengths over each dimension of the dataset (i.e. the range for each nested for-loop). + + In order to load a data set with $\text{selection\_shape} == (P, Q, R, S)$ the QROM bloq + needs four selection registers with bitsizes $(p, q, r, s)$ where each of + $p,q,r,s \geq \log_2{P}, \log_2{Q}, \log_2{R}, \log_2{S}$. + + In general, to load $K$ dimensional data, we use $K$ named selection registers + $(\text{selection}_0, \text{selection}_1, ..., \text{selection}_k)$ to index and + load the data. For the $i$'th selection register, its size is configured using + attribute $\text{selection\_bitsizes[i]}$ and the iteration range is configued + using $\text{data\_or\_shape[0].shape[i]}$. + + ### Number of target registers, their quantum datatype and shape + QROM bloq uses one target register for each entry corresponding to classical dataset in the + tuple `data_or_shape`. Thus, to load $L$ classical datsets, we use $L$ names target registers + $(\text{target}_0, \text{target}_1, ..., \text{target}_L)$ + + Each named target register has a bitsize $b_{i}=\text{target\_bitsizes[i]}$ that represents + the size of the register and depends upon the maximum value of individual elements in the + $i$'th classical dataset. + + Each named target register has a shape that can be configured using attribute + $\text{target\_shape[i]}$ that represents the number of target registers if the output to load + is multidimensional. + + Args: + data_or_shape: List of numpy ndarrays specifying the data to load. If the length + of this list ($L$) is greater than one then we use the same selection indices + to load each dataset. The shape of a classical dataset is a concatenation of + selection_shape and target_shape[i]; i.e. `data_or_shape[i].shape = + selection_shape + target_shape[i]`. Thus, each data set is required to have the + same selection shape $(S_1, S_2, ..., S_K)$ and can have a different + target shape given by `target_shapes[i]`. For symbolic QROMs, pass a list of + `Shaped` objects instead with shape $(S_1, S_2, ..., S_K) + target_shape[i]$. + selection_bitsizes: The number of bits used to represent each selection register + corresponding to the size of each dimension of the selection_shape + $(S_1, S_2, ..., S_K)$. Should be the same length as the selection shape of + each of the datasets and $2**\text{selection\_bitsizes[i]} >= S_i$ + target_shapes: Shape of target registers for each classical dataset to be loaded. + Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_shapes)` + and `data_or_shape[-len(target_shapes[i]):] == target_shapes[i]`. + target_bitsizes: Bitsize (or qdtype) of the target registers for each classical + dataset to be loaded. This can be deduced from the maximum element of each + of the datasets. Must be consistent with `data_or_shape` s.t. + `len(data_or_shape) == len(target_bitsizes)` and + `target_bitsizes[i] >= max(data[i]).bitsize`. + num_controls: The number of controls to instanstiate a controlled version of this bloq. + """ + + data_or_shape: Tuple[Union[NDArray, Shaped], ...] = attrs.field( + converter=lambda x: tuple(np.array(y) if isinstance(y, (list, tuple)) else y for y in 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) + ) + target_shapes: Tuple[Tuple[SymbolicInt, ...], ...] = attrs.field( + converter=lambda x: tuple(tuple(y) for y in x) + ) + num_controls: SymbolicInt = 0 + + @target_shapes.default + def _default_target_shapes(self): + return ((),) * len(self.data_or_shape) + + @cached_property + def data_shape(self) -> Tuple[SymbolicInt, ...]: + ret: Tuple[SymbolicInt, ...] = () + for data_or_shape, target_shape in zip(self.data_or_shape, self.target_shapes): + data_shape = shape(data_or_shape) + if target_shape: + data_shape = data_shape[: -len(target_shape)] + if ret == (): + ret = data_shape + if tuple(ret) != tuple(data_shape): + raise ValueError("All datasets must have same shape") + return ret + + def has_data(self) -> bool: + return all(isinstance(d, np.ndarray) for d in self.data_or_shape) + + @property + def data(self) -> Tuple[np.ndarray, ...]: + if not self.has_data(): + raise ValueError(f"Data not available for symbolic QROM {self}") + assert all(isinstance(d, np.ndarray) for d in self.data_or_shape) + return cast(Tuple[np.ndarray, ...], 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) == len(self.data_or_shape), ( + f"len(self.target_bitsizes)={len(self.target_bitsizes)} should be same as " + f"len(self.data)={len(self.data_or_shape)}" + ) + assert len(self.target_bitsizes) == len(self.target_shapes) + 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: Type[QROM_T], + *data: ArrayLike, + target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + num_controls: SymbolicInt = 0, + ) -> QROM_T: + _data = [np.array(d, dtype=int) for d in data] + selection_bitsizes = tuple((s - 1).bit_length() for s in _data[0].shape) + if target_bitsizes is None: + target_bitsizes = tuple(max(int(np.max(d)).bit_length(), 1) for d in data) + return cls( + data_or_shape=_data, + selection_bitsizes=selection_bitsizes, + target_bitsizes=target_bitsizes, + num_controls=num_controls, + ) + + def with_data(self: QROM_T, *data: ArrayLike) -> QROM_T: + _data = tuple([np.array(d, dtype=int) for d in data]) + assert all(shape(d1) == shape(d2) for d1, d2 in zip(_data, self.data_or_shape)) + assert len(_data) == len(self.target_bitsizes) + assert all(t >= int(np.max(d)).bit_length() for t, d in zip(self.target_bitsizes, _data)) + return attrs.evolve(self, data_or_shape=_data) + + @classmethod + def _build_from_bitsize( + cls: Type[QROM_T], + 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_T: + data_shape: Tuple[SymbolicInt, ...] = ( + (data_len_or_shape,) + if isinstance(data_len_or_shape, (int, numbers.Number, sympy.Basic)) + else data_len_or_shape + ) + if not isinstance(target_bitsizes, tuple): + target_bitsizes = (target_bitsizes,) + if target_shapes == (): + target_shapes = ((),) * len(target_bitsizes) + _data = [Shaped(data_shape + sh) for sh in target_shapes] + if selection_bitsizes == (): + selection_bitsizes = tuple(bit_length(s - 1) for s in data_shape) + assert len(selection_bitsizes) == len(data_shape) + return cls( + data_or_shape=_data, + selection_bitsizes=selection_bitsizes, + target_bitsizes=target_bitsizes, + target_shapes=target_shapes, + 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_shape, 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), shape=sh) + for i, (l, sh) in enumerate(zip(self.target_bitsizes, self.target_shapes)) + if is_symbolic(l) or l + ) + + def on_classical_vals( + self, **vals: Union['sympy.Symbol', 'ClassicalValT'] + ) -> Dict[str, 'ClassicalValT']: + if not self.has_data(): + raise NotImplementedError(f'Symbolic {self} does not support classical simulation') + vals = cast(Dict[str, 'ClassicalValT'], vals) + 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 _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 __pow__(self, power: int): + if power in [1, -1]: + return self + return NotImplemented # pragma: no cover + + +_QROM_BASE_DOC = BloqDocSpec(bloq_cls=QROMBase, import_line='', examples=[]) diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.ipynb b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb new file mode 100644 index 000000000..5fb2b5c42 --- /dev/null +++ b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb @@ -0,0 +1,368 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8f717aa7", + "metadata": { + "cq.autogen": "title_cell" + }, + "source": [ + "# SelectSwapQROM" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d0e99ddc", + "metadata": { + "cq.autogen": "top_imports" + }, + "outputs": [], + "source": [ + "from qualtran import Bloq, CompositeBloq, BloqBuilder, Signature, Register\n", + "from qualtran import QBit, QInt, QUInt, QAny\n", + "from qualtran.drawing import show_bloq, show_call_graph, show_counts_sigma\n", + "from typing import *\n", + "import numpy as np\n", + "import sympy\n", + "import cirq" + ] + }, + { + "cell_type": "markdown", + "id": "fc5770d6", + "metadata": { + "cq.autogen": "QROMBase.bloq_doc.md" + }, + "source": [ + "## `QROMBase`\n", + "Interface for Bloqs to load `data[l]` when the selection register stores index `l`.\n", + "\n", + "## Overview\n", + "The action of a QROM can be described as\n", + "$$\n", + " \\text{QROM}_{s_1, s_2, \\dots, s_K}^{d_1, d_2, \\dots, d_L}\n", + " |s_1\\rangle |s_2\\rangle \\dots |s_K\\rangle\n", + " |0\\rangle^{\\otimes b_1} |0\\rangle^{\\otimes b_2} \\dots |0\\rangle^{\\otimes b_L}\n", + " \\rightarrow\n", + " |s_1\\rangle |s_2\\rangle \\dots |s_K\\rangle\n", + " |d_1[s_1, s_2, \\dots, s_k]\\rangle\n", + " |d_2[s_1, s_2, \\dots, s_k]\\rangle \\dots\n", + " |d_L[s_1, s_2, \\dots, s_k]\\rangle\n", + "$$\n", + "\n", + "A behavior of a QROM can be understood in terms of its classical analogue, where a for-loop\n", + "over one or more (selection) indices can be used to load one or more classical datasets, where\n", + "each of the classical dataset can be multidimensional.\n", + "\n", + "```\n", + ">>> # N, M, P, Q, R, S, T are pre-initialized integer parameters.\n", + ">>> output = [np.zeros((P, Q)), np.zeros((R, S, T))]\n", + ">>> # Load two different classical datasets; each of different shape.\n", + ">>> data = [np.random.rand(N, M, P, Q), np.random.rand(N, M, R, S, T)]\n", + ">>> for i in range(N): # For loop over two selection indices i and j.\n", + ">>> for j in range(M):\n", + ">>> # Load two multidimensional classical datasets data[0] and data[1] s.t.\n", + ">>> # |i, j⟩|0⟩ -> |i, j⟩|data[0][i, j, :]⟩|data[1][i, j, :]⟩\n", + ">>> output[0] = data[0][i, j, :]\n", + ">>> output[1] = data[1][i, j, :]\n", + "```\n", + "\n", + "The parameters that control the behavior and costs of a QROM are -\n", + "\n", + "1. Number of selection registers (eg: $i$, $j$) and their iteration lengths (eg: $N$, $M$).\n", + "2. Number of target registers, their quantum datatype and shape.\n", + " - Number of target registers: One for each classical dataset to load (eg: $\\text{data}[0]$\n", + " and $\\text{data}[1]$)\n", + " - QDType of target registers: Depends on `dtype` of the $i$'th classical dataset\n", + " - Shape of target registers: Depends on shape of classical data (eg: $(P, Q)$ and\n", + " $(R, S, T)$ above)\n", + "\n", + "### Specification of classical data via `data_or_shape`\n", + "Users can specify the classical data to load via QROM by passing in an appropriate value\n", + "for `data_or_shape` attribute. This is a list of numpy arrays or `Shaped` objects, where\n", + "each item of the list corresponds to a classical dataset to load.\n", + "\n", + "Each classical dataset to load can be specified as a numpy array (or a `Shaped` object for\n", + "symbolic bloqs). The shape of the dataset is a union of the selection shape and target shape,\n", + "s.t.\n", + "$$\n", + " \\text{data[i].shape} = \\text{selection\\_shape} + \\text{target\\_shape[i]}\n", + "$$\n", + "\n", + "Note that the $\\text{selection\\_shape}$ should be same across all classical datasets to be\n", + "loaded and correspond to a tuple of iteration lengths of selection indices (i.e. $(N, M)$\n", + "in the example above).\n", + "\n", + "The target shape of each classical dataset can be different and parameterizes the size of\n", + "the desired output that should be loaded in a target register.\n", + "\n", + "### Number of selection registers and their iteration lengths\n", + "As describe in the previous section, the number of selection registers and their iteration\n", + "lengths can be inferred from the shape of the classical dataset. All classical datasets\n", + "to be loaded must have the same $\\text{selection\\_shape}$, which is a tuple of iteration\n", + "lengths over each dimension of the dataset (i.e. the range for each nested for-loop).\n", + "\n", + "In order to load a data set with $\\text{selection\\_shape} == (P, Q, R, S)$ the QROM bloq\n", + "needs four selection registers with bitsizes $(p, q, r, s)$ where each of\n", + "$p,q,r,s \\geq \\log_2{P}, \\log_2{Q}, \\log_2{R}, \\log_2{S}$.\n", + "\n", + "In general, to load $K$ dimensional data, we use $K$ named selection registers\n", + "$(\\text{selection}_0, \\text{selection}_1, ..., \\text{selection}_k)$ to index and\n", + "load the data. For the $i$'th selection register, its size is configured using\n", + "attribute $\\text{selection\\_bitsizes[i]}$ and the iteration range is configued\n", + "using $\\text{data\\_or\\_shape[0].shape[i]}$.\n", + "\n", + "### Number of target registers, their quantum datatype and shape\n", + "QROM bloq uses one target register for each entry corresponding to classical dataset in the\n", + "tuple `data_or_shape`. Thus, to load $L$ classical datsets, we use $L$ names target registers\n", + "$(\\text{target}_0, \\text{target}_1, ..., \\text{target}_L)$\n", + "\n", + "Each named target register has a bitsize $b_{i}=\\text{target\\_bitsizes[i]}$ that represents\n", + "the size of the register and depends upon the maximum value of individual elements in the\n", + "$i$'th classical dataset.\n", + "\n", + "Each named target register has a shape that can be configured using attribute\n", + "$\\text{target\\_shape[i]}$ that represents the number of target registers if the output to load\n", + "is multidimensional.\n", + "\n", + "#### Parameters\n", + " - `data_or_shape`: List of numpy ndarrays specifying the data to load. If the length of this list ($L$) is greater than one then we use the same selection indices to load each dataset. The shape of a classical dataset is a concatenation of selection_shape and target_shape[i]; i.e. `data_or_shape[i].shape = selection_shape + target_shape[i]`. Thus, each data set is required to have the same selection shape $(S_1, S_2, ..., S_K)$ and can have a different target shape given by `target_shapes[i]`. For symbolic QROMs, pass a list of `Shaped` objects instead with shape $(S_1, S_2, ..., S_K) + target_shape[i]$.\n", + " - `selection_bitsizes`: The number of bits used to represent each selection register corresponding to the size of each dimension of the selection_shape $(S_1, S_2, ..., S_K)$. Should be the same length as the selection shape of each of the datasets and $2**\\text{selection\\_bitsizes[i]} >= S_i$\n", + " - `target_shapes`: Shape of target registers for each classical dataset to be loaded. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_shapes)` and `data_or_shape[-len(target_shapes[i]):] == target_shapes[i]`.\n", + " - `target_bitsizes`: Bitsize (or qdtype) of the target registers for each classical dataset to be loaded. This can be deduced from the maximum element of each of the datasets. Must be consistent with `data_or_shape` s.t. `len(data_or_shape) == len(target_bitsizes)` and `target_bitsizes[i] >= max(data[i]).bitsize`.\n", + " - `num_controls`: The number of controls to instanstiate a controlled version of this bloq.\n" + ] + }, + { + "cell_type": "markdown", + "id": "818551d0", + "metadata": { + "cq.autogen": "SelectSwapQROM.bloq_doc.md" + }, + "source": [ + "## `SelectSwapQROM`\n", + "Gate to load data[l] in the target register when the selection register stores integer l.\n", + "\n", + "Let\n", + " N:= Number of data elements to load.\n", + " b:= Bit-length of the target register in which data elements should be loaded.\n", + "\n", + "The `SelectSwapQROM` is a hybrid of the following two existing primitives:\n", + "\n", + "- Unary Iteration based `QROM` requires O(N) T-gates to load `N` data\n", + "elements into a b-bit target register. Note that the T-complexity is independent of `b`.\n", + "- `SwapWithZeroGate` can swap a `b` bit register indexed `x` with a `b`\n", + "bit register at index `0` using O(b) T-gates, if the selection register stores integer `x`.\n", + "Note that the swap complexity is independent of the iteration length `N`.\n", + "\n", + "The `SelectSwapQROM` uses square root decomposition by combining the above two approaches to\n", + "further optimize the T-gate complexity of loading `N` data elements, each into a `b` bit\n", + "target register as follows:\n", + "\n", + "- Divide the `N` data elements into batches of size `B` (a variable) and\n", + "load each batch simultaneously into `B` distinct target signature using the conventional\n", + "QROM. This has T-complexity `O(N / B)`.\n", + "- Use `SwapWithZeroGate` to swap the `i % B`'th target register in batch number `i / B`\n", + "to load `data[i]` in the 0'th target register. This has T-complexity `O(B * b)`.\n", + "\n", + "This, the final T-complexity of `SelectSwapQROM` is `O(B * b + N / B)` T-gates; where `B` is\n", + "the block-size with an optimal value of `O(sqrt(N / b))`.\n", + "\n", + "This improvement in T-complexity is achieved at the cost of using an additional `O(B * b)`\n", + "ancilla qubits, with a nice property that these additional ancillas can be `dirty`; i.e.\n", + "they don't need to start in the |0> state and thus can be borrowed from other parts of the\n", + "algorithm. The state of these dirty ancillas would be unaffected after the operation has\n", + "finished.\n", + "\n", + "For more details, see the reference below:\n", + "\n", + "#### References\n", + " - [Trading T-gates for dirty qubits in state preparation and unitary synthesis](https://arxiv.org/abs/1812.00954). Low, Kliuchnikov, Schaeffer. 2018.\n", + " - [Qubitization of Arbitrary Basis Quantum Chemistry Leveraging Sparsity and Low Rank Factorization](https://arxiv.org/abs/1902.02134). Berry et. al. (2019). Appendix A. and B.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b91df7f9", + "metadata": { + "cq.autogen": "SelectSwapQROM.bloq_doc.py" + }, + "outputs": [], + "source": [ + "from qualtran.bloqs.data_loading.select_swap_qrom import SelectSwapQROM" + ] + }, + { + "cell_type": "markdown", + "id": "514bbe62", + "metadata": { + "cq.autogen": "SelectSwapQROM.example_instances.md" + }, + "source": [ + "### Example Instances" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17f62143", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_multi_data" + }, + "outputs": [], + "source": [ + "data1 = np.arange(5)\n", + "data2 = np.arange(5) + 1\n", + "qroam_multi_data = SelectSwapQROM.build_from_data([data1, data2])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "757edd48", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_multi_dim" + }, + "outputs": [], + "source": [ + "data1 = np.arange(25).reshape((5, 5))\n", + "data2 = (np.arange(25) + 1).reshape((5, 5))\n", + "qroam_multi_dim = SelectSwapQROM.build_from_data([data1, data2])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "655da4ad", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_symb_dirty_1d" + }, + "outputs": [], + "source": [ + "N, b, k, c = sympy.symbols('N b k c')\n", + "qroam_symb_dirty_1d = SelectSwapQROM.build_from_bitsize(\n", + " (N,), (b,), log_block_sizes=(k,), num_controls=c\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ebd4b95", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_symb_dirty_2d" + }, + "outputs": [], + "source": [ + "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c')\n", + "log_block_sizes = (k1, k2)\n", + "qroam_symb_dirty_2d = SelectSwapQROM.build_from_bitsize(\n", + " (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "792cda31", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_symb_clean_1d" + }, + "outputs": [], + "source": [ + "N, b, k, c = sympy.symbols('N b k c')\n", + "qroam_symb_clean_1d = SelectSwapQROM.build_from_bitsize(\n", + " (N,), (b,), log_block_sizes=(k,), num_controls=c, use_dirty_ancilla=False\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05eac1c8", + "metadata": { + "cq.autogen": "SelectSwapQROM.qroam_symb_clean_2d" + }, + "outputs": [], + "source": [ + "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c')\n", + "log_block_sizes = (k1, k2)\n", + "qroam_symb_clean_2d = SelectSwapQROM.build_from_bitsize(\n", + " (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c, use_dirty_ancilla=False\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "21aac4da", + "metadata": { + "cq.autogen": "SelectSwapQROM.graphical_signature.md" + }, + "source": [ + "#### Graphical Signature" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e86e1c78", + "metadata": { + "cq.autogen": "SelectSwapQROM.graphical_signature.py" + }, + "outputs": [], + "source": [ + "from qualtran.drawing import show_bloqs\n", + "show_bloqs([qroam_multi_data, qroam_multi_dim, qroam_symb_dirty_1d, qroam_symb_dirty_2d, qroam_symb_clean_1d, qroam_symb_clean_2d],\n", + " ['`qroam_multi_data`', '`qroam_multi_dim`', '`qroam_symb_dirty_1d`', '`qroam_symb_dirty_2d`', '`qroam_symb_clean_1d`', '`qroam_symb_clean_2d`'])" + ] + }, + { + "cell_type": "markdown", + "id": "3f1de46e", + "metadata": { + "cq.autogen": "SelectSwapQROM.call_graph.md" + }, + "source": [ + "### Call Graph" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66358a0c", + "metadata": { + "cq.autogen": "SelectSwapQROM.call_graph.py" + }, + "outputs": [], + "source": [ + "from qualtran.resource_counting.generalizers import ignore_split_join\n", + "qroam_multi_data_g, qroam_multi_data_sigma = qroam_multi_data.call_graph(max_depth=1, generalizer=ignore_split_join)\n", + "show_call_graph(qroam_multi_data_g)\n", + "show_counts_sigma(qroam_multi_data_sigma)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.py b/qualtran/bloqs/data_loading/select_swap_qrom.py index 67008785f..99f299987 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom.py @@ -11,22 +11,34 @@ # 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 numbers +from collections import defaultdict from functools import cached_property -from typing import Iterator, List, Optional, Sequence, Tuple +from typing import cast, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING, Union +import attrs import cirq import numpy as np -from numpy.typing import NDArray +import sympy +from numpy.typing import ArrayLike, NDArray -from qualtran import BoundedQUInt, GateWithRegisters, QAny, Register, Signature -from qualtran._infra.gate_with_registers import merge_qubits, split_qubits, total_bits +from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Register, Signature +from qualtran._infra.gate_with_registers import total_bits +from qualtran.bloqs.basic_gates import CNOT from qualtran.bloqs.data_loading.qrom import QROM +from qualtran.bloqs.data_loading.qrom_base import QROMBase from qualtran.bloqs.swap_network import SwapWithZero from qualtran.drawing import Circle, Text, TextBox, WireSymbol +from qualtran.symbolics import ceil, is_symbolic, log2, SymbolicInt + +if TYPE_CHECKING: + from qualtran import Bloq + from qualtran.resource_counting import BloqCountT, SympySymbolAllocator -def find_optimal_log_block_size(iteration_length: int, target_bitsize: int) -> int: +def find_optimal_log_block_size( + iteration_length: SymbolicInt, target_bitsize: SymbolicInt +) -> SymbolicInt: """Find optimal block size, which is a power of 2, for SelectSwapQROM. This functions returns the optimal `k` s.t. @@ -34,7 +46,10 @@ def find_optimal_log_block_size(iteration_length: int, target_bitsize: int) -> i * iteration_length/2^k + target_bitsize*(2^k - 1) is minimized. The corresponding block size for SelectSwapQROM would be 2^k. """ - k = 0.5 * np.log2(iteration_length / target_bitsize) + k = 0.5 * log2(iteration_length / target_bitsize) + if is_symbolic(k): + return ceil(k) + if k < 0: return 1 @@ -45,8 +60,9 @@ def value(kk: List[int]): return int(k_int[np.argmin(value(k_int))]) # obtain optimal k -@cirq.value_equality() -class SelectSwapQROM(GateWithRegisters): +@cirq.value_equality(distinct_child_types=True) +@attrs.frozen +class SelectSwapQROM(QROMBase, GateWithRegisters): # type: ignore[misc] """Gate to load data[l] in the target register when the selection register stores integer l. Let @@ -55,21 +71,21 @@ class SelectSwapQROM(GateWithRegisters): The `SelectSwapQROM` is a hybrid of the following two existing primitives: - * Unary Iteration based `QROM` requires O(N) T-gates to load `N` data - elements into a b-bit target register. Note that the T-complexity is independent of `b`. - * `SwapWithZeroGate` can swap a `b` bit register indexed `x` with a `b` - bit register at index `0` using O(b) T-gates, if the selection register stores integer `x`. - Note that the swap complexity is independent of the iteration length `N`. + - Unary Iteration based `QROM` requires O(N) T-gates to load `N` data + elements into a b-bit target register. Note that the T-complexity is independent of `b`. + - `SwapWithZeroGate` can swap a `b` bit register indexed `x` with a `b` + bit register at index `0` using O(b) T-gates, if the selection register stores integer `x`. + Note that the swap complexity is independent of the iteration length `N`. The `SelectSwapQROM` uses square root decomposition by combining the above two approaches to further optimize the T-gate complexity of loading `N` data elements, each into a `b` bit target register as follows: - * Divide the `N` data elements into batches of size `B` (a variable) and - load each batch simultaneously into `B` distinct target signature using the conventional - QROM. This has T-complexity `O(N / B)`. - * Use `SwapWithZeroGate` to swap the `i % B`'th target register in batch number `i / B` - to load `data[i]` in the 0'th target register. This has T-complexity `O(B * b)`. + - Divide the `N` data elements into batches of size `B` (a variable) and + load each batch simultaneously into `B` distinct target signature using the conventional + QROM. This has T-complexity `O(N / B)`. + - Use `SwapWithZeroGate` to swap the `i % B`'th target register in batch number `i / B` + to load `data[i]` in the 0'th target register. This has T-complexity `O(B * b)`. This, the final T-complexity of `SelectSwapQROM` is `O(B * b + N / B)` T-gates; where `B` is the block-size with an optimal value of `O(sqrt(N / b))`. @@ -85,162 +101,219 @@ class SelectSwapQROM(GateWithRegisters): References: [Trading T-gates for dirty qubits in state preparation and unitary synthesis](https://arxiv.org/abs/1812.00954). Low, Kliuchnikov, Schaeffer. 2018. + + [Qubitization of Arbitrary Basis Quantum Chemistry Leveraging Sparsity and Low Rank Factorization](https://arxiv.org/abs/1902.02134). + Berry et. al. (2019). Appendix A. and B. """ - def __init__( - self, - *data: Sequence[int], - target_bitsizes: Optional[Sequence[int]] = None, - block_size: Optional[int] = None, - ): - """Initializes SelectSwapQROM - - For a single data sequence of length `N`, maximum target bitsize `b` and block size `B`; - SelectSwapQROM requires: - - Selection register & ancilla of size `logN` for QROM data load. - - 1 clean target register of size `b`. - - `B` dirty target signature, each of size `b`. - - Similarly, to load `M` such data sequences, `SelectSwapQROM` requires: - - Selection register & ancilla of size `logN` for QROM data load. - - 1 clean target register of size `sum(target_bitsizes)`. - - `B` dirty target signature, each of size `sum(target_bitsizes)`. - - Args: - data: Sequence of integers to load in the target register. If more than one sequence - is provided, each sequence must be of the same length. - target_bitsizes: Sequence of integers describing the size of target register for each - data sequence to load. Defaults to `max(data[i]).bit_length()` for each i. - block_size(B): Load batches of `B` data elements in each iteration of traditional QROM - (N/B iterations required). Complexity of SelectSwap QROAM scales as - `O(B * b + N / B)`, where `B` is the block_size. Defaults to optimal value of - `O(sqrt(N / b))`. - - Raises: - ValueError: If all target data sequences to load do not have the same length. - """ - # Validate input. - if len(set(len(d) for d in data)) != 1: - raise ValueError("All data sequences to load must be of equal length.") - if target_bitsizes is None: - target_bitsizes = [max(d).bit_length() for d in data] - assert len(target_bitsizes) == len(data) - assert all(t >= max(d).bit_length() for t, d in zip(target_bitsizes, data)) - self._num_sequences = len(data) - self._target_bitsizes = tuple(target_bitsizes) - self._iteration_length = len(data[0]) - if block_size is None: - # Figure out optimal value of block_size - block_size = 2 ** find_optimal_log_block_size(len(data[0]), sum(target_bitsizes)) - assert 0 < block_size <= self._iteration_length - self._block_size = block_size - self._num_blocks = int(np.ceil(self._iteration_length / self.block_size)) - self.selection_q, self.selection_r = tuple( - (L - 1).bit_length() for L in [self.num_blocks, self.block_size] - ) - self._data = tuple(tuple(d) for d in data) + log_block_sizes: Tuple[SymbolicInt, ...] = attrs.field( + converter=lambda x: tuple(x.tolist() if isinstance(x, np.ndarray) else x) + ) + use_dirty_ancilla: bool = True @cached_property - def selection_registers(self) -> Tuple[Register, ...]: - return ( - Register( - 'selection', - BoundedQUInt(self.selection_q + self.selection_r, self._iteration_length), - ), + def signature(self) -> Signature: + return Signature( + [*self.control_registers, *self.selection_registers, *self.target_registers] ) - @cached_property - def target_registers(self) -> Tuple[Register, ...]: - # See https://github.com/quantumlib/Qualtran/issues/556 for unusual placement of underscore. + # Builder methods and helpers. + @log_block_sizes.default + def _default_block_sizes(self) -> Tuple[SymbolicInt, ...]: return tuple( - Register(f'target{sequence_id}_', QAny(self._target_bitsizes[sequence_id])) - for sequence_id in range(self._num_sequences) + find_optimal_log_block_size(ilen, sbitsize) + for ilen, sbitsize in zip(self.data_shape, self.selection_bitsizes) + ) + + @classmethod + def build_from_data( + cls: Type['SelectSwapQROM'], + *data: ArrayLike, + target_bitsizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + use_dirty_ancilla: bool = True, + ) -> 'SelectSwapQROM': + qroam: 'SelectSwapQROM' = cls._build_from_data( + *data, target_bitsizes=target_bitsizes, num_controls=num_controls + ) + qroam = attrs.evolve(qroam, use_dirty_ancilla=use_dirty_ancilla) + return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) + + @classmethod + def build_from_bitsize( + cls: Type['SelectSwapQROM'], + data_len_or_shape: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + target_bitsizes: Union[SymbolicInt, Tuple[SymbolicInt, ...]], + *, + selection_bitsizes: Tuple[SymbolicInt, ...] = (), + num_controls: SymbolicInt = 0, + log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None, + use_dirty_ancilla: bool = True, + ) -> 'SelectSwapQROM': + qroam: 'SelectSwapQROM' = cls._build_from_bitsize( + data_len_or_shape, + target_bitsizes, + selection_bitsizes=selection_bitsizes, + num_controls=num_controls, ) + qroam = attrs.evolve(qroam, use_dirty_ancilla=use_dirty_ancilla) + return qroam.with_log_block_sizes(log_block_sizes=log_block_sizes) + + def with_log_block_sizes( + self, log_block_sizes: Optional[Union[SymbolicInt, Tuple[SymbolicInt, ...]]] = None + ) -> 'SelectSwapQROM': + if log_block_sizes is None: + return self + if isinstance(log_block_sizes, (int, sympy.Basic, numbers.Number)): + log_block_sizes = (log_block_sizes,) + if not is_symbolic(*log_block_sizes): + assert all(1 <= 2**bs <= ilen for bs, ilen in zip(log_block_sizes, self.data_shape)) + return attrs.evolve(self, log_block_sizes=log_block_sizes) @cached_property - def signature(self) -> Signature: - return Signature([*self.selection_registers, *self.target_registers]) + def block_sizes(self) -> Tuple[SymbolicInt, ...]: + return tuple(2**log_K for log_K in self.log_block_sizes) - @property - def data(self) -> Tuple[Tuple[int, ...], ...]: - return self._data + @cached_property + def batched_qrom_shape(self) -> Tuple[SymbolicInt, ...]: + return tuple(ceil(N / K) for N, K in zip(self.data_shape, self.block_sizes)) - @property - def block_size(self) -> int: - return self._block_size + @cached_property + def batched_qrom_selection_bitsizes(self) -> Tuple[SymbolicInt, ...]: + return tuple(s - log_K for s, log_K in zip(self.selection_bitsizes, self.log_block_sizes)) - @property - def num_blocks(self) -> int: - return self._num_blocks + @cached_property + def padded_data(self) -> List[np.ndarray]: + pad_width = tuple( + (0, ceil(N / K) * K - N) for N, K in zip(self.data_shape, self.block_sizes) + ) + return [np.pad(d, pad_width) for d in self.data] - def decompose_from_registers( + @cached_property + def batched_data_shape(self) -> Tuple[int, ...]: + return cast(Tuple[int, ...], self.batched_qrom_shape + self.block_sizes) + + @cached_property + def batched_data(self) -> List[np.ndarray]: + # In SelectSwapQROM, for N-dimensional data (one or more datasets), you pick block sizes for + # each dimension and load a batched N-dimensional output "at-once" using a traditional QROM read + # followed by an N-dimensional SwapWithZero swap. + # + # For data[N1][N2] with block sizes (k1, k2), you load batches of size `(k1, k2)` at once. + # Thus, you load batch[N1/k1][N2/k2] where batch[i][j] = data[i*k1:(i + 1)*k1][j*k2:(j + 1)*k2] + batched_data = [np.empty(self.batched_data_shape) for _ in range(len(self.target_bitsizes))] + block_slices = [slice(0, k) for k in self.block_sizes] + for i, data in enumerate(self.padded_data): + for batch_idx in np.ndindex(cast(Tuple[int, ...], self.batched_qrom_shape)): + data_idx = [slice(x * k, (x + 1) * k) for x, k in zip(batch_idx, self.block_sizes)] + batched_data[i][(*batch_idx, *block_slices)] = data[tuple(data_idx)] + return batched_data + + @cached_property + def qrom_bloq(self) -> QROM: + return QROM.build_from_bitsize( + self.batched_qrom_shape, + self.target_bitsizes, + target_shapes=(self.block_sizes,) * len(self.target_bitsizes), + selection_bitsizes=self.batched_qrom_selection_bitsizes, + num_controls=self.num_controls, + ) + + @cached_property + def swap_with_zero_bloqs(self) -> List[SwapWithZero]: + return [ + SwapWithZero( + self.log_block_sizes, + target_bitsize=target_bitsize, + n_target_registers=self.block_sizes, + ) + for target_bitsize in self.target_bitsizes + ] + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + ret: Dict[Bloq, SymbolicInt] = defaultdict(lambda: 0) + toggle_overhead = 2 if self.use_dirty_ancilla else 1 + ret[self.qrom_bloq] += 1 + ret[self.qrom_bloq.adjoint()] += 1 + ret[CNOT()] += toggle_overhead * total_bits(self.target_registers) + for swz in self.swap_with_zero_bloqs: + if any(is_symbolic(s) or s > 0 for s in swz.selection_bitsizes): + ret[swz] += toggle_overhead + ret[swz.adjoint()] += toggle_overhead + return set(ret.items()) + + def _alloc_qrom_target_qubits( + self, reg: Register, qm: cirq.QubitManager + ) -> NDArray[cirq.Qid]: # type:ignore[type-var] + qubits = ( + qm.qborrow(total_bits([reg])) + if self.use_dirty_ancilla + else qm.qalloc(total_bits([reg])) + ) + return np.array(qubits).reshape(reg.shape + (reg.bitsize,)) + + def decompose_from_registers( # type: ignore[return] self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid], # type:ignore[type-var] - ) -> Iterator[cirq.Operation]: - # Divide each data sequence and corresponding target registers into - # `self.num_blocks` batches of size `self.block_size`. - selection, targets = quregs.pop('selection'), quregs - qrom_data: List[NDArray] = [] - qrom_target_bitsizes: List[int] = [] - ordered_target_qubits: List[cirq.Qid] = [] - for block_id in range(self.block_size): - for sequence_id in range(self._num_sequences): - data = self.data[sequence_id] - target_bitsize = self._target_bitsizes[sequence_id] - ordered_target_qubits.extend(context.qubit_manager.qborrow(target_bitsize)) - data_for_current_block = data[block_id :: self.block_size] - if len(data_for_current_block) < self.num_blocks: - zero_pad = (0,) * (self.num_blocks - len(data_for_current_block)) - data_for_current_block = data_for_current_block + zero_pad - qrom_data.append(np.array(data_for_current_block)) - qrom_target_bitsizes.append(target_bitsize) - # Construct QROM, SwapWithZero and CX operations using the batched data and qubits. - k = (self.block_size - 1).bit_length() - q, r = selection[: self.selection_q], selection[self.selection_q :] - qrom_gate = QROM( - qrom_data, - selection_bitsizes=(self.selection_q,), - target_bitsizes=tuple(qrom_target_bitsizes), - ) - qrom_op = qrom_gate.on_registers( - selection=q, **split_qubits(qrom_gate.target_registers, ordered_target_qubits) - ) - if self.block_size > 1: - swap_with_zero_gate = SwapWithZero( - k, total_bits(self.target_registers), self.block_size - ) - swap_with_zero_op = swap_with_zero_gate.on_registers( - selection=r, - **split_qubits(swap_with_zero_gate.target_registers, ordered_target_qubits), - ) - clean_targets = merge_qubits(self.target_registers, **targets) - cnot_op = cirq.Moment(cirq.CNOT(s, t) for s, t in zip(ordered_target_qubits, clean_targets)) + ) -> cirq.OP_TREE: + # 1. Construct QROM to load the batched data. + qrom = self.qrom_bloq.with_data(*self.batched_data) + qrom_ctrls = {reg.name: quregs[reg.name] for reg in qrom.control_registers} + qrom_selection = { + qrom_reg.name: quregs[sel_reg.name][: qrom_reg.bitsize] + for sel_reg, qrom_reg in zip(self.selection_registers, qrom.selection_registers) + } + qrom_targets = { + reg.name: self._alloc_qrom_target_qubits(reg, context.qubit_manager) + for reg in qrom.target_registers + } + qrom_op = qrom.on_registers(**qrom_ctrls, **qrom_selection, **qrom_targets) + # 2. Construct SwapWithZero. + swz_ops = [] + assert len(qrom_targets) == len(self.swap_with_zero_bloqs) + for targets, swz in zip(qrom_targets.values(), self.swap_with_zero_bloqs): + if len(targets) <= 1: + continue + swz_selection = { + swz_reg.name: quregs[sel_reg.name][-swz_reg.bitsize :] + for sel_reg, swz_reg in zip(self.selection_registers, swz.selection_registers) + } + swz_ops.append(swz.on_registers(**swz_selection, targets=targets)) + # 3. Construct CNOTs from 0th borrowed register to clean target registers. + cnot_ops = [] + for qrom_batched_target, target_reg in zip(qrom_targets.values(), self.target_registers): + cnot_ops += [ + [cirq.CNOT(a, b) for a, b in zip(qrom_batched_target[0], quregs[target_reg.name])] + ] + # Yield the operations in correct order. - if self.block_size > 1: + if any(b > 0 for b in self.block_sizes): yield qrom_op - yield swap_with_zero_op - yield from cnot_op - yield cirq.inverse(swap_with_zero_op) + yield swz_ops + yield cnot_ops + yield cirq.inverse(swz_ops) yield cirq.inverse(qrom_op) - yield swap_with_zero_op - yield from cnot_op - yield cirq.inverse(swap_with_zero_op) + if self.use_dirty_ancilla: + yield swz_ops + yield cnot_ops + yield cirq.inverse(swz_ops) else: yield qrom_op - yield from cnot_op + yield cnot_ops yield cirq.inverse(qrom_op) - yield from cnot_op + yield cnot_ops + + context.qubit_manager.qfree( + [q for targets in qrom_targets.values() for q in targets.flatten()] + ) - context.qubit_manager.qfree(ordered_target_qubits) + def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo: + from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info - def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo: - wire_symbols = ["In_q"] * self.selection_q - wire_symbols += ["In_r"] * self.selection_r - for i, target in enumerate(self.target_registers): - wire_symbols += [f"QROAM_{i}"] * target.total_bits() - return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols) + return _wire_symbol_to_cirq_diagram_info(self, args) def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: @@ -252,10 +325,78 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - trg_indx = int(name.replace('target', '').replace('_', '')) # match the sel index subscript = chr(ord('a') + trg_indx) - return TextBox(f'data_{subscript}') + return TextBox(f'QROAM_{subscript}') elif name == 'control': return Circle() raise ValueError(f'Unknown register name {name}') def _value_equality_values_(self): - return self.block_size, self._target_bitsizes, self.data + return self.log_block_sizes, *super()._value_equality_values_() + + +@bloq_example +def _qroam_multi_data() -> SelectSwapQROM: + data1 = np.arange(5) + data2 = np.arange(5) + 1 + qroam_multi_data = SelectSwapQROM.build_from_data([data1, data2]) + return qroam_multi_data + + +@bloq_example +def _qroam_multi_dim() -> SelectSwapQROM: + data1 = np.arange(25).reshape((5, 5)) + data2 = (np.arange(25) + 1).reshape((5, 5)) + qroam_multi_dim = SelectSwapQROM.build_from_data([data1, data2]) + return qroam_multi_dim + + +@bloq_example +def _qroam_symb_dirty_1d() -> SelectSwapQROM: + N, b, k, c = sympy.symbols('N b k c') + qroam_symb_dirty_1d = SelectSwapQROM.build_from_bitsize( + (N,), (b,), log_block_sizes=(k,), num_controls=c + ) + return qroam_symb_dirty_1d + + +@bloq_example +def _qroam_symb_dirty_2d() -> SelectSwapQROM: + N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c') + log_block_sizes = (k1, k2) + qroam_symb_dirty_2d = SelectSwapQROM.build_from_bitsize( + (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c + ) + return qroam_symb_dirty_2d + + +@bloq_example +def _qroam_symb_clean_1d() -> SelectSwapQROM: + N, b, k, c = sympy.symbols('N b k c') + qroam_symb_clean_1d = SelectSwapQROM.build_from_bitsize( + (N,), (b,), log_block_sizes=(k,), num_controls=c, use_dirty_ancilla=False + ) + return qroam_symb_clean_1d + + +@bloq_example +def _qroam_symb_clean_2d() -> SelectSwapQROM: + N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c') + log_block_sizes = (k1, k2) + qroam_symb_clean_2d = SelectSwapQROM.build_from_bitsize( + (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c, use_dirty_ancilla=False + ) + return qroam_symb_clean_2d + + +_SELECT_SWAP_QROM_DOC = BloqDocSpec( + bloq_cls=SelectSwapQROM, + import_line='from qualtran.bloqs.data_loading.select_swap_qrom import SelectSwapQROM', + examples=[ + _qroam_multi_data, + _qroam_multi_dim, + _qroam_symb_dirty_1d, + _qroam_symb_dirty_2d, + _qroam_symb_clean_1d, + _qroam_symb_clean_2d, + ], +) diff --git a/qualtran/bloqs/data_loading/select_swap_qrom_test.py b/qualtran/bloqs/data_loading/select_swap_qrom_test.py index 1425a250a..3f050027b 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom_test.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom_test.py @@ -34,25 +34,28 @@ data, block_size, id=f"{block_size}-data{didx}", - marks=pytest.mark.slow if block_size == 3 and didx == 0 else (), + marks=pytest.mark.slow if block_size == 2 and didx == 0 else (), ) for didx, data in enumerate([[[1, 2, 3, 4, 5]], [[1, 2, 3], [3, 2, 1]]]) - for block_size in [None, 1, 2, 3] + for block_size in [None, 0, 1] ], ) def test_select_swap_qrom(data, block_size): - qrom = SelectSwapQROM(*data, block_size=block_size) - + qrom = SelectSwapQROM.build_from_data(*data, log_block_sizes=block_size) assert_valid_bloq_decomposition(qrom) qubit_regs = get_named_qubits(qrom.signature) selection = qubit_regs["selection"] - selection_q, selection_r = selection[: qrom.selection_q], selection[qrom.selection_q :] + q_len = qrom.batched_qrom_selection_bitsizes[0] + assert isinstance(q_len, int) + selection_q, selection_r = (selection[:q_len], selection[q_len:]) targets = [qubit_regs[f"target{i}_"] for i in range(len(data))] greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True) context = cirq.DecompositionContext(greedy_mm) - qrom_circuit = cirq.Circuit(cirq.decompose(qrom.on_registers(**qubit_regs), context=context)) + qrom_circuit = cirq.Circuit( + cirq.decompose_once(qrom.on_registers(**qubit_regs), context=context) + ) dirty_target_ancilla = [ q for q in qrom_circuit.all_qubits() if isinstance(q, cirq.ops.BorrowableQubit) @@ -72,8 +75,8 @@ def test_select_swap_qrom(data, block_size): dtype = qrom.selection_registers[0].dtype assert isinstance(dtype, BoundedQUInt) for selection_integer in range(int(dtype.iteration_length)): - svals_q = list(iter_bits(selection_integer // qrom.block_size, len(selection_q))) - svals_r = list(iter_bits(selection_integer % qrom.block_size, len(selection_r))) + svals_q = list(iter_bits(selection_integer // qrom.block_sizes[0], len(selection_q))) + svals_r = list(iter_bits(selection_integer % qrom.block_sizes[0], len(selection_r))) qubit_vals = {x: 0 for x in all_qubits} qubit_vals.update({s: sval for s, sval in zip(selection_q, svals_q)}) qubit_vals.update({s: sval for s, sval in zip(selection_r, svals_r)}) @@ -92,42 +95,44 @@ def test_select_swap_qrom(data, block_size): def test_qroam_diagram(): data = [[1, 2, 3], [4, 5, 6]] blocksize = 2 - qrom = SelectSwapQROM(*data, block_size=blocksize) + qrom = SelectSwapQROM.build_from_data(*data, log_block_sizes=(1,)) q = cirq.LineQubit.range(cirq.num_qubits(qrom)) circuit = cirq.Circuit(qrom.on_registers(**split_qubits(qrom.signature, q))) cirq.testing.assert_has_diagram( circuit, """ -0: ───In_q────── +0: ───In──────── │ -1: ───In_r────── +1: ───In──────── │ -2: ───QROAM_0─── +2: ───QROAM_a─── │ -3: ───QROAM_0─── +3: ───QROAM_a─── │ -4: ───QROAM_1─── +4: ───QROAM_b─── │ -5: ───QROAM_1─── +5: ───QROAM_b─── │ -6: ───QROAM_1─── +6: ───QROAM_b─── """, ) def test_qroam_raises(): - with pytest.raises(ValueError, match="must be of equal length"): - _ = SelectSwapQROM([1, 2], [1, 2, 3]) + with pytest.raises(ValueError, match="must have same shape"): + _ = SelectSwapQROM.build_from_data([1, 2], [1, 2, 3]) def test_qroam_hashable(): - qrom = SelectSwapQROM([1, 2, 5, 6, 7, 8]) + qrom = SelectSwapQROM.build_from_data([1, 2, 5, 6, 7, 8]) assert hash(qrom) is not None assert t_complexity(qrom) == TComplexity(32, 160, 0) def test_qrom_t_complexity(): - qrom = SelectSwapQROM([1, 2, 3, 4, 5, 6, 7, 8], target_bitsizes=(4,), block_size=4) + qrom = SelectSwapQROM.build_from_data( + [1, 2, 3, 4, 5, 6, 7, 8], target_bitsizes=(4,), log_block_sizes=(2,) + ) _, sigma = qrom.call_graph() assert t_counts_from_sigma(sigma) == qrom.t_complexity().t == 192 @@ -135,8 +140,8 @@ def test_qrom_t_complexity(): def test_qroam_many_registers(): # Test > 10 registers which resulted in https://github.com/quantumlib/Qualtran/issues/556 target_bitsizes = (3,) * 10 + (1,) * 2 + (3,) - block_size = 2 ** find_optimal_log_block_size(10, sum(target_bitsizes)) - qrom = SelectSwapQROM( + log_block_size = find_optimal_log_block_size(10, sum(target_bitsizes)) + qrom = SelectSwapQROM.build_from_data( (1,) * 10, (1,) * 10, (1,) * 10, @@ -150,7 +155,6 @@ def test_qroam_many_registers(): (0,) * 10, (1,) * 10, (3,) * 10, - target_bitsizes=target_bitsizes, - block_size=block_size, + log_block_sizes=(log_block_size,), ) qrom.call_graph()