Symbolic Shaped (#867)

* support `Shaped` for symbolic analysis

fix expression for Jacobi-Anger degree

add test using 1d ising walk op, fix call graph

fix test

fix hamsim nb

`slen`, docstrings, generalize `is_symbolic` (to remove dep. on cirq)

test shape with symbolic entries

docstring refs: remove extra indent

skip serialize test for symbolic hamsim

* remove `_is_parameterized_`

* unittest gqsp symbolic call graph

* lint

* fix type hint

Author: anurudhp

qualtran/bloqs/basic_gates/su2_rotation.py

@@ -117,13 +117,13 @@ def add_my_tensors(
         )
 
     def _unitary_(self):
-        if self._is_parameterized_():
+        if self.is_symbolic():
             return None
         return self.rotation_matrix
 
     def build_composite_bloq(self, bb: 'BloqBuilder', q: 'SoquetT') -> Dict[str, 'SoquetT']:
-        pi = sympy.pi if self._is_parameterized_() else np.pi
-        exp = sympy.exp if self._is_parameterized_() else np.exp
+        pi = sympy.pi if self.is_symbolic() else np.pi
+        exp = sympy.exp if self.is_symbolic() else np.exp
 
         bb.add(GlobalPhase(coefficient=-exp(1j * self.global_shift), eps=self.eps / 4))
         q = bb.add(ZPowGate(exponent=1 - self.lambd / pi, global_shift=-1, eps=self.eps / 4), q=q)
@@ -151,7 +151,7 @@ def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
     def _t_complexity_(self) -> TComplexity:
         return TComplexity(rotations=3)
 
-    def _is_parameterized_(self) -> bool:
+    def is_symbolic(self) -> bool:
         return is_symbolic(self.theta, self.phi, self.lambd, self.global_shift)
 
     @classmethod

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.ipynb

@@ -80,7 +80,7 @@
     "This therefore block-encodes $e^{-iHt}$ in the block where the signal qubit and selection registers are all $|0\\rangle$.\n",
     "\n",
     "#### References\n",
-    " - [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501).     Motlagh and Wiebe. (2023). Theorem 7, Corollary 8. \n",
+    " - [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501). Motlagh and Wiebe. (2023). Theorem 7, Corollary 8. \n",
     "\n",
     "#### Parameters\n",
     " - `walk_operator`: qubitization walk operator of $H$ constructed from SELECT and PREPARE oracles.\n",
@@ -128,6 +128,26 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11baa846",
+   "metadata": {
+    "cq.autogen": "HamiltonianSimulationByGQSP.symbolic_hamsim_by_gqsp"
+   },
+   "outputs": [],
+   "source": [
+    "import sympy\n",
+    "\n",
+    "from qualtran.bloqs.hubbard_model import get_walk_operator_for_hubbard_model\n",
+    "\n",
+    "walk_op = get_walk_operator_for_hubbard_model(2, 2, 1, 1)\n",
+    "t, alpha, inv_eps = sympy.symbols(\"t alpha N\")\n",
+    "symbolic_hamsim_by_gqsp = HamiltonianSimulationByGQSP(\n",
+    "    walk_op, t=t, alpha=alpha, precision=1 / inv_eps\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "a6e29295",
@@ -148,8 +168,8 @@
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([hubbard_time_evolution_by_gqsp],\n",
-    "           ['`hubbard_time_evolution_by_gqsp`'])"
+    "show_bloqs([hubbard_time_evolution_by_gqsp, symbolic_hamsim_by_gqsp],\n",
+    "           ['`hubbard_time_evolution_by_gqsp`', '`symbolic_hamsim_by_gqsp`'])"
    ]
   },
   {
@@ -171,21 +191,42 @@
    },
    "outputs": [],
    "source": [
-    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
-    "hubbard_time_evolution_by_gqsp_g, hubbard_time_evolution_by_gqsp_sigma = hubbard_time_evolution_by_gqsp.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "from qualtran.resource_counting.generalizers import generalize_rotation_angle, ignore_split_join, ignore_alloc_free\n",
+    "hubbard_time_evolution_by_gqsp_g, hubbard_time_evolution_by_gqsp_sigma = hubbard_time_evolution_by_gqsp.call_graph(generalizer=[ignore_split_join, ignore_alloc_free, generalize_rotation_angle])\n",
     "show_call_graph(hubbard_time_evolution_by_gqsp_g)\n",
     "show_counts_sigma(hubbard_time_evolution_by_gqsp_sigma)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7045d204b9fcd6f1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import generalize_rotation_angle, ignore_split_join, ignore_alloc_free\n",
+    "_, symbolic_hamsim_by_gqsp_sigma = symbolic_hamsim_by_gqsp.call_graph(max_depth=2, generalizer=[ignore_split_join, ignore_alloc_free, generalize_rotation_angle])\n",
+    "show_counts_sigma(symbolic_hamsim_by_gqsp_sigma)"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py

@@ -12,22 +12,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Dict, Set, Tuple, TYPE_CHECKING
+from typing import Dict, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import field, frozen
 from numpy.typing import NDArray
 
-from qualtran import (
-    bloq_example,
-    BloqDocSpec,
-    Controlled,
-    CtrlSpec,
-    GateWithRegisters,
-    Signature,
-    Soquet,
-)
-from qualtran.bloqs.basic_gates import SU2RotationGate
+from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, Signature, Soquet
 from qualtran.bloqs.qsp.generalized_qsp import GeneralizedQSP, scale_down_to_qsp_polynomial
 from qualtran.bloqs.qubitization_walk_operator import QubitizationWalkOperator
 from qualtran.linalg.jacobi_anger_approximations import (
@@ -36,13 +27,13 @@
 )
 from qualtran.resource_counting.symbolic_counting_utils import (
     is_symbolic,
+    Shaped,
     SymbolicFloat,
     SymbolicInt,
 )
 
 if TYPE_CHECKING:
-    from qualtran import BloqBuilder, Soquet, SoquetT
-    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran import BloqBuilder, SoquetT
 
 
 @frozen
@@ -92,7 +83,7 @@ class HamiltonianSimulationByGQSP(GateWithRegisters):
 
     References:
         [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501)
-            Motlagh and Wiebe. (2023). Theorem 7, Corollary 8.
+        Motlagh and Wiebe. (2023). Theorem 7, Corollary 8.
 
     Args:
         walk_operator: qubitization walk operator of $H$ constructed from SELECT and PREPARE oracles.
@@ -106,7 +97,7 @@ class HamiltonianSimulationByGQSP(GateWithRegisters):
     alpha: SymbolicFloat = field(kw_only=True)
     precision: SymbolicFloat = field(kw_only=True)
 
-    def _parameterized_(self):
+    def is_symbolic(self):
         return is_symbolic(self.t, self.alpha, self.precision)
 
     @cached_property
@@ -115,10 +106,11 @@ def degree(self) -> SymbolicInt:
         return degree_jacobi_anger_approximation(self.t * self.alpha, precision=self.precision)
 
     @cached_property
-    def approx_cos(self) -> NDArray[np.complex_]:
+    def approx_cos(self) -> Union[NDArray[np.complex_], Shaped]:
         r"""polynomial approximation for $$e^{i\theta} \mapsto e^{it\cos(\theta)}$$"""
-        if self._parameterized_():
-            raise ValueError(f"cannot compute `cos` approximation for parameterized Bloq {self}")
+        if self.is_symbolic():
+            return Shaped((2 * self.degree + 1,))
+
         poly = approx_exp_cos_by_jacobi_anger(-self.t * self.alpha, degree=self.degree)
         # TODO(#860) current scaling method does not compute true maximum, so we scale down a bit more by (1 - 2\eps)
         poly = scale_down_to_qsp_polynomial(poly) * (1 - 2 * self.precision)
@@ -129,7 +121,7 @@ def gqsp(self) -> GeneralizedQSP:
         return GeneralizedQSP.from_qsp_polynomial(
             self.walk_operator,
             self.approx_cos,
-            negative_power=int(self.degree),
+            negative_power=self.degree,
             verify=True,
             verify_precision=1e-4,
         )
@@ -180,18 +172,6 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
 
         return soqs
 
-    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
-        if self._parameterized_():
-            d = self.degree
-            return {
-                (Controlled(self.walk_operator.adjoint(), CtrlSpec()), d),
-                (Controlled(self.walk_operator, CtrlSpec(cvs=0)), d),
-                (self.walk_operator.prepare, 1),
-                (self.walk_operator.prepare.adjoint(), 1),
-                (SU2RotationGate.arbitrary(ssa), 2 * d + 1),
-            }
-        return self.decompose_bloq().build_call_graph(ssa)
-
 
 @bloq_example
 def _hubbard_time_evolution_by_gqsp() -> HamiltonianSimulationByGQSP:
@@ -204,8 +184,22 @@ def _hubbard_time_evolution_by_gqsp() -> HamiltonianSimulationByGQSP:
     return hubbard_time_evolution_by_gqsp
 
 
+@bloq_example
+def _symbolic_hamsim_by_gqsp() -> HamiltonianSimulationByGQSP:
+    import sympy
+
+    from qualtran.bloqs.hubbard_model import get_walk_operator_for_hubbard_model
+
+    walk_op = get_walk_operator_for_hubbard_model(2, 2, 1, 1)
+    t, alpha, inv_eps = sympy.symbols("t alpha N")
+    symbolic_hamsim_by_gqsp = HamiltonianSimulationByGQSP(
+        walk_op, t=t, alpha=alpha, precision=1 / inv_eps
+    )
+    return symbolic_hamsim_by_gqsp
+
+
 _Hamiltonian_Simulation_by_GQSP_DOC = BloqDocSpec(
     bloq_cls=HamiltonianSimulationByGQSP,
     import_line='from qualtran.bloqs.hamiltonian_simulation.hamiltonian_simulation_by_gqsp import HamiltonianSimulationByGQSP',
-    examples=[_hubbard_time_evolution_by_gqsp],
+    examples=[_hubbard_time_evolution_by_gqsp, _symbolic_hamsim_by_gqsp],
 )

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py

@@ -29,6 +29,7 @@
 
 from .hamiltonian_simulation_by_gqsp import (
     _hubbard_time_evolution_by_gqsp,
+    _symbolic_hamsim_by_gqsp,
     HamiltonianSimulationByGQSP,
 )
 
@@ -37,6 +38,12 @@ def test_examples(bloq_autotester):
     bloq_autotester(_hubbard_time_evolution_by_gqsp)
 
 
+def test_symbolic_examples(bloq_autotester):
+    if bloq_autotester.check_name == 'serialization':
+        pytest.xfail(f"serialization for {_symbolic_hamsim_by_gqsp.name} not yet supported")
+    bloq_autotester(_symbolic_hamsim_by_gqsp)
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize("bitsize", [1, 2])
 @pytest.mark.parametrize("t", [2, 3, 5, 10])

qualtran/bloqs/qsp/generalized_qsp.ipynb

@@ -87,7 +87,7 @@
     " - `negative_power`: value of $k$, which effectively applies $z^{-k} P(z)$. defaults to 0. \n",
     "\n",
     "#### References\n",
-    " - [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501).     Motlagh and Wiebe. (2023). Theorem 3; Figure 2; Theorem 6.\n"
+    " - [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501). Motlagh and Wiebe. (2023). Theorem 3; Figure 2; Theorem 6.\n"
    ]
   },
   {

qualtran/bloqs/qsp/generalized_qsp.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 from collections import Counter
 from functools import cached_property
-from typing import Iterable, Sequence, Set, Tuple, TYPE_CHECKING
+from typing import Iterable, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import field, frozen
@@ -22,6 +22,14 @@
 
 from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, QBit, Register, Signature
 from qualtran.bloqs.basic_gates.su2_rotation import SU2RotationGate
+from qualtran.resource_counting.symbolic_counting_utils import (
+    is_symbolic,
+    Shaped,
+    slen,
+    smax,
+    smin,
+    SymbolicInt,
+)
 
 if TYPE_CHECKING:
     import cirq
@@ -253,8 +261,10 @@ def assert_is_qsp_polynomial(P: Sequence[complex], *, n_points: int = 2**17):
     ), f"Not a QSP polynomial! maximum absolute value {max_value} is greater than 1."
 
 
-def _to_tuple(x: Iterable[complex]) -> Sequence[complex]:
+def _to_tuple(x: Iterable[complex]) -> Union[Tuple[complex, ...], Shaped]:
     """mypy-compatible attrs converter for GeneralizedQSP.P and Q"""
+    if isinstance(x, Shaped):
+        return x
     return tuple(x)
 
 
@@ -312,18 +322,19 @@ class GeneralizedQSP(GateWithRegisters):
 
     References:
         [Generalized Quantum Signal Processing](https://arxiv.org/abs/2308.01501)
-            Motlagh and Wiebe. (2023). Theorem 3; Figure 2; Theorem 6.
+        Motlagh and Wiebe. (2023). Theorem 3; Figure 2; Theorem 6.
     """
 
     U: GateWithRegisters
-    P: Tuple[complex, ...] = field(converter=_to_tuple)
-    Q: Tuple[complex, ...] = field(converter=_to_tuple)
-    negative_power: int = field(default=0, kw_only=True)
+    P: Union[Tuple[complex, ...], Shaped] = field(converter=_to_tuple)
+    Q: Union[Tuple[complex, ...], Shaped] = field(converter=_to_tuple)
+    negative_power: SymbolicInt = field(default=0, kw_only=True)
 
     @P.validator
     @Q.validator
     def _check_polynomial(self, attribute, value):
-        if len(value) <= 1:
+        degree = slen(value) - 1
+        if not is_symbolic(degree) and degree <= 0:
             raise ValueError("GQSP Polynomial must have degree at least 1")
 
     @cached_property
@@ -334,12 +345,15 @@ def signature(self) -> Signature:
     def from_qsp_polynomial(
         cls,
         U: GateWithRegisters,
-        P: Sequence[complex],
+        P: Union[Sequence[complex], Shaped],
         *,
         negative_power: int = 0,
         verify: bool = False,
         verify_precision=1e-7,
     ) -> 'GeneralizedQSP':
+        if is_symbolic(P):
+            return GeneralizedQSP(U, P, P, negative_power=negative_power)
+
         if verify:
             assert_is_qsp_polynomial(P)
         Q = qsp_complementary_polynomial(P, verify=verify, verify_precision=verify_precision)
@@ -391,7 +405,20 @@ def decompose_from_registers(
         for _ in range(num_inverse_applications):
             yield self.U.adjoint().on_registers(**quregs)
 
+    def is_symbolic(self) -> bool:
+        return is_symbolic(self.P, self.Q, self.negative_power)
+
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        if self.is_symbolic():
+            degree = slen(self.P) - 1
+
+            return {
+                (self.U.controlled(control_values=[0]), smax(0, degree - self.negative_power)),
+                (self.U.adjoint(), smax(0, self.negative_power - degree)),
+                (self.U.adjoint().controlled(), smin(degree, self.negative_power)),
+                (SU2RotationGate.arbitrary(ssa), degree + 1),
+            }
+
         degree = len(self.P) - 1
 
         counts = set(Counter(self.signal_rotations).items())