Consistent Toffoli.controlled() (#1771)

This fixes the inconsistency raised in #1768 

This is a step towards #1556 -- specifically, we now always return
`ControlledViaAnd(XGate(), ...)` for all Toffoli cases (as well as for
most cases with XGate(); but this PR does not change the behavior of
`CNOT`)

Author: mpharrigan

qualtran/_infra/controlled.py

@@ -33,7 +33,8 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from ..symbolics import is_symbolic, prod, Shaped, SymbolicInt
+from qualtran.symbolics import is_symbolic, prod, Shaped, SymbolicInt
+
 from .bloq import Bloq, DecomposeNotImplementedError, DecomposeTypeError
 from .data_types import CDType, QBit, QCDType, QDType
 from .gate_with_registers import GateWithRegisters
@@ -52,18 +53,11 @@
 ControlBit: TypeAlias = int
 """A control bit, either 0 or 1."""
 
+_CVInLeafT: TypeAlias = Union[int, np.integer, NDArray[np.integer], Shaped]
+_CVInType: TypeAlias = Union[_CVInLeafT, Sequence['_CVInType']]
+
 
-def _cvs_convert(
-    cvs: Union[
-        int,
-        np.integer,
-        NDArray[np.integer],
-        Shaped,
-        Sequence[Union[int, np.integer]],
-        Sequence[Sequence[Union[int, np.integer]]],
-        Sequence[Union[NDArray[np.integer], Shaped]],
-    ]
-) -> Tuple[Union[NDArray[np.integer], Shaped], ...]:
+def _cvs_convert(cvs: _CVInType) -> Tuple[Union[NDArray[np.integer], Shaped], ...]:
     if isinstance(cvs, Shaped):
         return (cvs,)
     if isinstance(cvs, (int, np.integer)):

qualtran/bloqs/basic_gates/toffoli.ipynb

@@ -141,13 +141,109 @@
     "show_counts_sigma(toffoli_sigma)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "dc13c6f1-86b7-4a0b-9b4a-e7f3ed67e3a6",
+   "metadata": {},
+   "source": [
+    "## Controlled\n",
+    "\n",
+    "We use `Toffoli` for specifically two control bits of the `XGate`. Requesting `Toffoli().controlled(ctrl_spec)` will give an instance of the generic `ControlledViaAnd(XGate(), ...)` bloq, which will control the `XGate` on an arbitrary boolean function supported by `CtrlSpec`. The set of control registers present on the resulting bloq will 1) the control registers needed for the provided `ctrl_spec` and 2) one additional `QBit()[2]` register for the original Toffoli control bits."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4e0647b-ce06-47eb-9e18-7d49e375a9ed",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran import CtrlSpec\n",
+    "print(repr(Toffoli().controlled()))\n",
+    "show_bloq(Toffoli().controlled(), 'musical_score')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7fffec98-d1fe-42d5-8878-4117ecc881bc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(repr(Toffoli().controlled(CtrlSpec(cvs=0))))\n",
+    "show_bloq(Toffoli().controlled(CtrlSpec(cvs=0)), 'musical_score')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4a58a0e7-1219-4bc6-ae2b-8df71d7d85d5",
+   "metadata": {},
+   "source": [
+    "### Complex control specs\n",
+    "\n",
+    "Per above, `Toffoli.controlled()` can support arbitrary control specs."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ad6ba003",
+   "id": "59d97e1f-b957-432f-b0de-edeee84e5b71",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "active_six = CtrlSpec(qdtypes=QInt(8), cvs=6)\n",
+    "six_and_toffoli = Toffoli().controlled(active_six)\n",
+    "print(repr(six_and_toffoli))\n",
+    "show_bloq(six_and_toffoli, 'musical_score')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "afcbb7aa-c098-4c56-8961-e0cfc9661a5c",
+   "metadata": {},
+   "source": [
+    "### Calling controlled gates automatically\n",
+    "\n",
+    "The `Toffoli().controlled()` method (and in general, the `.controlled()` method for any bloq object) returns a bloq that implements a controlled version of the original gate. The returned bloq can structure its signature however it chooses. You can see that `Toffoli.controlled()` returns a `ControlledViaAnd(XGate(), ...)` which names its control registers `ctrl{n}` with the final one corresponding to the original Toffoli control bits. \n",
+    "\n",
+    "When writing meta-bloqs that need to support adding arbitrarily-controlled versions of Toffoli to a decomposition, you should not rely on any particular signature. `Bloq.get_ctrl_system()` is the method for automatically `bb.add`-ing controlled versions of arbitrary bloqs. We demonstrate that below."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "afc14f80-44cd-4fc8-a626-427ef8fd4a99",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.basic_gates import IntState, IntEffect\n",
+    "\n",
+    "bb = BloqBuilder()\n",
+    "# q0, q1, target will be original wires for calling Toffoli\n",
+    "q0 = bb.add_register('q0', QBit())\n",
+    "q1 = bb.add_register('q1', QBit())\n",
+    "target = bb.add_register('target', QBit())\n",
+    "\n",
+    "# \"six\" is a new wire that we additionally want to control the Toffoli with\n",
+    "six = bb.add(IntState(val=6, bitsize=8))\n",
+    "ctrl_spec = CtrlSpec(qdtypes=QUInt(8), cvs=6)\n",
+    "\n",
+    "# Instead of using Toffoli().controlled() directly, we use the ctrl_system\n",
+    "ctof, ctof_adder = Toffoli().get_ctrl_system(ctrl_spec)\n",
+    "\n",
+    "# The new control variables are passed in their own list and returned in their own list\n",
+    "# The existing Toffoli inputs are provided in a dictionary and returned in their own list\n",
+    "(six,), ((q0, q1), target) = ctof_adder(bb, ctrl_soqs=[six], in_soqs=dict(ctrl=[q0,q1], target=target))\n",
+    "\n",
+    "# # below is a brittle way to achieve the same effect for this particular ctrl_spec\n",
+    "# # it assumes control registers named ctrl1 and ctrl2\n",
+    "# six_and_toffoli = Toffoli().controlled(ctrl_spec)\n",
+    "# six, (q0, q1), target = bb.add(six_and_toffoli, ctrl1=six, ctrl2=[q0, q1], q=target)\n",
+    "\n",
+    "bb.add(IntEffect(val=6, bitsize=8), val=six)\n",
+    "program = bb.finalize(q0=q0, q1=q1, target=target)\n",
+    "show_bloq(program, 'musical_score')"
+   ]
   }
  ],
  "metadata": {
@@ -166,7 +262,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.13.9"
   }
  },
  "nbformat": 4,

qualtran/bloqs/basic_gates/toffoli.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 import itertools
 from functools import cached_property
-from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import frozen
@@ -22,6 +22,7 @@
 from qualtran import (
     Bloq,
     bloq_example,
+    BloqBuilder,
     BloqDocSpec,
     CompositeBloq,
     Connection,
@@ -30,6 +31,7 @@
     QBit,
     Register,
     Signature,
+    SoquetT,
 )
 
 if TYPE_CHECKING:
@@ -38,7 +40,7 @@
     from pennylane.operation import Operation
     from pennylane.wires import Wires
 
-    from qualtran import AddControlledT, BloqBuilder, SoquetT
+    from qualtran import AddControlledT
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
     from qualtran.simulation.classical_sim import ClassicalValT
@@ -137,27 +139,33 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.basic_gates import CNOT
-        from qualtran.bloqs.mcmt import ControlledViaAnd
-
-        if ctrl_spec != CtrlSpec():
-            return super().get_ctrl_system(ctrl_spec)
-
-        cc_cnot = ControlledViaAnd(CNOT(), CtrlSpec(cvs=[1, 1]))
-
-        def add_controlled(
-            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
-        ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
-            (new_ctrl,) = ctrl_soqs
-            ctrl0, ctrl1 = cast(NDArray, in_soqs.pop('ctrl'))
-
-            (new_ctrl, ctrl0), ctrl1, target = bb.add(
-                cc_cnot, ctrl2=np.array([new_ctrl, ctrl0]), ctrl=ctrl1, target=in_soqs.pop('target')
+        from qualtran.bloqs.basic_gates import XGate
+
+        # Perform the Logical-AND of `ctrl_spec` with Toffoli's implicit XGate ctrl spec
+        # of cvs=(1, 1).
+        anded_ctrl_spec = CtrlSpec(
+            qdtypes=ctrl_spec.qdtypes + (QBit(),), cvs=ctrl_spec.cvs + (np.array([1, 1]),)
+        )
+        ctrl_via_and, adder1 = XGate().get_ctrl_system(anded_ctrl_spec)
+
+        # We have to wire up the *new* ctrl registers vs Toffoli's *existing* ctrl register
+        # while also translating between Toffoli() and ControlledViaAnd(XGate(), cvs=(1,1)).
+        def adder2(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            # Note: in adder2, `ctrl_soqs` matches the requested `ctrl_spec`
+            #       in adder1, `ctrl_soqs` matches `anded_ctrl_spec`
+            # Note: in adder2 `in_soqs` matches the signature of Toffoli
+            #       in adder1, `in_soqs` matches the signature of XGate
+            tof_ctrl = in_soqs['ctrl']
+            tof_target = in_soqs['target']
+            (*ctrl_soqs, tof_ctrl), (tof_target,) = adder1(
+                bb, list(ctrl_soqs) + [tof_ctrl], {'q': tof_target}
             )
 
-            return [new_ctrl], [np.array([ctrl0, ctrl1]), target]
+            return ctrl_soqs, [tof_ctrl, tof_target]
 
-        return cc_cnot, add_controlled
+        return ctrl_via_and, adder2
 
 
 @bloq_example

qualtran/bloqs/basic_gates/toffoli_test.py

@@ -17,10 +17,10 @@
 import numpy as np
 
 import qualtran.testing as qlt_testing
-from qualtran import BloqBuilder, CtrlSpec
-from qualtran.bloqs.basic_gates import Toffoli, ZeroState
+from qualtran import BloqBuilder, CtrlSpec, QBit, Register, Signature
+from qualtran.bloqs.basic_gates import CNOT, Toffoli, XGate, ZeroState
 from qualtran.bloqs.basic_gates.toffoli import _toffoli
-from qualtran.bloqs.mcmt import And
+from qualtran.bloqs.mcmt import And, ControlledViaAnd
 from qualtran.drawing.musical_score import Circle, ModPlus
 from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
 
@@ -106,15 +106,15 @@ def test_ctrl_toffoli_cost():
     ctrl_tof = Toffoli().controlled()
 
     _, sigma = ctrl_tof.call_graph()
-    assert sigma == {Toffoli(): 1, And(): 1, And().adjoint(): 1}
+    assert sigma == {CNOT(): 1, And(): 2, And().adjoint(): 2}
 
     gc = get_cost_value(ctrl_tof, QECGatesCost())
-    assert gc == GateCounts(and_bloq=1, toffoli=1, clifford=1, measurement=1)
+    assert gc == GateCounts(and_bloq=2, clifford=3, measurement=2)
 
     cc_tof = Toffoli().controlled(CtrlSpec(cvs=[1, 1]))
 
     _, sigma = cc_tof.call_graph()
-    assert sigma == {Toffoli(): 1, And(): 2, And().adjoint(): 2}
+    assert sigma == {CNOT(): 1, And(): 3, And().adjoint(): 3}
 
 
 def test_toffoli_controlled():
@@ -123,3 +123,18 @@ def test_toffoli_controlled():
 
     bloq = Controlled(Toffoli().as_composite_bloq(), CtrlSpec())
     qlt_testing.assert_valid_bloq_decomposition(bloq)
+
+
+def test_toffoli_controlled_2():
+    # https://github.com/quantumlib/Qualtran/issues/1768
+
+    c0t = Toffoli().controlled(CtrlSpec(QBit(), cvs=0))
+    c1t = Toffoli().controlled(CtrlSpec(QBit(), cvs=1))
+
+    assert c0t == ControlledViaAnd(XGate(), CtrlSpec(qdtypes=(QBit(), QBit()), cvs=(0, [1, 1])))
+    assert c1t == ControlledViaAnd(XGate(), CtrlSpec(qdtypes=(QBit(), QBit()), cvs=(1, [1, 1])))
+
+    assert c0t.signature == c1t.signature
+    assert c0t.signature == Signature(
+        [Register('ctrl1', QBit()), Register('ctrl2', QBit(), shape=(2,)), Register('q', QBit())]
+    )