Simplify QubitizationWalkOperator: remove specializations for control/adjoint and use defaults instead (#1481)

* simplify walk operator

* add ctrl systems for various bloqs

* special adjoint wrapper for bloqs with specialized ctrl

* fix ctrl sys for `CZ` and `CY`: target register names of uncontrolled and controlled versions do not match

* nb

* fix bug in select hubbard (`0` ctrl not supported by decomp, so fallback to default)

* move new adjoint bloq to `mcmt.specialized_ctrl`

* add tests for adjoint

* imports and mypy

* `Y, Z`: explicitly implement ctrl sys (avoids AutoPartition in decomp)

* docstring `AdjointWithSpecializedCtrl`

* add notebook

* format

* simplify ctrl sys for CZ, CY

* hubbard: not impl error on cv=0

Author: anurudhp

qualtran/bloqs/basic_gates/swap.py

@@ -105,19 +105,15 @@ def adjoint(self) -> 'Bloq':
         return self
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:
-        if ctrl_spec != CtrlSpec():
-            return super().get_ctrl_system(ctrl_spec=ctrl_spec)
-
-        cswap = TwoBitCSwap()
-
-        def adder(
-            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
-        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
-            (ctrl,) = ctrl_soqs
-            ctrl, x, y = bb.add(cswap, ctrl=ctrl, x=in_soqs['x'], y=in_soqs['y'])
-            return [ctrl], [x, y]
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
-        return cswap, adder
+        return get_ctrl_system_1bit_cv_from_bloqs(
+            self,
+            ctrl_spec,
+            current_ctrl_bit=None,
+            bloq_with_ctrl=TwoBitCSwap(),
+            ctrl_reg_name='ctrl',
+        )
 
 
 @bloq_example
@@ -201,6 +197,13 @@ def wire_symbol(self, reg: Optional['Register'], idx: Tuple[int, ...] = ()) -> '
         else:
             return TextBox('×')
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+
+        return get_ctrl_system_1bit_cv_from_bloqs(
+            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='ctrl'
+        )
+
 
 @bloq_example
 def _cswap_bit() -> TwoBitCSwap:

qualtran/bloqs/basic_gates/y_gate.py

@@ -173,6 +173,13 @@ def wire_symbol(
             return TextBox('Y')
         raise ValueError(f"Unknown register {reg}.")
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+
+        return get_ctrl_system_1bit_cv_from_bloqs(
+            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='ctrl'
+        )
+
 
 @bloq_example
 def _cy_gate() -> CYGate:

qualtran/bloqs/basic_gates/z_basis.py

@@ -340,6 +340,13 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return Circle()
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+
+        return get_ctrl_system_1bit_cv_from_bloqs(
+            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='q1'
+        )
+
 
 @bloq_example
 def _cz() -> CZ:

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

@@ -90,6 +90,10 @@ class SelectHubbard(SelectOracle):
     def __attrs_post_init__(self):
         if self.x_dim != self.y_dim:
             raise NotImplementedError("Currently only supports the case where x_dim=y_dim.")
+        if self.control_val == 0:
+            raise NotImplementedError(
+                "control_val=0 not supported, use `SelectHubbard(x, y).controlled(CtrlSpec(cvs=0))` instead"
+            )
 
     @cached_property
     def control_registers(self) -> Tuple[Register, ...]:
@@ -191,18 +195,24 @@ def __str__(self):
         return s
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
-        return get_ctrl_system_1bit_cv(
+        return get_ctrl_system_1bit_cv_from_bloqs(
             self,
             ctrl_spec=ctrl_spec,
             current_ctrl_bit=self.control_val,
-            get_ctrl_bloq_and_ctrl_reg_name=lambda cv: (
-                attrs.evolve(self, control_val=cv),
-                'control',
-            ),
+            bloq_with_ctrl=attrs.evolve(self, control_val=1),
+            ctrl_reg_name='control',
+        )
+
+    def adjoint(self) -> 'Bloq':
+        from qualtran.bloqs.mcmt.specialized_ctrl import (
+            AdjointWithSpecializedCtrl,
+            SpecializeOnCtrlBit,
         )
 
+        return AdjointWithSpecializedCtrl(self, specialize_on_ctrl=SpecializeOnCtrlBit.ONE)
+
 
 @bloq_example
 def _sel_hubb() -> SelectHubbard:

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

@@ -11,14 +11,15 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from unittest.mock import ANY
 
 import pytest
 
 from qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard import (
     _sel_hubb,
     SelectHubbard,
 )
-from qualtran.cirq_interop.t_complexity_protocol import t_complexity
+from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
 
 
 def test_sel_hubb_auto(bloq_autotester):
@@ -28,8 +29,19 @@ def test_sel_hubb_auto(bloq_autotester):
 @pytest.mark.parametrize('dim', [*range(2, 10)])
 def test_select_t_complexity(dim):
     select = SelectHubbard(x_dim=dim, y_dim=dim, control_val=1)
-    cost = t_complexity(select)
+    cost = get_cost_value(select, QECGatesCost())
     N = 2 * dim * dim
     logN = 2 * (dim - 1).bit_length() + 1
-    assert cost.t == 10 * N + 14 * logN - 8
-    assert cost.rotations == 0
+    assert cost == GateCounts(
+        cswap=2 * logN, and_bloq=5 * (N // 2) - 2, measurement=5 * (N // 2) - 2, clifford=ANY
+    )
+    assert cost.total_t_count() == 10 * N + 14 * logN - 8
+
+
+def test_adjoint_controlled():
+    bloq = _sel_hubb()
+
+    adj_ctrl_bloq = bloq.controlled().adjoint()
+    ctrl_adj_bloq = bloq.adjoint().controlled()
+
+    assert adj_ctrl_bloq == ctrl_adj_bloq

qualtran/bloqs/mcmt/specialized_ctrl.ipynb

@@ -0,0 +1,201 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "## Bloqs with specialized controlled implementations\n",
+    "\n",
+    "In some cases, a bloq may have a specialized singly-controlled version (e.g. `LCUBlockEncoding`).\n",
+    "Qualtran provides a convenience methods `get_ctrl_system_1bit_cv` and `get_ctrl_system_1bit_cv_from_bloqs` to override the `get_ctrl_system`. These methods ensure that multiply-controlled bloqs are correctly reduced to the provided singly-controlled variants.\n",
+    "\n",
+    "- `get_ctrl_system_1bit_cv_from_bloqs` - Override when a specialized controlled-by-1 implementation is available.\n",
+    "- `get_ctrl_system_1bit_cv` - Override when both specialized controlled-by-1 and controlled-by-0 implementations are available.\n",
+    "\n",
+    "The following demonstrates an example for a bloq implementing $T^\\dagger X T$, where the controlled version only needs to control the $X$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import attrs\n",
+    "from qualtran import Bloq, BloqBuilder, Soquet, SoquetT, Signature, CtrlSpec, AddControlledT\n",
+    "from qualtran.bloqs.basic_gates import TGate, XGate, CNOT\n",
+    "\n",
+    "\n",
+    "@attrs.frozen\n",
+    "class BloqWithSpecializedCtrl(Bloq):\n",
+    "    \"\"\"Bloq implementing $T^\\dagger X T$\"\"\"\n",
+    "    is_controlled: bool = False\n",
+    "\n",
+    "    @property\n",
+    "    def signature(self) -> 'Signature':\n",
+    "        n_ctrls = 1 if self.is_controlled else 0\n",
+    "        return Signature.build(ctrl=n_ctrls, q=1)\n",
+    "    \n",
+    "    def build_composite_bloq(self, bb: 'BloqBuilder', q: 'Soquet', **soqs) -> dict[str, 'SoquetT']:\n",
+    "        ctrl = soqs.pop('ctrl', None)\n",
+    "        \n",
+    "        q = bb.add(TGate(), q=q)\n",
+    "        if self.is_controlled:\n",
+    "            ctrl, q = bb.add(CNOT(), ctrl=ctrl, target=q)\n",
+    "        else:\n",
+    "            ctrl, q = bb.add(XGate(), ctrl=ctrl, target=q)\n",
+    "        q = bb.add(TGate().adjoint(), q=q)\n",
+    "        \n",
+    "        out_soqs = {'q': q}\n",
+    "        if ctrl:\n",
+    "            out_soqs |= {'ctrl': ctrl}\n",
+    "        return out_soqs\n",
+    "    \n",
+    "    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:\n",
+    "        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs\n",
+    "\n",
+    "        return get_ctrl_system_1bit_cv_from_bloqs(\n",
+    "            self,\n",
+    "            ctrl_spec,\n",
+    "            current_ctrl_bit=1 if self.is_controlled else None,\n",
+    "            bloq_with_ctrl=attrs.evolve(self, is_controlled=True),\n",
+    "            ctrl_reg_name='ctrl',\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloq, show_call_graph\n",
+    "\n",
+    "bloq = BloqWithSpecializedCtrl().controlled().controlled()\n",
+    "show_bloq(bloq.decompose_bloq().flatten())\n",
+    "show_call_graph(bloq)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "## Propagating the Adjoint\n",
+    "\n",
+    "In the above bloq, calling controlled on the adjoint does not push the controls into the bloq, and therefore does not use the specialized implementation provided."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BloqWithSpecializedCtrl().adjoint().controlled()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "This can be fixed by overriding the adjoint using a special wrapper for this case - `AdjointWithSpecializedCtrl`. This is a subclass of the default `Adjoint` metabloq, and ensures that single-qubit controls are pushed into the underlying bloq."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@attrs.frozen\n",
+    "class BloqWithSpecializedCtrlWithAdjoint(Bloq):\n",
+    "    \"\"\"Bloq implementing $T^\\dagger X T$\"\"\"\n",
+    "    is_controlled: bool = False\n",
+    "\n",
+    "    @property\n",
+    "    def signature(self) -> 'Signature':\n",
+    "        n_ctrls = 1 if self.is_controlled else 0\n",
+    "        return Signature.build(ctrl=n_ctrls, q=1)\n",
+    "    \n",
+    "    def build_composite_bloq(self, bb: 'BloqBuilder', q: 'Soquet', **soqs) -> dict[str, 'SoquetT']:\n",
+    "        ctrl = soqs.pop('ctrl', None)\n",
+    "        \n",
+    "        q = bb.add(TGate(), q=q)\n",
+    "        if self.is_controlled:\n",
+    "            ctrl, q = bb.add(CNOT(), ctrl=ctrl, target=q)\n",
+    "        else:\n",
+    "            ctrl, q = bb.add(XGate(), ctrl=ctrl, target=q)\n",
+    "        q = bb.add(TGate().adjoint(), q=q)\n",
+    "        \n",
+    "        out_soqs = {'q': q}\n",
+    "        if ctrl:\n",
+    "            out_soqs |= {'ctrl': ctrl}\n",
+    "        return out_soqs\n",
+    "    \n",
+    "    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:\n",
+    "        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs\n",
+    "\n",
+    "        return get_ctrl_system_1bit_cv_from_bloqs(\n",
+    "            self,\n",
+    "            ctrl_spec,\n",
+    "            current_ctrl_bit=1 if self.is_controlled else None,\n",
+    "            bloq_with_ctrl=attrs.evolve(self, is_controlled=True),\n",
+    "            ctrl_reg_name='ctrl',\n",
+    "        )\n",
+    "\n",
+    "    def adjoint(self):\n",
+    "        from qualtran.bloqs.mcmt.specialized_ctrl import AdjointWithSpecializedCtrl, SpecializeOnCtrlBit\n",
+    "        \n",
+    "        return AdjointWithSpecializedCtrl(self, SpecializeOnCtrlBit.ONE)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BloqWithSpecializedCtrlWithAdjoint().adjoint().controlled()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert BloqWithSpecializedCtrlWithAdjoint().adjoint().controlled() == BloqWithSpecializedCtrlWithAdjoint(is_controlled=True).adjoint()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}