Use symbolics for computing register bitsizes (#1353)

* use symbolics for computing register bitsizes

* symbolic `n` in `Partition`

Author: anurudhp

qualtran/_infra/registers.py

@@ -19,11 +19,10 @@
 from typing import cast, Dict, Iterable, Iterator, List, overload, Tuple, Union
 
 import attrs
-import numpy as np
 import sympy
 from attrs import field, frozen
 
-from qualtran.symbolics import is_symbolic, smax, SymbolicInt
+from qualtran.symbolics import is_symbolic, prod, smax, ssum, SymbolicInt
 
 from .data_types import QAny, QBit, QDType
 
@@ -99,7 +98,7 @@ def total_bits(self) -> int:
 
         This is the product of bitsize and each of the dimensions in `shape`.
         """
-        return self.bitsize * int(np.prod(self.shape))
+        return self.bitsize * prod(self.shape_symbolic)
 
     def adjoint(self) -> 'Register':
         """Return the 'adjoint' of this register by switching RIGHT and LEFT registers."""
@@ -202,8 +201,8 @@ def n_qubits(self) -> int:
         is taken to be the greater of the number of left or right qubits. A bloq with this
         signature uses at least this many qubits.
         """
-        left_size = sum(reg.total_bits() for reg in self.lefts())
-        right_size = sum(reg.total_bits() for reg in self.rights())
+        left_size = ssum(reg.total_bits() for reg in self.lefts())
+        right_size = ssum(reg.total_bits() for reg in self.rights())
         return smax(left_size, right_size)
 
     def __repr__(self):

qualtran/bloqs/bookkeeping/partition.py

@@ -32,6 +32,7 @@
 )
 from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
 from qualtran.drawing import directional_text_box, Text, WireSymbol
+from qualtran.symbolics import is_symbolic, ssum, SymbolicInt
 
 if TYPE_CHECKING:
     import quimb.tensor as qtn
@@ -54,14 +55,14 @@ class Partition(_BookkeepingBloq):
         [user spec]: The registers provided by the `regs` argument. RIGHT by default.
     """
 
-    n: int
+    n: SymbolicInt
     regs: Tuple[Register, ...] = field(
         converter=lambda x: x if isinstance(x, tuple) else tuple(x), validator=validators.min_len(1)
     )
     partition: bool = True
 
     def __attrs_post_init__(self):
-        if self.n != sum(r.total_bits() for r in self.regs):
+        if self.n != ssum(r.total_bits() for r in self.regs):
             raise ValueError("Total bitsize not equal to sum of registers to partition into")
         if len(set(r.name for r in self.regs)) != len(self.regs):
             raise ValueError("Duplicate register names")
@@ -104,6 +105,9 @@ def my_tensors(
     ) -> List['qtn.Tensor']:
         import quimb.tensor as qtn
 
+        if is_symbolic(self.n):
+            raise DecomposeTypeError(f"cannot compute tensors for symbolic {self}")
+
         grouped = incoming['x'] if self.partition else outgoing['x']
         partitioned = outgoing if self.partition else incoming