Fix invocation of PrepareIdentity in block encodings (#1345)

* Fix invocation of `PrepareIdentity` in block encodings

* Fix `TestBlockEncoding`

Author: charlesyuan314

qualtran/bloqs/block_encoding/chebyshev_polynomial.py

@@ -114,7 +114,7 @@ def epsilon(self) -> SymbolicFloat:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
 
     @cached_property
     def reflection_bloq(self):
@@ -249,7 +249,7 @@ def selection_registers(self) -> Tuple[Register, ...]:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
 
     @cached_property
     def linear_combination(self) -> Union[LinearCombination, ChebyshevPolynomial]:

qualtran/bloqs/block_encoding/chebyshev_polynomial_test.py

@@ -20,7 +20,7 @@
 from attr import field, frozen
 from numpy.typing import NDArray
 
-from qualtran import BloqBuilder, QAny, Signature, Soquet, SoquetT
+from qualtran import BloqBuilder, Signature, Soquet, SoquetT
 from qualtran.bloqs.basic_gates import Hadamard, Identity, IntEffect, IntState, XGate
 from qualtran.bloqs.block_encoding import BlockEncoding, Unitary
 from qualtran.bloqs.block_encoding.chebyshev_polynomial import (
@@ -200,7 +200,7 @@ def signature(self) -> Signature:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(1),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([1]))
 
     def build_composite_bloq(
         self, bb: BloqBuilder, system: Soquet, ancilla: Soquet, resource: Soquet
@@ -221,6 +221,7 @@ def test_chebyshev_matrix():
 
 def test_chebyshev_poly_signal_state():
     assert isinstance(_chebyshev_poly_even().signal_state.prepare, PrepareIdentity)
+    _ = _chebyshev_poly_even().signal_state.decompose_bloq()
 
 
 @pytest.mark.slow

qualtran/bloqs/block_encoding/linear_combination.py

@@ -178,7 +178,7 @@ def epsilon(self) -> SymbolicFloat:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
 
     @cached_property
     def prepare(self) -> BlackBoxPrepare:

qualtran/bloqs/block_encoding/linear_combination_test.py

@@ -215,6 +215,7 @@ def test_linear_combination_approx_random():
 
 def test_linear_combination_signal_state():
     assert isinstance(_linear_combination_block_encoding().signal_state.prepare, PrepareIdentity)
+    _ = _linear_combination_block_encoding().signal_state.decompose_bloq()
 
 
 @pytest.mark.notebook

qualtran/bloqs/block_encoding/product.py

@@ -130,7 +130,7 @@ def epsilon(self) -> SymbolicFloat:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
 
     def build_composite_bloq(
         self, bb: BloqBuilder, system: SoquetT, **soqs: SoquetT

qualtran/bloqs/block_encoding/product_test.py

@@ -212,6 +212,7 @@ def test_product_random():
 
 def test_product_signal_state():
     assert isinstance(_product_block_encoding().signal_state.prepare, PrepareIdentity)
+    _ = _product_block_encoding().signal_state.decompose_bloq()
 
 
 @pytest.mark.notebook

qualtran/bloqs/block_encoding/sparse_matrix.py

@@ -211,7 +211,7 @@ def epsilon(self) -> SymbolicFloat:
 
     @property
     def signal_state(self) -> BlackBoxPrepare:
-        return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+        return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
 
     @cached_property
     def diffusion(self):

qualtran/bloqs/block_encoding/sparse_matrix_test.py

@@ -315,6 +315,7 @@ def test_symmetric_banded_counts():
 
 def test_sparse_matrix_signal_state():
     assert isinstance(_sparse_matrix_block_encoding().signal_state.prepare, PrepareIdentity)
+    _ = _sparse_matrix_block_encoding().signal_state.decompose_bloq()
 
 
 @pytest.mark.notebook

qualtran/bloqs/block_encoding/tensor_product.py

@@ -104,7 +104,7 @@ def epsilon(self) -> SymbolicFloat:
     @property
     def signal_state(self) -> BlackBoxPrepare:
         if all(isinstance(u.signal_state.prepare, PrepareIdentity) for u in self.block_encodings):
-            return BlackBoxPrepare(PrepareIdentity((QAny(self.ancilla_bitsize),)))
+            return BlackBoxPrepare(PrepareIdentity.from_bitsizes([self.ancilla_bitsize]))
         else:
             # TODO: implement by taking tensor product of component signal states
             raise NotImplementedError

qualtran/bloqs/block_encoding/tensor_product_test.py

@@ -131,6 +131,7 @@ def test_tensor_product_cirq():
 
 def test_tensor_product_signal_state():
     assert isinstance(_tensor_product_block_encoding().signal_state.prepare, PrepareIdentity)
+    _ = _tensor_product_block_encoding().signal_state.decompose_bloq()
 
 
 @pytest.mark.notebook