bloqade-circuit 0.4.4__py3-none-any.whl → 0.5.0__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/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/squin/__init__.py

@@ -4,6 +4,7 @@ from . import (
     noise as noise,
     qubit as qubit,
     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/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/__init__.py

@@ -9,8 +9,9 @@ from . import lowering as lowering
 from .. import kernel
 
 # NOTE: just to register methods
-from .emit import op as op, qubit as qubit
+from .emit import op as op, noise as noise, qubit as qubit
 from .lowering import Squin
+from ..noise.rewrite import RewriteNoiseStmts
 from .emit.emit_circuit import EmitCirq
 
 
@@ -18,6 +19,9 @@ def load_circuit(
     circuit: cirq.Circuit,
     kernel_name: str = "main",
     dialects: ir.DialectGroup = kernel,
+    register_as_argument: bool = False,
+    return_register: bool = False,
+    register_argument_name: str = "q",
     globals: dict[str, Any] | None = None,
     file: str | None = None,
     lineno_offset: int = 0,
@@ -32,13 +36,23 @@ def load_circuit(
     Keyword Args:
         kernel_name (str): The name of the kernel to load. Defaults to "main".
         dialects (ir.DialectGroup | None): The dialects to use. Defaults to `squin.kernel`.
+        register_as_argument (bool): Determine whether the resulting kernel function should accept
+            a single `ilist.IList[Qubit, Any]` argument that is a list of qubits used within the
+            function. This allows you to compose kernel functions generated from circuits.
+            Defaults to `False`.
+        return_register (bool): Determine whether the resulting kernel functionr returns a
+            single value of type `ilist.IList[Qubit, Any]` that is the list of qubits used
+            in the kernel function. Useful when you want to compose multiple kernel functions
+            generated from circuits. Defaults to `False`.
+        register_argument_name (str): The name of the argument that represents the qubit register.
+            Only used when `register_as_argument=True`. Defaults to "q".
         globals (dict[str, Any] | None): The global variables to use. Defaults to None.
         file (str | None): The file name for error reporting. Defaults to None.
         lineno_offset (int): The line number offset for error reporting. Defaults to 0.
         col_offset (int): The column number offset for error reporting. Defaults to 0.
         compactify (bool): Whether to compactify the output. Defaults to True.
 
-    Example:
+    ## Usage Examples:
 
     ```python
     # from cirq's "hello qubit" example
@@ -60,6 +74,30 @@ def load_circuit(
     # print the resulting IR
     main.print()
     ```
+
+    You can also compose kernel functions generated from circuits by passing in
+    and / or returning the respective quantum registers:
+
+    ```python
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.H(q[0]), cirq.CX(*q))
+
+    get_entangled_qubits = squin.cirq.load_circuit(
+        circuit, return_register=True, kernel_name="get_entangled_qubits"
+    )
+    get_entangled_qubits.print()
+
+    entangle_qubits = squin.cirq.load_circuit(
+        circuit, register_as_argument=True, kernel_name="entangle_qubits"
+    )
+
+    @squin.kernel
+    def main():
+        qreg = get_entangled_qubits()
+        qreg2 = squin.qubit.new(1)
+        entangle_qubits([qreg[1], qreg2[0]])
+        return squin.qubit.measure(qreg2)
+    ```
     """
 
     target = Squin(dialects=dialects, circuit=circuit)
@@ -71,16 +109,38 @@ def load_circuit(
         lineno_offset=lineno_offset,
         col_offset=col_offset,
         compactify=compactify,
+        register_as_argument=register_as_argument,
+        register_argument_name=register_argument_name,
     )
 
-    # NOTE: no return value
-    return_value = func.ConstantNone()
-    body.blocks[0].stmts.append(return_value)
-    body.blocks[0].stmts.append(func.Return(value_or_stmt=return_value))
+    if return_register:
+        return_value = target.qreg
+    else:
+        return_value = func.ConstantNone()
+        body.blocks[0].stmts.append(return_value)
+
+    return_node = func.Return(value_or_stmt=return_value)
+    body.blocks[0].stmts.append(return_node)
+
+    self_arg_name = kernel_name + "_self"
+    arg_names = [self_arg_name]
+    if register_as_argument:
+        args = (target.qreg.type,)
+        arg_names.append(register_argument_name)
+    else:
+        args = ()
+
+    # NOTE: add _self as argument; need to know signature before so do it after lowering
+    signature = func.Signature(args, return_node.value.type)
+    body.blocks[0].args.insert_from(
+        0,
+        types.Generic(ir.Method, types.Tuple.where(signature.inputs), signature.output),
+        self_arg_name,
+    )
 
     code = func.Function(
         sym_name=kernel_name,
-        signature=func.Signature((), types.NoneType),
+        signature=signature,
         body=body,
     )
 
@@ -88,7 +148,7 @@ def load_circuit(
         mod=None,
         py_func=None,
         sym_name=kernel_name,
-        arg_names=[],
+        arg_names=arg_names,
         dialects=dialects,
         code=code,
     )
@@ -176,7 +236,12 @@ def emit_circuit(
         )
 
     emitter = EmitCirq(qubits=qubits)
-    return emitter.run(mt, args=())
+
+    # Rewrite noise statements
+    mt_ = mt.similar(mt.dialects)
+    RewriteNoiseStmts(mt_.dialects)(mt_)
+
+    return emitter.run(mt_, args=())
 
 
 def dump_circuit(mt: ir.Method, qubits: Sequence[cirq.Qid] | None = None, **kwargs):

bloqade/squin/cirq/emit/noise.py

@@ -0,0 +1,49 @@
+import cirq
+from kirin.emit import EmitError
+from kirin.interp import MethodTable, impl
+
+from ... import noise
+from .runtime import (
+    KronRuntime,
+    BasicOpRuntime,
+    OperatorRuntimeABC,
+    PauliStringRuntime,
+)
+from .emit_circuit import EmitCirq, EmitCirqFrame
+
+
+@noise.dialect.register(key="emit.cirq")
+class EmitCirqNoiseMethods(MethodTable):
+
+    @impl(noise.stmts.StochasticUnitaryChannel)
+    def stochastic_unitary_channel(
+        self,
+        emit: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: noise.stmts.StochasticUnitaryChannel,
+    ):
+        ops = frame.get(stmt.operators)
+        ps = frame.get(stmt.probabilities)
+
+        error_probabilities = {self._op_to_key(op_): p for op_, p in zip(ops, ps)}
+        cirq_op = cirq.asymmetric_depolarize(error_probabilities=error_probabilities)
+        return (BasicOpRuntime(cirq_op),)
+
+    @staticmethod
+    def _op_to_key(operator: OperatorRuntimeABC) -> str:
+        match operator:
+            case KronRuntime():
+                key_lhs = EmitCirqNoiseMethods._op_to_key(operator.lhs)
+                key_rhs = EmitCirqNoiseMethods._op_to_key(operator.rhs)
+                return key_lhs + key_rhs
+
+            case BasicOpRuntime():
+                return str(operator.gate)
+
+            case PauliStringRuntime():
+                return operator.string
+
+            case _:
+                raise EmitError(
+                    f"Unexpected operator runtime in StochasticUnitaryChannel of type {type(operator).__name__} encountered!"
+                )

bloqade/squin/cirq/emit/runtime.py

@@ -21,7 +21,10 @@ class OperatorRuntimeABC:
 
     def unsafe_apply(
         self, qubits: Sequence[cirq.Qid], adjoint: bool = False
-    ) -> list[cirq.Operation]: ...
+    ) -> list[cirq.Operation]:
+        raise NotImplementedError(
+            f"Apply method needs to be implemented in {self.__class__.__name__}"
+        )
 
 
 @dataclass
@@ -38,6 +41,11 @@ class BasicOpRuntime(UnsafeOperatorRuntimeABC):
     def num_qubits(self) -> int:
         return self.gate.num_qubits()
 
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        return [self.gate(*qubits)]
+
 
 @dataclass
 class UnitaryRuntime(BasicOpRuntime):

bloqade/squin/cirq/lowering.py

@@ -3,7 +3,7 @@ from typing import Any
 from dataclasses import field, dataclass
 
 import cirq
-from kirin import ir, lowering
+from kirin import ir, types, lowering
 from kirin.rewrite import Walk, CFGCompactify
 from kirin.dialects import py, scf, ilist
 
@@ -25,27 +25,26 @@ class Squin(lowering.LoweringABC[CirqNode]):
     """Lower a cirq.Circuit object to a squin kernel"""
 
     circuit: cirq.Circuit
-    qreg: qubit.New = field(init=False)
+    qreg: ir.SSAValue = field(init=False)
     qreg_index: dict[cirq.Qid, int] = field(init=False, default_factory=dict)
     next_qreg_index: int = field(init=False, default=0)
 
-    def lower_qubit_getindex(self, state: lowering.State[CirqNode], qid: cirq.Qid):
-        index = self.qreg_index.get(qid)
-
-        if index is None:
-            index = self.next_qreg_index
-            self.qreg_index[qid] = index
-            self.next_qreg_index += 1
+    def __post_init__(self):
+        # TODO: sort by cirq ordering
+        qbits = sorted(self.circuit.all_qubits())
+        self.qreg_index = {qid: idx for (idx, qid) in enumerate(qbits)}
 
+    def lower_qubit_getindex(self, state: lowering.State[CirqNode], qid: cirq.Qid):
+        index = self.qreg_index[qid]
         index_ssa = state.current_frame.push(py.Constant(index)).result
-        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg.result, index_ssa))
+        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg, index_ssa))
         return qbit_getitem.result
 
     def lower_qubit_getindices(
         self, state: lowering.State[CirqNode], qids: list[cirq.Qid]
     ):
         qbits_getitem = [self.lower_qubit_getindex(state, qid) for qid in qids]
-        qbits_stmt = ilist.New(values=qbits_getitem)
+        qbits_stmt = ilist.New(values=qbits_getitem, elem_type=qubit.QubitType)
         qbits_result = state.current_frame.get(qbits_stmt.name)
 
         if qbits_result is not None:
@@ -64,6 +63,8 @@ class Squin(lowering.LoweringABC[CirqNode]):
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
+        register_as_argument: bool = False,
+        register_argument_name: str = "q",
     ) -> ir.Region:
 
         state = lowering.State(
@@ -73,16 +74,21 @@ class Squin(lowering.LoweringABC[CirqNode]):
             col_offset=col_offset,
         )
 
-        with state.frame(
-            [stmt],
-            globals=globals,
-            finalize_next=False,
-        ) as frame:
-            # NOTE: create a global register of qubits first
-            # TODO: can there be a circuit without qubits?
-            n_qubits = cirq.num_qubits(self.circuit)
-            n = frame.push(py.Constant(n_qubits))
-            self.qreg = frame.push(qubit.New(n_qubits=n.result))
+        with state.frame([stmt], globals=globals, finalize_next=False) as frame:
+
+            # NOTE: need a register of qubits before lowering statements
+            if register_as_argument:
+                # NOTE: register as argument to the kernel; we have freedom of choice for the name here
+                frame.curr_block.args.append_from(
+                    ilist.IListType[qubit.QubitType, types.Any],
+                    name=register_argument_name,
+                )
+                self.qreg = frame.curr_block.args[0]
+            else:
+                # NOTE: create a new register of appropriate size
+                n_qubits = len(self.qreg_index)
+                n = frame.push(py.Constant(n_qubits))
+                self.qreg = frame.push(qubit.New(n_qubits=n.result)).result
 
             self.visit(state, stmt)
 
@@ -362,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/_wrapper.py

@@ -1,3 +1,6 @@
+from typing import Literal
+
+from kirin.dialects import ilist
 from kirin.lowering import wraps
 
 from bloqade.squin.op.types import Op
@@ -18,11 +21,15 @@ def depolarize(p: float) -> Op: ...
 
 
 @wraps(stmts.SingleQubitPauliChannel)
-def single_qubit_pauli_channel(params: tuple[float, float, float]) -> Op: ...
+def single_qubit_pauli_channel(
+    params: ilist.IList[float, Literal[3]] | list[float] | tuple[float, float, float],
+) -> Op: ...
 
 
 @wraps(stmts.TwoQubitPauliChannel)
-def two_qubit_pauli_channel(params: tuple[float, ...]) -> Op: ...
+def two_qubit_pauli_channel(
+    params: ilist.IList[float, Literal[15]] | list[float] | tuple[float, ...],
+) -> Op: ...
 
 
 @wraps(stmts.QubitLoss)

bloqade/squin/noise/rewrite.py

@@ -58,12 +58,12 @@ class _RewriteNoiseStmts(RewriteRule):
     def rewrite_two_qubit_pauli_channel(
         self, node: TwoQubitPauliChannel
     ) -> RewriteResult:
-        paulis = (X(), Y(), Z(), Identity(sites=1))
+        paulis = (Identity(sites=1), X(), Y(), Z())
         for op in paulis:
             op.insert_before(node)
 
-        # NOTE: collect list so we can skip the last entry, which will be two identities
-        combinations = list(itertools.product(paulis, repeat=2))[:-1]
+        # NOTE: collect list so we can skip the first entry, which will be two identities
+        combinations = list(itertools.product(paulis, repeat=2))[1:]
         operators: list[ir.SSAValue] = []
         for pauli_1, pauli_2 in combinations:
             op = Kron(pauli_1.result, pauli_2.result)

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: ...