bloqade-circuit 0.4.5__py3-none-any.whl → 0.5.1__py3-none-any.whl

bloqade/pyqrack/squin/op.py

@@ -96,10 +96,15 @@ class PyQrackMethods(interp.MethodTable):
         return (PhaseOpRuntime(theta, global_=global_),)
 
     @interp.impl(op.stmts.Reset)
+    @interp.impl(op.stmts.ResetToOne)
     def reset(
-        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Reset
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.Reset | op.stmts.ResetToOne,
     ) -> tuple[OperatorRuntimeABC]:
-        return (ResetRuntime(),)
+        target_state = isinstance(stmt, op.stmts.ResetToOne)
+        return (ResetRuntime(target_state=target_state),)
 
     @interp.impl(op.stmts.X)
     @interp.impl(op.stmts.Y)

bloqade/pyqrack/squin/runtime.py

@@ -43,7 +43,9 @@ class OperatorRuntimeABC:
 
 @dataclass(frozen=True)
 class ResetRuntime(OperatorRuntimeABC):
-    """Reset the qubit to |0> state"""
+    """Reset the qubit to the target state"""
+
+    target_state: bool
 
     @property
     def n_sites(self) -> int:
@@ -55,7 +57,7 @@ class ResetRuntime(OperatorRuntimeABC):
                 continue
 
             res: bool = qubit.sim_reg.m(qubit.addr)
-            if res:
+            if res != self.target_state:
                 qubit.sim_reg.x(qubit.addr)
 
 

bloqade/qasm2/dialects/expr/stmts.py

@@ -1,34 +1,16 @@
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
+from kirin.dialects import func
 from kirin.print.printer import Printer
-from kirin.dialects.func.attrs import Signature
 
 from ._dialect import dialect
 
 
-class GateFuncOpCallableInterface(ir.CallableStmtInterface["GateFunction"]):
-
-    @classmethod
-    def get_callable_region(cls, stmt: "GateFunction") -> ir.Region:
-        return stmt.body
-
-
 @statement(dialect=dialect)
-class GateFunction(ir.Statement):
+class GateFunction(func.Function):
     """Special Function for qasm2 gate subroutine."""
 
     name = "gate.func"
-    traits = frozenset(
-        {
-            ir.IsolatedFromAbove(),
-            ir.SymbolOpInterface(),
-            ir.HasSignature(),
-            GateFuncOpCallableInterface(),
-        }
-    )
-    sym_name: str = info.attribute()
-    signature: Signature = info.attribute()
-    body: ir.Region = info.region(multi=True)
 
     def print_impl(self, printer: Printer) -> None:
         with printer.rich(style="red"):

bloqade/qasm2/parse/lowering.py

@@ -36,6 +36,7 @@ class QASM2(lowering.LoweringABC[ast.Node]):
             file=file,
             lineno_offset=lineno_offset,
             col_offset=col_offset,
+            compactify=compactify,
         )
 
         return frame.curr_region

bloqade/qasm2/passes/parallel.py

@@ -26,6 +26,7 @@ from bloqade.qasm2.rewrite import (
     ParallelToUOpRule,
     RaiseRegisterRule,
     UOpToParallelRule,
+    ParallelToGlobalRule,
     SimpleOptimalMergePolicy,
     RydbergGateSetRewriteRule,
 )
@@ -183,3 +184,20 @@ class UOpToParallel(Pass):
             CommonSubexpressionElimination(),
         )
         return Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+
+@dataclass
+class ParallelToGlobal(Pass):
+
+    def generate_rule(self, mt: ir.Method) -> ParallelToGlobalRule:
+        address_analysis = address.AddressAnalysis(mt.dialects)
+        frame, _ = address_analysis.run_analysis(mt)
+        return ParallelToGlobalRule(frame.entries)
+
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
+        rule = self.generate_rule(mt)
+
+        result = Walk(rule).rewrite(mt.code)
+        result = Walk(DeadCodeElimination()).rewrite(mt.code).join(result)
+
+        return result

bloqade/qasm2/passes/unroll_if.py

@@ -7,15 +7,22 @@ from kirin.rewrite import (
     ConstantFold,
     CommonSubexpressionElimination,
 )
+from kirin.dialects import scf, func
 
-from ..rewrite.split_ifs import LiftThenBody, SplitIfStmts
+from bloqade.rewrite.rules import LiftThenBody, SplitIfStmts
+
+from ..dialects.uop.stmts import SingleQubitGate, TwoQubitCtrlGate
+from ..dialects.core.stmts import Reset, Measure
+
+AllowedThenType = (SingleQubitGate, TwoQubitCtrlGate, Measure, Reset)
+DontLiftType = AllowedThenType + (scf.Yield, func.Return, func.Invoke)
 
 
 class UnrollIfs(Pass):
     """This pass lifts statements that are not UOP out of the if body and then splits whatever is left into multiple if statements so you obtain valid QASM2"""
 
     def unsafe_run(self, mt: ir.Method):
-        result = Walk(LiftThenBody()).rewrite(mt.code)
+        result = Walk(LiftThenBody(exclude_stmts=DontLiftType)).rewrite(mt.code)
         result = Walk(SplitIfStmts()).rewrite(mt.code).join(result)
         result = (
             Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))

bloqade/qasm2/rewrite/__init__.py

@@ -11,5 +11,6 @@ from .uop_to_parallel import (
     SimpleGreedyMergePolicy as SimpleGreedyMergePolicy,
     SimpleOptimalMergePolicy as SimpleOptimalMergePolicy,
 )
+from .parallel_to_glob import ParallelToGlobalRule as ParallelToGlobalRule
 from .noise.remove_noise import RemoveNoisePass as RemoveNoisePass
 from .noise.heuristic_noise import NoiseRewriteRule as NoiseRewriteRule

bloqade/qasm2/rewrite/parallel_to_glob.py

@@ -0,0 +1,82 @@
+from typing import Dict
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.rewrite import abc
+from kirin.analysis import const
+from kirin.dialects import ilist
+
+from bloqade.analysis import address
+
+from ..dialects import core, glob, parallel
+
+
+@dataclass
+class ParallelToGlobalRule(abc.RewriteRule):
+    address_analysis: Dict[ir.SSAValue, address.Address]
+
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
+        if not isinstance(node, parallel.UGate):
+            return abc.RewriteResult()
+
+        qargs = node.qargs
+        qarg_addresses = self.address_analysis.get(qargs, None)
+
+        if isinstance(qarg_addresses, address.AddressReg):
+            # NOTE: we only have an AddressReg if it's an entire register, definitely rewrite that
+            return self._rewrite_parallel_to_glob(node)
+
+        if not isinstance(qarg_addresses, address.AddressTuple):
+            return abc.RewriteResult()
+
+        idxs, qreg = self._find_qreg(qargs.owner, set())
+
+        if qreg is None:
+            # NOTE: no unique register found
+            return abc.RewriteResult()
+
+        if not isinstance(hint := qreg.n_qubits.hints.get("const"), const.Value):
+            # NOTE: non-constant number of qubits
+            return abc.RewriteResult()
+
+        n = hint.data
+        if len(idxs) != n:
+            # NOTE: not all qubits of the register are there
+            return abc.RewriteResult()
+
+        return self._rewrite_parallel_to_glob(node)
+
+    @staticmethod
+    def _rewrite_parallel_to_glob(node: parallel.UGate) -> abc.RewriteResult:
+        theta, phi, lam = node.theta, node.phi, node.lam
+        global_u = glob.UGate(node.qargs, theta=theta, phi=phi, lam=lam)
+        node.replace_by(global_u)
+        return abc.RewriteResult(has_done_something=True)
+
+    @staticmethod
+    def _find_qreg(
+        qargs_owner: ir.Statement | ir.Block, idxs: set
+    ) -> tuple[set, core.stmts.QRegNew | None]:
+
+        if isinstance(qargs_owner, core.stmts.QRegGet):
+            idxs.add(qargs_owner.idx)
+            qreg = qargs_owner.reg.owner
+            if not isinstance(qreg, core.stmts.QRegNew):
+                # NOTE: this could potentially be casted
+                qreg = None
+            return idxs, qreg
+
+        if isinstance(qargs_owner, ilist.New):
+            vals = qargs_owner.values
+            if len(vals) == 0:
+                return idxs, None
+
+            idxs, first_qreg = ParallelToGlobalRule._find_qreg(vals[0].owner, idxs)
+            for val in vals[1:]:
+                idxs, qreg = ParallelToGlobalRule._find_qreg(val.owner, idxs)
+                if qreg != first_qreg:
+                    return idxs, None
+
+            return idxs, first_qreg
+
+        return idxs, None

bloqade/rewrite/passes/__init__.py

@@ -0,0 +1 @@
+from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList

bloqade/rewrite/passes/canonicalize_ilist.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+)
+from kirin.analysis import const
+
+from ..rules.flatten_ilist import FlattenAddOpIList
+from ..rules.inline_getitem_ilist import InlineGetItemFromIList
+
+
+@dataclass
+class CanonicalizeIList(Pass):
+
+    def unsafe_run(self, mt: ir.Method):
+
+        cp_result_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
+
+        return Fixpoint(
+            Chain(
+                Walk(InlineGetItemFromIList(constprop_result=cp_result_frame.entries)),
+                Walk(FlattenAddOpIList()),
+            )
+        ).rewrite(mt.code)

bloqade/rewrite/rules/__init__.py

@@ -0,0 +1 @@
+from .split_ifs import LiftThenBody as LiftThenBody, SplitIfStmts as SplitIfStmts

bloqade/rewrite/rules/flatten_ilist.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import py, ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class FlattenAddOpIList(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.binop.Add):
+            return RewriteResult()
+
+        # check if we are adding two ilist.New objects
+        new_data = ()
+
+        # lhs:
+        if not isinstance(node.lhs.owner, ilist.New):
+            if not (
+                isinstance(node.lhs.owner, py.Constant)
+                and isinstance(
+                    const_ilist := node.lhs.owner.value.unwrap(), ilist.IList
+                )
+                and len(const_ilist.data) == 0
+            ):
+                return RewriteResult()
+
+        else:
+            new_data += node.lhs.owner.values
+
+        # rhs:
+        if not isinstance(node.rhs.owner, ilist.New):
+            if not (
+                isinstance(node.rhs.owner, py.Constant)
+                and isinstance(
+                    const_ilist := node.rhs.owner.value.unwrap(), ilist.IList
+                )
+                and len(const_ilist.data) == 0
+            ):
+                return RewriteResult()
+
+        else:
+            new_data += node.rhs.owner.values
+
+        new_stmt = ilist.New(values=new_data)
+        node.replace_by(new_stmt)
+
+        return RewriteResult(
+            has_done_something=True,
+        )

bloqade/rewrite/rules/inline_getitem_ilist.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.analysis import const
+from kirin.dialects import py, ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class InlineGetItemFromIList(RewriteRule):
+    constprop_result: dict[ir.SSAValue, const.Result]
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.indexing.GetItem):
+            return RewriteResult()
+
+        if not isinstance(node.obj.owner, ilist.New):
+            return RewriteResult()
+
+        if not isinstance(
+            index_value := self.constprop_result.get(node.index), const.Value
+        ):
+            return RewriteResult()
+
+        elem_ssa = node.obj.owner.values[index_value.data]
+
+        node.result.replace_by(elem_ssa)
+
+        return RewriteResult(
+            has_done_something=True,
+        )

bloqade/{qasm2/rewrite → rewrite/rules}/split_ifs.py

@@ -1,18 +1,23 @@
+from dataclasses import field, dataclass
+
 from kirin import ir
 from kirin.dialects import scf, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from ..dialects.uop.stmts import SingleQubitGate, TwoQubitCtrlGate
-from ..dialects.core.stmts import Reset, Measure
 
-# TODO: unify with PR #248
-AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | Reset
+@dataclass
+class LiftThenBody(RewriteRule):
+    """
+    Lifts anything that's not in the `exclude_stmts` in the *then* body
+
 
-DontLiftType = AllowedThenType | scf.Yield | func.Return | func.Invoke
+    Args:
+        exclude_stmts: A tuple of statement types that should not be lifted from the then body.
+            Defaults to an empty tuple, meaning all statements are lifted.
 
+    """
 
-class LiftThenBody(RewriteRule):
-    """Lifts anything that's not a UOP or a yield/return out of the then body"""
+    exclude_stmts: tuple[type[ir.Statement], ...] = field(default_factory=tuple)
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         if not isinstance(node, scf.IfElse):
@@ -20,7 +25,9 @@ class LiftThenBody(RewriteRule):
 
         then_stmts = node.then_body.stmts()
 
-        lift_stmts = [stmt for stmt in then_stmts if not isinstance(stmt, DontLiftType)]
+        lift_stmts = [
+            stmt for stmt in then_stmts if not isinstance(stmt, self.exclude_stmts)
+        ]
 
         if len(lift_stmts) == 0:
             return RewriteResult()

bloqade/squin/__init__.py

@@ -3,7 +3,9 @@ from . import (
     wire as wire,
     noise as noise,
     qubit as qubit,
+    analysis as analysis,
     lowering as lowering,
+    _typeinfer as _typeinfer,
 )
 from .groups import wired as wired, kernel as kernel
 

bloqade/squin/_typeinfer.py

@@ -0,0 +1,20 @@
+from kirin import types, interp
+from kirin.analysis import TypeInference, const
+from kirin.dialects import ilist
+
+from bloqade import squin
+
+
+@squin.qubit.dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+    @interp.impl(squin.qubit.New)
+    def _call(self, interp: TypeInference, frame: interp.Frame, stmt: squin.qubit.New):
+        # based on Xiu-zhe (Roger) Luo's get_const_value function
+
+        if (hint := stmt.n_qubits.hints.get("const")) is None:
+            return (ilist.IListType[squin.qubit.QubitType, types.Any],)
+
+        if isinstance(hint, const.Value) and isinstance(hint.data, int):
+            return (ilist.IListType[squin.qubit.QubitType, types.Literal(hint.data)],)
+
+        return (ilist.IListType[squin.qubit.QubitType, types.Any],)

bloqade/squin/analysis/__init__.py

@@ -0,0 +1 @@
+from . import address_impl as address_impl

bloqade/squin/analysis/address_impl.py

@@ -0,0 +1,71 @@
+from kirin import interp
+from kirin.analysis import ForwardFrame
+
+from bloqade.analysis.address.lattice import (
+    Address,
+    AddressReg,
+    AddressWire,
+    AddressQubit,
+)
+from bloqade.analysis.address.analysis import AddressAnalysis
+
+from .. import wire, qubit
+
+# Address lattice elements we can work with:
+## NotQubit (bottom), AnyAddress (top)
+
+## AddressTuple -> data: tuple[Address, ...]
+### Recursive type, could contain itself or other variants
+### This pops up in cases where you can have an IList/Tuple
+### That contains elements that could be other Address types
+
+## AddressReg -> data: Sequence[int]
+### specific to creation of a register of qubits
+
+## AddressQubit -> data: int
+### Base qubit address type
+
+
+@wire.dialect.register(key="qubit.address")
+class SquinWireMethodTable(interp.MethodTable):
+
+    @interp.impl(wire.Unwrap)
+    def unwrap(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: wire.Unwrap,
+    ):
+
+        origin_qubit = frame.get(stmt.qubit)
+
+        if isinstance(origin_qubit, AddressQubit):
+            return (AddressWire(origin_qubit=origin_qubit),)
+        else:
+            return (Address.top(),)
+
+    @interp.impl(wire.Apply)
+    def apply(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: wire.Apply,
+    ):
+        return frame.get_values(stmt.inputs)
+
+
+@qubit.dialect.register(key="qubit.address")
+class SquinQubitMethodTable(interp.MethodTable):
+
+    # This can be treated like a QRegNew impl
+    @interp.impl(qubit.New)
+    def new(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: qubit.New,
+    ):
+        n_qubits = interp_.get_const_value(int, stmt.n_qubits)
+        addr = AddressReg(range(interp_.next_address, interp_.next_address + n_qubits))
+        interp_.next_address += n_qubits
+        return (addr,)

bloqade/squin/analysis/nsites/impls.py

@@ -1,5 +1,5 @@
 from kirin import interp
-from kirin.dialects import scf
+from kirin.dialects import scf, func
 from kirin.dialects.scf.typeinfer import TypeInfer as ScfTypeInfer
 
 from bloqade.squin import op, wire
@@ -85,3 +85,8 @@ class SquinOp(interp.MethodTable):
 @scf.dialect.register(key="op.nsites")
 class ScfSquinOp(ScfTypeInfer):
     pass
+
+
+@func.dialect.register(key="op.nsites")
+class FuncSquinOp(func.typeinfer.TypeInfer):
+    pass

bloqade/squin/cirq/lowering.py

@@ -368,11 +368,24 @@ class Squin(lowering.LoweringABC[CirqNode]):
         state: lowering.State[CirqNode],
         node: cirq.GeneralizedAmplitudeDampingChannel,
     ):
-        raise NotImplementedError("TODO: needs a new operator statement")
-        # p = state.current_frame.push(py.Constant(node.p))
-        # gamma = state.current_frame.push(py.Constant(node.gamma))
+        p = state.current_frame.push(py.Constant(node.p)).result
+        gamma = state.current_frame.push(py.Constant(node.gamma)).result
 
-        # p1 =
+        # NOTE: cirq has a weird convention here: if p == 1, we have AmplitudeDampingChannel,
+        # which basically means p is the probability of the environment being in the vacuum state
+        prob0 = state.current_frame.push(py.binop.Mult(p, gamma)).result
+        one_ = state.current_frame.push(py.Constant(1)).result
+        p_minus_1 = state.current_frame.push(py.binop.Sub(one_, p)).result
+        prob1 = state.current_frame.push(py.binop.Mult(p_minus_1, gamma)).result
 
-        # x = state.current_frame.push(op.stmts.X())
-        # noise_channel1 = noise.stmts.PauliError(basis=x.result, p=)
+        r0 = state.current_frame.push(op.stmts.Reset()).result
+        r1 = state.current_frame.push(op.stmts.ResetToOne()).result
+
+        probs = state.current_frame.push(ilist.New(values=(prob0, prob1))).result
+        ops = state.current_frame.push(ilist.New(values=(r0, r1))).result
+
+        noise_channel = state.current_frame.push(
+            noise.stmts.StochasticUnitaryChannel(probabilities=probs, operators=ops)
+        )
+
+        return noise_channel

bloqade/squin/noise/stmts.py

@@ -10,7 +10,7 @@ from ..op.types import NumOperators
 
 @statement
 class NoiseChannel(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     result: ir.ResultValue = info.result(OpType)
 
 

bloqade/squin/op/__init__.py

@@ -37,4 +37,5 @@ from ._wrapper import (
     control as control,
     identity as identity,
     pauli_string as pauli_string,
+    reset_to_one as reset_to_one,
 )

bloqade/squin/op/_wrapper.py

@@ -41,6 +41,10 @@ def control(op: types.Op, *, n_controls: int) -> types.Op:
 def reset() -> types.Op: ...
 
 
+@wraps(stmts.ResetToOne)
+def reset_to_one() -> types.Op: ...
+
+
 @wraps(stmts.Identity)
 def identity(*, sites: int) -> types.Op: ...
 

bloqade/squin/op/stmts.py

@@ -98,6 +98,15 @@ class ConstantUnitary(ConstantOp):
 
 @statement(dialect=dialect)
 class U3(PrimitiveOp):
+    """
+    The rotation operator U3(theta, phi, lam).
+    Note that we use the convention from the QASM2 specification, namely
+
+    $$
+    U_3(\theta, \phi, \lambda) = R_z(\phi) R_y(\theta) R_z(\lambda)
+    $$
+    """
+
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
     phi: ir.SSAValue = info.argument(types.Float)
@@ -110,7 +119,7 @@ class PhaseOp(PrimitiveOp):
     A phase operator.
 
     $$
-    PhaseOp(theta) = e^{i \theta} I
+    PhaseOp(\theta) = e^{i \theta} I
     $$
     """
 
@@ -124,7 +133,7 @@ class ShiftOp(PrimitiveOp):
     A phase shift operator.
 
     $$
-    Shift(theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
+    Shift(\theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
     $$
     """
 
@@ -141,6 +150,15 @@ class Reset(PrimitiveOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
 
 
+@statement(dialect=dialect)
+class ResetToOne(PrimitiveOp):
+    """
+    Reset qubits to the one state. Mainly needed to accommodate cirq's GeneralizedAmplitudeDampingChannel
+    """
+
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
+
+
 @statement
 class CliffordOp(ConstantUnitary):
     pass

bloqade/squin/qubit.py

@@ -17,6 +17,7 @@ from kirin.lowering import wraps
 from bloqade.types import Qubit, QubitType
 from bloqade.squin.op.types import Op, OpType
 
+from .types import MeasurementResult, MeasurementResultType
 from .lowering import ApplyAnyCallLowering
 
 dialect = ir.Dialect("squin.qubit")
@@ -65,8 +66,8 @@ class MeasureQubit(ir.Statement):
     name = "measure.qubit"
 
     traits = frozenset({lowering.FromPythonCall()})
-    qubit: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
+    qubit: ir.SSAValue = info.argument(QubitType)
+    result: ir.ResultValue = info.result(MeasurementResultType)
 
 
 @statement(dialect=dialect)
@@ -75,7 +76,7 @@ class MeasureQubitList(ir.Statement):
 
     traits = frozenset({lowering.FromPythonCall()})
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
+    result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType])
 
 
 # NOTE: no dependent types in Python, so we have to mark it Any...
@@ -131,9 +132,11 @@ def apply(operator: Op, *qubits) -> None: ...
 
 
 @overload
-def measure(input: Qubit) -> bool: ...
+def measure(input: Qubit) -> MeasurementResult: ...
 @overload
-def measure(input: ilist.IList[Qubit, Any] | list[Qubit]) -> ilist.IList[bool, Any]: ...
+def measure(
+    input: ilist.IList[Qubit, Any] | list[Qubit],
+) -> ilist.IList[MeasurementResult, Any]: ...
 
 
 @wraps(MeasureAny)

bloqade/squin/rewrite/__init__.py

@@ -4,4 +4,5 @@ from .wrap_analysis import (
     WrapOpSiteAnalysis as WrapOpSiteAnalysis,
     WrapAddressAnalysis as WrapAddressAnalysis,
 )
+from .U3_to_clifford import SquinU3ToClifford as SquinU3ToClifford
 from .remove_dangling_qubits import RemoveDeadRegister as RemoveDeadRegister