fix bug in CtrlSpecAnd decomposition, add more bloq examples

quantumlib · Nov 11, 2024 · 304da47 · 304da47
1 parent ee153f4
commit 304da47
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 22 deletions.
diff --git a/qualtran/bloqs/mcmt/ctrl_spec_and.ipynb b/qualtran/bloqs/mcmt/ctrl_spec_and.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "0beec492",
+   "id": "209f4989",
    "metadata": {
     "cq.autogen": "title_cell"
    },
@@ -13,7 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4d31094e",
+   "id": "373927c3",
    "metadata": {
     "cq.autogen": "top_imports"
    },
@@ -30,7 +30,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "12aba876",
+   "id": "d7dd05fa",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.bloq_doc.md"
    },
@@ -66,7 +66,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "842307af",
+   "id": "837cc8f9",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.bloq_doc.py"
    },
@@ -77,7 +77,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "76f2965f",
+   "id": "760fb1b0",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.example_instances.md"
    },
@@ -88,7 +88,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "68a43214",
+   "id": "8486e383",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.ctrl_on_int"
    },
@@ -102,7 +102,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "945fa4a4",
+   "id": "1510fd98",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.ctrl_on_bits"
    },
@@ -116,7 +116,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5e6374b2",
+   "id": "489cf199",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.ctrl_on_nd_bits"
    },
@@ -132,7 +132,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2047fccf",
+   "id": "bc189810",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.ctrl_on_multiple_values"
    },
@@ -145,9 +145,39 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fa2c166",
+   "metadata": {
+    "cq.autogen": "CtrlSpecAnd.ctrl_on_symbolic_cv"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import CtrlSpec\n",
+    "from qualtran.symbolics import Shaped\n",
+    "\n",
+    "ctrl_on_symbolic_cv = CtrlSpecAnd(CtrlSpec(cvs=Shaped((2,))))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0d78907",
+   "metadata": {
+    "cq.autogen": "CtrlSpecAnd.ctrl_on_symbolic_cv_multi"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import CtrlSpec\n",
+    "from qualtran.symbolics import Shaped\n",
+    "\n",
+    "ctrl_on_symbolic_cv_multi = CtrlSpecAnd(CtrlSpec(cvs=Shaped((3,))))"
+   ]
+  },
   {
    "cell_type": "markdown",
-   "id": "55581cc0",
+   "id": "60a9fcc0",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.graphical_signature.md"
    },
@@ -158,20 +188,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5ea9e6e5",
+   "id": "f1b9274c",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.graphical_signature.py"
    },
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([ctrl_on_int, ctrl_on_bits, ctrl_on_nd_bits, ctrl_on_multiple_values],\n",
-    "           ['`ctrl_on_int`', '`ctrl_on_bits`', '`ctrl_on_nd_bits`', '`ctrl_on_multiple_values`'])"
+    "show_bloqs([ctrl_on_int, ctrl_on_bits, ctrl_on_nd_bits, ctrl_on_multiple_values, ctrl_on_symbolic_cv, ctrl_on_symbolic_cv_multi],\n",
+    "           ['`ctrl_on_int`', '`ctrl_on_bits`', '`ctrl_on_nd_bits`', '`ctrl_on_multiple_values`', '`ctrl_on_symbolic_cv`', '`ctrl_on_symbolic_cv_multi`'])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "3f5bb7d6",
+   "id": "704a2863",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.call_graph.md"
    },
@@ -182,7 +212,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f83e1715",
+   "id": "1c2ab375",
    "metadata": {
     "cq.autogen": "CtrlSpecAnd.call_graph.py"
    },

diff --git a/qualtran/bloqs/mcmt/ctrl_spec_and.py b/qualtran/bloqs/mcmt/ctrl_spec_and.py
@@ -15,6 +15,7 @@
 from typing import Optional, TYPE_CHECKING, Union
 
 import numpy as np
+import sympy
 from attrs import frozen
 
 from qualtran import (
@@ -91,7 +92,7 @@ def signature(self) -> Signature:
         return Signature(
             [
                 *self.control_registers,
-                *self.junk_registers(),
+                *self.junk_registers,
                 Register('target', QBit(), side=Side.RIGHT),
             ]
         )
@@ -103,6 +104,7 @@ def control_registers(self) -> tuple[Register, ...]:
             for i, (dtype, shape) in enumerate(self.ctrl_spec.activation_function_dtypes())
         )
 
+    @cached_property
     def junk_registers(self) -> tuple[Register, ...]:
         if not is_symbolic(self.n_ctrl_qubits) and self.n_ctrl_qubits == 2:
             return ()
@@ -140,8 +142,10 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str
 
         # Compute the single control qubit `target`
         if self.n_ctrl_qubits == 2:
-            assert isinstance(self._flat_cvs, tuple)
-            cv1, cv2 = self._flat_cvs
+            if isinstance(self._flat_cvs, tuple):
+                cv1, cv2 = self._flat_cvs
+            else:
+                cv1, cv2 = sympy.symbols("cv1, cv2")
             ctrl_qubits, target = bb.add(And(cv1, cv2), ctrl=ctrl_qubits)
             junk = None
         else:
@@ -177,8 +181,10 @@ def wire_symbol(self, reg: Optional[Register], idx: tuple[int, ...] = tuple()) -
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         if not is_symbolic(self.n_ctrl_qubits) and self.n_ctrl_qubits == 2:
-            assert isinstance(self._flat_cvs, tuple)
-            cv1, cv2 = self._flat_cvs
+            if isinstance(self._flat_cvs, tuple):
+                cv1, cv2 = self._flat_cvs
+            else:
+                cv1, cv2 = sympy.symbols("cv1, cv2")
             return {And(cv1, cv2): 1}
 
         return {MultiAnd(self._flat_cvs): 1}
@@ -220,7 +226,42 @@ def _ctrl_on_multiple_values() -> CtrlSpecAnd:
     return ctrl_on_multiple_values
 
 
+@bloq_example(generalizer=ignore_split_join)
+def _ctrl_on_symbolic_cv() -> CtrlSpecAnd:
+    from qualtran import CtrlSpec
+    from qualtran.symbolics import Shaped
+
+    ctrl_on_symbolic_cv = CtrlSpecAnd(CtrlSpec(cvs=Shaped((2,))))
+    return ctrl_on_symbolic_cv
+
+
+@bloq_example(generalizer=ignore_split_join)
+def _ctrl_on_symbolic_cv_multi() -> CtrlSpecAnd:
+    from qualtran import CtrlSpec
+    from qualtran.symbolics import Shaped
+
+    ctrl_on_symbolic_cv_multi = CtrlSpecAnd(CtrlSpec(cvs=Shaped((3,))))
+    return ctrl_on_symbolic_cv_multi
+
+
+@bloq_example(generalizer=ignore_split_join)
+def _ctrl_on_symbolic_n_ctrls() -> CtrlSpecAnd:
+    from qualtran import CtrlSpec
+    from qualtran.symbolics import Shaped
+
+    n = sympy.Symbol("n")
+    ctrl_on_symbolic_cv_multi = CtrlSpecAnd(CtrlSpec(cvs=Shaped((n,))))
+    return ctrl_on_symbolic_cv_multi
+
+
 _CTRLSPEC_AND_DOC = BloqDocSpec(
     bloq_cls=CtrlSpecAnd,
-    examples=(_ctrl_on_int, _ctrl_on_bits, _ctrl_on_nd_bits, _ctrl_on_multiple_values),
+    examples=(
+        _ctrl_on_int,
+        _ctrl_on_bits,
+        _ctrl_on_nd_bits,
+        _ctrl_on_multiple_values,
+        _ctrl_on_symbolic_cv,
+        _ctrl_on_symbolic_cv_multi,
+    ),
 )
diff --git a/qualtran/bloqs/mcmt/ctrl_spec_and_test.py b/qualtran/bloqs/mcmt/ctrl_spec_and_test.py
@@ -14,19 +14,33 @@
 import numpy as np
 import pytest
 
+import qualtran.testing as qlt_testing
 from qualtran import CtrlSpec, QUInt
 from qualtran.bloqs.mcmt.ctrl_spec_and import (
     _ctrl_on_bits,
     _ctrl_on_int,
     _ctrl_on_multiple_values,
     _ctrl_on_nd_bits,
+    _ctrl_on_symbolic_cv,
+    _ctrl_on_symbolic_cv_multi,
+    _ctrl_on_symbolic_n_ctrls,
     CtrlSpecAnd,
 )
 from qualtran.simulation.classical_sim import get_classical_truth_table
 
 
 @pytest.mark.parametrize(
-    "example", [_ctrl_on_bits, _ctrl_on_nd_bits, _ctrl_on_int, _ctrl_on_multiple_values]
+    "example",
+    [
+        _ctrl_on_bits,
+        _ctrl_on_nd_bits,
+        _ctrl_on_int,
+        _ctrl_on_multiple_values,
+        _ctrl_on_symbolic_cv,
+        _ctrl_on_symbolic_cv_multi,
+        _ctrl_on_symbolic_n_ctrls,
+    ],
+    ids=lambda ex: ex.name,
 )
 def test_examples(bloq_autotester, example):
     bloq_autotester(example)
@@ -51,3 +65,8 @@ def test_truth_table_using_classical_sim(ctrl_spec: CtrlSpec):
 
         # check: target bit (last output value) matches `is_active`
         assert out_vals[-1] == ctrl_spec.is_active(*in_vals)
+
+
+@pytest.mark.notebook
+def test_notebook():
+    qlt_testing.execute_notebook('ctrl_spec_and')
diff --git a/qualtran/conftest.py b/qualtran/conftest.py
@@ -138,6 +138,9 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
         'black_box_select',  # cannot serialize AutoPartition
         'black_box_prepare',  # cannot serialize AutoPartition
         'kaiser_window_state_symbolic',  # Split cannot have a symbolic data type.
+        'ctrl_on_symbolic_cv',  # cannot serialize Shaped
+        'ctrl_on_symbolic_cv_multi',  # cannot serialize Shaped
+        'ctrl_on_symbolic_n_ctrls',  # cannot serialize Shaped
     ]:
         pytest.xfail("Skipping serialization test for bloq examples that cannot yet be serialized.")