From b334f47ed3daedf29e5ff0b71456428eb351100d Mon Sep 17 00:00:00 2001 From: Matthew Harrigan Date: Wed, 8 Apr 2026 18:14:53 -0700 Subject: [PATCH] Bloqify syntax basics --- qualtran/__init__.py | 2 + qualtran/_infra/composite_bloq.py | 102 ++++++++++-- qualtran/_infra/composite_bloq_test.py | 4 +- qualtran/bloqify_syntax/__init__.py | 16 ++ qualtran/bloqify_syntax/_infra.py | 210 +++++++++++++++++++++++++ qualtran/bloqify_syntax/_infra_test.py | 115 ++++++++++++++ 6 files changed, 432 insertions(+), 17 deletions(-) create mode 100644 qualtran/bloqify_syntax/__init__.py create mode 100644 qualtran/bloqify_syntax/_infra.py create mode 100644 qualtran/bloqify_syntax/_infra_test.py diff --git a/qualtran/__init__.py b/qualtran/__init__.py index 9ef7e03b2..b41f8e6e7 100644 --- a/qualtran/__init__.py +++ b/qualtran/__init__.py @@ -105,4 +105,6 @@ from ._infra.bloq_example import BloqExample, bloq_example, BloqDocSpec +from .bloqify_syntax import bloqify + # -------------------------------------------------------------------------------------------------- diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py index 5297456a0..924ea6f43 100644 --- a/qualtran/_infra/composite_bloq.py +++ b/qualtran/_infra/composite_bloq.py @@ -62,6 +62,8 @@ if TYPE_CHECKING: import cirq + import qualtran as qlt + import qualtran.dtype as qdt from qualtran.bloqs.bookkeeping.auto_partition import Unused from qualtran.cirq_interop._cirq_to_bloq import CirqQuregInT, CirqQuregT from qualtran.drawing import WireSymbol @@ -467,7 +469,7 @@ def final_soqs(self) -> Dict[str, _SoquetT]: def copy(self) -> 'CompositeBloq': """Create a copy of this composite bloq by re-building it.""" - bb, _ = BloqBuilder.from_signature(self.signature) + bb, _ = BloqBuilder.from_signature(self.signature, bloq_key=self.bloq_key) soq_map = bb.initial_soq_map(self.signature.lefts()) for binst, in_soqs, old_out_soqs in self.iter_bloqsoqs(): @@ -641,6 +643,11 @@ def __str__(self): return self.bloq_key return f'CompositeBloq([{len(self.bloq_instances)} subbloqs...])' + def __repr__(self): + if self.bloq_key is not None: + return f'CompositeBloq(..., bloq_key={self.bloq_key!r})' + return f'CompositeBloq([{len(self.bloq_instances)} subbloqs...])' + def _create_binst_graph( cxns: Iterable[Connection], nodes: Iterable[BloqInstance] = () @@ -1098,7 +1105,7 @@ def build_composite_bloq(self, bb: BloqBuilder, q0, q1): by the framework or by the `BloqBuilder.from_signature(s)` factory method. """ - def __init__(self, add_registers_allowed: bool = True): + def __init__(self, add_registers_allowed: bool = True, *, bloq_key: Optional[str] = None): # To be appended to: self._cxns: List[Connection] = [] self._regs: List[Register] = [] @@ -1113,6 +1120,8 @@ def __init__(self, add_registers_allowed: bool = True): # Whether we can call `add_register` and do non-strict `finalize()`. self.add_register_allowed = add_registers_allowed + self._bloq_key = bloq_key + def add_register_from_dtype( self, reg: Union[str, Register], dtype: Optional[QCDType] = None ) -> Union[None, QVarT]: @@ -1135,7 +1144,7 @@ def add_register_from_dtype( from qualtran.symbolics import is_symbolic if not self.add_register_allowed: - raise ValueError( + raise BloqError( "This BloqBuilder was constructed from pre-specified registers. " "Ad hoc addition of more registers is not allowed." ) @@ -1208,7 +1217,11 @@ def add_register( @classmethod def from_signature( - cls, signature: Signature, add_registers_allowed: bool = False + cls, + signature: Signature, + add_registers_allowed: bool = False, + *, + bloq_key: Optional[str] = None, ) -> Tuple['BloqBuilder', Dict[str, QVarT]]: """Construct a BloqBuilder with a pre-specified signature. @@ -1216,7 +1229,7 @@ def from_signature( to match. This constructor is used by `Bloq.decompose_bloq()`. """ # Initial construction: allow register addition for the following loop. - bb = cls(add_registers_allowed=True) + bb = cls(add_registers_allowed=True, bloq_key=bloq_key) initial_soqs: Dict[str, QVarT] = {} for reg in signature: @@ -1454,7 +1467,12 @@ def add(self, bloq: Bloq, **in_soqs: SoquetInT): be unpacked with tuple unpacking. In this final case, the ordering is according to `bloq.signature` and irrespective of the order of `**in_soqs`. """ - outs = self.add_t(bloq, **in_soqs) + try: + outs = self.add_t(bloq, **in_soqs) + except BloqError as be: + # Error source shown as `bb.add(...)` + raise BloqError(*be.args) from None + if len(outs) == 0: return None if len(outs) == 1: @@ -1539,6 +1557,38 @@ def add_from(self, bloq: Bloq, **in_soqs: SoquetInT) -> Tuple['QVarT', ...]: fsoqs = _map_soqs(cbloq.final_soqs(), soq_map) return tuple(fsoqs[reg.name] for reg in cbloq.signature.rights()) + def _change_THRU_to_LEFT(self, reg_name: str): + """Used during loose `finalize` to force LEFT registers.""" + for reg_i, reg in enumerate(self._regs): + if reg.name == reg_name: + break + else: + raise AssertionError(f"{reg_name} doesn't exist in the registers.") + + if reg.side != Side.THRU: + raise ValueError(f"{reg} is supposed to be a THRU register.") + + new_reg = attrs.evolve(reg, side=Side.LEFT) + + # Replace in `self._available` + soqs_to_replace = [] + for soq in self._available: + if soq.binst is LeftDangle and soq.reg == reg: + soqs_to_replace.append(soq) + for soq in soqs_to_replace: + self._available.remove(soq) + self._available.add(attrs.evolve(soq, reg=new_reg)) + + # Replace in `self._cxns` + for j in range(len(self._cxns)): + cxn = self._cxns[j] + if cxn.left.reg == reg: + new_cxn = attrs.evolve(cxn, left=attrs.evolve(cxn.left, reg=new_reg)) + self._cxns[j] = new_cxn + + # Replace in `self._regs` + self._regs[reg_i] = new_reg + def finalize(self, **final_soqs: SoquetInT) -> CompositeBloq: """Finish building a CompositeBloq and return the immutable CompositeBloq. @@ -1546,10 +1596,14 @@ def finalize(self, **final_soqs: SoquetInT) -> CompositeBloq: it configures the final "dangling" soquets that serve as the outputs for the composite bloq as a whole. - If `self.add_registers_allowed` is set to `True`, additional register - names passed to this function will be added as RIGHT registers. Otherwise, - this method validates the provided `final_soqs` against our list of RIGHT - (and THRU) registers. + If `self.add_registers_allowed` is set to `False`, the kwqargs to this method must + exactly match the signature configured for this bloq builder. + Otherwise: + - surplus arguments (with register names not in our signature) will be interpreted as + output quantum variables and we'll add a corresponding RIGHT register. + - missing arguments (i.e. a register name was introduced somewhere but + no quantum variable was provided for it), we will set that register to be a + LEFT register. Args: **final_soqs: Keyword arguments mapping the composite bloq's register names to @@ -1577,6 +1631,12 @@ def _infer_shaped_dtype(soq: SoquetT) -> Tuple['QCDType', Tuple[int, ...]]: dtype, shape = _infer_shaped_dtype(np.asarray(soq)) self._regs.append(Register(name=name, dtype=dtype, shape=shape, side=Side.RIGHT)) + # If a soquet is missing, we're charitable and consider it de-allocated, see docstring. + deletable_right_reg_names = [reg.name for reg in self._regs if reg.side == Side.THRU] + for name in deletable_right_reg_names: + if name not in final_soqs: + self._change_THRU_to_LEFT(reg_name=name) + return self._finalize_strict(**final_soqs) def _finalize_strict(self, **final_soqs: SoquetInT) -> CompositeBloq: @@ -1592,16 +1652,25 @@ def _fin(idxed_soq: _QVar, reg: Register, idx: Tuple[int, ...]): # close over `RightDangle` return self._add_cxn(RightDangle, idxed_soq, reg, idx) + if self._bloq_key: + debug_str = f'finalization of {self._bloq_key}' + else: + debug_str = 'finalization' _process_soquets( - registers=signature.rights(), debug_str='Finalizing', in_soqs=final_soqs, func=_fin + registers=signature.rights(), debug_str=debug_str, in_soqs=final_soqs, func=_fin ) if self._available: - raise BloqError( - f"During finalization, {self._available} Soquets were not used." - ) from None + if self._bloq_key is None: + ctx = '' + else: + ctx = f' of {self._bloq_key}' + raise BloqError(f"During finalization{ctx}, {self._available} Soquets were not used.") return CompositeBloq( - connections=self._cxns, signature=signature, bloq_instances=self._binsts + connections=self._cxns, + signature=signature, + bloq_instances=self._binsts, + bloq_key=self._bloq_key, ) def allocate( @@ -1653,3 +1722,6 @@ def join(self, soqs: SoquetInT, dtype: Optional[QDType] = None) -> 'Soquet': dtype = QAny(n) return self.add(Join(dtype=dtype), reg=soqs) + + def in_register(self, name: str, dtype: QCDType, shape=()) -> Union[None, QVarT]: + return self.add_register_from_dtype(Register(name=name, dtype=dtype, shape=shape)) diff --git a/qualtran/_infra/composite_bloq_test.py b/qualtran/_infra/composite_bloq_test.py index f9f3953f7..65cbda8c0 100644 --- a/qualtran/_infra/composite_bloq_test.py +++ b/qualtran/_infra/composite_bloq_test.py @@ -312,7 +312,7 @@ def test_finalize_missing_args(): x2, y2 = bb.add(TestTwoBitOp(), ctrl=x, target=y) bb.add_register_allowed = False - with pytest.raises(BloqError, match=r"During Finalizing, we expected a value for 'x'\."): + with pytest.raises(BloqError, match=r"During finalization, we expected a value for 'x'\."): bb.finalize(y=y2) @@ -321,7 +321,7 @@ def test_finalize_strict_too_many_args(): x2, y2 = bb.add(TestTwoBitOp(), ctrl=x, target=y) bb.add_register_allowed = False - with pytest.raises(BloqError, match=r'Finalizing does not accept Soquets.*z.*'): + with pytest.raises(BloqError, match=r'finalization does not accept Soquets.*z.*'): bb.finalize(x=x2, y=y2, z=_Soquet(RightDangle, Register('asdf', QBit()))) diff --git a/qualtran/bloqify_syntax/__init__.py b/qualtran/bloqify_syntax/__init__.py new file mode 100644 index 000000000..f85e40ab4 --- /dev/null +++ b/qualtran/bloqify_syntax/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 ._infra import bloqify diff --git a/qualtran/bloqify_syntax/_infra.py b/qualtran/bloqify_syntax/_infra.py new file mode 100644 index 000000000..257219e00 --- /dev/null +++ b/qualtran/bloqify_syntax/_infra.py @@ -0,0 +1,210 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +# pylint: disable=keyword-arg-before-vararg +import inspect +from typing import Any, Dict, Optional, Protocol, Sequence, Set, TYPE_CHECKING + +import attrs +import numpy as np + +from qualtran import BloqBuilder, CompositeBloq, Signature +from qualtran._infra.quantum_graph import _QVar + +if TYPE_CHECKING: + import qualtran as qlt + + +class _TracingBloqFuncT(Protocol): + """Protocol for functions compatible with `@bloqify`.""" + + __name__: str + + def __call__(self, bb: 'BloqBuilder', *args: Any, **kwargs: Any) -> Dict[str, Any]: + """the structure of the function. + + During normal operation + - a `bb: BloqBuilder` will be passed positionally as the first argument. + - all other quantum and classical args and kwargs will be bound according to the Python + rules and passed to the function by keyword. There are no additional positional-only + arguments allowed (other than optionally `bb`). + + During `make` + - a `bb: BloqBuilder` will be passed positionally as the first argument. + - positional arguments will the be passed + - any provided keyword arguments will be merged with a dictionary of initial quantum + variables constructed by this method, and passed by keyword. An error is raised + if a keyworkd argument is provided that interferes with this. + + - During normal operation, all quantum and classical arguments will be + - Quantum arguments will always be passed by keyword, so you should probably put them last. + """ + + +@attrs.mutable +class _BloqifyPrepResult: + """Container for the results of tracing a function to build a Bloq.""" + + cbloq: 'CompositeBloq' + in_qargnames: Set[str] + out_qargnames: Set[str] + explicit_bb: Optional['BloqBuilder'] = None + found_bb: Optional['BloqBuilder'] = None + + @property + def bb(self) -> 'BloqBuilder': + if self.explicit_bb is not None: + return self.explicit_bb + if self.found_bb is not None: + return self.found_bb + + raise ValueError( + "Could not find a valid bloq builder object to qcall this function. " + "Please add an explicit `bb: BloqBuilder` argument to your `@bloqify` function." + ) + + +_BB_PLACEHOLDER = object() + + +class _TracingBloqIntermediate: + """Wrapper for a tracing function that can be used to construct or call a Bloq. + + This class provides methods to finalize the composition into a CompositeBloq + or to add the traced operations to another BloqBuilder. + """ + + def __init__(self, func: _TracingBloqFuncT): + sig = inspect.signature(func) + params = list(sig.parameters.values()) + if not params or params[0].name != 'bb': + raise ValueError( + f"Bloqified function '{func.__name__}' must take 'bb' as its first argument." + ) + self.func: _TracingBloqFuncT = func + + @property + def name(self) -> str: + return self.func.__name__ + + @property + def pkg(self) -> str: + return self.func.__module__ + + def _bound_kvs(self, *args, **kwargs): + yield from inspect.signature(self.func).bind( + _BB_PLACEHOLDER, *args, **kwargs + ).arguments.items() + + def _is_qvar_array(self, v): + if not isinstance(v, (Sequence, np.ndarray)): + return False + + try: + x = np.asarray(v).reshape(-1).item(0) + return isinstance(x, _QVar) + except (ValueError, IndexError): + return False + + def _prep_qstackframe(self, kv_iter): + bb = BloqBuilder(bloq_key=self.name, add_registers_allowed=True) + qkwargs = {} + ckwargs = {} + + for k, v in kv_iter: + if v is _BB_PLACEHOLDER: + continue + elif isinstance(v, _QVar): + qkwargs[k] = bb.in_register(name=k, dtype=v.dtype) + elif self._is_qvar_array(v): + v = np.asarray(v) + qkwargs[k] = bb.in_register( + name=k, dtype=v.reshape(-1).item(0).dtype, shape=v.shape + ) + else: + ckwargs[k] = v + + out_qvars = self.func(bb, **ckwargs, **qkwargs) + if not isinstance(out_qvars, dict): + raise ValueError( + f"{self.name} is expected to return a dictionary mapping " + f"output register name to output quantum variable." + ) + cbloq = bb.finalize(**out_qvars) + return _BloqifyPrepResult( + cbloq=cbloq, in_qargnames=set(qkwargs.keys()), out_qargnames=set(out_qvars.keys()) + ) + + def make( + self, signature: Optional['qlt.Signature'] = None, *classical_args, **classical_kwargs + ): + """Trace the function and finalize it into a CompositeBloq.""" + if signature is None: + signature = Signature([]) + bb, soqs = BloqBuilder.from_signature( + signature, bloq_key=self.name, add_registers_allowed=True + ) + + dupes = set(classical_kwargs.keys()) & set(soqs.keys()) + if dupes: + raise ValueError( + f"`make` called with keyword arguments that shadow quantum " + f"register names: {dupes}. Please do not provide quantum variables " + f"when calling `make`." + ) + + kwargs = classical_kwargs | soqs + soqs = self.func(bb, *classical_args, **kwargs) + return bb.finalize(**soqs) + + def __call__(self, bb: 'BloqBuilder', /, *args, **kwargs): + """Trace the function and add it as a sub-bloq to the provided builder.""" + f = self._prep_qstackframe(self._bound_kvs(*args, **kwargs)) + return bb.add( + f.cbloq, **{k: v for k, v in self._bound_kvs(*args, **kwargs) if k in f.in_qargnames} + ) + + def inline(self, bb: 'BloqBuilder', /, *args, **kwargs): + """Run the function directly on the provided builder without creating a sub-bloq.""" + ret_dict = self.func(bb, *args, **kwargs) + return tuple(ret_dict.values()) + + def dump_l1( + self, signature: Optional['qlt.Signature'] = None, *classical_args, **classical_kwargs + ): + """Trace the function and return its L1 representation.""" + from qualtran.l1 import dump_root_l1 + + return dump_root_l1(self.make(signature, *classical_args, **classical_kwargs)) + + def print_l1( + self, signature: Optional['qlt.Signature'] = None, *classical_args, **classical_kwargs + ): + """Trace the function and print its L1 representation.""" + print(self.dump_l1(signature, *classical_args, **classical_kwargs)) + + def draw( + self, signature: Optional['qlt.Signature'] = None, *classical_args, **classical_kwargs + ): + """Trace the function and draw the resulting CompositeBloq.""" + draw_type = classical_kwargs.pop("type", 'graph') + self.make(signature, *classical_args, **classical_kwargs).draw(type=draw_type) + + +def bloqify(func: _TracingBloqFuncT) -> _TracingBloqIntermediate: + """Decorator to enable tracing of a function that builds a Bloq. + + The decorated function must take a `BloqBuilder` as its first argument + and return a dictionary mapping output register names to quantum variables. + """ + return _TracingBloqIntermediate(func) diff --git a/qualtran/bloqify_syntax/_infra_test.py b/qualtran/bloqify_syntax/_infra_test.py new file mode 100644 index 000000000..4e4a0ea6a --- /dev/null +++ b/qualtran/bloqify_syntax/_infra_test.py @@ -0,0 +1,115 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +import inspect +from typing import Dict + +import pytest + +import qualtran as qlt +from qualtran import BloqBuilder, QVar + + +def func1(x, k, n=1): + pass + + +def call_func1_with_kwargs(*args, **kwargs): + stuff = inspect.signature(func1).bind(*args, **kwargs).arguments + return dict(stuff) + + +def test_inspect(): + a = call_func1_with_kwargs('xv', 'kv', n=2) + assert a == {'x': 'xv', 'k': 'kv', 'n': 2} + + # Note: we could use `.apply_defaults` if we wanted the default values + a = call_func1_with_kwargs('xv', 'kv') + assert a == {'x': 'xv', 'k': 'kv'} + + a = call_func1_with_kwargs('xv', n=2, k=1) + assert a == {'x': 'xv', 'n': 2, 'k': 1} + + +@qlt.bloqify +def minimal_bloq(bb: 'BloqBuilder', x: 'QVar') -> Dict[str, 'QVar']: + return {'x': x} + + +def test_bloqify_decorator(): + assert minimal_bloq.name == 'minimal_bloq' + assert minimal_bloq.pkg == 'qualtran.bloqify_syntax._infra_test' + + +def test_bloqify_make(): + sig = qlt.Signature.build(x=1) + cbloq = minimal_bloq.make(sig) + assert isinstance(cbloq, qlt.CompositeBloq) + assert cbloq.signature == sig + + +def test_bloqify_call(): + @qlt.bloqify + def outer_program(bb: 'BloqBuilder', x: 'QVar') -> Dict[str, 'QVar']: + x = minimal_bloq(bb, x=x) + return {'x': x} + + sig = qlt.Signature.build(x=1) + cbloq = outer_program.make(sig) + assert isinstance(cbloq, qlt.CompositeBloq) + + +def test_bloqify_inline(): + @qlt.bloqify + def outer_program(bb: 'BloqBuilder', x: 'QVar') -> Dict[str, 'QVar']: + x_out = minimal_bloq.inline(bb, x=x) + return {'x': x_out[0]} + + sig = qlt.Signature.build(x=1) + cbloq = outer_program.make(sig) + assert isinstance(cbloq, qlt.CompositeBloq) + + +def test_bloqify_l1(): + sig = qlt.Signature.build(x=1) + l1_str = minimal_bloq.dump_l1(sig) + assert isinstance(l1_str, str) + + +def test_bloqify_non_dict_return(): + @qlt.bloqify + def bad_bloq(bb: 'BloqBuilder', x: 'QVar'): + return x # Not a dict + + @qlt.bloqify + def outer_program(bb: 'BloqBuilder', x: 'QVar') -> Dict[str, 'QVar']: + x = bad_bloq(bb, x=x) + return {'x': x} + + sig = qlt.Signature.build(x=1) + with pytest.raises(ValueError, match="bad_bloq is expected to return a dictionary"): + outer_program.make(sig) + + +def test_bloqify_make_shadowing(): + sig = qlt.Signature.build(x=1) + with pytest.raises(ValueError, match="shadow quantum register names"): + minimal_bloq.make(sig, x=1) + + +def test_bloqify_missing_bb(): + with pytest.raises(ValueError, match="must take 'bb' as its first argument"): + + @qlt.bloqify # type: ignore[arg-type] + def no_bb_func(x: 'QVar'): + return {'x': x}