Soquet to _Soquet, _QVar, Soquet and QVar (#1831)

In Qualtran, we compose quantum operations into larger operations by
connecting the outputs of a predecessor bloq to the input "ports" of a
successor bloq. We've used the `Soquet` frozen attrs class to represent
these nodes in the compute graph, which address the "ports" of a
particular instance of the bloq.

Previously, the immutable `Soquet` object represented both the "quantum
variables" being passed around during bloq building as well as the nodes
of the compute graph. With this PR, these concerns are separated. The
compute graph nodes are frozen dataclasses of type `_Soquet`, and the
quantum variables are *mutable* objects of type `_QVar`. We will put
additional helper attributes and methods onto `_QVar` to assist in bloq
building.

For backwards compatibility, the `Soquet` name is now a
`typing.Protocol` that encapsulates the duck-typing behavior of
`_Soquet` and `_QVar`. Bloqs in the wild should not have to update the
type annotations in `build_composite_bloq`
with this backwards compatibilty typing shim.

`isinstance(..., Soquet)` checks will emit a deprecation warning and
return True for *either* `_Soquet` or `_QVar`.

If you're using isinstance(soq, Soquet) to determine whether an item is
a single object
or an ndarray of those objects, use `BloqBuilder.is_single(x)` or
`BloqBuilder.is_ndarray(x)`. See the documentation in `QVarT` for an
example.

If you're developing library functionality, you can port isinstance
checks to either `_Soquet` or `QVar` as appropriate.

----

*Backwards compatibility note:* Despite the fundamental change to the
data model, this PR is designed to be runtime-backwards compatible with
deprecation warnings. The way I've accomplished this is through some
dark Python incantations, so no guarantees. Please send reports if this
breaks something.

Note that type checkers like `mypy` will likely still complain about
things even if the code runs fine without modification following this
PR.

----

### High-priority deprecations

- Please avoid constructing `Soquet`s yourself. You really shouldn't
have been doing this. `_Soquet` is now "private"
- If you're copying or modifying `CompositeBloq` by iterating over
connections and using `map_soqs`, please urgently upgrade your code to
use the new idiom where you must initialize the mapping with
`bb.initial_soq_map`. The shim I put in is brittle. #1742

Author: mpharrigan

qualtran/__init__.py

@@ -46,8 +46,11 @@
     CompositeBloq,
     BloqBuilder,
     DidNotFlattenAnythingError,
+    Soquet,
     SoquetT,
     ConnectionT,
+    QVar,
+    QVarT,
 )
 
 from ._infra.data_types import (
@@ -84,14 +87,7 @@
 
 # Internal imports: none
 # External imports: none
-from ._infra.quantum_graph import (
-    BloqInstance,
-    Connection,
-    DanglingT,
-    LeftDangle,
-    RightDangle,
-    Soquet,
-)
+from ._infra.quantum_graph import BloqInstance, Connection, DanglingT, LeftDangle, RightDangle
 
 from ._infra.gate_with_registers import GateWithRegisters
 

qualtran/_infra/adjoint.py

@@ -18,20 +18,28 @@
 
 from attrs import frozen
 
-from .composite_bloq import _binst_to_cxns, _cxns_to_soq_dict, _map_soqs, _reg_to_soq, BloqBuilder
+from .composite_bloq import (
+    _binst_to_cxns,
+    _cxns_to_soq_dict,
+    _map_soqs,
+    _reg_to_soq,
+    _SoquetT,
+    BloqBuilder,
+    QVarT,
+)
 from .gate_with_registers import GateWithRegisters
 from .quantum_graph import LeftDangle, RightDangle
 from .registers import Signature
 
 if TYPE_CHECKING:
     import cirq
 
-    from qualtran import Bloq, CompositeBloq, Register, Signature, SoquetT
+    from qualtran import Bloq, CompositeBloq, Register, Signature
     from qualtran.drawing import WireSymbol
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 
 
-def _adjoint_final_soqs(cbloq: 'CompositeBloq', new_signature: Signature) -> Dict[str, 'SoquetT']:
+def _adjoint_final_soqs(cbloq: 'CompositeBloq', new_signature: Signature) -> Dict[str, '_SoquetT']:
     """`CompositeBloq.final_soqs()` but backwards."""
     if LeftDangle not in cbloq._binst_graph:
         return {}
@@ -57,15 +65,15 @@ def _adjoint_cbloq(cbloq: 'CompositeBloq') -> 'CompositeBloq':
     # First, we reverse the registers to initialize the BloqBuilder.
     old_signature = cbloq.signature
     new_signature = cbloq.signature.adjoint()
-    old_i_soqs = [_reg_to_soq(RightDangle, reg) for reg in old_signature.rights()]
-    new_i_soqs = [_reg_to_soq(LeftDangle, reg) for reg in new_signature.lefts()]
-    soq_map: List[Tuple[SoquetT, SoquetT]] = list(zip(old_i_soqs, new_i_soqs))
 
     # Then we reverse the order of subbloqs
     bloqnections = reversed(list(cbloq.iter_bloqnections()))
 
     # And add subbloq.adjoint() back in for each subbloq.
     bb, _ = BloqBuilder.from_signature(new_signature)
+    old_i_soqs = [_reg_to_soq(RightDangle, reg) for reg in old_signature.rights()]
+    new_i_soqs = [bb._reg_to_qvar(LeftDangle, reg) for reg in new_signature.lefts()]
+    soq_map: List[Tuple[_SoquetT, QVarT]] = list(zip(old_i_soqs, new_i_soqs))
     for binst, preds, succs in bloqnections:
         # Instead of get_me returning the right element of a predecessor connection,
         # it's the left element of a successor connection.

qualtran/_infra/composite_bloq.ipynb

@@ -343,7 +343,7 @@
     "bb, _ = BloqBuilder.from_signature(cbloq.signature)\n",
     "\n",
     "# We'll have to \"map\" the soquets from our template cbloq to our new one\n",
-    "soq_map: List[Tuple[SoquetT, SoquetT]] = []\n",
+    "soq_map = bb.initial_soq_map(cbloq.signature.lefts())\n",
     "    \n",
     "# Iteration yields each bloq instance as well as its input and output soquets.\n",
     "for binst, in_soqs, old_out_soqs in cbloq.iter_bloqsoqs():\n",
@@ -512,8 +512,8 @@
     "# Go through and decompose each subbloq\n",
     "# We'll manually code this up in this notebook since this isn't a useful operation.\n",
     "bb, _ = BloqBuilder.from_signature(flat_three_p.signature)\n",
-    "soq_map: List[Tuple[SoquetT, SoquetT]] = []\n",
-    "    \n",
+    "soq_map: List[Tuple[SoquetT, SoquetT]] = bb.initial_soq_map(flat_three_p.signature.lefts())\n",
+    "\n",
     "for binst, in_soqs, old_out_soqs in flat_three_p.iter_bloqsoqs():\n",
     "    in_soqs = bb.map_soqs(in_soqs, soq_map)\n",
     "    \n",