Skip to content

Commit 0c3da85

Browse files
anurudhpanurudhp
authored
fix slow import: only import quimb.tensor when needed (#949)
* only import `quimb` when needed * one more lazy import
1 parent c10d386 commit 0c3da85

9 files changed

Lines changed: 80 additions & 28 deletions

File tree

‎qualtran/_infra/controlled.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import attrs
3030
import cirq
3131
import numpy as np
32-
import quimb.tensor as qtn
3332
from numpy.typing import NDArray
3433

3534
from .bloq import Bloq
@@ -38,6 +37,8 @@
3837
from .registers import Register, Side, Signature
3938

4039
if TYPE_CHECKING:
40+
import quimb.tensor as qtn
41+
4142
from qualtran import BloqBuilder, CompositeBloq, Soquet, SoquetT
4243
from qualtran.cirq_interop import CirqQuregT
4344
from qualtran.drawing import WireSymbol
@@ -399,6 +400,8 @@ def add_my_tensors(
399400
incoming: Dict[str, 'SoquetT'],
400401
outgoing: Dict[str, 'SoquetT'],
401402
):
403+
import quimb.tensor as qtn
404+
402405
from qualtran._infra.composite_bloq import _flatten_soquet_collection
403406
from qualtran.simulation.tensor._tensor_data_manipulation import (
404407
active_space_for_ctrl_spec,

‎qualtran/bloqs/basic_gates/cnot.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING
1818

1919
import numpy as np
20-
import quimb.tensor as qtn
2120
from attrs import frozen
2221

2322
from qualtran import (
@@ -38,11 +37,11 @@
3837

3938
if TYPE_CHECKING:
4039
import cirq
40+
import quimb.tensor as qtn
4141

4242
from qualtran.cirq_interop import CirqQuregT
4343
from qualtran.simulation.classical_sim import ClassicalValT
4444

45-
4645
COPY = [1, 0, 0, 0, 0, 0, 0, 1]
4746
COPY = np.array(COPY, dtype=np.complex128).reshape((2, 2, 2))
4847

@@ -72,7 +71,7 @@ def adjoint(self) -> 'Bloq':
7271

7372
def add_my_tensors(
7473
self,
75-
tn: qtn.TensorNetwork,
74+
tn: 'qtn.TensorNetwork',
7675
tag: Any,
7776
*,
7877
incoming: Dict[str, SoquetT],
@@ -86,6 +85,8 @@ def add_my_tensors(
8685
References:
8786
[Lectures on Quantum Tensor Networks](https://arxiv.org/abs/1912.10049). Biamonte 2019.
8887
"""
88+
import quimb.tensor as qtn
89+
8990
internal = qtn.rand_uuid()
9091
tn.add(
9192
qtn.Tensor(

‎qualtran/bloqs/basic_gates/swap.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
import cirq
1919
import numpy as np
20-
import quimb.tensor as qtn
2120
import sympy
2221
from attrs import frozen
2322
from numpy.typing import NDArray
@@ -43,6 +42,8 @@
4342
from .t_gate import TGate
4443

4544
if TYPE_CHECKING:
45+
import quimb.tensor as qtn
46+
4647
from qualtran import CompositeBloq
4748
from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
4849
from qualtran.simulation.classical_sim import ClassicalValT
@@ -90,6 +91,8 @@ def add_my_tensors(
9091
incoming: Dict[str, 'SoquetT'],
9192
outgoing: Dict[str, 'SoquetT'],
9293
):
94+
import quimb.tensor as qtn
95+
9396
matrix = _swap_matrix()
9497
out_inds = [outgoing['x'], outgoing['y']]
9598
in_inds = [incoming['x'], incoming['y']]
@@ -157,6 +160,8 @@ def add_my_tensors(
157160
incoming: Dict[str, 'SoquetT'],
158161
outgoing: Dict[str, 'SoquetT'],
159162
):
163+
import quimb.tensor as qtn
164+
160165
matrix = _controlled_swap_matrix()
161166
out_inds = [outgoing['ctrl'], outgoing['x'], outgoing['y']]
162167
in_inds = [incoming['ctrl'], incoming['x'], incoming['y']]

‎qualtran/bloqs/basic_gates/x_basis.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING, Union
1717

1818
import numpy as np
19-
import quimb.tensor as qtn
2019
from attrs import frozen
2120

2221
from qualtran import (
@@ -37,6 +36,7 @@
3736

3837
if TYPE_CHECKING:
3938
import cirq
39+
import quimb.tensor as qtn
4040

4141
from qualtran.cirq_interop import CirqQuregT
4242
from qualtran.simulation.classical_sim import ClassicalValT
@@ -74,12 +74,14 @@ def signature(self) -> 'Signature':
7474

7575
def add_my_tensors(
7676
self,
77-
tn: qtn.TensorNetwork,
77+
tn: 'qtn.TensorNetwork',
7878
tag: Any,
7979
*,
8080
incoming: Dict[str, SoquetT],
8181
outgoing: Dict[str, SoquetT],
8282
):
83+
import quimb.tensor as qtn
84+
8385
side = outgoing if self.state else incoming
8486
tn.add(
8587
qtn.Tensor(
@@ -221,12 +223,14 @@ def adjoint(self) -> 'Bloq':
221223

222224
def add_my_tensors(
223225
self,
224-
tn: qtn.TensorNetwork,
226+
tn: 'qtn.TensorNetwork',
225227
tag: Any,
226228
*,
227229
incoming: Dict[str, SoquetT],
228230
outgoing: Dict[str, SoquetT],
229231
):
232+
import quimb.tensor as qtn
233+
230234
tn.add(
231235
qtn.Tensor(
232236
data=_PAULIX, inds=(outgoing['q'], incoming['q']), tags=[self.pretty_name(), tag]

‎qualtran/bloqs/basic_gates/y_gate.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@
1616
from typing import Any, Dict, Tuple, TYPE_CHECKING
1717

1818
import numpy as np
19-
import quimb.tensor as qtn
2019
from attrs import frozen
2120

2221
from qualtran import Bloq, Signature, SoquetT
2322
from qualtran.cirq_interop.t_complexity_protocol import TComplexity
2423

2524
if TYPE_CHECKING:
2625
import cirq
26+
import quimb.tensor as qtn
2727

2828
from qualtran.cirq_interop import CirqQuregT
2929

@@ -46,12 +46,14 @@ def adjoint(self) -> 'Bloq':
4646

4747
def add_my_tensors(
4848
self,
49-
tn: qtn.TensorNetwork,
49+
tn: 'qtn.TensorNetwork',
5050
tag: Any,
5151
*,
5252
incoming: Dict[str, SoquetT],
5353
outgoing: Dict[str, SoquetT],
5454
):
55+
import quimb.tensor as qtn
56+
5557
tn.add(qtn.Tensor(data=_PAULIY, inds=(outgoing['q'], incoming['q']), tags=["Y", tag]))
5658

5759
def as_cirq_op(

‎qualtran/bloqs/basic_gates/z_basis.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
import attrs
1919
import numpy as np
20-
import quimb.tensor as qtn
2120
import sympy
2221
from attrs import frozen
2322
from numpy.typing import NDArray
@@ -43,6 +42,7 @@
4342

4443
if TYPE_CHECKING:
4544
import cirq
45+
import quimb.tensor as qtn
4646

4747
from qualtran.cirq_interop import CirqQuregT
4848
from qualtran.drawing import WireSymbol
@@ -84,12 +84,14 @@ def decompose_bloq(self) -> CompositeBloq:
8484

8585
def add_my_tensors(
8686
self,
87-
tn: qtn.TensorNetwork,
87+
tn: 'qtn.TensorNetwork',
8888
tag: Any,
8989
*,
9090
incoming: Dict[str, SoquetT],
9191
outgoing: Dict[str, SoquetT],
9292
):
93+
import quimb.tensor as qtn
94+
9395
side = outgoing if self.state else incoming
9496
tn.add(
9597
qtn.Tensor(
@@ -240,12 +242,14 @@ def decompose_bloq(self) -> CompositeBloq:
240242

241243
def add_my_tensors(
242244
self,
243-
tn: qtn.TensorNetwork,
245+
tn: 'qtn.TensorNetwork',
244246
tag: Any,
245247
*,
246248
incoming: Dict[str, SoquetT],
247249
outgoing: Dict[str, SoquetT],
248250
):
251+
import quimb.tensor as qtn
252+
249253
tn.add(qtn.Tensor(data=_PAULIZ, inds=(outgoing['q'], incoming['q']), tags=["Z", tag]))
250254

251255
def as_cirq_op(
@@ -334,12 +338,14 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **val: 'SoquetT') -> Dict[str,
334338

335339
def add_my_tensors(
336340
self,
337-
tn: qtn.TensorNetwork,
341+
tn: 'qtn.TensorNetwork',
338342
tag: Any,
339343
*,
340344
incoming: Dict[str, SoquetT],
341345
outgoing: Dict[str, SoquetT],
342346
):
347+
import quimb.tensor as qtn
348+
343349
if isinstance(self.bitsize, sympy.Expr):
344350
raise ValueError(f'Symbolic bitsize {self.bitsize} not supported')
345351
data = np.zeros(2**self.bitsize).reshape((2,) * self.bitsize)

‎qualtran/bloqs/mcmt/and_bloq.py‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,22 @@
2424

2525
import itertools
2626
from functools import cached_property
27-
from typing import Any, Dict, Iterable, Iterator, Optional, Sequence, Set, Tuple, Union
27+
from typing import (
28+
Any,
29+
Dict,
30+
Iterable,
31+
Iterator,
32+
Optional,
33+
Sequence,
34+
Set,
35+
Tuple,
36+
TYPE_CHECKING,
37+
Union,
38+
)
2839

2940
import attrs
3041
import cirq
3142
import numpy as np
32-
import quimb.tensor as qtn
3343
import sympy
3444
from attrs import field, frozen
3545
from numpy.typing import NDArray
@@ -61,6 +71,9 @@
6171
)
6272
from qualtran.simulation.classical_sim import ClassicalValT
6373

74+
if TYPE_CHECKING:
75+
import quimb.tensor as qtn
76+
6477

6578
@frozen
6679
class And(GateWithRegisters):
@@ -124,12 +137,14 @@ def on_classical_vals(
124137

125138
def add_my_tensors(
126139
self,
127-
tn: qtn.TensorNetwork,
140+
tn: 'qtn.TensorNetwork',
128141
tag: Any,
129142
*,
130143
incoming: Dict[str, NDArray[Soquet]], # type: ignore[type-var]
131144
outgoing: Dict[str, NDArray[Soquet]], # type: ignore[type-var]
132145
):
146+
import quimb.tensor as qtn
147+
133148
# Fill in our tensor using "and" logic.
134149
data = np.zeros((2, 2, 2, 2, 2), dtype=np.complex128)
135150
for c1, c2 in itertools.product((0, 1), repeat=2):

‎qualtran/bloqs/util_bloqs.py‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919

2020
import attrs
2121
import numpy as np
22-
import quimb.tensor as qtn
2322
from attrs import frozen
2423
from sympy import Expr
2524

@@ -42,6 +41,7 @@
4241

4342
if TYPE_CHECKING:
4443
import cirq
44+
import quimb.tensor as qtn
4545
from numpy.typing import NDArray
4646

4747
from qualtran import AddControlledT, CtrlSpec
@@ -83,12 +83,14 @@ def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
8383

8484
def add_my_tensors(
8585
self,
86-
tn: qtn.TensorNetwork,
86+
tn: 'qtn.TensorNetwork',
8787
tag: Any,
8888
*,
8989
incoming: Dict[str, 'SoquetT'],
9090
outgoing: Dict[str, 'SoquetT'],
9191
):
92+
import quimb.tensor as qtn
93+
9294
if not isinstance(outgoing['reg'], np.ndarray):
9395
raise ValueError('Outgoing register must be a numpy array')
9496
tn.add(
@@ -150,12 +152,14 @@ def _t_complexity_(self) -> 'TComplexity':
150152

151153
def add_my_tensors(
152154
self,
153-
tn: qtn.TensorNetwork,
155+
tn: 'qtn.TensorNetwork',
154156
tag: Any,
155157
*,
156158
incoming: Dict[str, 'SoquetT'],
157159
outgoing: Dict[str, 'SoquetT'],
158160
):
161+
import quimb.tensor as qtn
162+
159163
if not isinstance(incoming['reg'], np.ndarray):
160164
raise ValueError('Incoming register must be a numpy array')
161165
tn.add(
@@ -241,12 +245,14 @@ def _t_complexity_(self) -> 'TComplexity':
241245

242246
def add_my_tensors(
243247
self,
244-
tn: qtn.TensorNetwork,
248+
tn: 'qtn.TensorNetwork',
245249
tag: Any,
246250
*,
247251
incoming: Dict[str, 'SoquetT'],
248252
outgoing: Dict[str, 'SoquetT'],
249253
):
254+
import quimb.tensor as qtn
255+
250256
unitary_shape = []
251257
soquets = []
252258
_incoming = incoming if self.partition else outgoing
@@ -347,6 +353,8 @@ def add_my_tensors(
347353
incoming: Dict[str, 'SoquetT'],
348354
outgoing: Dict[str, 'SoquetT'],
349355
):
356+
import quimb.tensor as qtn
357+
350358
data = np.zeros(1 << self.dtype.num_qubits)
351359
data[0] = 1
352360
tn.add(qtn.Tensor(data=data, inds=(outgoing['reg'],), tags=['Allocate', tag]))
@@ -395,6 +403,8 @@ def add_my_tensors(
395403
incoming: Dict[str, 'SoquetT'],
396404
outgoing: Dict[str, 'SoquetT'],
397405
):
406+
import quimb.tensor as qtn
407+
398408
data = np.zeros(1 << self.dtype.num_qubits)
399409
data[0] = 1
400410
tn.add(qtn.Tensor(data=data, inds=(incoming['reg'],), tags=['Free', tag]))
@@ -468,12 +478,14 @@ def signature(self) -> Signature:
468478

469479
def add_my_tensors(
470480
self,
471-
tn: qtn.TensorNetwork,
481+
tn: 'qtn.TensorNetwork',
472482
tag: Any,
473483
*,
474484
incoming: Dict[str, 'SoquetT'],
475485
outgoing: Dict[str, 'SoquetT'],
476486
):
487+
import quimb.tensor as qtn
488+
477489
tn.add(
478490
qtn.Tensor(
479491
data=np.eye(2**self.inp_dtype.num_qubits, 2**self.inp_dtype.num_qubits),

0 commit comments

Comments
 (0)