Add is_symbolic to data types and registers. (#942)

* Move symbolic_counting_utils to a new module

* Add is_symbolic to data types and registers

* Fix mypy errors

* Fix tests

Author: tanujkhattar

qualtran/_infra/data_types.py

@@ -54,10 +54,11 @@
 
 import attrs
 import numpy as np
-import sympy
 from fxpmath import Fxp
 from numpy.typing import NDArray
 
+from qualtran.symbolics import is_symbolic, SymbolicInt
+
 
 class QDType(metaclass=abc.ABCMeta):
     """This defines the abstract interface for quantum data types."""
@@ -89,7 +90,11 @@ def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
             debug_str: Optional debugging information to use in exception messages.
         """
 
-    def iteration_length_or_zero(self) -> Union[int, sympy.Expr]:
+    @abc.abstractmethod
+    def is_symbolic(self) -> bool:
+        """Returns True if this qdtype is parameterized with symbolic objects."""
+
+    def iteration_length_or_zero(self) -> SymbolicInt:
         """Safe version of iteration length.
 
         Returns the iteration_length if the type has it or else zero.
@@ -130,6 +135,9 @@ def assert_valid_classical_val(self, val: int, debug_str: str = 'val'):
         if not (val == 0 or val == 1):
             raise ValueError(f"Bad {self} value {val} in {debug_str}")
 
+    def is_symbolic(self) -> bool:
+        return False
+
     def to_bits(self, x) -> List[int]:
         """Yields individual bits corresponding to binary representation of x"""
         self.assert_valid_classical_val(x)
@@ -154,7 +162,7 @@ def __str__(self):
 class QAny(QDType):
     """Opaque bag-of-qbits type."""
 
-    bitsize: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
 
     @property
     def num_qubits(self):
@@ -171,6 +179,9 @@ def from_bits(self, bits: Sequence[int]) -> int:
         # TODO: Raise an error once usage of `QAny` is minimized across the library
         return QUInt(self.bitsize).from_bits(bits)
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize)
+
     def assert_valid_classical_val(self, val, debug_str: str = 'val'):
         pass
 
@@ -188,12 +199,15 @@ class QInt(QDType):
         bitsize: The number of qubits used to represent the integer.
     """
 
-    bitsize: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
 
     @property
     def num_qubits(self):
         return self.bitsize
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize)
+
     def get_classical_domain(self) -> Iterable[int]:
         max_val = 1 << (self.bitsize - 1)
         return range(-max_val, max_val)
@@ -240,7 +254,7 @@ class QIntOnesComp(QDType):
         bitsize: The number of qubits used to represent the integer.
     """
 
-    bitsize: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
 
     def __attrs_post_init__(self):
         if isinstance(self.bitsize, int):
@@ -251,6 +265,9 @@ def __attrs_post_init__(self):
     def num_qubits(self):
         return self.bitsize
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize)
+
     def to_bits(self, x: int) -> List[int]:
         """Yields individual bits corresponding to binary representation of x"""
         self.assert_valid_classical_val(x)
@@ -286,12 +303,15 @@ class QUInt(QDType):
         bitsize: The number of qubits used to represent the integer.
     """
 
-    bitsize: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
 
     @property
     def num_qubits(self):
         return self.bitsize
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize)
+
     def get_classical_domain(self) -> Iterable[Any]:
         return range(2**self.bitsize)
 
@@ -371,11 +391,11 @@ class BoundedQUInt(QDType):
         iteration_length: The length of the iteration range.
     """
 
-    bitsize: Union[int, sympy.Expr]
-    iteration_length: Union[int, sympy.Expr] = attrs.field()
+    bitsize: SymbolicInt
+    iteration_length: SymbolicInt = attrs.field()
 
     def __attrs_post_init__(self):
-        if isinstance(self.bitsize, int):
+        if not self.is_symbolic():
             if self.iteration_length > 2**self.bitsize:
                 raise ValueError(
                     "BoundedQUInt iteration length is too large for given bitsize. "
@@ -386,6 +406,9 @@ def __attrs_post_init__(self):
     def _default_iteration_length(self):
         return 2**self.bitsize
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize, self.iteration_length)
+
     @property
     def num_qubits(self):
         return self.bitsize
@@ -446,16 +469,16 @@ class QFxp(QDType):
             number of integer bits is reduced by 1.
     """
 
-    bitsize: Union[int, sympy.Expr]
-    num_frac: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
+    num_frac: SymbolicInt
     signed: bool = False
 
     @property
     def num_qubits(self):
         return self.bitsize
 
     @property
-    def num_int(self) -> Union[int, sympy.Expr]:
+    def num_int(self) -> SymbolicInt:
         return self.bitsize - self.num_frac - int(self.signed)
 
     @property
@@ -466,6 +489,9 @@ def fxp_dtype_str(self) -> str:
     def _fxp_dtype(self) -> Fxp:
         return Fxp(None, dtype=self.fxp_dtype_str)
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize, self.num_frac)
+
     def to_bits(self, x: Union[float, Fxp]) -> List[int]:
         """Yields individual bits corresponding to binary representation of x"""
         self._assert_valid_classical_val(x)
@@ -539,12 +565,15 @@ class QMontgomeryUInt(QDType):
         [Montgomery modular multiplication](https://en.wikipedia.org/wiki/Montgomery_modular_multiplication)
     """
 
-    bitsize: Union[int, sympy.Expr]
+    bitsize: SymbolicInt
 
     @property
     def num_qubits(self):
         return self.bitsize
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.bitsize)
+
     def get_classical_domain(self) -> Iterable[Any]:
         return range(2**self.bitsize)
 

qualtran/_infra/data_types_test.py

@@ -16,6 +16,8 @@
 import pytest
 import sympy
 
+from qualtran.symbolics import is_symbolic
+
 from .data_types import (
     BoundedQUInt,
     check_dtypes_consistent,
@@ -39,6 +41,7 @@ def test_qint():
     qint_8 = QInt(n)
     assert qint_8.num_qubits == n
     assert str(qint_8) == 'QInt(x)'
+    assert is_symbolic(QInt(sympy.Symbol('x')))
 
 
 def test_qint_ones():
@@ -50,6 +53,7 @@ def test_qint_ones():
     n = sympy.symbols('x')
     qint_8 = QIntOnesComp(n)
     assert qint_8.num_qubits == n
+    assert is_symbolic(QIntOnesComp(sympy.Symbol('x')))
 
 
 def test_quint():
@@ -62,6 +66,7 @@ def test_quint():
     n = sympy.symbols('x')
     qint_8 = QUInt(n)
     assert qint_8.num_qubits == n
+    assert is_symbolic(QUInt(sympy.Symbol('x')))
 
 
 def test_bounded_quint():
@@ -77,6 +82,9 @@ def test_bounded_quint():
     qint_8 = BoundedQUInt(n, l)
     assert qint_8.num_qubits == n
     assert qint_8.iteration_length == l
+    assert is_symbolic(BoundedQUInt(sympy.Symbol('x'), 2))
+    assert is_symbolic(BoundedQUInt(2, sympy.Symbol('x')))
+    assert is_symbolic(BoundedQUInt(*sympy.symbols('x y')))
 
 
 def test_qfxp():
@@ -105,6 +113,7 @@ def test_qfxp():
     qfp = QFxp(b, f, True)
     assert qfp.num_qubits == b
     assert qfp.num_int == b - f - 1
+    assert is_symbolic(QFxp(*sympy.symbols('x y')))
 
 
 def test_qmontgomeryuint():
@@ -116,6 +125,7 @@ def test_qmontgomeryuint():
     n = sympy.symbols('x')
     qmontgomeryuint_8 = QMontgomeryUInt(n)
     assert qmontgomeryuint_8.num_qubits == n
+    assert is_symbolic(QMontgomeryUInt(sympy.Symbol('x')))
 
 
 @pytest.mark.parametrize('qdtype', [QBit(), QInt(4), QUInt(4), BoundedQUInt(3, 5)])

qualtran/_infra/registers.py

@@ -16,13 +16,15 @@
 import enum
 import itertools
 from collections import defaultdict
-from typing import Dict, Iterable, Iterator, List, overload, Tuple, Union
+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, SymbolicInt
+
 from .data_types import QAny, QBit, QDType
 
 
@@ -63,7 +65,7 @@ class Register:
 
     name: str
     dtype: QDType
-    shape: Tuple[int, ...] = field(
+    _shape: Tuple[SymbolicInt, ...] = field(
         default=tuple(), converter=lambda v: (v,) if isinstance(v, int) else tuple(v)
     )
     side: Side = Side.THRU
@@ -72,6 +74,19 @@ def __attrs_post_init__(self):
         if not isinstance(self.dtype, QDType):
             raise ValueError(f'dtype must be a QDType: found {type(self.dtype)}')
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.dtype, *self._shape)
+
+    @property
+    def shape_symbolic(self) -> Tuple[SymbolicInt, ...]:
+        return self._shape
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        if is_symbolic(*self._shape):
+            raise ValueError(f"{self} is symbolic. Cannot get real-valued shape.")
+        return cast(Tuple[int, ...], self._shape)
+
     @property
     def bitsize(self) -> int:
         return self.dtype.num_qubits

qualtran/_infra/registers_test.py

@@ -15,9 +15,11 @@
 import cirq
 import numpy as np
 import pytest
+import sympy
 
 from qualtran import BoundedQUInt, QAny, QBit, QInt, Register, Side, Signature
 from qualtran._infra.gate_with_registers import get_named_qubits
+from qualtran.symbolics import is_symbolic
 
 
 def test_register():
@@ -193,6 +195,13 @@ def test_dtypes_converter():
     r1 = Register("my_reg", QBit())
     r2 = Register("my_reg", QBit())
     assert r1 == r2
-    r2 = Register("my_reg", QAny(5))
+    r1 = Register("my_reg", QAny(5))
     r2 = Register("my_reg", QInt(5))
     assert r1 != r2
+
+
+def test_is_symbolic():
+    r = Register("my_reg", QAny(sympy.Symbol("x")))
+    assert is_symbolic(r)
+    r = Register("my_reg", QAny(2), shape=sympy.symbols("x y"))
+    assert is_symbolic(r)

qualtran/serialization/registers.py

@@ -32,7 +32,7 @@ def register_to_proto(register: Register) -> registers_pb2.Register:
     return registers_pb2.Register(
         name=register.name,
         dtype=data_type_to_proto(register.dtype),
-        shape=(args.int_or_sympy_to_proto(s) for s in register.shape),
+        shape=(args.int_or_sympy_to_proto(s) for s in register.shape_symbolic),
         side=_side_to_proto(register.side),
     )