Modifies QAny to not have an implicit _Unit encoding

pyproject.toml

@@ -15,6 +15,7 @@ skip_glob = ["qualtran/protos/*"]
 [tool.pytest.ini_options]
 filterwarnings = [
     'ignore::DeprecationWarning:quimb.linalg.approx_spectral:',
+    'ignore::qualtran.bloqs.bookkeeping.partition.LegacyPartitionWarning',
     'ignore:.*standard platformdirs.*:DeprecationWarning:jupyter_client.*'
 ]
 # we define classes like TestBloq etc. which pytest tries to collect,

qualtran/_infra/data_types.py

@@ -358,6 +358,76 @@ def __str__(self) -> str:
         return 'CBit()'
 
 
+@attrs.frozen
+class _Any(BitEncoding[int]):
+    """Bag of Qubits of a given bitsize.
+
+    Here (and throughout Qualtran), we use a big-endian bit convention. The most significant
+    bit is at index 0.
+    """
+
+    bitsize: SymbolicInt
+
+    def get_domain(self) -> Iterable[int]:
+        raise TypeError(f"Ambiguous domain for {self}. Please use a more specific type.")
+
+    def to_bits(self, x: int) -> List[int]:
+        if is_symbolic(self.bitsize):
+            raise ValueError(f"Cannot compute bits for symbolic {self.bitsize=}")
+        if x == 0:
+            return [0] * int(self.bitsize)
+
+        raise TypeError(
+            f"Ambiguous encoding for {self} when encoding non zero value {x=}. Please use a more specific type."
+        )
+
+    def to_bits_array(self, x_array: NDArray[np.integer]) -> NDArray[np.uint8]:
+        if is_symbolic(self.bitsize):
+            raise ValueError(f"Cannot compute bits for symbolic {self.bitsize=}")
+
+        values = np.atleast_1d(x_array)
+        if values.size == 0:
+            return np.zeros((values.shape[0], int(self.bitsize)), dtype=np.uint8)
+
+        if not np.all(values == 0):
+            raise TypeError(
+                f"Ambiguous encoding for {self} when encoding non zero values {values=}. Please use a more specific type."
+            )
+
+        return np.zeros((values.shape[0], int(self.bitsize)), dtype=np.uint8)
+
+    def from_bits(self, bits: Sequence[int]) -> int:
+        if all(x == 0 for x in bits):
+            return 0
+
+        raise TypeError(
+            f"Ambiguous value for {self} when bits ({bits}) are non zero. Please use a more specific type."
+        )
+
+    def from_bits_array(self, bits_array: NDArray[np.uint8]) -> NDArray[np.uint64]:
+        bitstrings = np.atleast_2d(bits_array)
+        if bitstrings.shape[1] != self.bitsize:
+            raise ValueError(f"Input bitsize {bitstrings.shape[1]} does not match {self.bitsize=}")
+
+        if bitstrings.size == 0:
+            return np.zeros(bitstrings.shape[0], dtype=np.uint64)
+
+        if not np.all(bitstrings == 0):
+            raise TypeError(
+                f"Ambiguous value for {self} when bits are non zero ({bits_array}). Please use a more specific type."
+            )
+
+        return np.zeros(bitstrings.shape[0], dtype=np.uint64)
+
+    def assert_valid_val(self, val: int, debug_str: str = 'val') -> None:
+        pass
+
+    def assert_valid_val_array(
+        self, val_array: NDArray[np.integer], debug_str: str = 'val'
+    ) -> None:
+        pass
+
+
 @attrs.frozen
 class QAny(QDType[Any]):
     """Opaque bag-of-qubits type."""
@@ -366,7 +436,7 @@ class QAny(QDType[Any]):
 
     @property
     def _bit_encoding(self) -> BitEncoding[Any]:
-        return _UInt(self.bitsize)
+        return _Any(self.bitsize)
 
     def __attrs_post_init__(self):
         if is_symbolic(self.bitsize):
@@ -375,15 +445,6 @@ def __attrs_post_init__(self):
         if not isinstance(self.bitsize, int):
             raise ValueError(f"Bad bitsize for QAny: {self.bitsize}")
 
-    def get_classical_domain(self) -> Iterable[Any]:
-        raise TypeError(f"Ambiguous domain for {self}. Please use a more specific type.")
-
-    def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
-        pass
-
-    def assert_valid_classical_val_array(self, val_array: NDArray, debug_str: str = 'val'):
-        pass
-
 
 @attrs.frozen
 class _Int(BitEncoding[int]):

qualtran/_infra/data_types_test.py

@@ -441,10 +441,32 @@ def test_qbit_to_and_from_bits():
     assert_to_and_from_bits_array_consistent(QBit(), [0, 1])
 
 
-def test_qany_to_and_from_bits():
-    assert list(QAny(4).to_bits(10)) == [1, 0, 1, 0]
+def test_qany_to_bits():
+    with pytest.raises(TypeError, match=r"Ambiguous encoding"):
+        QAny(4).to_bits(10)
 
-    assert_to_and_from_bits_array_consistent(QAny(4), range(16))
+
+def test_qany_from_bits_only_all_zeros():
+    assert QAny(4).from_bits([0, 0, 0, 0]) == 0
+
+    with pytest.raises(TypeError, match=r"Ambiguous value"):
+        QAny(4).from_bits([1, 0, 0, 0])
+
+
+def test_qany_to_bits_array():
+    enc = QAny(4)
+    assert np.all(enc.to_bits_array(np.array([0, 0])) == 0)
+
+    with pytest.raises(TypeError, match=r"Ambiguous encoding"):
+        enc.to_bits_array(np.array([1]))
+
+
+def test_qany_from_bits_array():
+    enc = QAny(4)
+    assert np.all(enc.from_bits_array(np.zeros((2, 4), dtype=np.uint8)) == 0)
+
+    with pytest.raises(TypeError, match=r"Ambiguous value"):
+        enc.from_bits_array(np.array([[1, 0, 0, 0]], dtype=np.uint8))
 
 
 def test_qintonescomp_to_and_from_bits():

qualtran/bloqs/basic_gates/z_basis.py

@@ -43,9 +43,9 @@
     ConnectionT,
     CtrlSpec,
     DecomposeTypeError,
-    QAny,
     QBit,
     QDType,
+    QUInt,
     Register,
     Side,
     Signature,
@@ -479,7 +479,7 @@ def check(self, attribute, val):
     def dtype(self) -> QDType:
         if self.bitsize == 1:
             return QBit()
-        return QAny(self.bitsize)
+        return QUInt(self.bitsize)
 
     @cached_property
     def signature(self) -> Signature:

qualtran/bloqs/bookkeeping/cast.py

@@ -11,6 +11,7 @@
 #  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 warnings
 from functools import cached_property
 from typing import Dict, List, Tuple, TYPE_CHECKING
 
@@ -26,13 +27,16 @@
     CompositeBloq,
     ConnectionT,
     DecomposeTypeError,
+    QAny,
     QCDType,
     QDType,
+    QUInt,
     Register,
     Side,
     Signature,
 )
 from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
+from qualtran.bloqs.bookkeeping.partition import LegacyPartitionWarning
 from qualtran.symbolics import is_symbolic
 
 if TYPE_CHECKING:
@@ -120,7 +124,23 @@ def my_tensors(
         ]
 
     def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
-        res = self.out_dtype.from_bits(self.inp_dtype.to_bits(reg))
+        if isinstance(self.inp_dtype, QAny) or isinstance(self.out_dtype, QAny):
+            warnings.warn(
+                "Doing classical casting with QAny is ambiguous, transforming it as QUInt for legacy purposes",
+                category=LegacyPartitionWarning,
+            )
+        match (self.inp_dtype, self.out_dtype):
+            case (QAny(), _):
+                res = self.out_dtype.from_bits(QUInt(self.inp_dtype.bitsize).to_bits(reg))
+            case (_, QAny()):
+                res = QUInt(self.out_dtype.bitsize).from_bits(self.inp_dtype.to_bits(reg))
+            case (QAny(), QAny()):
+                res = QUInt(self.out_dtype.bitsize).from_bits(
+                    QUInt(self.inp_dtype.bitsize).to_bits(reg)
+                )
+            case _:
+                res = self.out_dtype.from_bits(self.inp_dtype.to_bits(reg))
+
         return {'reg': res}
 
     def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]:

qualtran/bloqs/bookkeeping/join.py

@@ -11,6 +11,7 @@
 #  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 warnings
 from functools import cached_property
 from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING
 
@@ -24,6 +25,7 @@
     CompositeBloq,
     ConnectionT,
     DecomposeTypeError,
+    QAny,
     QBit,
     QDType,
     QUInt,
@@ -32,6 +34,7 @@
     Signature,
 )
 from qualtran.bloqs.bookkeeping._bookkeeping_bloq import _BookkeepingBloq
+from qualtran.bloqs.bookkeeping.partition import LegacyPartitionWarning
 from qualtran.drawing import directional_text_box, Text, WireSymbol
 
 if TYPE_CHECKING:
@@ -99,6 +102,13 @@ def my_tensors(
         ]
 
     def on_classical_vals(self, reg: 'NDArray[np.uint]') -> Dict[str, int]:
+        if isinstance(self.dtype, QAny):
+            warnings.warn(
+                "Doing classical operations with QAny is ambiguous, returning a QUInt for legacy purposes",
+                category=LegacyPartitionWarning,
+            )
+            return {'reg': QUInt(self.dtype.bitsize).from_bits(reg.tolist())}
+
         return {'reg': self.dtype.from_bits(reg.tolist())}
 
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':

qualtran/bloqs/bookkeeping/partition.ipynb

@@ -39,8 +39,9 @@
     "Partition a generic index into multiple registers.\n",
     "\n",
     "#### Parameters\n",
-    " - `n`: The total bitsize of the un-partitioned register\n",
+    " - `n`: The total bit-size of the un-partitioned register. Required if `dtype_in` is None. Deprecated. Kept for backward compatibility. Use `dtype_in` instead whenever possible.\n",
     " - `regs`: Registers to partition into. The `side` attribute is ignored.\n",
+    " - `dtype_in`: Type of the un-partitioned register. Required if `n` is None. If None, the type is inferred as `QUInt(n)`.\n",
     " - `partition`: `False` means un-partition instead. \n",
     "\n",
     "#### Registers\n",