revert changes to Partition

quantumlib · Jul 16, 2024 · b5567e1 · b5567e1
1 parent 0d7c985
commit b5567e1
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/qualtran/bloqs/bookkeeping/partition.py b/qualtran/bloqs/bookkeeping/partition.py
@@ -16,7 +16,6 @@
 
 import numpy as np
 from attrs import evolve, field, frozen, validators
-from numpy.typing import NDArray
 
 from qualtran import (
     bloq_example,
@@ -138,24 +137,22 @@ def _classical_partition(self, x: 'ClassicalValT') -> Dict[str, 'ClassicalValT']
             start += size
         return out_vals
 
-    def _classical_unpartition_to_bits(self, **vals: 'ClassicalValT') -> NDArray[np.uint]:
-        out_vals: list[list[int]] = []
+    def _classical_unpartition(self, **vals: 'ClassicalValT'):
+        out_vals = []
         for reg in self.regs:
             reg_val = vals[reg.name]
             if isinstance(reg_val, np.ndarray):
-                out_vals.extend(reg.dtype.to_bits(val) for val in reg_val.ravel())
+                out_vals.append(ints_to_bits(reg_val.ravel(), reg.bitsize).ravel())
             else:
-                out_vals.append(reg.dtype.to_bits(reg_val))
-        big_int_bits = np.concatenate(out_vals)
-        return big_int_bits
+                out_vals.append(ints_to_bits(reg_val, reg.bitsize)[0])
+        big_int = np.concatenate(out_vals)
+        return {'x': bits_to_ints(big_int)[0]}
 
     def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
         if self.partition:
             return self._classical_partition(vals['x'])
         else:
-            big_int_bits = self._classical_unpartition_to_bits(**vals)
-            (big_int,) = bits_to_ints(big_int_bits)
-            return {'x': big_int}
+            return self._classical_unpartition(**vals)
 
     def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None: