Add add_my_tensors implementation to Free and enable tensor contraction for PhaseGradientState adjoint Bloq

qualtran/bloqs/rotations/phase_gradient_test.py

@@ -15,6 +15,7 @@
 import numpy as np
 import pytest
 
+from qualtran import BloqBuilder
 from qualtran.bloqs.rotations.phase_gradient import (
     AddIntoPhaseGrad,
     AddScaledValIntoPhaseReg,
@@ -51,6 +52,12 @@ def test_phase_gradient_state_tensor_contract(n: int, t: float):
     bloq = PhaseGradientState(n, t)
     np.testing.assert_allclose(state_coefs, bloq.tensor_contract())
 
+    bb = BloqBuilder()
+    phase_reg = bb.add(bloq)
+    bb.add(PhaseGradientState(n, t, adjoint=True), phase_grad=phase_reg)
+    circuit = bb.finalize()
+    assert np.isclose(circuit.tensor_contract(), 1)
+
 
 @pytest.mark.parametrize('n', [6, 7, 8])
 @pytest.mark.parametrize('exponent', [-0.5, 1, 1 / 10])

qualtran/bloqs/util_bloqs.py

@@ -297,6 +297,10 @@ def add_my_tensors(
 class Free(Bloq):
     """Free (i.e. de-allocate) an `n` bit register.
 
+    The tensor decomposition assumes the `n` bit register is uncomputed and is in the $|0^{n}>$
+    state before getting freed. To verify that is the case, one can compute the resulting state
+    vector after freeing qubits and make sure it is normalized.
+
     Attributes:
         n: the bitsize of the register to be freed.
     """
@@ -318,6 +322,18 @@ def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
     def t_complexity(self) -> 'TComplexity':
         return TComplexity()
 
+    def add_my_tensors(
+        self,
+        tn: 'qtn.TensorNetwork',
+        tag: Any,
+        *,
+        incoming: Dict[str, 'SoquetT'],
+        outgoing: Dict[str, 'SoquetT'],
+    ):
+        data = np.zeros(1 << self.n)
+        data[0] = 1
+        tn.add(qtn.Tensor(data=data, inds=(incoming['reg'],), tags=['Free', tag]))
+
 
 @frozen
 class ArbitraryClifford(Bloq):

qualtran/bloqs/util_bloqs_test.py

@@ -58,6 +58,18 @@ def test_util_bloqs():
     assert isinstance(qs3, Soquet)
     no_return = bb.add(Free(10), reg=qs3)
     assert no_return is None
+    assert bb.finalize().tensor_contract() == 1.0
+
+
+def test_free_nonzero_state_vector_leads_to_unnormalized_state():
+    from qualtran.bloqs.basic_gates.hadamard import Hadamard
+    from qualtran.bloqs.on_each import OnEach
+
+    bb = BloqBuilder()
+    qs1 = bb.add(Allocate(10))
+    qs2 = bb.add(OnEach(10, Hadamard()), q=qs1)
+    no_return = bb.add(Free(10), reg=qs2)
+    assert np.allclose(bb.finalize().tensor_contract(), np.sqrt(1 / 2**10))
 
 
 def test_util_bloqs_tensor_contraction():