Restrict controlled from non-thru registers

qualtran/Controlled.ipynb

@@ -241,6 +241,14 @@
     "ccx3 = OnEach(n=3, gate=x).controlled(CtrlSpec(cvs=(1,1)))\n",
     "show_bloq(ccx3.decompose_bloq(), type='musical_score')"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24aa3d20-7803-4480-b717-5e3a2d5e6ba3",
+   "metadata": {},
+   "source": [
+    "Only bloqs with all-THRU registers can be controlled. Otherwise, it's not clear what the equivalent \"identity\" operation is."
+   ]
   }
  ],
  "metadata": {

qualtran/_infra/controlled.py

@@ -312,6 +312,13 @@ class Controlled(GateWithRegisters):
     subbloq: 'Bloq'
     ctrl_spec: 'CtrlSpec'
 
+    @cached_property
+    def _thru_registers_only(self) -> bool:
+        for reg in self.subbloq.signature:
+            if reg.side != Side.THRU:
+                return False
+        return True
+
     @classmethod
     def make_ctrl_system(cls, bloq: 'Bloq', ctrl_spec: 'CtrlSpec') -> Tuple[Bloq, AddControlledT]:
         """A factory method for creating both the Controlled and the adder function.
@@ -363,6 +370,9 @@ def decompose_bloq(self) -> 'CompositeBloq':
     def build_composite_bloq(
         self, bb: 'BloqBuilder', **initial_soqs: 'SoquetT'
     ) -> Dict[str, 'SoquetT']:
+        if not self._thru_registers_only:
+            raise DecomposeTypeError(f"Cannot handle non-thru registers in {self.subbloq}")
+
         # Use subbloq's decomposition but wire up the additional ctrl_soqs.
         from qualtran import CompositeBloq
 
@@ -406,6 +416,8 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         return counts
 
     def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
+        if not self._thru_registers_only:
+            raise ValueError(f"Cannot handle non-thru registers in {self}.")
         ctrl_vals = [vals[reg_name] for reg_name in self.ctrl_reg_names]
         other_vals = {reg.name: vals[reg.name] for reg in self.subbloq.signature}
         if self.ctrl_spec.is_active(*ctrl_vals):
@@ -418,6 +430,8 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT
         return vals
 
     def _tensor_data(self):
+        if not self._thru_registers_only:
+            raise ValueError(f"Cannot handle non-thru registers in {self}.")
         from qualtran.simulation.tensor._tensor_data_manipulation import (
             active_space_for_ctrl_spec,
             eye_tensor_for_signature,
@@ -447,7 +461,7 @@ def _unitary_(self):
             # to a unitary matrix.
             return self.tensor_contract()
         # Unable to determine the unitary effect.
-        return NotImplemented
+        raise ValueError(f"Cannot handle non-thru registers in {self}.")
 
     def my_tensors(
         self, incoming: Dict[str, 'ConnectionT'], outgoing: Dict[str, 'ConnectionT']

qualtran/_infra/controlled_test.py

@@ -11,51 +11,32 @@
 #  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 typing import Dict, List, Tuple, TYPE_CHECKING
+from typing import List, TYPE_CHECKING
 
-import attrs
 import cirq
 import numpy as np
 import pytest
 
 import qualtran.testing as qlt_testing
-from qualtran import (
-    Bloq,
-    BloqBuilder,
-    CompositeBloq,
-    Controlled,
-    CtrlSpec,
-    QBit,
-    QInt,
-    QUInt,
-    Register,
-    Side,
-    Signature,
-)
+from qualtran import Bloq, CompositeBloq, Controlled, CtrlSpec, QBit, QInt, QUInt
 from qualtran._infra.gate_with_registers import get_named_qubits, merge_qubits
 from qualtran.bloqs.basic_gates import (
     CSwap,
     GlobalPhase,
-    IntEffect,
-    IntState,
-    OneState,
     Swap,
     TwoBitCSwap,
     XGate,
     XPowGate,
     YGate,
-    ZeroState,
     ZGate,
 )
 from qualtran.bloqs.for_testing import TestAtom, TestParallelCombo, TestSerialCombo
-from qualtran.bloqs.mcmt import And
-from qualtran.cirq_interop.testing import GateHelper
 from qualtran.drawing import get_musical_score_data
 from qualtran.drawing.musical_score import Circle, SoqData, TextBox
 from qualtran.simulation.tensor import cbloq_to_quimb, get_right_and_left_inds
 
 if TYPE_CHECKING:
-    from qualtran import SoquetT
+    pass
 
 
 def test_ctrl_spec():
@@ -385,62 +366,6 @@ def test_controlled_global_phase_tensor():
     np.testing.assert_allclose(bloq.tensor_contract(), should_be)
 
 
-@attrs.frozen
-class TestCtrlStatePrepAnd(Bloq):
-    """Decomposes into a Controlled-AND gate + int effects & targets where ctrl is active.
-
-    Tensor contraction should give the output state vector corresponding to applying an
-    `And(and_ctrl)`; assuming all the control bits are active.
-    """
-
-    ctrl_spec: CtrlSpec
-    and_ctrl: Tuple[int, int]
-
-    @property
-    def signature(self) -> 'Signature':
-        return Signature([Register('x', QBit(), shape=(3,), side=Side.RIGHT)])
-
-    def build_composite_bloq(self, bb: 'BloqBuilder') -> Dict[str, 'SoquetT']:
-        one_or_zero = [ZeroState(), OneState()]
-        ctrl_bloq = Controlled(And(*self.and_ctrl), ctrl_spec=self.ctrl_spec)
-
-        ctrl_soqs = {}
-        for reg, cvs in zip(ctrl_bloq.ctrl_regs, self.ctrl_spec.cvs):
-            soqs = np.empty(shape=reg.shape, dtype=object)
-            for idx in reg.all_idxs():
-                soqs[idx] = bb.add(IntState(val=cvs[idx], bitsize=reg.dtype.num_qubits))
-            ctrl_soqs[reg.name] = soqs
-
-        and_ctrl = [bb.add(one_or_zero[cv]) for cv in self.and_ctrl]
-
-        ctrl_soqs = bb.add_d(ctrl_bloq, **ctrl_soqs, ctrl=and_ctrl)
-        out_soqs = np.asarray([*ctrl_soqs.pop('ctrl'), ctrl_soqs.pop('target')])  # type: ignore[misc]
-
-        for reg, cvs in zip(ctrl_bloq.ctrl_regs, self.ctrl_spec.cvs):
-            for idx in reg.all_idxs():
-                ctrl_soq = np.asarray(ctrl_soqs[reg.name])[idx]
-                bb.add(IntEffect(val=cvs[idx], bitsize=reg.dtype.num_qubits), val=ctrl_soq)
-        return {'x': out_soqs}
-
-
-def _verify_ctrl_tensor_for_and(ctrl_spec: CtrlSpec, and_ctrl: Tuple[int, int]):
-    bloq = TestCtrlStatePrepAnd(ctrl_spec, and_ctrl)
-    bloq_tensor = bloq.tensor_contract()
-    cirq_state_vector = GateHelper(And(*and_ctrl)).circuit.final_state_vector(
-        initial_state=and_ctrl + (0,)
-    )
-    np.testing.assert_allclose(bloq_tensor, cirq_state_vector, atol=1e-8)
-
-
-@pytest.mark.parametrize('ctrl_spec', interesting_ctrl_specs)
-def test_controlled_tensor_for_and_bloq(ctrl_spec: CtrlSpec):
-    # Test AND gate with one-sided signature (aka controlled state preparation).
-    _verify_ctrl_tensor_for_and(ctrl_spec, (1, 1))
-    _verify_ctrl_tensor_for_and(ctrl_spec, (1, 0))
-    _verify_ctrl_tensor_for_and(ctrl_spec, (0, 1))
-    _verify_ctrl_tensor_for_and(ctrl_spec, (0, 0))
-
-
 def test_controlled_diagrams():
     ctrl_gate = XPowGate(0.25).controlled()
     cirq.testing.assert_has_diagram(

qualtran/_infra/gate_with_registers_test.py

@@ -151,7 +151,7 @@ def test_gate_with_registers_decompose_from_context_auto_generated():
 
 
 def test_non_unitary_controlled():
-    bloq = BloqWithDecompose()
+    bloq = _TestGate()
     assert bloq.controlled(control_values=[0]) == Controlled(bloq, CtrlSpec(cvs=0))