diff --git a/qualtran/__init__.py b/qualtran/__init__.py index a4829c0dc..6b24feaca 100644 --- a/qualtran/__init__.py +++ b/qualtran/__init__.py @@ -40,6 +40,7 @@ BloqBuilder, DidNotFlattenAnythingError, SoquetT, + ConnectionT, ) from ._infra.data_types import ( diff --git a/qualtran/_infra/adjoint.py b/qualtran/_infra/adjoint.py index 16b9d2d3d..81c556ee9 100644 --- a/qualtran/_infra/adjoint.py +++ b/qualtran/_infra/adjoint.py @@ -19,7 +19,7 @@ from attrs import frozen from numpy.typing import NDArray -from .composite_bloq import _binst_to_cxns, _cxn_to_soq_dict, _map_soqs, _reg_to_soq, BloqBuilder +from .composite_bloq import _binst_to_cxns, _cxns_to_soq_dict, _map_soqs, _reg_to_soq, BloqBuilder from .gate_with_registers import GateWithRegisters from .quantum_graph import LeftDangle, RightDangle from .registers import Signature @@ -35,7 +35,7 @@ def _adjoint_final_soqs(cbloq: 'CompositeBloq', new_signature: Signature) -> Dic if LeftDangle not in cbloq._binst_graph: return {} _, init_succs = _binst_to_cxns(LeftDangle, binst_graph=cbloq._binst_graph) - return _cxn_to_soq_dict( + return _cxns_to_soq_dict( new_signature.rights(), init_succs, get_me=lambda x: x.left, get_assign=lambda x: x.right ) @@ -68,7 +68,7 @@ def _adjoint_cbloq(cbloq: 'CompositeBloq') -> 'CompositeBloq': for binst, preds, succs in bloqnections: # Instead of get_me returning the right element of a predecessor connection, # it's the left element of a successor connection. - soqs = _cxn_to_soq_dict( + soqs = _cxns_to_soq_dict( binst.bloq.signature.rights(), succs, get_me=lambda x: x.left, diff --git a/qualtran/_infra/bloq.py b/qualtran/_infra/bloq.py index d98a44565..55830a41b 100644 --- a/qualtran/_infra/bloq.py +++ b/qualtran/_infra/bloq.py @@ -16,7 +16,7 @@ """Contains the main interface for defining `Bloq`s.""" import abc -from typing import Any, Callable, Dict, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union if TYPE_CHECKING: import cirq @@ -30,11 +30,10 @@ Adjoint, BloqBuilder, CompositeBloq, + ConnectionT, CtrlSpec, - GateWithRegisters, Register, Signature, - Soquet, SoquetT, ) from qualtran.cirq_interop import CirqQuregT @@ -246,38 +245,39 @@ def tensor_contract(self) -> 'NDArray': return bloq_to_dense(self) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: """Override this method to support native quimb simulation of this Bloq. - This method is responsible for adding a tensor corresponding to the unitary, state, or - effect of the bloq to the provided tensor network `tn`. Often, this method will add - one tensor for a given Bloq, but some bloqs can be represented in a factorized form - requiring the addition of more than one tensor. - - If this method is not overriden, the default implementation will try to use the bloq's - decomposition to find a dense representation for this bloq. + This method is responsible for returning tensors corresponding to the unitary, state, or + effect of the bloq. Often, this method will return one tensor for a given Bloq, but + some bloqs can be represented in a factorized form using more than one tensor. + + By default, calls to `Bloq.tensor_contract()` will first decompose and flatten the bloq + before initiating the conversion to a tensor network. This has two consequences: + 1) Overriding this method is only necessary if this bloq does not define a decomposition + or if the fully-decomposed form contains a bloq that does not define its tensors. + 2) Even if you override this method to provide custom tensors, they may not be used + (by default) because we prefer the flat-decomposed version. This is usually desirable + for contraction performance; but for finer-grained control see + `qualtran.simulation.tensor.cbloq_to_quimb`. + + Quimb defines a connection between two tensors by a shared index. The returned tensors + from this method must use the Qualtran-Quimb index convention: + - Each tensor index is a tuple `(cxn, j)` + - The `cxn: qualtran.Connection` entry identifies the connection between bloq instances. + - The second integer `j` is the bit index within high-bitsize registers, + which is necessary due to technical restrictions. Args: - tn: The tensor network to which we add our tensor(s) - tag: An arbitrary tag that must be forwarded to `qtn.Tensor`'s `tag` attribute. - incoming: A mapping from register name to SoquetT to order left indices for - the tensor network. - outgoing: A mapping from register name to SoquetT to order right indices for - the tensor network. + incoming: A mapping from register name to Connection (or an array thereof) to use as + left indices for the tensor network. The shape of the array matches the register's + shape. + outgoing: A mapping from register name to Connection (or an array thereof) to use as + right indices for the tensor network. """ - from qualtran.simulation.tensor import cbloq_as_contracted_tensor - - cbloq = self.decompose_bloq() - tn.add( - cbloq_as_contracted_tensor(cbloq, incoming, outgoing, tags=[self.pretty_name(), tag]) - ) + raise NotImplementedError(f"{self} does not support tensor simulation.") def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: """Override this method to build the bloq call graph. diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py index e483774ee..b5b28f59b 100644 --- a/qualtran/_infra/composite_bloq.py +++ b/qualtran/_infra/composite_bloq.py @@ -66,6 +66,11 @@ canonicalize and return `SoquetT`. """ +_ConnectionType = TypeVar('_ConnectionType', bound=np.generic) + +ConnectionT = Union[Connection, NDArray[_ConnectionType]] +"""A `Connection` or array of connections.""" + def _to_tuple(x: Iterable[Connection]) -> Sequence[Connection]: """mypy-compatible attrs converter for CompositeBloq.connections""" @@ -280,7 +285,7 @@ def iter_bloqsoqs( """ for binst, preds, succs in self.iter_bloqnections(): - in_soqs = _cxn_to_soq_dict( + in_soqs = _cxns_to_soq_dict( binst.bloq.signature.lefts(), preds, get_me=lambda x: x.right, @@ -297,7 +302,7 @@ def final_soqs(self) -> Dict[str, SoquetT]: if RightDangle not in self._binst_graph: return {} final_preds, _ = _binst_to_cxns(RightDangle, binst_graph=self._binst_graph) - return _cxn_to_soq_dict( + return _cxns_to_soq_dict( self.signature.rights(), final_preds, get_me=lambda x: x.right, @@ -341,6 +346,9 @@ def flatten_once( the bloqs have decompositions. """ + if len(self.bloq_instances) == 0: + raise DidNotFlattenAnythingError() + bb, _ = BloqBuilder.from_signature(self.signature) # We take particular care during flattening to preserve the `binst.i` of bloq instances @@ -498,13 +506,13 @@ def _binst_to_cxns( return pred_cxns, succ_cxns -def _cxn_to_soq_dict( +def _cxns_to_soq_dict( regs: Iterable[Register], cxns: Iterable[Connection], get_me: Callable[[Connection], Soquet], get_assign: Callable[[Connection], Soquet], ) -> Dict[str, SoquetT]: - """Helper function to get a dictionary of incoming or outgoing soquets from a connection. + """Helper function to get a dictionary of soquets from a list of connections. Args: regs: Left or right registers (used as a reference to initialize multidimensional @@ -513,9 +521,12 @@ def _cxn_to_soq_dict( get_me: A function that says which soquet is used to derive keys for the returned dictionary. Generally: if `cxns` is predecessor connections, this will return the `right` element of the connection and opposite of successor connections. - get_assign: A function that says which soquet is used to dervice the values for the + get_assign: A function that says which soquet is used to derive the values for the returned dictionary. Generally, this is the opposite side vs. `get_me`, but we do something fancier in `cbloq_to_quimb`. + + Returns: + soqdict: A dictionary mapping register name to the selected soquets. """ soqdict: Dict[str, SoquetT] = {} @@ -538,7 +549,42 @@ def _cxn_to_soq_dict( return soqdict -def _get_dangling_soquets(signature: Signature, right=True) -> Dict[str, SoquetT]: +def _cxns_to_cxn_dict( + regs: Iterable[Register], cxns: Iterable[Connection], get_me: Callable[[Connection], Soquet] +) -> Dict[str, ConnectionT]: + """Helper function to get a dictionary of connections from a list of connections + + Args: + regs: Left or right registers (used as a reference to initialize multidimensional + registers correctly). + cxns: Predecessor or successor connections from which we get the connections of interest. + get_me: A function that says which soquet is used to derive keys for the returned + dictionary. Generally: if `cxns` is predecessor connections, this will return the + `right` element of the connection (opposite for successor connections). + + Returns: + cxndict: A dictionary mapping register name to the selected connections. + """ + cxndict: Dict[str, ConnectionT] = {} + + # Initialize multi-dimensional dictionary values. + for reg in regs: + if reg.shape: + cxndict[reg.name] = np.empty(reg.shape, dtype=object) + + # In the abstract: set `soqdict[me] = assign`. Specifically: use the register name as + # keys and handle multi-dimensional registers. + for cxn in cxns: + me = get_me(cxn) + if me.reg.shape: + cxndict[me.reg.name][me.idx] = cxn # type: ignore[index] + else: + cxndict[me.reg.name] = cxn + + return cxndict + + +def _get_dangling_soquets(signature: Signature, right: bool = True) -> Dict[str, SoquetT]: """Get instantiated dangling soquets from a `Signature`. Args: diff --git a/qualtran/_infra/controlled.py b/qualtran/_infra/controlled.py index 31161271d..c30b451b0 100644 --- a/qualtran/_infra/controlled.py +++ b/qualtran/_infra/controlled.py @@ -39,7 +39,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran import BloqBuilder, CompositeBloq, Soquet, SoquetT + from qualtran import BloqBuilder, CompositeBloq, ConnectionT, SoquetT from qualtran.cirq_interop import CirqQuregT from qualtran.drawing import WireSymbol from qualtran.resource_counting import BloqCountT, SympySymbolAllocator @@ -400,17 +400,7 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT return vals - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - import quimb.tensor as qtn - - from qualtran._infra.composite_bloq import _flatten_soquet_collection + def _tensor_data(self): from qualtran.simulation.tensor._tensor_data_manipulation import ( active_space_for_ctrl_spec, eye_tensor_for_signature, @@ -419,17 +409,15 @@ def add_my_tensors( # Create an identity tensor corresponding to the signature of current Bloq data = eye_tensor_for_signature(self.signature) - # Verify it has the right shape - in_ind = _flatten_soquet_collection(incoming[reg.name] for reg in self.signature.lefts()) - out_ind = _flatten_soquet_collection(outgoing[reg.name] for reg in self.signature.rights()) - assert data.shape == tuple(2**soq.reg.bitsize for ind in [out_ind, in_ind] for soq in ind) # Figure out the ctrl indexes for which the ctrl is "active" - active_idx = active_space_for_ctrl_spec(self.signature, self.ctrl_spec) - # Put the subbloq tensor at indices where ctrl is active. subbloq_shape = tensor_shape_from_signature(self.subbloq.signature) - data[active_idx] = self.subbloq.tensor_contract().reshape(subbloq_shape) - # Add the data to the tensor network. - tn.add(qtn.Tensor(data=data, inds=out_ind + in_ind, tags=[self.pretty_name(), tag])) + subbloq_tensor = self.subbloq.tensor_contract() + if subbloq_shape: + subbloq_tensor = subbloq_tensor.reshape(subbloq_shape) + # Put the subbloq tensor at indices where ctrl is active. + active_idx = active_space_for_ctrl_spec(self.signature, self.ctrl_spec) + data[active_idx] = subbloq_tensor + return data def _unitary_(self): if isinstance(self.subbloq, GateWithRegisters): @@ -444,11 +432,24 @@ def _unitary_(self): # Unable to determine the unitary effect. return NotImplemented + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + import quimb.tensor as qtn + + from qualtran.simulation.tensor._dense import _order_incoming_outgoing_indices + + inds = _order_incoming_outgoing_indices( + self.signature, incoming=incoming, outgoing=outgoing + ) + data = self._tensor_data().reshape((2,) * len(inds)) + return [qtn.Tensor(data=data, inds=inds, tags=[str(self)])] + def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': from qualtran.drawing import Text if reg is None: - return Text(f'C[{self.subbloq.wire_symbol(reg=None)}]') + return Text(f'C[{self.subbloq}]') if reg.name not in self.ctrl_reg_names: # Delegate to subbloq return self.subbloq.wire_symbol(reg, idx) diff --git a/qualtran/_infra/controlled_test.py b/qualtran/_infra/controlled_test.py index 55cd90292..19f52836b 100644 --- a/qualtran/_infra/controlled_test.py +++ b/qualtran/_infra/controlled_test.py @@ -35,10 +35,12 @@ from qualtran._infra.gate_with_registers import get_named_qubits, merge_qubits from qualtran.bloqs.basic_gates import ( CSwap, + GlobalPhase, IntEffect, IntState, OneState, Swap, + TwoBitCSwap, XGate, XPowGate, YGate, @@ -50,6 +52,7 @@ from qualtran.cirq_interop.testing import GateHelper from qualtran.drawing import get_musical_score_data from qualtran.drawing.musical_score import Circle, SoqData, TextBox +from qualtran.simulation.tensor import cbloq_to_quimb, get_right_and_left_inds if TYPE_CHECKING: from qualtran import SoquetT @@ -340,7 +343,7 @@ def test_notebook(): def _verify_ctrl_tensor_for_unitary(ctrl_spec: CtrlSpec, bloq: Bloq, gate: cirq.Gate): cbloq = Controlled(bloq, ctrl_spec) cgate = cirq.ControlledGate(gate, control_values=ctrl_spec.to_cirq_cv()) - np.testing.assert_array_equal(cbloq.tensor_contract(), cirq.unitary(cgate)) + np.testing.assert_allclose(cbloq.tensor_contract(), cirq.unitary(cgate), atol=1e-8) interesting_ctrl_specs = [ @@ -362,6 +365,26 @@ def test_controlled_tensor_for_unitary(ctrl_spec: CtrlSpec): _verify_ctrl_tensor_for_unitary(ctrl_spec, CSwap(3), CSwap(3)) +def test_controlled_tensor_without_decompose(): + ctrl_spec = CtrlSpec() + bloq = TwoBitCSwap() + cbloq = Controlled(bloq, ctrl_spec) + cgate = cirq.ControlledGate(cirq.CSWAP, control_values=ctrl_spec.to_cirq_cv()) + + tn = cbloq_to_quimb(cbloq.as_composite_bloq()) + # pylint: disable=unbalanced-tuple-unpacking + right, left = get_right_and_left_inds(tn, cbloq.signature) + # pylint: enable=unbalanced-tuple-unpacking + np.testing.assert_allclose(tn.to_dense(right, left), cirq.unitary(cgate), atol=1e-8) + np.testing.assert_allclose(cbloq.tensor_contract(), cirq.unitary(cgate), atol=1e-8) + + +def test_controlled_global_phase_tensor(): + bloq = GlobalPhase(1.0j).controlled() + should_be = np.diag([1, 1.0j]) + np.testing.assert_allclose(bloq.tensor_contract(), should_be) + + @attrs.frozen class TestCtrlStatePrepAnd(Bloq): """Decomposes into a Controlled-AND gate + int effects & targets where ctrl is active. diff --git a/qualtran/_infra/gate_with_registers.py b/qualtran/_infra/gate_with_registers.py index f13e9da28..1905ab443 100644 --- a/qualtran/_infra/gate_with_registers.py +++ b/qualtran/_infra/gate_with_registers.py @@ -14,7 +14,6 @@ import abc from typing import ( - Any, cast, Collection, Dict, @@ -40,7 +39,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran import AddControlledT, BloqBuilder, CtrlSpec, SoquetT + from qualtran import AddControlledT, BloqBuilder, ConnectionT, CtrlSpec, SoquetT from qualtran.cirq_interop import CirqQuregT from qualtran.drawing import WireSymbol @@ -507,22 +506,15 @@ def controlled( def _unitary_(self): return NotImplemented - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: 'Any', - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: if not self._unitary_.__qualname__.startswith('GateWithRegisters.'): - from qualtran.cirq_interop._cirq_to_bloq import _add_my_tensors_from_gate + from qualtran.cirq_interop._cirq_to_bloq import _my_tensors_from_gate - _add_my_tensors_from_gate( - self, self.signature, str(self), tn, tag, incoming=incoming, outgoing=outgoing - ) + return _my_tensors_from_gate(self, self.signature, incoming=incoming, outgoing=outgoing) else: - return super().add_my_tensors(tn, tag, incoming=incoming, outgoing=outgoing) + return super().my_tensors(incoming=incoming, outgoing=outgoing) def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo: """Default diagram info that uses register names to name the boxes in multi-qubit gates. diff --git a/qualtran/bloqs/arithmetic/addition.py b/qualtran/bloqs/arithmetic/addition.py index 88543c8b3..882722ec7 100644 --- a/qualtran/bloqs/arithmetic/addition.py +++ b/qualtran/bloqs/arithmetic/addition.py @@ -11,11 +11,9 @@ # 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 itertools import math from functools import cached_property from typing import ( - Any, Dict, Iterable, Iterator, @@ -63,8 +61,6 @@ from qualtran.drawing import directional_text_box, Text if TYPE_CHECKING: - import quimb.tensor as qtn - from qualtran.drawing import WireSymbol from qualtran.resource_counting import BloqCountT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -125,28 +121,6 @@ def dtype(self): def signature(self): return Signature([Register("a", self.a_dtype), Register("b", self.b_dtype)]) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - import quimb.tensor as qtn - - if isinstance(self.a_dtype, QInt) or isinstance(self.b_dtype, QInt): - raise TypeError("Tensor contraction for addition is only supported for unsigned ints.") - N_a = 2**self.a_dtype.bitsize - N_b = 2**self.b_dtype.bitsize - inds = (incoming['a'], incoming['b'], outgoing['a'], outgoing['b']) - unitary = np.zeros((N_a, N_b, N_a, N_b), dtype=np.complex128) - # TODO: Add a value-to-index method on dtype to make this easier. - for a, b in itertools.product(range(N_a), range(N_b)): - unitary[a, b, a, int(math.fmod(a + b, N_b))] = 1 - - tn.add(qtn.Tensor(data=unitary, inds=inds, tags=[self.pretty_name(), tag])) - def decompose_bloq(self) -> 'CompositeBloq': return decompose_from_cirq_style_method(self) diff --git a/qualtran/bloqs/arithmetic/addition_test.py b/qualtran/bloqs/arithmetic/addition_test.py index 039008cfb..858ea21f2 100644 --- a/qualtran/bloqs/arithmetic/addition_test.py +++ b/qualtran/bloqs/arithmetic/addition_test.py @@ -171,7 +171,7 @@ def test_addition_gate_counts(n: int): @pytest.mark.parametrize('a,b', itertools.product(range(2**3), repeat=2)) -def test_add_no_decompose(a, b): +def test_add_tensor_contract(a, b): num_bits = 5 bloq = Add(QUInt(num_bits)) @@ -184,7 +184,7 @@ def test_add_no_decompose(a, b): assert true_out_int == int(out_bin, 2) unitary = bloq.tensor_contract() - assert unitary[output_int, input_int] == 1 + np.testing.assert_allclose(unitary[output_int, input_int], 1) @pytest.mark.parametrize('a,b,num_bits', itertools.product(range(4), range(4), range(3, 5))) diff --git a/qualtran/bloqs/arithmetic/multiplication.py b/qualtran/bloqs/arithmetic/multiplication.py index e5ef7dd25..01ae9f44b 100644 --- a/qualtran/bloqs/arithmetic/multiplication.py +++ b/qualtran/bloqs/arithmetic/multiplication.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Iterable, Sequence, Set, TYPE_CHECKING, Union +from typing import Dict, Iterable, List, Sequence, Set, TYPE_CHECKING, Union import cirq import numpy as np @@ -22,6 +22,7 @@ Bloq, bloq_example, BloqDocSpec, + ConnectionT, GateWithRegisters, QFxp, QUInt, @@ -35,7 +36,6 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran import SoquetT from qualtran.resource_counting import BloqCountT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -90,19 +90,12 @@ def __pow__(self, power): ) raise NotImplementedError("PlusEqualProduct.__pow__ defined only for powers +1/-1.") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - from qualtran.cirq_interop._cirq_to_bloq import _add_my_tensors_from_gate - - _add_my_tensors_from_gate( - self, self.signature, self.pretty_name(), tn, tag, incoming=incoming, outgoing=outgoing - ) + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + from qualtran.cirq_interop._cirq_to_bloq import _my_tensors_from_gate + + return _my_tensors_from_gate(self, self.signature, incoming=incoming, outgoing=outgoing) def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: # TODO: The T-complexity here is approximate. @@ -169,25 +162,25 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: num_toff = self.bitsize * (self.bitsize - 1) return {(Toffoli(), num_toff)} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn + n = self.bitsize N = 2**self.bitsize data = np.zeros((N, N, N**2), dtype=np.complex128) for x in range(N): data[x, x, x**2] = 1 trg = incoming['result'] if self.uncompute else outgoing['result'] - tn.add( - qtn.Tensor(data=data, inds=(incoming['a'], outgoing['a'], trg), tags=['Square', tag]) + inds = ( + [(incoming['a'], j) for j in range(n)] + + [(outgoing['a'], j) for j in range(n)] + + [(trg, j) for j in range(2 * n)] ) + data = data.reshape((2,) * (4 * n)) + return [qtn.Tensor(data=data, inds=inds, tags=[str(self)])] def adjoint(self) -> 'Bloq': return Square(self.bitsize, not self.uncompute) diff --git a/qualtran/bloqs/arithmetic/subtraction_test.py b/qualtran/bloqs/arithmetic/subtraction_test.py index d878d77f8..e1da18665 100644 --- a/qualtran/bloqs/arithmetic/subtraction_test.py +++ b/qualtran/bloqs/arithmetic/subtraction_test.py @@ -31,7 +31,7 @@ def test_subtract_bloq_decomposition(): c = (a - b) % 32 want[(a << 5) | c][a_b] = 1 got = gate.tensor_contract() - np.testing.assert_equal(got, want) + np.testing.assert_allclose(got, want) def test_subtract_bloq_validation(): diff --git a/qualtran/bloqs/basic_gates/cnot.py b/qualtran/bloqs/basic_gates/cnot.py index c92d49478..62ae09258 100644 --- a/qualtran/bloqs/basic_gates/cnot.py +++ b/qualtran/bloqs/basic_gates/cnot.py @@ -14,7 +14,7 @@ import itertools from functools import cached_property -from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING +from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING import numpy as np from attrs import frozen @@ -26,6 +26,7 @@ BloqBuilder, BloqDocSpec, CompositeBloq, + ConnectionT, CtrlSpec, DecomposeTypeError, Register, @@ -69,35 +70,27 @@ def decompose_bloq(self) -> 'CompositeBloq': def adjoint(self) -> 'Bloq': return self - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): - """Append tensors to `tn` that represent this operation. - - This bloq uses the factored form of CNOT composed of a COPY and XOR tensor joined - by an internal index. - - References: - [Lectures on Quantum Tensor Networks](https://arxiv.org/abs/1912.10049). Biamonte 2019. - """ + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + # This bloq uses the factored form of CNOT composed of a COPY and XOR tensor joined + # by an internal index. + # [Lectures on Quantum Tensor Networks](https://arxiv.org/abs/1912.10049). Biamonte 2019. import quimb.tensor as qtn internal = qtn.rand_uuid() - tn.add( + return [ qtn.Tensor( - data=COPY, inds=(incoming['ctrl'], outgoing['ctrl'], internal), tags=['COPY', tag] - ) - ) - tn.add( + data=COPY, + inds=[(incoming['ctrl'], 0), (outgoing['ctrl'], 0), internal], + tags=['COPY'], + ), qtn.Tensor( - data=XOR, inds=(incoming['target'], outgoing['target'], internal), tags=['XOR'] - ) - ) + data=XOR, + inds=[(incoming['target'], 0), (outgoing['target'], 0), internal], + tags=['XOR'], + ), + ] def on_classical_vals(self, ctrl: int, target: int) -> Dict[str, 'ClassicalValT']: return {'ctrl': ctrl, 'target': (ctrl + target) % 2} diff --git a/qualtran/bloqs/basic_gates/global_phase.py b/qualtran/bloqs/basic_gates/global_phase.py index 5b23bd88c..4c0a87a05 100644 --- a/qualtran/bloqs/basic_gates/global_phase.py +++ b/qualtran/bloqs/basic_gates/global_phase.py @@ -12,18 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING +from typing import Dict, List, TYPE_CHECKING import attrs import cirq from attrs import frozen -from qualtran import bloq_example, BloqDocSpec, DecomposeTypeError +from qualtran import bloq_example, BloqDocSpec, ConnectionT, DecomposeTypeError from qualtran.cirq_interop import CirqGateAsBloqBase from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.symbolics import sconj, SymbolicComplex if TYPE_CHECKING: + import quimb.tensor as qtn + from qualtran import CompositeBloq @@ -49,6 +51,13 @@ def decompose_bloq(self) -> 'CompositeBloq': def adjoint(self) -> 'GlobalPhase': return attrs.evolve(self, coefficient=sconj(self.coefficient)) + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + import quimb.tensor as qtn + + return [qtn.Tensor(data=self.coefficient, inds=[], tags=[str(self)])] + def pretty_name(self) -> str: return 'GPhase' diff --git a/qualtran/bloqs/basic_gates/hadamard.py b/qualtran/bloqs/basic_gates/hadamard.py index 6a834d780..02a3bf324 100644 --- a/qualtran/bloqs/basic_gates/hadamard.py +++ b/qualtran/bloqs/basic_gates/hadamard.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import frozen @@ -23,10 +23,10 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, Register, Signature, - SoquetT, ) from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.drawing import Text, TextBox, WireSymbol @@ -65,17 +65,16 @@ def adjoint(self) -> 'Bloq': def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add(qtn.Tensor(data=_HADAMARD, inds=(outgoing['q'], incoming['q']), tags=["H", tag])) + return [ + qtn.Tensor( + data=_HADAMARD, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)] + ) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT' # type: ignore[type-var] diff --git a/qualtran/bloqs/basic_gates/identity.py b/qualtran/bloqs/basic_gates/identity.py index ae3e26fcf..86f03ed1f 100644 --- a/qualtran/bloqs/basic_gates/identity.py +++ b/qualtran/bloqs/basic_gates/identity.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import frozen @@ -23,10 +23,10 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, Register, Signature, - SoquetT, ) from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.drawing import Text, TextBox, WireSymbol @@ -57,17 +57,16 @@ def adjoint(self) -> 'Bloq': def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add(qtn.Tensor(data=np.eye(2), inds=(outgoing['q'], incoming['q']), tags=["I", tag])) + return [ + qtn.Tensor( + data=np.eye(2), inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)] + ) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT' # type: ignore[type-var] diff --git a/qualtran/bloqs/basic_gates/s_gate.py b/qualtran/bloqs/basic_gates/s_gate.py index 3dbb7020f..a716d06fe 100644 --- a/qualtran/bloqs/basic_gates/s_gate.py +++ b/qualtran/bloqs/basic_gates/s_gate.py @@ -13,13 +13,13 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING import attrs import numpy as np from attrs import frozen -from qualtran import Bloq, bloq_example, BloqDocSpec, Register, Signature, SoquetT +from qualtran import Bloq, bloq_example, BloqDocSpec, ConnectionT, Register, Signature from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.drawing import Text, TextBox, WireSymbol @@ -57,22 +57,15 @@ def signature(self) -> 'Signature': def _t_complexity_(self) -> 'TComplexity': return TComplexity(clifford=1) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn data = _SMATRIX.conj().T if self.is_adjoint else _SMATRIX - tn.add( - qtn.Tensor( - data=data, inds=(outgoing['q'], incoming['q']), tags=[self.pretty_name(), tag] - ) - ) + return [ + qtn.Tensor(data=data, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)]) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT' # type:ignore[type-var] diff --git a/qualtran/bloqs/basic_gates/su2_rotation.py b/qualtran/bloqs/basic_gates/su2_rotation.py index 79022f659..3e7e2f51f 100644 --- a/qualtran/bloqs/basic_gates/su2_rotation.py +++ b/qualtran/bloqs/basic_gates/su2_rotation.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np import sympy from attrs import frozen from numpy.typing import NDArray -from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Register, Signature +from qualtran import bloq_example, BloqDocSpec, ConnectionT, GateWithRegisters, Register, Signature from qualtran.bloqs.basic_gates import GlobalPhase, Ry, ZPowGate from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.drawing import Text, TextBox @@ -106,23 +106,18 @@ def from_matrix(mat: NDArray[np.complex_]) -> 'SU2RotationGate': return SU2RotationGate(theta, phi, lambd, alpha) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add( + return [ qtn.Tensor( data=self.rotation_matrix, - inds=(outgoing['q'], incoming['q']), - tags=[self.pretty_name(), tag], + inds=[(outgoing['q'], 0), (incoming['q'], 0)], + tags=[str(self)], ) - ) + ] def _unitary_(self): if self.is_symbolic(): diff --git a/qualtran/bloqs/basic_gates/swap.py b/qualtran/bloqs/basic_gates/swap.py index 592b4378c..5f65ac2c2 100644 --- a/qualtran/bloqs/basic_gates/swap.py +++ b/qualtran/bloqs/basic_gates/swap.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Iterator, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union import cirq import numpy as np @@ -26,6 +26,7 @@ bloq_example, BloqBuilder, BloqDocSpec, + ConnectionT, DecomposeTypeError, GateWithRegisters, Register, @@ -83,20 +84,15 @@ def as_cirq_op( def _t_complexity_(self) -> 'TComplexity': return TComplexity(clifford=1) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn matrix = _swap_matrix() - out_inds = [outgoing['x'], outgoing['y']] - in_inds = [incoming['x'], incoming['y']] - tn.add(qtn.Tensor(data=matrix, inds=out_inds + in_inds, tags=["swap", tag])) + out_inds = [(outgoing['x'], 0), (outgoing['y'], 0)] + in_inds = [(incoming['x'], 0), (incoming['y'], 0)] + return [qtn.Tensor(data=matrix, inds=out_inds + in_inds, tags=[str(self)])] def on_classical_vals( self, x: 'ClassicalValT', y: 'ClassicalValT' @@ -152,20 +148,20 @@ def decompose_from_registers( yield [cirq.T(x), cirq.H(y)] yield [cirq.CNOT(y, x)] - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn + # TODO: https://github.com/quantumlib/Qualtran/issues/873. Since this bloq + # has a decomposition, it will be used (by default) whenever `bloq.tensor_contract()` + # is called on a bloq containing a TwoBitCSwap instead of this implementation. + # When this becomes a leaf bloq, this explicit tensor will be used. + matrix = _controlled_swap_matrix() - out_inds = [outgoing['ctrl'], outgoing['x'], outgoing['y']] - in_inds = [incoming['ctrl'], incoming['x'], incoming['y']] - tn.add(qtn.Tensor(data=matrix, inds=out_inds + in_inds, tags=["swap", tag])) + out_inds = [(outgoing['ctrl'], 0), (outgoing['x'], 0), (outgoing['y'], 0)] + in_inds = [(incoming['ctrl'], 0), (incoming['x'], 0), (incoming['y'], 0)] + return [qtn.Tensor(data=matrix, inds=out_inds + in_inds, tags=[str(self)])] def on_classical_vals( self, ctrl: 'ClassicalValT', x: 'ClassicalValT', y: 'ClassicalValT' diff --git a/qualtran/bloqs/basic_gates/swap_test.py b/qualtran/bloqs/basic_gates/swap_test.py index 26edaba0c..7a46ce621 100644 --- a/qualtran/bloqs/basic_gates/swap_test.py +++ b/qualtran/bloqs/basic_gates/swap_test.py @@ -62,6 +62,11 @@ def test_two_bit_swap_unitary_vs_cirq(): np.testing.assert_array_equal(swap.tensor_contract(), cirq.unitary(cirq.SWAP)) +def test_two_bit_cswap_unitary_vs_cirq(): + swap = TwoBitCSwap() + np.testing.assert_allclose(swap.tensor_contract(), cirq.unitary(cirq.CSWAP), atol=1e-8) + + def test_two_bit_swap_as_cirq_op(): q = cirq.LineQubit.range(2) expected_circuit = cirq.Circuit(cirq.SWAP(*q)) @@ -93,16 +98,16 @@ def _set_ctrl_two_bit_swap(ctrl_bit): def test_two_bit_cswap(): cswap = TwoBitCSwap() - np.testing.assert_array_equal(cswap.tensor_contract(), cirq.unitary(cirq.CSWAP)) + np.testing.assert_allclose(cswap.tensor_contract(), cirq.unitary(cirq.CSWAP), atol=1e-8) np.testing.assert_allclose( cswap.decompose_bloq().tensor_contract(), cirq.unitary(cirq.CSWAP), atol=1e-8 ) # Zero ctrl -- it's identity - np.testing.assert_array_equal(np.eye(4), _set_ctrl_two_bit_swap(0).tensor_contract()) + np.testing.assert_allclose(np.eye(4), _set_ctrl_two_bit_swap(0).tensor_contract(), atol=1e-8) # One ctrl -- it's swap - np.testing.assert_array_equal( - _swap_matrix().reshape(4, 4), _set_ctrl_two_bit_swap(1).tensor_contract() + np.testing.assert_allclose( + _swap_matrix().reshape(4, 4), _set_ctrl_two_bit_swap(1).tensor_contract(), atol=1e-8 ) # classical logic @@ -159,13 +164,15 @@ def test_cswap_unitary(): cswap = CSwap(bitsize=4) # Zero ctrl -- it's identity - np.testing.assert_array_equal(np.eye(2 ** (4 * 2)), _set_ctrl_swap(0, cswap).tensor_contract()) + np.testing.assert_allclose( + np.eye(2 ** (4 * 2)), _set_ctrl_swap(0, cswap).tensor_contract(), atol=1e-8 + ) # One ctrl -- it's multi-swap qubits = cirq.LineQubit.range(8) q_x, q_y = qubits[:4], qubits[4:] unitary = cirq.unitary(cirq.Circuit(cirq.SWAP(x, y) for x, y in zip(q_x, q_y))) - np.testing.assert_array_equal(unitary, _set_ctrl_swap(1, cswap).tensor_contract()) + np.testing.assert_allclose(unitary, _set_ctrl_swap(1, cswap).tensor_contract(), atol=1e-8) def test_cswap_classical(): diff --git a/qualtran/bloqs/basic_gates/t_gate.py b/qualtran/bloqs/basic_gates/t_gate.py index 2a4616fca..b80f752ff 100644 --- a/qualtran/bloqs/basic_gates/t_gate.py +++ b/qualtran/bloqs/basic_gates/t_gate.py @@ -13,13 +13,13 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING import attrs import numpy as np from attrs import frozen -from qualtran import Bloq, bloq_example, BloqDocSpec, Register, Signature, SoquetT +from qualtran import Bloq, bloq_example, BloqDocSpec, ConnectionT, Register, Signature from qualtran.drawing import Text, TextBox, WireSymbol if TYPE_CHECKING: @@ -74,22 +74,15 @@ class TGate(Bloq): def signature(self) -> 'Signature': return Signature.build(q=1) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn data = _TMATRIX.conj().T if self.is_adjoint else _TMATRIX - tn.add( - qtn.Tensor( - data=data, inds=(outgoing['q'], incoming['q']), tags=[self.pretty_name(), tag] - ) - ) + return [ + qtn.Tensor(data=data, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)]) + ] def adjoint(self) -> 'Bloq': return attrs.evolve(self, is_adjoint=not self.is_adjoint) diff --git a/qualtran/bloqs/basic_gates/toffoli.py b/qualtran/bloqs/basic_gates/toffoli.py index 3c8599c31..7c9d48fee 100644 --- a/qualtran/bloqs/basic_gates/toffoli.py +++ b/qualtran/bloqs/basic_gates/toffoli.py @@ -13,13 +13,13 @@ # limitations under the License. import itertools from functools import cached_property -from typing import Any, Dict, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING, Union import numpy as np from attrs import frozen from numpy.typing import NDArray -from qualtran import Bloq, bloq_example, BloqDocSpec, QBit, Register, Signature, Soquet +from qualtran import Bloq, bloq_example, BloqDocSpec, Connection, QBit, Register, Signature from qualtran.bloqs.basic_gates import TGate from qualtran.cirq_interop.t_complexity_protocol import TComplexity from qualtran.resource_counting import SympySymbolAllocator @@ -62,14 +62,11 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: def _t_complexity_(self): return TComplexity(t=4) - def add_my_tensors( + def my_tensors( self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, NDArray['Soquet']], # type: ignore[type-var] - outgoing: Dict[str, NDArray['Soquet']], # type: ignore[type-var] - ): + incoming: Dict[str, NDArray[Connection]], # type: ignore[type-var] + outgoing: Dict[str, NDArray[Connection]], # type: ignore[type-var] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn from qualtran.bloqs.basic_gates.cnot import XOR @@ -77,24 +74,26 @@ def add_my_tensors( # Set up the CTRL tensor which copies inputs to outputs and activates # when a==1 and b==1 internal = qtn.rand_uuid() - inds = ( - incoming['ctrl'][0], - incoming['ctrl'][1], - outgoing['ctrl'][0], - outgoing['ctrl'][1], + inds = [ + (incoming['ctrl'][0], 0), + (incoming['ctrl'][1], 0), + (outgoing['ctrl'][0], 0), + (outgoing['ctrl'][1], 0), internal, - ) + ] CTRL = np.zeros((2,) * 5, dtype=np.complex128) for a, b in itertools.product([0, 1], repeat=2): CTRL[a, b, a, b, int(a == 1 and b == 1)] = 1.0 # Wire up the CTRL tensor to XOR to flip `target` when active. - tn.add(qtn.Tensor(data=CTRL, inds=inds, tags=['COPY', tag])) - tn.add( + return [ + qtn.Tensor(data=CTRL, inds=inds, tags=['COPY']), qtn.Tensor( - data=XOR, inds=(incoming['target'], outgoing['target'], internal), tags=['XOR'] - ) - ) + data=XOR, + inds=[(incoming['target'], 0), (outgoing['target'], 0), internal], + tags=['XOR'], + ), + ] def on_classical_vals( self, ctrl: NDArray[np.integer], target: 'ClassicalValT' diff --git a/qualtran/bloqs/basic_gates/x_basis.py b/qualtran/bloqs/basic_gates/x_basis.py index 086b0bd41..88e8ba7de 100644 --- a/qualtran/bloqs/basic_gates/x_basis.py +++ b/qualtran/bloqs/basic_gates/x_basis.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING, Union +from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union import numpy as np from attrs import frozen @@ -24,6 +24,7 @@ bloq_example, BloqBuilder, BloqDocSpec, + ConnectionT, CtrlSpec, QBit, Register, @@ -72,24 +73,15 @@ def __attrs_post_init__(self): def signature(self) -> 'Signature': return Signature([Register('q', QBit(), side=Side.RIGHT if self.state else Side.LEFT)]) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn side = outgoing if self.state else incoming - tn.add( - qtn.Tensor( - data=_MINUS if self.bit else _PLUS, - inds=(side['q'],), - tags=[self.pretty_name(), tag], - ) - ) + return [ + qtn.Tensor(data=_MINUS if self.bit else _PLUS, inds=[(side['q'], 0)], tags=[str(self)]) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', **cirq_quregs: 'CirqQuregT' # type: ignore[type-var] @@ -222,21 +214,16 @@ def signature(self) -> 'Signature': def adjoint(self) -> 'Bloq': return self - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add( + return [ qtn.Tensor( - data=_PAULIX, inds=(outgoing['q'], incoming['q']), tags=[self.pretty_name(), tag] + data=_PAULIX, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)] ) - ) + ] def get_ctrl_system( self, ctrl_spec: Optional['CtrlSpec'] = None diff --git a/qualtran/bloqs/basic_gates/y_gate.py b/qualtran/bloqs/basic_gates/y_gate.py index 9eb676edf..b3ee7255e 100644 --- a/qualtran/bloqs/basic_gates/y_gate.py +++ b/qualtran/bloqs/basic_gates/y_gate.py @@ -13,12 +13,12 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING import numpy as np from attrs import frozen -from qualtran import Bloq, Signature, SoquetT +from qualtran import Bloq, ConnectionT, Signature from qualtran.cirq_interop.t_complexity_protocol import TComplexity if TYPE_CHECKING: @@ -44,17 +44,16 @@ def signature(self) -> 'Signature': def adjoint(self) -> 'Bloq': return self - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add(qtn.Tensor(data=_PAULIY, inds=(outgoing['q'], incoming['q']), tags=["Y", tag])) + return [ + qtn.Tensor( + data=_PAULIY, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)] + ) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT' diff --git a/qualtran/bloqs/basic_gates/z_basis.py b/qualtran/bloqs/basic_gates/z_basis.py index 589330062..48b69b021 100644 --- a/qualtran/bloqs/basic_gates/z_basis.py +++ b/qualtran/bloqs/basic_gates/z_basis.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, cast, Dict, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, List, Optional, Set, Tuple, TYPE_CHECKING, Union import attrs import numpy as np @@ -27,6 +27,7 @@ BloqBuilder, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QAny, QBit, @@ -82,24 +83,15 @@ def signature(self) -> 'Signature': def decompose_bloq(self) -> CompositeBloq: raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn side = outgoing if self.state else incoming - tn.add( - qtn.Tensor( - data=_ONE if self.bit else _ZERO, - inds=(side['q'],), - tags=['1' if self.bit else '0', tag], - ) - ) + return [ + qtn.Tensor(data=_ONE if self.bit else _ZERO, inds=[(side['q'], 0)], tags=[str(self)]) + ] def on_classical_vals(self, *, q: Optional[int] = None) -> Dict[str, int]: """Return or consume 1 or 0 depending on `self.state` and `self.bit`. @@ -248,17 +240,16 @@ def adjoint(self) -> 'Bloq': def decompose_bloq(self) -> CompositeBloq: raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add(qtn.Tensor(data=_PAULIZ, inds=(outgoing['q'], incoming['q']), tags=["Z", tag])) + return [ + qtn.Tensor( + data=_PAULIZ, inds=[(outgoing['q'], 0), (incoming['q'], 0)], tags=[str(self)] + ) + ] def as_cirq_op( self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT' @@ -344,30 +335,6 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **val: 'SoquetT') -> Dict[str, else: return self._build_composite_effect(bb, cast(Soquet, val['val']), bits) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): - import quimb.tensor as qtn - - if isinstance(self.bitsize, sympy.Expr): - raise ValueError(f'Symbolic bitsize {self.bitsize} not supported') - data = np.zeros(2**self.bitsize).reshape((2,) * self.bitsize) - bitstring = ints_to_bits(np.array([self.val]), w=self.bitsize)[0] - data[tuple(bitstring)] = 1 - data = data.reshape(-1) - - if self.state: - inds = (outgoing['val'],) - else: - inds = (incoming['val'],) - - tn.add(qtn.Tensor(data=data, inds=inds, tags=[f'{self.val}', tag])) - def on_classical_vals(self, *, val: Optional[int] = None) -> Dict[str, Union[int, sympy.Expr]]: if self.state: assert val is None @@ -405,6 +372,9 @@ class IntState(_IntVector): def __init__(self, val: Union[int, sympy.Expr], bitsize: Union[int, sympy.Expr]): self.__attrs_init__(val=val, bitsize=bitsize, state=True) + def __str__(self): + return f'IntState({self.val})' + @bloq_example def _int_state() -> IntState: @@ -430,6 +400,9 @@ class IntEffect(_IntVector): def __init__(self, val: int, bitsize: int): self.__attrs_init__(val=val, bitsize=bitsize, state=False) + def __str__(self): + return f'IntEffect({self.val})' + @bloq_example def _int_effect() -> IntEffect: diff --git a/qualtran/bloqs/block_encoding/lcu_block_encoding.py b/qualtran/bloqs/block_encoding/lcu_block_encoding.py index 107dc39ef..a0dbc5c8e 100644 --- a/qualtran/bloqs/block_encoding/lcu_block_encoding.py +++ b/qualtran/bloqs/block_encoding/lcu_block_encoding.py @@ -354,7 +354,6 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: SoquetT) -> Dict[str, def _extract_soqs(bloq: Bloq) -> Dict[str, 'SoquetT']: return {reg.name: soqs.pop(reg.name) for reg in bloq.signature.lefts()} - print('xxxx') soqs |= bb.add_d(self.prepare, **_extract_soqs(self.prepare)) soqs |= bb.add_d(self.select, **_extract_soqs(self.select)) soqs |= bb.add_d(self.prepare.adjoint(), **_extract_soqs(self.prepare.adjoint())) diff --git a/qualtran/bloqs/bookkeeping/allocate.py b/qualtran/bloqs/bookkeeping/allocate.py index 953807739..f7775a8d7 100644 --- a/qualtran/bloqs/bookkeeping/allocate.py +++ b/qualtran/bloqs/bookkeeping/allocate.py @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING -import numpy as np import sympy from attrs import frozen @@ -23,13 +22,13 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QDType, QUInt, Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq from qualtran.drawing import directional_text_box, Text, WireSymbol @@ -66,19 +65,17 @@ def adjoint(self) -> 'Bloq': def on_classical_vals(self) -> Dict[str, int]: return {'reg': 0} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - data = np.zeros(1 << self.dtype.num_qubits) - data[0] = 1 - tn.add(qtn.Tensor(data=data, inds=(outgoing['reg'],), tags=['Allocate', tag])) + from qualtran.bloqs.basic_gates.z_basis import _ZERO + + return [ + qtn.Tensor(data=_ZERO, inds=[(outgoing['reg'], i)], tags=[str(self)]) + for i in range(self.dtype.num_qubits) + ] def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: diff --git a/qualtran/bloqs/bookkeeping/cast.py b/qualtran/bloqs/bookkeeping/cast.py index f7a4f3409..ae574953f 100644 --- a/qualtran/bloqs/bookkeeping/cast.py +++ b/qualtran/bloqs/bookkeeping/cast.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING import attrs import numpy as np @@ -23,13 +23,13 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QDType, QFxp, Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq @@ -82,23 +82,17 @@ def signature(self) -> Signature: def adjoint(self) -> 'Bloq': return Cast(inp_dtype=self.out_dtype, out_dtype=self.inp_dtype) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add( + return [ qtn.Tensor( - data=np.eye(2**self.inp_dtype.num_qubits, 2**self.inp_dtype.num_qubits), - inds=[outgoing['reg']] + [incoming['reg']], - tags=['Cast', tag], + data=np.eye(2), inds=[(outgoing['reg'], j), (incoming['reg'], j)], tags=[str(self)] ) - ) + for j in range(self.inp_dtype.num_qubits) + ] def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']: if isinstance(self.out_dtype, QFxp): diff --git a/qualtran/bloqs/bookkeeping/cast_test.py b/qualtran/bloqs/bookkeeping/cast_test.py index d34da5699..7390258e9 100644 --- a/qualtran/bloqs/bookkeeping/cast_test.py +++ b/qualtran/bloqs/bookkeeping/cast_test.py @@ -26,9 +26,8 @@ def test_cast(bloq_autotester): def test_cast_tensor_contraction(): bloq = TestCastToFrom() - tn, _ = cbloq_to_quimb(bloq.decompose_bloq()) - assert len(tn.tensors) == 3 - assert tn.shape == (2**4,) * 4 + tn = cbloq_to_quimb(bloq.as_composite_bloq().flatten()) + assert tn.shape == (2,) * 4 * 4 def test_cast_classical_sim(): diff --git a/qualtran/bloqs/bookkeeping/free.py b/qualtran/bloqs/bookkeeping/free.py index 2aea1c11b..f00851cd3 100644 --- a/qualtran/bloqs/bookkeeping/free.py +++ b/qualtran/bloqs/bookkeeping/free.py @@ -13,9 +13,8 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING -import numpy as np import sympy from attrs import frozen @@ -24,13 +23,13 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QDType, QUInt, Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq from qualtran.drawing import directional_text_box, Text, WireSymbol @@ -75,19 +74,17 @@ def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']: raise ValueError(f"Tried to free a non-zero register: {reg}.") return {} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - data = np.zeros(1 << self.dtype.num_qubits) - data[0] = 1 - tn.add(qtn.Tensor(data=data, inds=(incoming['reg'],), tags=['Free', tag])) + from qualtran.bloqs.basic_gates.z_basis import _ZERO + + return [ + qtn.Tensor(data=_ZERO, inds=[(incoming['reg'], i)], tags=[str(self)]) + for i in range(self.dtype.num_qubits) + ] def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: diff --git a/qualtran/bloqs/bookkeeping/join.py b/qualtran/bloqs/bookkeeping/join.py index 743c112d3..49d2674cb 100644 --- a/qualtran/bloqs/bookkeeping/join.py +++ b/qualtran/bloqs/bookkeeping/join.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import field, frozen @@ -22,6 +22,7 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QBit, QDType, @@ -29,7 +30,6 @@ Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq from qualtran.drawing import directional_text_box, Text, WireSymbol @@ -81,27 +81,18 @@ def adjoint(self) -> 'Bloq': def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]: return None, {'reg': reg.reshape(self.dtype.num_qubits)} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - if not isinstance(incoming['reg'], np.ndarray): - raise ValueError('Incoming register must be a numpy array') - tn.add( - qtn.Tensor( - data=np.eye(2**self.dtype.num_qubits, 2**self.dtype.num_qubits).reshape( - (2,) * self.dtype.num_qubits + (2**self.dtype.num_qubits,) - ), - inds=incoming['reg'].tolist() + [outgoing['reg']], - tags=['Join', tag], - ) - ) + eye = np.eye(2) + incoming = cast('NDArray', incoming['reg']) + outgoing = outgoing['reg'] + return [ + qtn.Tensor(data=eye, inds=[(outgoing, i), (incoming[i], 0)], tags=[str(self)]) + for i in range(self.dtype.num_qubits) + ] def on_classical_vals(self, reg: 'NDArray[np.uint]') -> Dict[str, int]: return {'reg': bits_to_ints(reg)[0]} diff --git a/qualtran/bloqs/bookkeeping/partition.py b/qualtran/bloqs/bookkeeping/partition.py index 7be28353a..fdfca1e75 100644 --- a/qualtran/bloqs/bookkeeping/partition.py +++ b/qualtran/bloqs/bookkeeping/partition.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING import numpy as np from attrs import evolve, field, frozen, validators @@ -21,12 +21,12 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QAny, Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq from qualtran.drawing import directional_text_box, Text, WireSymbol @@ -88,38 +88,28 @@ def as_cirq_op(self, qubit_manager, **cirq_quregs) -> Tuple[None, Dict[str, 'Cir else: return None, {'x': np.concatenate([v.ravel() for _, v in cirq_quregs.items()])} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - unitary_shape = [] - soquets = [] - _incoming = incoming if self.partition else outgoing - _outgoing = outgoing if self.partition else incoming + grouped = incoming['x'] if self.partition else outgoing['x'] + partitioned = outgoing if self.partition else incoming + + partitioned_inds = [] for reg in self.regs: - for i in range(int(np.prod(reg.shape))): - unitary_shape.append(2**reg.bitsize) - outgoing_reg = _outgoing[reg.name] - if isinstance(outgoing_reg, np.ndarray): - soquets.append(outgoing_reg.ravel()[i]) - else: - soquets.append(outgoing_reg) - - tn.add( - qtn.Tensor( - data=np.eye(2**self.n, 2**self.n).reshape( - tuple(unitary_shape) + (2**self.n,) - ), - inds=soquets + [_incoming['x']], - tags=['Partition', tag], - ) - ) + part_ind = partitioned[reg.name] + for idx in reg.all_idxs(): + for j in range(reg.bitsize): + if isinstance(part_ind, np.ndarray): + partitioned_inds.append((part_ind[idx], j)) + else: + partitioned_inds.append((part_ind, j)) + + return [ + qtn.Tensor(data=np.eye(2), inds=[partitioned_inds[j], (grouped, j)], tags=[str(self)]) + for j in range(self.n) + ] def _classical_partition(self, x: 'ClassicalValT') -> Dict[str, 'ClassicalValT']: out_vals = {} diff --git a/qualtran/bloqs/bookkeeping/partition_test.py b/qualtran/bloqs/bookkeeping/partition_test.py index 8e7f54813..3c97cc465 100644 --- a/qualtran/bloqs/bookkeeping/partition_test.py +++ b/qualtran/bloqs/bookkeeping/partition_test.py @@ -66,9 +66,12 @@ def test_partition_wrapper(): def test_partition_wrapper_tensor_contract(): bloq = TestPartition(test_bloq=TestMultiRegister()) - tn, _ = cbloq_to_quimb(bloq.decompose_bloq()) - assert len(tn.tensors) == 3 - assert tn.shape == (4096, 4096) + tn = cbloq_to_quimb(bloq.as_composite_bloq().flatten()) + assert tn.shape == (2,) * (2 * int(np.log2(4096))) + + tn = tn.rank_simplify() + for tensor in tn.tensors: + assert np.prod(tensor.shape) <= 2**4 def test_partition_wrapper_as_cirq_op(): diff --git a/qualtran/bloqs/bookkeeping/split.py b/qualtran/bloqs/bookkeeping/split.py index 14fef7a15..f73a8adc8 100644 --- a/qualtran/bloqs/bookkeeping/split.py +++ b/qualtran/bloqs/bookkeeping/split.py @@ -13,16 +13,18 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np from attrs import field, frozen +from numpy.typing import NDArray from qualtran import ( Bloq, bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, DecomposeTypeError, QBit, QDType, @@ -30,7 +32,6 @@ Register, Side, Signature, - SoquetT, ) from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq from qualtran.drawing import directional_text_box, Text, WireSymbol @@ -89,27 +90,18 @@ def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']: return {'reg': ints_to_bits(np.array([reg]), self.dtype.num_qubits)[0]} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - if not isinstance(outgoing['reg'], np.ndarray): - raise ValueError('Outgoing register must be a numpy array') - tn.add( - qtn.Tensor( - data=np.eye(2**self.dtype.num_qubits, 2**self.dtype.num_qubits).reshape( - (2,) * self.dtype.num_qubits + (2**self.dtype.num_qubits,) - ), - inds=outgoing['reg'].tolist() + [incoming['reg']], - tags=['Split', tag], - ) - ) + eye = np.eye(2) + incoming = incoming['reg'] + outgoing = cast(NDArray, outgoing['reg']) + return [ + qtn.Tensor(data=eye, inds=[(outgoing[i], 0), (incoming, i)], tags=[str(self)]) + for i in range(self.dtype.num_qubits) + ] def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: diff --git a/qualtran/bloqs/for_testing/atom.py b/qualtran/bloqs/for_testing/atom.py index 2a6515dac..c7781083b 100644 --- a/qualtran/bloqs/for_testing/atom.py +++ b/qualtran/bloqs/for_testing/atom.py @@ -13,13 +13,20 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import Dict, List, Optional, TYPE_CHECKING import attrs import numpy as np from attrs import frozen -from qualtran import Bloq, CompositeBloq, DecomposeTypeError, GateWithRegisters, Signature, SoquetT +from qualtran import ( + Bloq, + CompositeBloq, + ConnectionT, + DecomposeTypeError, + GateWithRegisters, + Signature, +) from qualtran.cirq_interop.t_complexity_protocol import TComplexity if TYPE_CHECKING: @@ -46,23 +53,18 @@ def signature(self) -> Signature: def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - tn.add( + return [ qtn.Tensor( data=np.array([[0, 1], [1, 0]]), - inds=(outgoing['q'], incoming['q']), - tags=[self.pretty_name(), tag], + inds=[(outgoing['q'], 0), (incoming['q'], 0)], + tags=[str(self)], ) - ) + ] def _t_complexity_(self) -> 'TComplexity': return TComplexity(100) @@ -92,26 +94,26 @@ def signature(self) -> Signature: def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn _I = [[1, 0], [0, 1]] _NULL = [[0, 0], [0, 0]] _X = [[0, 1], [1, 0]] - tn.add( + return [ qtn.Tensor( data=np.array([[_I, _NULL], [_NULL, _X]]), - inds=(outgoing['ctrl'], incoming['ctrl'], outgoing['target'], incoming['target']), - tags=[self.pretty_name(), tag], + inds=[ + (outgoing['ctrl'], 0), + (incoming['ctrl'], 0), + (outgoing['target'], 0), + (incoming['target'], 0), + ], + tags=[str(self)], ) - ) + ] def pretty_name(self) -> str: return 'TestTwoBitOp' @@ -140,23 +142,12 @@ def signature(self) -> Signature: def decompose_bloq(self) -> 'CompositeBloq': raise DecomposeTypeError(f"{self} is atomic") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - import quimb.tensor as qtn + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + from qualtran.cirq_interop._cirq_to_bloq import _my_tensors_from_gate - tn.add( - qtn.Tensor( - data=self._unitary_(), - inds=(outgoing['q'], incoming['q']), - tags=[self.pretty_name(), tag], - ) - ) + return _my_tensors_from_gate(self, self.signature, incoming=incoming, outgoing=outgoing) def _unitary_(self): return np.eye(2) diff --git a/qualtran/bloqs/for_testing/matrix_gate.py b/qualtran/bloqs/for_testing/matrix_gate.py index 459ee246b..8fc7b6d40 100644 --- a/qualtran/bloqs/for_testing/matrix_gate.py +++ b/qualtran/bloqs/for_testing/matrix_gate.py @@ -11,12 +11,15 @@ # 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. -from typing import Tuple +from typing import Dict, List, Tuple, TYPE_CHECKING import numpy as np from attrs import field, frozen -from qualtran import GateWithRegisters, Signature +from qualtran import ConnectionT, GateWithRegisters, Signature + +if TYPE_CHECKING: + import quimb.tensor as qtn @frozen @@ -54,6 +57,21 @@ def random(cls, bitsize: int, *, random_state=None) -> 'MatrixGate': matrix = random_unitary(2**bitsize, random_state=random_state) return cls(bitsize, matrix) + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + import quimb.tensor as qtn + + data = np.array(self.matrix).reshape((2,) * (self.bitsize * 2)) + return [ + qtn.Tensor( + data=data, + inds=[(outgoing['q'], j) for j in range(self.bitsize)] + + [(incoming['q'], j) for j in range(self.bitsize)], + tags=[str(self)], + ) + ] + def _unitary_(self): return np.array(self.matrix) diff --git a/qualtran/bloqs/for_testing/matrix_gate_test.py b/qualtran/bloqs/for_testing/matrix_gate_test.py index 3b04e8766..7aa95ce11 100644 --- a/qualtran/bloqs/for_testing/matrix_gate_test.py +++ b/qualtran/bloqs/for_testing/matrix_gate_test.py @@ -25,3 +25,4 @@ def test_create_and_unitary(bitsize: int): gate = MatrixGate.random(bitsize, random_state=random_state) np.testing.assert_allclose(cirq.unitary(gate), gate.matrix) np.testing.assert_allclose(cirq.unitary(gate**-1), np.conj(gate.matrix).T) + np.testing.assert_allclose(gate.tensor_contract(), gate.matrix) diff --git a/qualtran/bloqs/mcmt/and_bloq.py b/qualtran/bloqs/mcmt/and_bloq.py index 3b93ce241..1f9cbb572 100644 --- a/qualtran/bloqs/mcmt/and_bloq.py +++ b/qualtran/bloqs/mcmt/and_bloq.py @@ -23,7 +23,7 @@ """ import itertools from functools import cached_property -from typing import Any, Dict, Iterable, Iterator, Optional, Set, Tuple, TYPE_CHECKING, Union +from typing import cast, Dict, Iterable, Iterator, List, Optional, Set, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -37,12 +37,12 @@ bloq_example, BloqDocSpec, CompositeBloq, + ConnectionT, GateWithRegisters, QBit, Register, Side, Signature, - Soquet, ) from qualtran.bloqs.basic_gates import TGate, XGate from qualtran.bloqs.bookkeeping import ArbitraryClifford @@ -119,14 +119,9 @@ def on_classical_vals( assert target == out return {'ctrl': ctrl} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, NDArray[Soquet]], # type: ignore[type-var] - outgoing: Dict[str, NDArray[Soquet]], # type: ignore[type-var] - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn # Fill in our tensor using "and" logic. @@ -137,25 +132,24 @@ def add_my_tensors( else: data[c1, c2, c1, c2, 0] = 1 - # Here: uncompute just switches the direction of the target index. - if self.uncompute: - trg = incoming['target'] - else: - trg = outgoing['target'] + # uncompute just switches the direction of the target index. + trg = incoming['target'] if self.uncompute else outgoing['target'] - tn.add( + in_ctrls = cast(NDArray, incoming['ctrl']) + out_ctrls = cast(NDArray, outgoing['ctrl']) + return [ qtn.Tensor( data=data, - inds=( - incoming['ctrl'][0], - incoming['ctrl'][1], - outgoing['ctrl'][0], - outgoing['ctrl'][1], - trg, - ), - tags=['And', tag], + inds=[ + (in_ctrls[0], 0), + (in_ctrls[1], 0), + (out_ctrls[0], 0), + (out_ctrls[1], 0), + (trg, 0), + ], + tags=[str(self)], ) - ) + ] def pretty_name(self) -> str: dag = '†' if self.uncompute else '' diff --git a/qualtran/bloqs/mcmt/and_bloq_test.py b/qualtran/bloqs/mcmt/and_bloq_test.py index 654c42d96..70d853579 100644 --- a/qualtran/bloqs/mcmt/and_bloq_test.py +++ b/qualtran/bloqs/mcmt/and_bloq_test.py @@ -66,9 +66,9 @@ def test_truth_table(cv1, cv2): for cbloq, a, b in _iter_and_truth_table(cv1, cv2): vec = cbloq.tensor_contract() if (a == cv1) and (b == cv2): - np.testing.assert_allclose([0, 1], vec) + np.testing.assert_allclose([0, 1], vec, atol=1e-8) else: - np.testing.assert_allclose([1, 0], vec) + np.testing.assert_allclose([1, 0], vec, atol=1e-8) @pytest.mark.parametrize('cv2', [0, 1]) diff --git a/qualtran/bloqs/mcmt/multi_control_multi_target_pauli.py b/qualtran/bloqs/mcmt/multi_control_multi_target_pauli.py index 823e7affe..4501a9b2d 100644 --- a/qualtran/bloqs/mcmt/multi_control_multi_target_pauli.py +++ b/qualtran/bloqs/mcmt/multi_control_multi_target_pauli.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Iterator, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Iterator, Set, Tuple, TYPE_CHECKING, Union import cirq import numpy as np @@ -38,8 +38,6 @@ from qualtran.symbolics import HasLength, SymbolicInt if TYPE_CHECKING: - import quimb.tensor as qtn - from qualtran.resource_counting import BloqCountT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT @@ -216,20 +214,6 @@ def _apply_unitary_(self, args: 'cirq.ApplyUnitaryArgs') -> np.ndarray: ) return cirq.apply_unitary(cpauli, args) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - from qualtran.cirq_interop._cirq_to_bloq import _add_my_tensors_from_gate - - _add_my_tensors_from_gate( - self, self.signature, self.pretty_name(), tn, tag, incoming=incoming, outgoing=outgoing - ) - def _has_unitary_(self) -> bool: return not is_symbolic(self.n_ctrls) diff --git a/qualtran/bloqs/phase_estimation/lp_resource_state_test.py b/qualtran/bloqs/phase_estimation/lp_resource_state_test.py index 96cf7420b..2e374d929 100644 --- a/qualtran/bloqs/phase_estimation/lp_resource_state_test.py +++ b/qualtran/bloqs/phase_estimation/lp_resource_state_test.py @@ -29,6 +29,7 @@ ignore_alloc_free, ignore_split_join, ) +from qualtran.simulation.tensor import initialize_from_zero def test_lprs_interim_auto(bloq_autotester): @@ -57,20 +58,72 @@ def get_resource_state(m: int) -> np.ndarray: return np.sqrt(2 / (1 + N)) * np.sin(np.pi * (1 + np.arange(N)) / (1 + N)) +def test_intermediate_resource_state_cirq_quick(): + n = 3 + bloq = LPRSInterimPrep(n) + state = GateHelper(bloq).circuit.final_state_vector() + np.testing.assert_allclose(state, get_interim_resource_state(n)) + + +def test_intermediate_resource_state_tensor_quick(): + n = 3 + bloq = LPRSInterimPrep(n) + state_prep = initialize_from_zero(bloq) + state_vec = state_prep.tensor_contract() + pytest.xfail("https://github.com/quantumlib/Qualtran/issues/1068") + np.testing.assert_allclose(state_vec, get_interim_resource_state(n)) + + +@pytest.mark.slow @pytest.mark.parametrize('n', [*range(1, 14, 2)]) -def test_intermediate_resource_state(n): +def test_intermediate_resource_state_cirq(n): bloq = LPRSInterimPrep(n) state = GateHelper(bloq).circuit.final_state_vector() np.testing.assert_allclose(state, get_interim_resource_state(n)) +@pytest.mark.slow +@pytest.mark.parametrize('n', [*range(1, 14, 2)]) +def test_intermediate_resource_state_tensor(n): + bloq = LPRSInterimPrep(n) + state_prep = initialize_from_zero(bloq) + state_vec = state_prep.tensor_contract() + pytest.xfail("https://github.com/quantumlib/Qualtran/issues/1068") + np.testing.assert_allclose(state_vec, get_interim_resource_state(n)) + + +def test_prepares_resource_state_cirq_quick(): + n = 3 + bloq = LPResourceState(n) + state = GateHelper(bloq).circuit.final_state_vector() + np.testing.assert_allclose(state, get_resource_state(n)) + + +def test_prepares_resource_state_tensor_quick(): + n = 3 + bloq = LPResourceState(n) + state_prep = initialize_from_zero(bloq) + state_vec = state_prep.tensor_contract() + np.testing.assert_allclose(state_vec, get_resource_state(n)) + + +@pytest.mark.slow @pytest.mark.parametrize('n', [*range(1, 14, 2)]) -def test_prepares_resource_state(n): +def test_prepares_resource_state_cirq(n): bloq = LPResourceState(n) state = GateHelper(bloq).circuit.final_state_vector() np.testing.assert_allclose(state, get_resource_state(n)) +@pytest.mark.slow +@pytest.mark.parametrize('n', [*range(1, 14, 2)]) +def test_prepares_resource_state_tensor(n): + bloq = LPResourceState(n) + state_prep = initialize_from_zero(bloq) + state_vec = state_prep.tensor_contract() + np.testing.assert_allclose(state_vec, get_resource_state(n)) + + @pytest.mark.parametrize('n', [*range(1, 14, 2)]) def test_t_complexity(n): bloq = LPResourceState(n) diff --git a/qualtran/bloqs/qft/two_bit_ffft.py b/qualtran/bloqs/qft/two_bit_ffft.py index e9db23e48..08a4d008d 100644 --- a/qualtran/bloqs/qft/two_bit_ffft.py +++ b/qualtran/bloqs/qft/two_bit_ffft.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Set, TYPE_CHECKING +from typing import Dict, List, Set, TYPE_CHECKING import numpy as np from attrs import frozen @@ -23,6 +23,7 @@ bloq_example, BloqBuilder, BloqDocSpec, + ConnectionT, QBit, Register, Signature, @@ -88,27 +89,21 @@ def signature(self) -> Signature: def pretty_name(self) -> str: return 'F(k, n)' - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: import quimb.tensor as qtn - out_inds = [outgoing['x'], outgoing['y']] - in_inds = [incoming['x'], incoming['y']] + # TODO: https://github.com/quantumlib/Qualtran/issues/873. This tensor definition + # isn't used by default since this isn't (yet) a "leaf bloq". + + out_inds = [(outgoing['x'], 0), (outgoing['y'], 0)] + in_inds = [(incoming['x'], 0), (incoming['y'], 0)] matrix = _fkn_matrix(self.k, self.n) matrix = matrix.conj().T if self.is_adjoint else matrix - tn.add( - qtn.Tensor( - data=matrix.reshape((2,) * 4), - inds=out_inds + in_inds, - tags=[self.pretty_name(), tag], - ) - ) + return [ + qtn.Tensor(data=matrix.reshape((2,) * 4), inds=out_inds + in_inds, tags=[str(self)]) + ] def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: return { diff --git a/qualtran/bloqs/qft/two_bit_ffft_test.py b/qualtran/bloqs/qft/two_bit_ffft_test.py index b45bf726e..1fd5edf1b 100644 --- a/qualtran/bloqs/qft/two_bit_ffft_test.py +++ b/qualtran/bloqs/qft/two_bit_ffft_test.py @@ -14,7 +14,6 @@ import itertools import cirq -import numpy as np import pytest from qualtran import Bloq @@ -41,7 +40,5 @@ def all_t(bloq) -> Bloq: @pytest.mark.parametrize('k,n', itertools.product(range(0, 4), range(1, 4))) def test_tensors(k, n): # Eq. E11 in https://arxiv.org/pdf/2012.09238.pdf - from_tensors = TwoBitFFFT(k, n).tensor_contract() - np.testing.assert_allclose(from_tensors, _fkn_matrix(k, n)) from_decomp = TwoBitFFFT(k, n).decompose_bloq().tensor_contract() cirq.testing.assert_allclose_up_to_global_phase(from_decomp, _fkn_matrix(k, n), atol=1e-12) diff --git a/qualtran/bloqs/rotations/phase_gradient.py b/qualtran/bloqs/rotations/phase_gradient.py index 593bf8c14..01947d6f0 100644 --- a/qualtran/bloqs/rotations/phase_gradient.py +++ b/qualtran/bloqs/rotations/phase_gradient.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union +from typing import Dict, Iterator, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union import attrs import cirq @@ -25,6 +25,7 @@ from qualtran import ( bloq_example, BloqDocSpec, + ConnectionT, GateWithRegisters, QBit, QFxp, @@ -39,7 +40,7 @@ if TYPE_CHECKING: import quimb.tensor as qtn - from qualtran import Bloq, SoquetT + from qualtran import Bloq from qualtran.resource_counting import BloqCountT, SympySymbolAllocator from qualtran.simulation.classical_sim import ClassicalValT from qualtran.symbolics import SymbolicFloat, SymbolicInt @@ -308,19 +309,12 @@ def __pow__(self, power): return self.adjoint() raise NotImplementedError("AddIntoPhaseGrad.__pow__ defined only for powers +1/-1.") - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - from qualtran.cirq_interop._cirq_to_bloq import _add_my_tensors_from_gate - - _add_my_tensors_from_gate( - self, self.signature, self.pretty_name(), tn, tag, incoming=incoming, outgoing=outgoing - ) + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + from qualtran.cirq_interop._cirq_to_bloq import _my_tensors_from_gate + + return _my_tensors_from_gate(self, self.signature, incoming=incoming, outgoing=outgoing) @bloq_example @@ -508,20 +502,6 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: num_additions = min(num_additions_naive, num_additions) return {(AddIntoPhaseGrad(self.x_dtype.bitsize, self.phase_bitsize), num_additions)} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - from qualtran.cirq_interop._cirq_to_bloq import _add_my_tensors_from_gate - - _add_my_tensors_from_gate( - self, self.signature, self.pretty_name(), tn, tag, incoming=incoming, outgoing=outgoing - ) - @bloq_example def _add_scaled_val_into_phase_reg() -> AddScaledValIntoPhaseReg: diff --git a/qualtran/bloqs/rotations/phasing_via_cost_function_test.py b/qualtran/bloqs/rotations/phasing_via_cost_function_test.py index dbd71c47c..ebb8bdc18 100644 --- a/qualtran/bloqs/rotations/phasing_via_cost_function_test.py +++ b/qualtran/bloqs/rotations/phasing_via_cost_function_test.py @@ -18,14 +18,17 @@ import numpy as np import pytest -from qualtran import Bloq, BloqBuilder, GateWithRegisters, Register, Signature, SoquetT +from qualtran import Bloq, BloqBuilder, BloqError, GateWithRegisters, Register, Signature, SoquetT from qualtran._infra.data_types import QFxp from qualtran.bloqs.arithmetic.hamming_weight import HammingWeightCompute from qualtran.bloqs.arithmetic.multiplication import Square from qualtran.bloqs.basic_gates import Hadamard from qualtran.bloqs.basic_gates.on_each import OnEach from qualtran.bloqs.rotations.phase_gradient import PhaseGradientState -from qualtran.bloqs.rotations.phasing_via_cost_function import PhasingViaCostFunction +from qualtran.bloqs.rotations.phasing_via_cost_function import ( + _square_via_zpow_phasing, + PhasingViaCostFunction, +) from qualtran.bloqs.rotations.quantum_variable_rotation import ( QvrInterface, QvrPhaseGradient, @@ -35,6 +38,10 @@ from qualtran.testing import assert_valid_bloq_decomposition +def test_square_via_zpow_phasing(bloq_autotester): + bloq_autotester(_square_via_zpow_phasing) + + @attrs.frozen class TestHammingWeightPhasing(GateWithRegisters): bitsize: int @@ -191,5 +198,9 @@ def test_square_phasing_via_phase_gradient( a = bb.add(OnEach(n, Hadamard()), q=a) a = bb.add(test_bloq, a=a) cbloq = bb.finalize(a=a) - hw_final_state = cbloq.tensor_contract() + try: + flat_cbloq = cbloq.flatten() + except BloqError: + pytest.xfail("https://github.com/quantumlib/Qualtran/issues/1069") + hw_final_state = flat_cbloq.tensor_contract() np.testing.assert_allclose(expected_final_state, hw_final_state, atol=eps) diff --git a/qualtran/bloqs/swap_network/swap_with_zero_test.py b/qualtran/bloqs/swap_network/swap_with_zero_test.py index d683a826c..87fcdfee8 100644 --- a/qualtran/bloqs/swap_network/swap_with_zero_test.py +++ b/qualtran/bloqs/swap_network/swap_with_zero_test.py @@ -26,7 +26,6 @@ from qualtran.bloqs.bookkeeping import ArbitraryClifford from qualtran.bloqs.swap_network.swap_with_zero import _swz, _swz_small, SwapWithZero from qualtran.cirq_interop.t_complexity_protocol import TComplexity -from qualtran.simulation.tensor import flatten_for_tensor_contraction from qualtran.testing import assert_valid_bloq_decomposition random.seed(12345) @@ -57,8 +56,7 @@ def test_swap_with_zero_bloq(selection_bitsize, target_bitsize, n_target_registe trgs.append(trg) sel, trgs = bb.add(swz, selection=sel, targets=np.array(trgs)) circuit = bb.finalize(sel=sel, trgs=trgs) - flat_circuit = flatten_for_tensor_contraction(circuit) - full_state_vector = flat_circuit.tensor_contract() + full_state_vector = circuit.tensor_contract() result_state_vector = cirq.sub_state_vector( full_state_vector, keep_indices=list(range(selection_bitsize, selection_bitsize + target_bitsize)), diff --git a/qualtran/cirq_interop/_bloq_to_cirq_test.py b/qualtran/cirq_interop/_bloq_to_cirq_test.py index 9d9ab0f2e..c149bdfb2 100644 --- a/qualtran/cirq_interop/_bloq_to_cirq_test.py +++ b/qualtran/cirq_interop/_bloq_to_cirq_test.py @@ -11,14 +11,14 @@ # 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. -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING import cirq import numpy as np import pytest from attrs import frozen -from qualtran import Bloq, BloqBuilder, Signature, Soquet, SoquetT +from qualtran import Bloq, BloqBuilder, ConnectionT, Signature, Soquet, SoquetT from qualtran._infra.gate_with_registers import get_named_qubits from qualtran.bloqs.basic_gates import Toffoli, XGate from qualtran.bloqs.factoring import ModExp @@ -44,17 +44,12 @@ def as_cirq_op( (y,) = y return cirq.SWAP(x, y), {'x': np.array([x]), 'y': np.array([y])} - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: from qualtran.bloqs.basic_gates import TwoBitSwap - TwoBitSwap().add_my_tensors(tn, tag, incoming=incoming, outgoing=outgoing) + return TwoBitSwap().my_tensors(incoming=incoming, outgoing=outgoing) def test_swap_two_bits_to_cirq(): @@ -94,15 +89,19 @@ class SwapTestWithOnlyTensorData(Bloq): def signature(self): return Signature.build(x=self.n, y=self.n) - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - SwapTest(self.n).add_my_tensors(tn, tag, incoming=incoming, outgoing=outgoing) + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + import quimb.tensor as qtn + + from qualtran.simulation.tensor._dense import _order_incoming_outgoing_indices + + inds = _order_incoming_outgoing_indices( + self.signature, incoming=incoming, outgoing=outgoing + ) + data = SwapTest(self.n).tensor_contract().reshape((2,) * len(inds)) + + return [qtn.Tensor(data=data, inds=inds)] @pytest.mark.parametrize('n', [1, 2, 3, 4]) diff --git a/qualtran/cirq_interop/_cirq_to_bloq.py b/qualtran/cirq_interop/_cirq_to_bloq.py index 004e1f950..a97b44a92 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq.py +++ b/qualtran/cirq_interop/_cirq_to_bloq.py @@ -28,6 +28,7 @@ Bloq, BloqBuilder, CompositeBloq, + ConnectionT, Controlled, CtrlSpec, DecomposeNotImplementedError, @@ -40,7 +41,6 @@ Side, Signature, Soquet, - SoquetT, ) from qualtran._infra.gate_with_registers import ( _get_all_and_output_quregs_from_input, @@ -102,22 +102,11 @@ def decompose_from_registers( except TypeError as e: raise DecomposeNotImplementedError(f"{self} does not declare a decomposition.") from e - def add_my_tensors( - self, - tn: 'qtn.TensorNetwork', - tag: Any, - *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], - ): - _add_my_tensors_from_gate( - self.cirq_gate, - self.signature, - str(self.cirq_gate), - tn=tn, - tag=tag, - incoming=incoming, - outgoing=outgoing, + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + return _my_tensors_from_gate( + self.cirq_gate, self.signature, incoming=incoming, outgoing=outgoing ) def _t_complexity_(self) -> 'TComplexity': @@ -204,39 +193,27 @@ def _wire_symbol_from_gate( return _cirq_wire_symbol_to_qualtran_wire_symbol(symbol, wire_reg.side) -def _add_my_tensors_from_gate( +def _my_tensors_from_gate( gate: cirq.Gate, signature: Signature, - short_name: str, - tn: 'qtn.TensorNetwork', - tag: Any, *, - incoming: Dict[str, 'SoquetT'], - outgoing: Dict[str, 'SoquetT'], -): + incoming: Dict[str, 'ConnectionT'], + outgoing: Dict[str, 'ConnectionT'], +) -> List['qtn.Tensor']: import quimb.tensor as qtn + from qualtran.simulation.tensor._dense import _order_incoming_outgoing_indices from qualtran.simulation.tensor._tensor_data_manipulation import ( tensor_data_from_unitary_and_signature, ) if not cirq.has_unitary(gate): - raise NotImplementedError( - f"CirqGateAsBloq.add_my_tensors is currently supported only for unitary gates. " - f"Found {gate}." - ) + raise NotImplementedError(f"Tensors are only supported for unitary gates, not {gate}.") + unitary = tensor_data_from_unitary_and_signature(cirq.unitary(gate), signature) - incoming_list = [ - *itertools.chain.from_iterable( - [np.array(incoming[reg.name]).flatten() for reg in signature.lefts()] - ) - ] - outgoing_list = [ - *itertools.chain.from_iterable( - [np.array(outgoing[reg.name]).flatten() for reg in signature.rights()] - ) - ] - tn.add(qtn.Tensor(data=unitary, inds=outgoing_list + incoming_list, tags=[short_name, tag])) + inds = _order_incoming_outgoing_indices(signature, incoming=incoming, outgoing=outgoing) + unitary = unitary.reshape((2,) * len(inds)) + return [qtn.Tensor(data=unitary, inds=inds, tags=[str(gate)])] @frozen(eq=False) diff --git a/qualtran/cirq_interop/_cirq_to_bloq_test.py b/qualtran/cirq_interop/_cirq_to_bloq_test.py index 3e7c2d2a3..64a7c3087 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq_test.py +++ b/qualtran/cirq_interop/_cirq_to_bloq_test.py @@ -89,7 +89,7 @@ def test_cirq_gate_as_bloq_tensor_contract_for_and_gate(): state_vector = cbloq.tensor_contract() assert np.isclose(state_vector[7], 1) - with pytest.raises(NotImplementedError, match="supported only for unitary gates"): + with pytest.raises(NotImplementedError, match=r".*only supported for unitary gates.*"): _ = CirqGateAsBloq(And(uncompute=True)).as_composite_bloq().tensor_contract() diff --git a/qualtran/simulation/tensor.ipynb b/qualtran/simulation/tensor.ipynb index 2f3165d0a..014193182 100644 --- a/qualtran/simulation/tensor.ipynb +++ b/qualtran/simulation/tensor.ipynb @@ -99,7 +99,9 @@ "source": [ "## Additional functionality\n", "\n", - "A composite bloq has a 1-to-1 mapping with a tensor network. We use [Quimb](https://quimb.readthedocs.io/) to handle efficient contraction of such networks.\n", + "### Direct manipulation of the `qtn.TensorNetwork`\n", + "\n", + "A composite bloq can be easily transformed into a tensor network. We use [Quimb](https://quimb.readthedocs.io/) to handle efficient contraction of such networks.\n", "\n", "The most important library function is `qualtran.simulation.tensor.cbloq_to_quimb`. This will build a quimb `qtn.TensorNetwork` tensor network representation of the composite bloq. You may want to manipulate this object directly using the full Quimb API. Otherwise, this function is used as the workhorse behind the public functions and methods like `Bloq.tensor_contract()`. \n", "\n", @@ -138,7 +140,7 @@ "source": [ "from qualtran.simulation.tensor import cbloq_to_quimb\n", "\n", - "tn, fix = cbloq_to_quimb(cbloq)\n", + "tn = cbloq_to_quimb(cbloq)\n", "tn.draw(show_inds=False)" ] }, @@ -147,7 +149,7 @@ "id": "f18e1f2d", "metadata": {}, "source": [ - "The entire suite of Quimb tools are now available." + "With this `qtn.TensorNetwork` in hand, the entire suite of Quimb tools are available." ] }, { @@ -160,6 +162,73 @@ "tn.contraction_info()" ] }, + { + "cell_type": "markdown", + "id": "fc6c12e0-6c89-4f7b-8814-a9d5724e73d9", + "metadata": {}, + "source": [ + "### Quimb index format\n", + "\n", + "In `CompositeBloq`, we form the compute graph by storing a list of nodes and edges. Quimb uses a different strategy for representing the tensor network graphs. To form a tensor network in Quimb, we provide a list of `qtn.Tensor` which contain not only the tensor data but also a list of \"indices\" that can form connections to other tensors. Similar to the Einstein summation convention, if two tensors each have an index with the same name: an edge is created in the tensor network graph and this shared index is summed over. These indices are traditionally short strings like `\"k\"`, but can be any hashable Python object. In `CompositeBloq`, the unique object that identifies a connection between bloqs is `qualtran.Connection`, so we use these connection objects as our indices.\n", + "\n", + "Qualtran and Quimb both support \"wires\" with arbitrary bitsize. Qualtran uses bookkeeping bloqs like `Join` and `Split` to fuse and un-fuse indices. In theory, these operations should be free in the tensor contraction, as they are essentially an identity tensor. In our understanding, Quimb does not have special support for supporting these re-shaping operations within the tensor network graph. In versions of Qualtran prior to v0.5, split and join operations were tensored up to `n`-sized identity tensors. This would form a bottleneck in any contraction ordering. Therefore, we keep all the indices un-fused in the tensor network representation and use tuples of `(cxn, i)` for our tensor indices, where the second integer `i` indexes the individual bits in a register with `reg.dtype.num_qubits` > 1.\n", + "\n", + "**In summary:**\n", + " - Each tensor index is a tuple `(cxn, i)`\n", + " - The `cxn: qualtran.Connection` entry identifies the connection between soquets in a Qualtran compute graph.\n", + " - The second integer `i` is the bit index within high-bitsize registers, which is necessary due to technical restrictions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d9c0ed47-275e-417b-bbf5-b8dff2f84099", + "metadata": {}, + "outputs": [], + "source": [ + "# Use `get_right_and_left_inds` to get the quimb indices ordered according to\n", + "# the bloq's signature.\n", + "\n", + "from qualtran.simulation.tensor import get_right_and_left_inds\n", + "left, right = get_right_and_left_inds(tn, bloq.signature)\n", + "\n", + "print(\"Left outer inds:\")\n", + "for cxn, i in left:\n", + " print(' ', cxn.left)\n", + "\n", + "print(\"Right outer inds\")\n", + "for cxn, i in right:\n", + " print(' ', cxn.right)" + ] + }, + { + "cell_type": "markdown", + "id": "05279b9b-b7db-4b4e-ad27-adfe2d88a4d7", + "metadata": {}, + "source": [ + "### Flattening\n", + "\n", + "A call to `Bloq.tensor_contract` will first \"flatten\" the bloq by doing `bloq.as_composite_bloq().flatten()`. This is a sensible default default for constructing tensor networks, as the best contraction performance can generally be achieved by keeping tensors as small as possible in the network. \n", + "\n", + "In Qualtran, we usually avoid flattening bloqs and strongly to prefer to work with one level of decomposition at a time. This is to avoid performance issues with large, abstract algorithms. But typically if the full circuit is large enough to cause performance issues with flattening it is also too large to simulate numerically; so an exception to the general advice is made here.\n", + "\n", + "All bloqs in the flattened circuit must provide their explicit tensors. If your bloq's tensors ought to be derived from its decomposition: this is achieved by the previously mentioned flattening operation. If a bloq provides tensors through overriding `Bloq.my_tensors` _and also_ defines a decomposition, the explicit tensors will not be used (by default). This is because any bloq with a decomposition will be flattened. If you would like to control the flattening operation, use the free functions to control the tensor network construction and contraction." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7d9983b2-e241-4cab-a518-fe2fa7a5ac92", + "metadata": {}, + "outputs": [], + "source": [ + "# Stop flattening if e.g. a bloq supports explicit tensors _and_ a decomposition.\n", + "# Use the flatten predicate to control.\n", + "custom_flat = bloq.as_composite_bloq().flatten(lambda binst: binst.i != 2)\n", + "tn = cbloq_to_quimb(custom_flat)\n", + "len(tn.tensors)" + ] + }, { "cell_type": "markdown", "id": "85fb5d5f", @@ -167,11 +236,13 @@ "source": [ "## Implementation\n", "\n", - "The `qualtran.simulation.tensor` functions rely on the `Bloq.add_my_tensors(...)` method to implement the protocol. This is where a bloq's tensor information is actually encoded.\n", + "The `qualtran.simulation.tensor` functions rely on the `Bloq.my_tensors(...)` method to implement the protocol. This is where a bloq's tensor information is actually encoded.\n", + "\n", + "Usually, the most efficient way of supporting tensor simulation is by providing a decomposition for your bloq. However, bloq authors may want to override `my_tensors` if the bloq can't or shouldn't define a decomposition. The method takes dictionaries of incoming and outgoing indices (keyed by register name) to asist the author in matching up dimensions of their `np.ndarray` to the incoming and outgoing wires when constructing `qtn.Tensor`s.\n", "\n", - "Bloq authors may want to override this method. The library will provide a (partial) `qtn.TensorNetwork` as well as dictionaries of incoming and outgoing indices (keyed by register name) to asist the author in matching up dimensions of their `np.ndarray` to the incoming and outgoing wires. Bloq authors are encouraged to read the docstring for this method for more details.\n", + "The docstring for `Bloq.my_tensors` provides a complete, technical description of how to successfully override this method. In brief, the method must return one or more `qtn.Tensor`s that get added to the tensor network. The indices used to construct these **must** be of the correct form. Each tensor index is a tuple `(cxn, j)`. The connection entry comes from the `incoming` and `outgoing` arguments to the method. The `j` integer is the bit index within multi-bit registers.\n", "\n", - "Below, we write our own `CNOT` bloq with custom tensors." + "New, we write our own `CNOT` bloq with custom tensors." ] }, { @@ -182,7 +253,7 @@ "outputs": [], "source": [ "from functools import cached_property\n", - "from typing import Any, Dict, Tuple\n", + "from typing import Any, Dict, Tuple, List\n", "\n", "import numpy as np\n", "import quimb.tensor as qtn\n", @@ -196,10 +267,9 @@ " def signature(self) -> 'Signature':\n", " return Signature.build(ctrl=1, target=1)\n", "\n", - " def add_my_tensors(\n", - " self, tn: qtn.TensorNetwork, tag: Any,\n", - " *, incoming: Dict[str, SoquetT], outgoing: Dict[str, SoquetT],\n", - " ):\n", + " def my_tensors(\n", + " self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT']\n", + " ) -> List['qtn.Tensor']:\n", " # The familiar CNOT matrix. We make sure to\n", " # cast this to np.complex128 so we don't accidentally\n", " # lose precision anywhere else in the contraction.\n", @@ -217,12 +287,21 @@ " tensor = matrix.reshape((2,2,2,2))\n", " \n", "\n", - " tn.add(qtn.Tensor(\n", + " # This is a simple case: we only need one tensor and\n", + " # each register is one bit.\n", + " outgoing_inds = [\n", + " (outgoing['ctrl'], 0),\n", + " (outgoing['target'], 0),\n", + " ]\n", + " incoming_inds = [\n", + " (incoming['ctrl'], 0),\n", + " (incoming['target'], 0),\n", + " ]\n", + " \n", + " return [qtn.Tensor(\n", " data=tensor, \n", - " inds=(outgoing['ctrl'], outgoing['target'], \n", - " incoming['ctrl'], incoming['target']),\n", - " tags = ['cnot', tag],\n", - " ))" + " inds=outgoing_inds + incoming_inds\n", + " )]" ] }, { @@ -243,17 +322,12 @@ "source": [ "## Default Fallback\n", "\n", - "If a bloq does not override `add_my_tensors(...)`, the default fallback will be used by `qualtran` to support the tensor protocol.\n", - "\n", - "By default, qualtran will fall back on the tensor contraction of the decomposition. This recursion will continue until each leaf bloq defines `add_my_tensors` or a bloq cannot be further decomposed.\n", + "By default, the flattening operation in a given tensor contraction should mean that only a finite number of small, target-gateset bloqs\n", + "should need to explicitly override `my_tensors()`. \n", "\n", - "Specifically, the system will:\n", - "\n", - " - decompose the bloq with `bloq.decompose_bloq()` into a composite bloq.\n", - " - use the result of `cbloq_as_contracted_tensor(...)` as the bloq's tensor.\n", - " \n", + "If a bloq does not override `my_tensors(...)` and doesn't provide a decomposition, the tensor protocol will throw an error when trying to construct a tensor network.\n", " \n", - "For example, below we author a `BellState` bloq. We define a decomposition but do not explicitly provide tensor information." + "For example, we author a `BellState` bloq. We define a decomposition but do not explicitly provide tensor information." ] }, { @@ -288,7 +362,7 @@ "id": "d0eebfb4", "metadata": {}, "source": [ - "Nevertheless, the system can recursively determine the tensor form:" + "The system can still contract the tensor network implied by this bloq because it will automatically flatten the bloq." ] }, { @@ -306,7 +380,20 @@ "id": "ebbe9a2a", "metadata": {}, "source": [ - "Note that the composite bloq is fully contracted to a dense tensor at each level of decomposition, which likely will prevent quimb from finding the best contraction ordering. See `flatten_for_tensor_contraction` if this is an issue." + "If you don't flatten first, an error will be raised" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aeb3096d-8745-49ec-8809-61bd6d17ae67", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " cbloq_to_quimb(BellState().as_composite_bloq())\n", + "except NotImplementedError as e:\n", + " print(\"Expected error because we didn't flatten first:\", repr(e))" ] }, { @@ -324,7 +411,7 @@ "source": [ "### Gates with factorized tensors\n", "\n", - "The `add_my_tensors` method can add multiple `Tensor` objects to the method if there is a known factorization of the bloq's tensors. For example: CNOT can be written as a dense 4x4 matrix or by contracting the so-called COPY and XOR tensors. " + "The `my_tensors` method can return multiple `Tensor` objects if there is a known factorization of the bloq's tensors. For example: CNOT can be written as a dense 4x4 matrix or by contracting the so-called COPY and XOR tensors. " ] }, { @@ -340,18 +427,15 @@ ")\n", "\n", "cbloq = CNOT().as_composite_bloq()\n", - "tn, _ = cbloq_to_quimb(cbloq)\n", + "tn = cbloq_to_quimb(cbloq)\n", "\n", "# Rename the indices to something less verbose\n", - "from qualtran._infra.composite_bloq import _get_dangling_soquets\n", - "lsoqs = _get_dangling_soquets(cbloq.signature, right=False)\n", - "rsoqs = _get_dangling_soquets(cbloq.signature, right=True)\n", - "\n", - "rename = {lsoqs[k]: f'{k}_in' for k in lsoqs.keys()}\n", - "rename |= {rsoqs[k]: f'{k}_out' for k in rsoqs.keys()}\n", + "lefts, rights = get_right_and_left_inds(tn, cbloq.signature)\n", + "rename = {left: f'{left[0].left.reg.name}_in' for left in lefts}\n", + "rename |= {right: f'{right[0].right.reg.name}_out' for right in rights}\n", "tn = tn.reindex(rename)\n", "\n", - "tn.draw(color=['COPY', 'XOR'], show_tags=False, initial_layout='spectral')\n", + "tn.draw(color=['COPY', 'XOR'], show_tags=True, initial_layout='spectral')\n", "for tensor in tn:\n", " print(tensor.tags)\n", " print(tensor.data)\n", diff --git a/qualtran/simulation/tensor/__init__.py b/qualtran/simulation/tensor/__init__.py index 6f5b1f814..33e340000 100644 --- a/qualtran/simulation/tensor/__init__.py +++ b/qualtran/simulation/tensor/__init__.py @@ -13,9 +13,9 @@ # limitations under the License. -from ._dense import bloq_to_dense, get_right_and_left_inds +from ._dense import bloq_to_dense, get_right_and_left_inds, quimb_to_dense from ._flattening import bloq_has_custom_tensors, flatten_for_tensor_contraction -from ._quimb import cbloq_as_contracted_tensor, cbloq_to_quimb +from ._quimb import cbloq_to_quimb, initialize_from_zero from ._tensor_data_manipulation import ( active_space_for_ctrl_spec, eye_tensor_for_signature, diff --git a/qualtran/simulation/tensor/_dense.py b/qualtran/simulation/tensor/_dense.py index c9ccb5f55..bcfa9105a 100644 --- a/qualtran/simulation/tensor/_dense.py +++ b/qualtran/simulation/tensor/_dense.py @@ -12,55 +12,147 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List +import logging +from typing import Dict, List, Tuple, TYPE_CHECKING from numpy.typing import NDArray -from qualtran import Bloq, Signature, Soquet -from qualtran._infra.composite_bloq import _get_flat_dangling_soqs +from qualtran import Bloq, Connection, ConnectionT, LeftDangle, RightDangle, Signature, Soquet +from ._flattening import flatten_for_tensor_contraction from ._quimb import cbloq_to_quimb +if TYPE_CHECKING: + import quimb.tensor as qtn -def get_right_and_left_inds(signature: Signature) -> List[List[Soquet]]: +logger = logging.getLogger(__name__) + + +def _order_incoming_outgoing_indices( + signature: Signature, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] +) -> List[Tuple[Connection, int]]: + """Order incoming and outgoing indices provided by the tensor protocol according to `signature`. + + This can be used if you have a well-ordered, dense numpy array. + + >>> inds = _order_incoming_outgoing_indices(signature, incoming, outgoing) + >>> data = unitary.reshape((2,) * len(inds)) + >>> return [qtn.Tensor(data=data, inds=inds)] + """ + + inds: List[Tuple[Connection, int]] = [] + + # Nested for loops: + # reg: each register in the signature + # idx: each index into a shaped register, or () for a non-shaped register + # j: each qubit (sub-)index for a given data type + for reg in signature.rights(): + for idx in reg.all_idxs(): + for j in range(reg.dtype.num_qubits): + if idx: + inds.append((outgoing[reg.name][idx], j)) # type: ignore[index] + else: + inds.append((outgoing[reg.name], j)) # type: ignore[arg-type] + for reg in signature.lefts(): + for idx in reg.all_idxs(): + for j in range(reg.dtype.num_qubits): + if idx: + inds.append((incoming[reg.name][idx], j)) # type: ignore[index] + else: + inds.append((incoming[reg.name], j)) # type: ignore[arg-type] + + return inds + + +def get_right_and_left_inds(tn: 'qtn.TensorNetwork', signature: Signature) -> List[List[Soquet]]: """Return right and left tensor indices. In general, this will be returned as a list of length-2 corresponding - to the right and left indices, respectively. If there *are* no right + to the right and left indices, respectively. If there *are no* right or left indices, that entry will be omitted from the returned list. Right indices come first to match the quantum computing / matrix multiplication convention where U_tot = U_n ... U_2 U_1. + + Args: + tn: The tensor network to fetch the outer indices, which won't necessarily be ordered. + signature: The signature of the bloq used to order the indices. """ + left_inds = {} + right_inds = {} + cxn: Connection + j: int + for ind in tn.outer_inds(): + cxn, j = ind + if cxn.left.binst is LeftDangle: + soq = cxn.left + left_inds[soq.reg, soq.idx, j] = ind + elif cxn.right.binst is RightDangle: + soq = cxn.right + right_inds[soq.reg, soq.idx, j] = ind + else: + raise ValueError( + "Outer indices of a tensor network should be " + "connections to LeftDangle or RightDangle" + ) + + left_ordered_inds = [] + for reg in signature.lefts(): + for idx in reg.all_idxs(): + for j in range(reg.dtype.num_qubits): + left_ordered_inds.append(left_inds[reg, idx, j]) + + right_ordered_inds = [] + for reg in signature.rights(): + for idx in reg.all_idxs(): + for j in range(reg.dtype.num_qubits): + right_ordered_inds.append(right_inds[reg, idx, j]) + inds = [] - rsoqs = _get_flat_dangling_soqs(signature, right=True) - if rsoqs: - inds.append(rsoqs) - lsoqs = _get_flat_dangling_soqs(signature, right=False) - if lsoqs: - inds.append(lsoqs) + if right_ordered_inds: + inds.append(right_ordered_inds) + if left_ordered_inds: + inds.append(left_ordered_inds) + return inds -def bloq_to_dense(bloq: Bloq) -> NDArray: +def quimb_to_dense(tn: 'qtn.TensorNetwork', signature: Signature) -> NDArray: + """Contract a quimb tensor network `tn` to a dense matrix consistent with `signature`.""" + inds = get_right_and_left_inds(tn, signature) + if tn.contraction_width() > 8: + tn.full_simplify(inplace=True) + + if inds: + data = tn.to_dense(*inds) + else: + data = tn.contract() + + return data + + +def bloq_to_dense(bloq: Bloq, full_flatten: bool = True) -> NDArray: """Return a contracted, dense ndarray representing the composite bloq. - The public version of this function is available as the `Bloq.tensor_contract()` - method. + This function is also available as the `Bloq.tensor_contract()` method. - This constructs a tensor network and then contracts it according to the cbloq's registers, - i.e. the dangling indices. The returned array will be 0-, 1- or 2- dimensional. If it is - a 2-dimensional matrix, we follow the quantum computing / matrix multiplication convention - of (right, left) indices. + This function decomposes and flattens a given bloq into a factorized CompositeBloq, + turns that composite bloq into a Quimb tensor network, and contracts it into a dense + matrix. - For more fine grained control over the final shape of the tensor, use - `cbloq_to_quimb` and `TensorNetwork.to_dense` directly. - """ - cbloq = bloq.as_composite_bloq() - tn, _ = cbloq_to_quimb(cbloq) - inds = get_right_and_left_inds(cbloq.signature) + The returned array will be 0-, 1- or 2- dimensional with indices arranged according to the + bloq's signature. In the case of a 2-dimensional matrix, we follow the + quantum computing / matrix multiplication convention of (right, left) order of dimensions. - if inds: - return tn.to_dense(*inds) + For fine-grained control over the tensor contraction, use + `cbloq_to_quimb` and `TensorNetwork.to_dense` directly. - return tn.contract() + Args: + bloq: The bloq + full_flatten: Whether to completely flatten the bloq into the smallest possible + bloqs. Otherwise, stop flattening if custom tensors are encountered. + """ + logging.info("bloq_to_dense() on %s", bloq) + flat_cbloq = flatten_for_tensor_contraction(bloq, full_flatten=full_flatten) + tn = cbloq_to_quimb(flat_cbloq) + return quimb_to_dense(tn, bloq.signature) diff --git a/qualtran/simulation/tensor/_dense_test.py b/qualtran/simulation/tensor/_dense_test.py index 61940e39d..991d68df9 100644 --- a/qualtran/simulation/tensor/_dense_test.py +++ b/qualtran/simulation/tensor/_dense_test.py @@ -13,18 +13,28 @@ # limitations under the License. from functools import cached_property -from typing import Dict +from typing import Dict, List import cirq import numpy as np import quimb.tensor as qtn from attrs import frozen -from numpy.typing import NDArray -from qualtran import Bloq, BloqBuilder, QAny, QBit, Register, Side, Signature, Soquet, SoquetT -from qualtran._infra.composite_bloq import _get_dangling_soquets +from qualtran import ( + Bloq, + BloqBuilder, + Connection, + ConnectionT, + QAny, + QBit, + Register, + Side, + Signature, + Soquet, + SoquetT, +) from qualtran.bloqs.basic_gates import CNOT, XGate, ZGate -from qualtran.simulation.tensor import bloq_to_dense, get_right_and_left_inds +from qualtran.simulation.tensor import bloq_to_dense from qualtran.testing import assert_valid_bloq_decomposition @@ -40,62 +50,44 @@ def signature(self) -> 'Signature': ] ) - def add_my_tensors( - self, - tn: qtn.TensorNetwork, - tag, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: assert sorted(incoming.keys()) == ['qubits', 'x'] in_qubits = incoming['qubits'] assert isinstance(in_qubits, np.ndarray) assert in_qubits.shape == (2,) - assert isinstance(incoming['x'], Soquet) - assert incoming['x'].reg.bitsize == 2 + assert isinstance(incoming['x'], Connection) + assert incoming['x'].right.reg.bitsize == 2 assert sorted(outgoing.keys()) == ['qubits', 'y'] out_qubits = outgoing['qubits'] assert isinstance(out_qubits, np.ndarray) assert out_qubits.shape == (2,) - assert isinstance(outgoing['y'], Soquet) - assert outgoing['y'].reg.bitsize == 1 + assert isinstance(outgoing['y'], Connection) + assert outgoing['y'].left.reg.bitsize == 1 data = np.zeros((2**2, 2, 2, 2, 2, 2)) data[3, 0, 1, 0, 1, 0] = 1 - tn.add( + data = data.reshape((2,) * 7) + return [ qtn.Tensor( data=data, - inds=( - incoming['x'], - in_qubits[0], - in_qubits[1], - outgoing['y'], - out_qubits[0], - out_qubits[1], - ), - tags=[tag], + inds=[ + (incoming['x'], 0), + (incoming['x'], 1), + (in_qubits[0], 0), + (in_qubits[1], 0), + (outgoing['y'], 0), + (out_qubits[0], 0), + (out_qubits[1], 0), + ], ) - ) - - -def _old_bloq_to_dense(bloq: Bloq) -> NDArray: - """Old code for tensor-contracting a bloq without wrapping it in length-1 composite bloq.""" - tn = qtn.TensorNetwork([]) - lsoqs = _get_dangling_soquets(bloq.signature, right=False) - rsoqs = _get_dangling_soquets(bloq.signature, right=True) - bloq.add_my_tensors(tn, None, incoming=lsoqs, outgoing=rsoqs) - - inds = get_right_and_left_inds(bloq.signature) - matrix = tn.to_dense(*inds) - return matrix + ] def test_bloq_to_dense(): - mat1 = _old_bloq_to_dense(TensorAdderTester()) mat2 = bloq_to_dense(TensorAdderTester()) - np.testing.assert_allclose(mat1, mat2, atol=1e-8) # Right inds: qubits=(1,0), y=0 right = 1 * 2**2 + 0 * 2**1 + 0 * 2**0 @@ -166,3 +158,49 @@ def test_bloq_with_non_trivial_inds(): cirq_unitary = cirq_circuit.unitary(qubit_order=cirq_qubits) np.testing.assert_allclose(cirq_unitary, bloq.decompose_bloq().tensor_contract()) np.testing.assert_allclose(cirq_unitary, bloq.tensor_contract()) + + +class BloqWithTensorsAndDecomp(Bloq): + def __init__(self): + self.called_build_composite_bloq = False + self.called_my_tensors = False + + @cached_property + def signature(self) -> 'Signature': + return Signature.build(a=1, b=1) + + def build_composite_bloq(self, bb: 'BloqBuilder', a: Soquet, b: Soquet) -> Dict[str, 'SoquetT']: + self.called_build_composite_bloq = True + a, b = bb.add(CNOT(), ctrl=a, target=b) + a, b = bb.add(CNOT(), ctrl=a, target=b) + return {'a': a, 'b': b} + + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: + self.called_my_tensors = True + return [ + qtn.Tensor( + data=np.eye(4).reshape((2, 2, 2, 2)), + inds=[ + (outgoing['a'], 0), + (outgoing['b'], 0), + (incoming['a'], 0), + (incoming['b'], 0), + ], + ) + ] + + +def test_bloq_stop_flattening(): + bloq = BloqWithTensorsAndDecomp() + u2 = bloq_to_dense(bloq, full_flatten=True) + assert bloq.called_build_composite_bloq + assert not bloq.called_my_tensors + + bloq = BloqWithTensorsAndDecomp() + u1 = bloq_to_dense(bloq, full_flatten=False) + assert not bloq.called_build_composite_bloq + assert bloq.called_my_tensors + + np.testing.assert_allclose(u1, u2) diff --git a/qualtran/simulation/tensor/_flattening.py b/qualtran/simulation/tensor/_flattening.py index 45f6ac0a0..e31a3d327 100644 --- a/qualtran/simulation/tensor/_flattening.py +++ b/qualtran/simulation/tensor/_flattening.py @@ -16,23 +16,28 @@ def bloq_has_custom_tensors(bloq: Bloq) -> bool: - """Whether this bloq declares custom tensors by overriding `.add_my_tensors(...)`. + """Whether this bloq declares custom tensors by overriding `.my_tensors(...)`. - This is a heuristic that checks that the method is overriden. This is used as the - flattening predicate in `flatten_for_tensor_contraction`. + This is a heuristic that checks that the method is overriden. This is used as + an optional predicate in `flatten_for_tensor_contraction`. """ - return not bloq.add_my_tensors.__qualname__.startswith( + return not bloq.my_tensors.__qualname__.startswith( 'Bloq.' - ) and not bloq.add_my_tensors.__qualname__.startswith('GateWithRegisters.') + ) and not bloq.my_tensors.__qualname__.startswith('GateWithRegisters.') -def flatten_for_tensor_contraction(bloq: Bloq, max_depth: int = 1_000) -> CompositeBloq: +def flatten_for_tensor_contraction(bloq: Bloq, full_flatten: bool = True) -> CompositeBloq: """Flatten a (composite) bloq as much as possible to enable efficient tensor contraction. - Without this function, bloqs without custom tensors will be contracted to a dense tensor using - their decomposition and then that dense tensor will be used in the enclosing tensor network. - To allow a more efficient contraction ordering, use this function to decompose-and-flatten - as much as possible before starting the tensor contraction. + Args: + bloq: The bloq to flatten. + full_flatten: Whether to completely flatten the bloq into the smallest possible + bloqs. Otherwise, stop flattening if custom tensors are encountered. """ cbloq = bloq.as_composite_bloq() - return cbloq.flatten(lambda binst: not bloq_has_custom_tensors(binst.bloq), max_depth=max_depth) + if full_flatten: + pred = lambda b: True + else: + pred = lambda binst: not bloq_has_custom_tensors(binst.bloq) + + return cbloq.flatten(pred) diff --git a/qualtran/simulation/tensor/_quimb.py b/qualtran/simulation/tensor/_quimb.py index ee78ac406..d4957e477 100644 --- a/qualtran/simulation/tensor/_quimb.py +++ b/qualtran/simulation/tensor/_quimb.py @@ -11,136 +11,111 @@ # 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 itertools -from typing import Dict, Optional, Set, Tuple +import logging +from typing import cast, Dict, Iterable import numpy as np import quimb.tensor as qtn from qualtran import ( Bloq, - BloqInstance, CompositeBloq, Connection, - DanglingT, LeftDangle, + QBit, + Register, RightDangle, Soquet, SoquetT, ) -from qualtran._infra.composite_bloq import ( - _cxn_to_soq_dict, - _flatten_soquet_collection, - _get_flat_dangling_soqs, -) - +from qualtran._infra.composite_bloq import _cxns_to_cxn_dict, BloqBuilder -def cbloq_to_quimb( - cbloq: CompositeBloq, pos: Optional[Dict[BloqInstance, Tuple[float, float]]] = None -) -> Tuple[qtn.TensorNetwork, Dict]: - """Convert a composite bloq into a Quimb tensor network. +logger = logging.getLogger(__name__) - External indices are the dangling soquets of the compute graph. - Args: - cbloq: The composite bloq. A composite bloq is a container class analogous to a - `TensorNetwork`. This function will simply add the tensor(s) for each Bloq - that constitutes the `CompositeBloq`. - pos: Optional mapping of each `binst` to (x, y) coordinates which will be converted - into a `fix` dictionary appropriate for `qtn.TensorNetwork.draw()`. +def cbloq_to_quimb(cbloq: CompositeBloq) -> qtn.TensorNetwork: + """Convert a composite bloq into a tensor network. - Returns: - tn: The `qtn.TensorNetwork` representing the quantum graph. This is constructed - by delegating to each bloq's `add_my_tensors` method. - fix: A version of `pos` suitable for `TensorNetwork.draw()` + This function will call `Bloq.my_tensors` on each subbloq in the composite bloq to add + tensors to a quimb tensor network. This method has no default fallback, so you likely want to + call `bloq.as_composite_bloq().flatten()` to decompose-and-flatten all bloqs down to their + smallest form first. The small bloqs that result from a flattening 1) likely already have + their `my_tensors` method implemented; and 2) can enable a more efficient tensor contraction + path. """ tn = qtn.TensorNetwork([]) - fix = {} - - def _assign_outgoing(cxn: Connection) -> Soquet: - """Logic for naming outgoing indices in quimb-land. - - In our representation, a `Connection` is a tuple of soquets. In quimb, connections are - made between nodes with indices having the same name. Conveniently, the indices - don't have to have string names, so we use a Soquet. - - Each binst makes a qtn.Tensor, and we use a soquet to name each index. We choose - the convention that each binst will respect its predecessors outgoing index names but - is in charge of its own outgoing index names. - - This convention breaks down at the end of our graph because we wish our external quimb - indices match the composite bloq's dangling soquets. Therefore: when the successor - is `RightDangle` the binst will respect the dangling soquets for naming its outgoing - indices. - """ - if isinstance(cxn.right.binst, DanglingT): - return cxn.right - return cxn.left - - visited_bloqs: Set[BloqInstance] = set() - for binst, incoming, outgoing in cbloq.iter_bloqnections(): + + logging.info( + "Constructing a tensor network for composite bloq of size %d", len(cbloq.bloq_instances) + ) + + for binst, pred_cxns, succ_cxns in cbloq.iter_bloqnections(): bloq = binst.bloq - visited_bloqs.add(binst) assert isinstance(bloq, Bloq) - inc_d = _cxn_to_soq_dict( - bloq.signature.lefts(), - incoming, - get_me=lambda cxn: cxn.right, - get_assign=lambda cxn: cxn.left, - ) - out_d = _cxn_to_soq_dict( - bloq.signature.rights(), - outgoing, - get_me=lambda cxn: cxn.left, - get_assign=_assign_outgoing, - ) - - bloq.add_my_tensors(tn, binst, incoming=inc_d, outgoing=out_d) - if pos is not None: - fix[tuple([binst])] = pos[binst] + inc_d = _cxns_to_cxn_dict(bloq.signature.lefts(), pred_cxns, get_me=lambda cxn: cxn.right) + out_d = _cxns_to_cxn_dict(bloq.signature.rights(), succ_cxns, get_me=lambda cxn: cxn.left) + + for tensor in bloq.my_tensors(inc_d, out_d): + tn.add(tensor) # Special case: Add variables corresponding to all registers that don't connect to any Bloq. # This is needed because `CompositeBloq.iter_bloqnections` ignores `LeftDangle/RightDangle` - # bloqs and therefore we never see connections the exist only b/w LeftDangle and + # bloqs, and therefore we never see connections that exist only b/w LeftDangle and # RightDangle bloqs. for cxn in cbloq.connections: if cxn.left.binst is LeftDangle and cxn.right.binst is RightDangle: - # This register has no Bloq acting on it, and thus it would not have a variable in the + # This register has no Bloq acting on it, and thus it would not have a variable in # the tensor network. Add an identity tensor acting on this register to make sure the # tensor network has variables corresponding to all input / output registers. - tn.add(qtn.Tensor(data=np.eye(2**cxn.left.reg.bitsize), inds=[cxn.right, cxn.left])) - return tn, fix + n = cxn.left.reg.bitsize + for j in range(cxn.left.reg.bitsize): + + placeholder = Soquet(None, Register('simulation_placeholder', QBit())) # type: ignore + Connection(cxn.left, placeholder) + tn.add( + qtn.Tensor( + data=np.eye(2), + inds=[ + (Connection(cxn.left, placeholder), j), + (Connection(placeholder, cxn.right), j), + ], + ) + ) + + return tn -def cbloq_as_contracted_tensor( - cbloq: CompositeBloq, incoming: Dict[str, SoquetT], outgoing: Dict[str, SoquetT], tags -) -> qtn.Tensor: - """`add_my_tensors` helper for contracting `cbloq` and adding it as a dense tensor. +def _add_classical_kets(bb: BloqBuilder, registers: Iterable[Register]) -> Dict[str, 'SoquetT']: + """Use `bb` to add `IntState(0)` for all the `vals`.""" - First, we turn the composite bloq into a TensorNetwork with `cbloq_to_quimb`. Then - we contract it to a dense ndarray. This function returns the dense array as well as - the indices munged from `incoming` and `outgoing` to match the structure of the ndarray. + from qualtran.bloqs.basic_gates import IntState + + soqs: Dict[str, 'SoquetT'] = {} + for reg in registers: + if reg.shape: + reg_vals = np.zeros(reg.shape, dtype=int) + soq = np.empty(reg.shape, dtype=object) + for idx in reg.all_idxs(): + soq[idx] = bb.add(IntState(val=cast(int, reg_vals[idx]), bitsize=reg.bitsize)) + else: + soq = bb.add(IntState(val=0, bitsize=reg.bitsize)) + + soqs[reg.name] = soq + return soqs + + +def initialize_from_zero(bloq: Bloq): + """Take `bloq` and compose it with initial zero states for each left register. + + This can be contracted to a state vector for a given unitary. """ + bb = BloqBuilder() + + # Add the initial 'kets' according to the provided values. + in_soqs = _add_classical_kets(bb, bloq.signature.lefts()) - # Turn into a dense ndarray, but instead of folding into a 1- or 2- - # dimensional state/effect or unitary; we keep all the indices as - # distinct dimensions. - signature = cbloq.signature - rsoqs = _get_flat_dangling_soqs(signature, right=True) - lsoqs = _get_flat_dangling_soqs(signature, right=False) - inds_for_contract = rsoqs + lsoqs - assert len(inds_for_contract) > 0 - tn, _ = cbloq_to_quimb(cbloq) - data = tn.to_dense(*([x] for x in inds_for_contract)) - assert data.ndim == len(inds_for_contract) - - # Now we just need to make sure the Soquets provided to us are in the correct - # order: namely the same order as how we got the indices to contract the composite bloq. - osoqs = (outgoing[reg.name] for reg in signature.rights()) - isoqs = (incoming[reg.name] for reg in signature.lefts()) - inds_for_adding = _flatten_soquet_collection(itertools.chain(osoqs, isoqs)) - assert len(inds_for_adding) == len(inds_for_contract) - - return qtn.Tensor(data=data, inds=inds_for_adding, tags=tags) + # Add the bloq itself + out_soqs = bb.add_d(bloq, **in_soqs) + return bb.finalize(**out_soqs) diff --git a/qualtran/simulation/tensor/_quimb_test.py b/qualtran/simulation/tensor/_quimb_test.py index a8595da6f..cde0e06e0 100644 --- a/qualtran/simulation/tensor/_quimb_test.py +++ b/qualtran/simulation/tensor/_quimb_test.py @@ -13,13 +13,13 @@ # limitations under the License. from functools import cached_property -from typing import Dict +from typing import Dict, List import numpy as np import quimb.tensor as qtn from attrs import frozen -from qualtran import Bloq, BloqBuilder, DanglingT, QAny, Signature, Soquet, SoquetT +from qualtran import Bloq, BloqBuilder, Connection, ConnectionT, DanglingT, QAny, Signature from qualtran.bloqs.bookkeeping import Join, Split from qualtran.simulation.tensor import cbloq_to_quimb @@ -30,17 +30,12 @@ class TensorAdderSimple(Bloq): def signature(self) -> 'Signature': return Signature.build(x=1) - def add_my_tensors( - self, - tn: qtn.TensorNetwork, - tag, - *, - incoming: Dict[str, SoquetT], - outgoing: Dict[str, SoquetT], - ): + def my_tensors( + self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT'] + ) -> List['qtn.Tensor']: assert list(incoming.keys()) == ['x'] assert list(outgoing.keys()) == ['x'] - tn.add(qtn.Tensor(data=np.eye(2), inds=(incoming['x'], outgoing['x']), tags=[tag])) + return [qtn.Tensor(data=np.eye(2), inds=[(incoming['x'], 0), (outgoing['x'], 0)])] def test_cbloq_to_quimb(): @@ -52,11 +47,12 @@ def test_cbloq_to_quimb(): x = bb.add(TensorAdderSimple(), x=x) cbloq = bb.finalize(x=x) - tn, _ = cbloq_to_quimb(cbloq) + tn = cbloq_to_quimb(cbloq) assert len(tn.tensors) == 4 - for oi in tn.outer_inds(): - assert isinstance(oi, Soquet) - assert isinstance(oi.binst, DanglingT) + for outer_ind in tn.outer_inds(): + cxn, j = outer_ind + assert isinstance(cxn, Connection) + assert isinstance(cxn.left.binst, DanglingT) or isinstance(cxn.right.binst, DanglingT) def test_cbloq_to_quimb_with_no_ops_on_register(): diff --git a/qualtran/simulation/xcheck_classical_quimb.ipynb b/qualtran/simulation/xcheck_classical_quimb.ipynb index 634f90416..9d20a4bf4 100644 --- a/qualtran/simulation/xcheck_classical_quimb.ipynb +++ b/qualtran/simulation/xcheck_classical_quimb.ipynb @@ -26,7 +26,9 @@ "outputs": [], "source": [ "from qualtran.drawing import show_bloq\n", - "from qualtran.simulation.xcheck_classical_quimb import flank_with_classical_vectors" + "from qualtran.simulation.xcheck_classical_quimb import flank_with_classical_vectors\n", + "\n", + "import numpy as np" ] }, { @@ -184,15 +186,73 @@ { "cell_type": "code", "execution_count": null, - "id": "82e374f9", + "id": "4f88f54d-abac-4545-87ef-ba1ff7f46048", "metadata": {}, "outputs": [], "source": [ "add_tt = flank_with_classical_vectors(add, {'a': 2, 'b': 3})\n", - "\n", - "assert add_tt.tensor_contract() == 1.0\n", "show_bloq(add_tt, type='musical_score')" ] + }, + { + "cell_type": "markdown", + "id": "0df5eced-4382-4c2c-af94-03d31118bd23", + "metadata": {}, + "source": [ + "Since `Add` is composed of other bloqs, we can contract the factorized network that comes from the \"flattened\" circuit." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84a00ba8-455e-434b-b618-8ba44183bc8d", + "metadata": {}, + "outputs": [], + "source": [ + "add_flat = add_tt.as_composite_bloq().flatten()\n", + "\n", + "from qualtran.drawing import get_musical_score_data, draw_musical_score\n", + "draw_musical_score(get_musical_score_data(add_flat), max_width=20, max_height=100)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "176355a9-6477-4932-b696-6dcc15dcd39d", + "metadata": {}, + "outputs": [], + "source": [ + "np.testing.assert_allclose(add_tt.tensor_contract(), 1.0)\n", + "np.testing.assert_allclose(add_flat.tensor_contract(), 1.0)" + ] + }, + { + "cell_type": "markdown", + "id": "bdb1f3da-0837-40cd-a810-be3bc9bd4638", + "metadata": {}, + "source": [ + "### Large bitsize\n", + "\n", + "Since we're using Quimb to find an efficient contraction ordering, we can handle large-width, bounded-depth tensor networks." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b569db6-c6cf-46ad-a625-7c0c72b307a5", + "metadata": {}, + "outputs": [], + "source": [ + "add_large = Add(QUInt(bitsize=32))\n", + "add_large_tt = flank_with_classical_vectors(add_large, {'a': 5555, 'b': 6666})\n", + "add_large_flat = add_large_tt.flatten()\n", + "\n", + "from qualtran.simulation.tensor import cbloq_to_quimb\n", + "tn = cbloq_to_quimb(add_large_flat)\n", + "\n", + "print(f\"We only use {tn.contraction_width()} qubit's worth of RAM for a {add_large.signature.n_qubits()}-qubit gate.\")\n", + "tn.contract()" + ] } ], "metadata": { @@ -211,7 +271,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.9" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/qualtran/simulation/xcheck_classical_quimb_test.py b/qualtran/simulation/xcheck_classical_quimb_test.py index 4e73b036e..cc2c07052 100644 --- a/qualtran/simulation/xcheck_classical_quimb_test.py +++ b/qualtran/simulation/xcheck_classical_quimb_test.py @@ -45,7 +45,7 @@ def test_toffoli(): def test_add(): add = Add(QUInt(bitsize=5)) add_tt = flank_with_classical_vectors(add, {'a': 2, 'b': 3}) - assert add_tt.tensor_contract() == 1.0 + np.testing.assert_allclose(add_tt.tensor_contract(), 1.0) @pytest.mark.notebook