stuff

Author: mpharrigan

.github/workflows/nightly.yaml

@@ -43,4 +43,4 @@ jobs:
           pip install -r dev_tools/requirements/envs/pytest.env.txt
           pip install --no-deps -e .
       - run: |
-          check/pytest
+          check/pytest --durations=10

qualtran/_infra/composite_bloq.py

@@ -397,11 +397,12 @@ def flatten_once(
         # pylint: disable=protected-access
         bb._i = max(binst.i for binst in self.bloq_instances) + 1
 
-        soq_map: List[Tuple[SoquetT, SoquetT]] = []
+        # soq_map: List[Tuple[SoquetT, SoquetT]] = []
+        flat_soq_map: Dict[Soquet, Soquet] = {}
         new_out_soqs: Tuple[SoquetT, ...]
         did_work = False
         for binst, in_soqs, old_out_soqs in self.iter_bloqsoqs():
-            in_soqs = _map_soqs(in_soqs, soq_map)  # update `in_soqs` from old to new.
+            in_soqs = _map_flat_soqs(in_soqs, flat_soq_map)  # update `in_soqs` from old to new.
             if pred(binst):
                 try:
                     new_out_soqs = bb.add_from(binst.bloq, **in_soqs)
@@ -416,12 +417,13 @@ def flatten_once(
                 # pylint: disable=protected-access
                 new_out_soqs = tuple(soq for _, soq in bb._add_binst(binst, in_soqs=in_soqs))
 
-            soq_map.extend(zip(old_out_soqs, new_out_soqs))
+            _update_flat_soq_map(zip(old_out_soqs, new_out_soqs), flat_soq_map)
+            # soq_map.extend(zip(old_out_soqs, new_out_soqs))
 
         if not did_work:
             raise DidNotFlattenAnythingError()
 
-        fsoqs = _map_soqs(self.final_soqs(), soq_map)
+        fsoqs = _map_flat_soqs(self.final_soqs(), flat_soq_map)
         return bb.finalize(**fsoqs)
 
     def flatten(
@@ -865,6 +867,40 @@ def _map_soqs(soqs: SoquetT) -> SoquetT:
     return {name: _map_soqs(soqs) for name, soqs in soqs.items()}
 
 
+def _map_flat_soqs(
+    soqs: Dict[str, SoquetT], flat_soq_map: Dict[Soquet, Soquet]
+) -> Dict[str, SoquetT]:
+
+    # use vectorize to use the flat mapping.
+    def _map_soq(soq: Soquet) -> Soquet:
+        # Helper function to map an individual soquet.
+        return flat_soq_map.get(soq, soq)
+
+    # Use `vectorize` to call `_map_soq` on each element of the array.
+    vmap = np.vectorize(_map_soq, otypes=[object])
+
+    def _map_soqs(soqs: SoquetT) -> SoquetT:
+        if isinstance(soqs, Soquet):
+            return _map_soq(soqs)
+        return vmap(soqs)
+
+    return {name: _map_soqs(soqs) for name, soqs in soqs.items()}
+
+
+def _update_flat_soq_map(soq_map: Iterable[Tuple[SoquetT, SoquetT]], flat_soq_map):
+    for old_soqs, new_soqs in soq_map:
+        if isinstance(old_soqs, Soquet):
+            assert isinstance(new_soqs, Soquet), new_soqs
+            flat_soq_map[old_soqs] = new_soqs
+            continue
+
+        assert isinstance(old_soqs, np.ndarray), old_soqs
+        assert isinstance(new_soqs, np.ndarray), new_soqs
+        assert old_soqs.shape == new_soqs.shape, (old_soqs.shape, new_soqs.shape)
+        for o, n in zip(old_soqs.reshape(-1), new_soqs.reshape(-1)):
+            flat_soq_map[o] = n
+
+
 class BloqBuilder:
     """A builder class for composing bloqs.
 

qualtran/_infra/gate_with_registers.py

@@ -508,13 +508,6 @@ def controlled(
         return controlled_bloq
 
     def _unitary_(self):
-        if all(reg.side == Side.THRU for reg in self.signature):
-            try:
-                tensor = self.tensor_contract()
-                assert tensor.ndim == 2, "All registers should have been checked to be THRU."
-                return tensor
-            except (DecomposeNotImplementedError, DecomposeTypeError, NotImplementedError):
-                return NotImplemented
         return NotImplemented
 
     def my_tensors(

qualtran/_infra/quantum_graph.py

@@ -43,6 +43,9 @@ def bloq_is(self, t) -> bool:
         `Union[BloqInstance, DanglingT]`"""
         return isinstance(self.bloq, t)
 
+    def __hash__(self):
+        return hash(self.i)
+
 
 class DanglingT:
     """The type of the singleton objects `LeftDangle` and `RightDangle`.

qualtran/bloqs/block_encoding/block_encoding_base.py

@@ -95,7 +95,6 @@ def signal_state(self) -> BlackBoxPrepare:
         r"""Returns the signal / ancilla flag state $|G\rangle."""
 
 
-
 _BLOCK_ENCODING_DOC = BloqDocSpec(
     bloq_cls=BlockEncoding, examples=[]  # type: ignore[type-abstract]
 )

qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard_test.py

@@ -73,20 +73,13 @@ def test_hubbard_majoranna_small(bloq_autotester):
     bloq_autotester(_hubb_majoranna_small)
 
 
-def test_hubb_majoranna_bloq_counts():
-    for n in [2, 3, 10, 11]:
-        bloq = HubbardMajorannaOperator(x_dim=n, y_dim=n, control_val=None)
-        qlt_testing.assert_equivalent_bloq_counts(
-            bloq,
-            generalizer=[ignore_cliffords, ignore_alloc_free, ignore_split_join, generalize_cvs],
-        )
-
-    for n in [2, 3, 10, 11]:
-        bloq = HubbardMajorannaOperator(x_dim=n, y_dim=n, control_val=1)
-        qlt_testing.assert_equivalent_bloq_counts(
-            bloq,
-            generalizer=[ignore_cliffords, ignore_alloc_free, ignore_split_join, generalize_cvs],
-        )
+@pytest.mark.parametrize('n', [2, 3, 10, 11])
+@pytest.mark.parametrize('cv', [None, 1])
+def test_hubb_majoranna_bloq_counts(cv, n):
+    bloq = HubbardMajorannaOperator(x_dim=n, y_dim=n, control_val=cv)
+    qlt_testing.assert_equivalent_bloq_counts(
+        bloq, generalizer=[ignore_cliffords, ignore_alloc_free, ignore_split_join, generalize_cvs]
+    )
 
 
 def test_hubbard_spin_up_z_symb(bloq_autotester):

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py

@@ -79,8 +79,9 @@ def verify_hamiltonian_simulation_by_gqsp(
     N = H.shape[0]
 
     W_e_iHt = HamiltonianSimulationByGQSP(W, t=t, precision=precision)
-    # TODO This cirq.unitary call is 4-5x faster than tensor_contract.
+    # TODO cirq.unitary was 4-5x faster than tensor_contract.
     #      https://github.com/quantumlib/Qualtran/issues/1336
+    #      TODO: note performance regression
     result_unitary = cirq.unitary(BloqAsCirqGate(W_e_iHt))
 
     expected_top_left = scipy.linalg.expm(-1j * H * t)
@@ -95,6 +96,7 @@ def verify_hamiltonian_simulation_by_gqsp(
 def test_hamiltonian_simulation_by_gqsp(
     select_bitsize: int, target_bitsize: int, t: float, precision: float
 ):
+    # TODO: this is slow again
     random_state = np.random.RandomState(42)
 
     for _ in range(5):

qualtran/bloqs/phase_estimation/qubitization_qpe_test.py

@@ -72,8 +72,9 @@ def test_qubitization_qpe_sparse_chem_bloq_autotester(bloq_autotester):
 )
 @pytest.mark.parametrize('use_resource_state', [True, False])
 def test_qubitization_phase_estimation_of_walk(num_terms: int, use_resource_state: bool):
-    if num_terms >= 2:
-        return pytest.xfail("Cirq regression") # TODO: Github Issue
+    # TODO: This is slow.
+    # if num_terms >= 2:
+    #     return pytest.xfail("Cirq regression") # TODO: Github Issue
 
     precision, eps = 5, 0.05
     ham, walk = get_uniform_pauli_qubitized_walk(num_terms)

qualtran/bloqs/phase_estimation/text_book_qpe_test.py

@@ -53,6 +53,7 @@ def test_textbook_phase_estimation_zpow_theta(theta):
 )
 @pytest.mark.parametrize('use_resource_state', [True, False])
 def test_textbook_phase_estimation_qubitized_walk(num_terms: int, use_resource_state: bool):
+    # TODO: broken
     precision, eps = 5, 0.05
     ham, walk = get_uniform_pauli_qubitized_walk(num_terms)
 
@@ -74,27 +75,23 @@ def test_textbook_phase_estimation_qubitized_walk(num_terms: int, use_resource_s
     # TODO cirq simulation seems to fail for controlled `QubitizationWalkOperator`.
     #      the following code decomposes a few levels till it gets only simulable bloqs.
     #      https://github.com/quantumlib/Qualtran/issues/1495
-    # def should_decompose(binst):
-    #     from qualtran import Adjoint, Controlled
-    #     from qualtran.bloqs.basic_gates import Power
-    #     from qualtran.bloqs.qubitization import QubitizationWalkOperator
-    #     from qualtran.bloqs.block_encoding import SelectBlockEncoding
-    #
-    #     bloqs_to_decompose = (TextbookQPE, QubitizationWalkOperator, Power, SelectBlockEncoding)
-    #
-    #     if binst.bloq_is(bloqs_to_decompose):
-    #         return True
-    #
-    #     if binst.bloq_is(Controlled) or binst.bloq_is(Adjoint):
-    #         return isinstance(binst.bloq.subbloq, bloqs_to_decompose)
-    #
-    #     return False
-
-    cbloq = qpe_bloq.as_composite_bloq().flatten()
-
-
-
-    quregs = get_named_qubits(qpe_bloq.signature.lefts())
+    def should_decompose(binst):
+        from qualtran import Adjoint, Controlled
+        from qualtran.bloqs.basic_gates import Power
+        from qualtran.bloqs.qubitization import QubitizationWalkOperator
+
+        bloqs_to_decompose = (TextbookQPE, QubitizationWalkOperator, Power)
+
+        if binst.bloq_is(bloqs_to_decompose):
+            return True
+
+        if binst.bloq_is(Controlled) or binst.bloq_is(Adjoint):
+            return isinstance(binst.bloq.subbloq, bloqs_to_decompose)
+
+        return False
+
+    cbloq = qpe_bloq.as_composite_bloq().flatten(pred=should_decompose)
+    quregs = get_named_qubits(cbloq.signature.lefts())
     qpe_circuit, quregs = cbloq.to_cirq_circuit_and_quregs(None, **quregs)
     for eig_idx, eig_val in enumerate(eigen_values):
         # Apply QPE to determine eigenvalue for walk operator W on initial state |L>|k>
@@ -107,7 +104,7 @@ def test_textbook_phase_estimation_qubitized_walk(num_terms: int, use_resource_s
 
         # 3. QPE circuit with state prep
         qpe_with_init = prep_L_K + qpe_circuit
-        # assert len(qpe_with_init.all_qubits()) < 23
+        assert len(qpe_with_init.all_qubits()) < 23
 
         # 4. Estimate theta
         theta = simulate_theta_estimate(qpe_with_init, quregs['qpe_reg'])

qualtran/bloqs/reflections/reflection_using_prepare_test.py

@@ -109,12 +109,13 @@ def get_3q_uniform_dirac_notation(signs, global_phase: complex = 1):
         ret = ret + f' {c} {term}'
     return ret
 
+
 @pytest.mark.slow
 @pytest.mark.parametrize('num_ones', [5])
 @pytest.mark.parametrize('eps', [0.05])
 @pytest.mark.parametrize('global_phase', [+1, -1j])
 def test_reflection_using_prepare(num_ones, eps, global_phase):
-    return pytest.xfail("Cirq regression") # TODO error message
+    return pytest.xfail("Cirq regression")  # TODO error message
 
     data = [1] * num_ones
     prepare_gate = StatePreparationAliasSampling.from_probabilities(data, precision=eps)