bloqade-circuit 0.7.12__py3-none-any.whl → 0.8.0__py3-none-any.whl

bloqade/squin/cirq/lowering.py

@@ -1,439 +0,0 @@
-import math
-from typing import Any
-from dataclasses import field, dataclass
-
-import cirq
-from kirin import ir, types, lowering
-from kirin.rewrite import Walk, CFGCompactify
-from kirin.dialects import py, scf, ilist
-
-from .. import op, noise, qubit
-
-CirqNode = cirq.Circuit | cirq.Moment | cirq.Gate | cirq.Qid | cirq.Operation
-
-DecomposeNode = (
-    cirq.SwapPowGate
-    | cirq.ISwapPowGate
-    | cirq.PhasedXPowGate
-    | cirq.PhasedXZGate
-    | cirq.CSwapGate
-)
-
-
-@dataclass
-class Squin(lowering.LoweringABC[CirqNode]):
-    """Lower a cirq.Circuit object to a squin kernel"""
-
-    circuit: cirq.Circuit
-    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 __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, 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]
-        return tuple(qbits_getitem)
-
-    def run(
-        self,
-        stmt: CirqNode,
-        *,
-        source: str | None = None,
-        globals: dict[str, Any] | None = None,
-        file: str | None = None,
-        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(
-            self,
-            file=file,
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-        )
-
-        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)
-
-            if compactify:
-                Walk(CFGCompactify()).rewrite(frame.curr_region)
-
-            region = frame.curr_region
-
-        return region
-
-    def visit(self, state: lowering.State[CirqNode], node: CirqNode) -> lowering.Result:
-        name = node.__class__.__name__
-        return getattr(self, f"visit_{name}", self.generic_visit)(state, node)
-
-    def generic_visit(self, state: lowering.State[CirqNode], node: CirqNode):
-        if isinstance(node, CirqNode):
-            raise lowering.BuildError(
-                f"Cannot lower {node.__class__.__name__} node: {node}"
-            )
-        raise lowering.BuildError(
-            f"Unexpected `{node.__class__.__name__}` node: {repr(node)} is not an AST node"
-        )
-
-    def lower_literal(self, state: lowering.State[CirqNode], value) -> ir.SSAValue:
-        raise lowering.BuildError("Literals not supported in cirq circuit")
-
-    def lower_global(
-        self, state: lowering.State[CirqNode], node: CirqNode
-    ) -> lowering.LoweringABC.Result:
-        raise lowering.BuildError("Literals not supported in cirq circuit")
-
-    def visit_Circuit(
-        self, state: lowering.State[CirqNode], node: cirq.Circuit
-    ) -> lowering.Result:
-        for moment in node:
-            state.lower(moment)
-
-    def visit_Moment(
-        self, state: lowering.State[CirqNode], node: cirq.Moment
-    ) -> lowering.Result:
-        for op_ in node.operations:
-            state.lower(op_)
-
-    def visit_GateOperation(
-        self, state: lowering.State[CirqNode], node: cirq.GateOperation
-    ):
-        if isinstance(node.gate, cirq.MeasurementGate):
-            # NOTE: special dispatch here, since measurement is a gate + a qubit in cirq,
-            # but a single statement in squin
-            return self.lower_measurement(state, node)
-
-        if isinstance(node.gate, DecomposeNode):
-            # NOTE: easier to decompose these, but for that we need the qubits too,
-            # so we need to do this within this method
-            for subnode in cirq.decompose_once(node):
-                state.lower(subnode)
-            return
-
-        op_ = state.lower(node.gate).expect_one()
-        qbits = self.lower_qubit_getindices(state, node.qubits)
-        return state.current_frame.push(qubit.Apply(operator=op_, qubits=qbits))
-
-    def visit_TaggedOperation(
-        self, state: lowering.State[CirqNode], node: cirq.TaggedOperation
-    ):
-        state.lower(node.untagged)
-
-    def lower_measurement(
-        self, state: lowering.State[CirqNode], node: cirq.GateOperation
-    ):
-        if len(node.qubits) == 1:
-            qbit = self.lower_qubit_getindex(state, node.qubits[0])
-            stmt = state.current_frame.push(qubit.MeasureQubit(qbit))
-        else:
-            qbits = self.lower_qubit_getindices(state, node.qubits)
-            qbits_list = state.current_frame.push(ilist.New(values=qbits))
-            stmt = state.current_frame.push(qubit.MeasureQubitList(qbits_list.result))
-
-        key = node.gate.key
-        if isinstance(key, cirq.MeasurementKey):
-            key = key.name
-
-        state.current_frame.defs[key] = stmt.result
-        return stmt
-
-    def visit_ClassicallyControlledOperation(
-        self, state: lowering.State[CirqNode], node: cirq.ClassicallyControlledOperation
-    ):
-        conditions: list[ir.SSAValue] = []
-        for outcome in node.classical_controls:
-            key = outcome.key
-            if isinstance(key, cirq.MeasurementKey):
-                key = key.name
-            measurement_outcome = state.current_frame.defs[key]
-
-            if measurement_outcome.type.is_subseteq(ilist.IListType):
-                # NOTE: there is currently no convenient ilist.any method, so we need to use foldl
-                # with a simple function that just does an or
-
-                def bool_op_or(x: bool, y: bool) -> bool:
-                    return x or y
-
-                f_code = state.current_frame.push(
-                    lowering.Python(self.dialects).python_function(bool_op_or)
-                )
-                fn = ir.Method(
-                    mod=None,
-                    py_func=bool_op_or,
-                    sym_name="bool_op_or",
-                    arg_names=[],
-                    dialects=self.dialects,
-                    code=f_code,
-                )
-                f_const = state.current_frame.push(py.constant.Constant(fn))
-                init_val = state.current_frame.push(py.Constant(False)).result
-                condition = state.current_frame.push(
-                    ilist.Foldl(f_const.result, measurement_outcome, init=init_val)
-                ).result
-            else:
-                condition = measurement_outcome
-
-            conditions.append(condition)
-
-        if len(conditions) == 1:
-            condition = conditions[0]
-        else:
-            condition = state.current_frame.push(
-                py.boolop.And(conditions[0], conditions[1])
-            ).result
-            for next_cond in conditions[2:]:
-                condition = state.current_frame.push(
-                    py.boolop.And(condition, next_cond)
-                ).result
-
-        then_stmt = self.visit(state, node.without_classical_controls())
-
-        assert isinstance(
-            then_stmt, ir.Statement
-        ), f"Expected operation of classically controlled node {node} to be lowered to a statement, got type {type(then_stmt)}. \
-        Please report this issue!"
-
-        # NOTE: remove stmt from parent block
-        then_stmt.detach()
-        then_body = ir.Block((then_stmt,))
-
-        return state.current_frame.push(scf.IfElse(condition, then_body=then_body))
-
-    def visit_SingleQubitPauliStringGateOperation(
-        self,
-        state: lowering.State[CirqNode],
-        node: cirq.SingleQubitPauliStringGateOperation,
-    ):
-
-        match node.pauli:
-            case cirq.X:
-                op_ = op.stmts.X()
-            case cirq.Y:
-                op_ = op.stmts.Y()
-            case cirq.Z:
-                op_ = op.stmts.Z()
-            case cirq.I:
-                op_ = op.stmts.Identity(sites=1)
-            case _:
-                raise lowering.BuildError(f"Unexpected Pauli operation {node.pauli}")
-
-        state.current_frame.push(op_)
-        qargs = self.lower_qubit_getindices(state, [node.qubit])
-        return state.current_frame.push(qubit.Apply(op_.result, qargs))
-
-    def visit_HPowGate(self, state: lowering.State[CirqNode], node: cirq.HPowGate):
-        if abs(node.exponent) == 1:
-            return state.current_frame.push(op.stmts.H())
-
-        # NOTE: decompose into products of paulis for arbitrary exponents according to _decompose_ method
-        # can't use decompose directly since that method requires qubits to be passed in for some reason
-        y_rhs = state.lower(cirq.YPowGate(exponent=0.25)).expect_one()
-        x = state.lower(
-            cirq.XPowGate(exponent=node.exponent, global_shift=node.global_shift)
-        ).expect_one()
-        y_lhs = state.lower(cirq.YPowGate(exponent=-0.25)).expect_one()
-
-        # NOTE: reversed order since we're creating a mult stmt
-        m_lhs = state.current_frame.push(op.stmts.Mult(y_lhs, x))
-        return state.current_frame.push(op.stmts.Mult(m_lhs.result, y_rhs))
-
-    def visit_XPowGate(self, state: lowering.State[CirqNode], node: cirq.XPowGate):
-        if abs(node.exponent == 1):
-            return state.current_frame.push(op.stmts.X())
-
-        return self.visit(state, node.in_su2())
-
-    def visit_YPowGate(self, state: lowering.State[CirqNode], node: cirq.YPowGate):
-        if abs(node.exponent == 1):
-            return state.current_frame.push(op.stmts.Y())
-
-        return self.visit(state, node.in_su2())
-
-    def visit_ZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZPowGate):
-        if node.exponent == 0.5:
-            return state.current_frame.push(op.stmts.S())
-
-        if node.exponent == 0.25:
-            return state.current_frame.push(op.stmts.T())
-
-        if abs(node.exponent == 1):
-            return state.current_frame.push(op.stmts.Z())
-
-        # NOTE: just for the Z gate, an arbitrary exponent is equivalent to the ShiftOp
-        # up to a minus sign!
-        t = -node.exponent
-        theta = state.current_frame.push(py.Constant(math.pi * t))
-        return state.current_frame.push(op.stmts.ShiftOp(theta=theta.result))
-
-    def visit_Rx(self, state: lowering.State[CirqNode], node: cirq.Rx):
-        x = state.current_frame.push(op.stmts.X())
-        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
-        return state.current_frame.push(op.stmts.Rot(axis=x.result, angle=angle.result))
-
-    def visit_Ry(self, state: lowering.State[CirqNode], node: cirq.Ry):
-        y = state.current_frame.push(op.stmts.Y())
-        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
-        return state.current_frame.push(op.stmts.Rot(axis=y.result, angle=angle.result))
-
-    def visit_Rz(self, state: lowering.State[CirqNode], node: cirq.Rz):
-        z = state.current_frame.push(op.stmts.Z())
-        angle = state.current_frame.push(py.Constant(value=math.pi * node.exponent))
-        return state.current_frame.push(op.stmts.Rot(axis=z.result, angle=angle.result))
-
-    def visit_CXPowGate(self, state: lowering.State[CirqNode], node: cirq.CXPowGate):
-        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Control(x, n_controls=1))
-
-    def visit_CZPowGate(self, state: lowering.State[CirqNode], node: cirq.CZPowGate):
-        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Control(z, n_controls=1))
-
-    def visit_ControlledOperation(
-        self, state: lowering.State[CirqNode], node: cirq.ControlledOperation
-    ):
-        return self.visit_GateOperation(state, node)
-
-    def visit_ControlledGate(
-        self, state: lowering.State[CirqNode], node: cirq.ControlledGate
-    ):
-        op_ = state.lower(node.sub_gate).expect_one()
-        n_controls = node.num_controls()
-        return state.current_frame.push(op.stmts.Control(op_, n_controls=n_controls))
-
-    def visit_XXPowGate(self, state: lowering.State[CirqNode], node: cirq.XXPowGate):
-        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Kron(x, x))
-
-    def visit_YYPowGate(self, state: lowering.State[CirqNode], node: cirq.YYPowGate):
-        y = state.lower(cirq.YPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Kron(y, y))
-
-    def visit_ZZPowGate(self, state: lowering.State[CirqNode], node: cirq.ZZPowGate):
-        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Kron(z, z))
-
-    def visit_CCXPowGate(self, state: lowering.State[CirqNode], node: cirq.CCXPowGate):
-        x = state.lower(cirq.XPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Control(x, n_controls=2))
-
-    def visit_CCZPowGate(self, state: lowering.State[CirqNode], node: cirq.CCZPowGate):
-        z = state.lower(cirq.ZPowGate(exponent=node.exponent)).expect_one()
-        return state.current_frame.push(op.stmts.Control(z, n_controls=2))
-
-    def visit_BitFlipChannel(
-        self, state: lowering.State[CirqNode], node: cirq.BitFlipChannel
-    ):
-        x = state.current_frame.push(op.stmts.X())
-        p = state.current_frame.push(py.Constant(node.p))
-        return state.current_frame.push(
-            noise.stmts.PauliError(basis=x.result, p=p.result)
-        )
-
-    def visit_AmplitudeDampingChannel(
-        self, state: lowering.State[CirqNode], node: cirq.AmplitudeDampingChannel
-    ):
-        r = state.current_frame.push(op.stmts.Reset())
-        p = state.current_frame.push(py.Constant(node.gamma))
-
-        # TODO: do we need a dedicated noise stmt for this? Using PauliError
-        # with this basis feels like a hack
-        noise_channel = state.current_frame.push(
-            noise.stmts.PauliError(basis=r.result, p=p.result)
-        )
-
-        return noise_channel
-
-    def visit_GeneralizedAmplitudeDampingChannel(
-        self,
-        state: lowering.State[CirqNode],
-        node: cirq.GeneralizedAmplitudeDampingChannel,
-    ):
-        p = state.current_frame.push(py.Constant(node.p)).result
-        gamma = state.current_frame.push(py.Constant(node.gamma)).result
-
-        # 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
-
-        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
-
-    def visit_DepolarizingChannel(
-        self, state: lowering.State[CirqNode], node: cirq.DepolarizingChannel
-    ):
-        p = state.current_frame.push(py.Constant(node.p)).result
-        return state.current_frame.push(noise.stmts.Depolarize(p))
-
-    def visit_AsymmetricDepolarizingChannel(
-        self, state: lowering.State[CirqNode], node: cirq.AsymmetricDepolarizingChannel
-    ):
-        nqubits = node.num_qubits()
-        if nqubits > 2:
-            raise lowering.BuildError(
-                "AsymmetricDepolarizingChannel applied to more than 2 qubits is not supported!"
-            )
-
-        if nqubits == 1:
-            p_x = state.current_frame.push(py.Constant(node.p_x)).result
-            p_y = state.current_frame.push(py.Constant(node.p_y)).result
-            p_z = state.current_frame.push(py.Constant(node.p_z)).result
-            params = state.current_frame.push(ilist.New(values=(p_x, p_y, p_z))).result
-            return state.current_frame.push(noise.stmts.SingleQubitPauliChannel(params))
-
-        # NOTE: nqubits == 2
-        error_probs = node.error_probabilities
-        paulis = ("I", "X", "Y", "Z")
-        values = []
-        for p1 in paulis:
-            for p2 in paulis:
-                if p1 == p2 == "I":
-                    continue
-
-                p = error_probs.get(p1 + p2, 0.0)
-                p_ssa = state.current_frame.push(py.Constant(p)).result
-                values.append(p_ssa)
-
-        params = state.current_frame.push(ilist.New(values=values)).result
-        return state.current_frame.push(noise.stmts.TwoQubitPauliChannel(params))

bloqade/squin/lowering.py

@@ -1,80 +0,0 @@
-import ast
-from dataclasses import dataclass
-
-from kirin import lowering
-
-from . import qubit
-
-
-@dataclass(frozen=True)
-class ApplyAnyCallLowering(lowering.FromPythonCall["qubit.ApplyAny"]):
-    """
-    Custom lowering for ApplyAny that collects vararg qubits into a single tuple argument
-    """
-
-    def lower(
-        self, stmt: type["qubit.ApplyAny"], state: lowering.State, node: ast.Call
-    ):
-        if len(node.args) + len(node.keywords) < 2:
-            raise lowering.BuildError(
-                "Apply requires at least one operator and one qubit as arguments!"
-            )
-
-        op, qubits = self.unpack_arguments(node)
-
-        op_ssa = state.lower(op).expect_one()
-        qubits_lowered = [state.lower(qbit).expect_one() for qbit in qubits]
-
-        s = stmt(op_ssa, tuple(qubits_lowered))
-        return state.current_frame.push(s)
-
-    def unpack_arguments(self, node: ast.Call) -> tuple[ast.expr, list[ast.expr]]:
-        if len(node.keywords) == 0:
-            op, *qubits = node.args
-            return op, qubits
-
-        kwargs = {kw.arg: kw.value for kw in node.keywords}
-        if len(kwargs) > 2 or "qubits" not in kwargs:
-            raise lowering.BuildError(f"Got unsupported keyword argument {kwargs}")
-
-        qubits = kwargs["qubits"]
-        if len(kwargs) == 1:
-            if len(node.args) != 1:
-                raise lowering.BuildError("Missing operator argument")
-            op = node.args[0]
-        else:
-            try:
-                op = kwargs["operator"]
-            except KeyError:
-                raise lowering.BuildError(f"Got unsupported keyword argument {kwargs}")
-
-        if isinstance(qubits, ast.List):
-            return op, qubits.elts
-
-        return op, [qubits]
-
-
-@dataclass(frozen=True)
-class BroadcastCallLowering(lowering.FromPythonCall["qubit.Broadcast"]):
-    """
-    Custom lowering for broadcast vararg call.
-
-    NOTE: we can re-use this to lower Apply too once we remove the deprecated syntax
-    """
-
-    def lower(
-        self, stmt: type["qubit.Broadcast"], state: lowering.State, node: ast.Call
-    ):
-        if len(node.args) < 2:
-            raise lowering.BuildError(
-                "Broadcast requires at least one operator and one qubit list argument"
-            )
-
-        op, *qubit_lists = node.args
-
-        op_lowered = state.lower(op).expect_one()
-        qubits_lists_lowered = [
-            state.lower(qubit_list).expect_one() for qubit_list in qubit_lists
-        ]
-
-        return state.current_frame.push(stmt(op_lowered, tuple(qubits_lists_lowered)))

bloqade/squin/noise/_wrapper.py

@@ -1,36 +0,0 @@
-from typing import Literal
-
-from kirin.dialects import ilist
-from kirin.lowering import wraps
-
-from bloqade.squin.op.types import Op, MultiQubitPauliOp
-
-from . import stmts
-
-
-@wraps(stmts.PauliError)
-def pauli_error(basis: MultiQubitPauliOp, p: float) -> Op: ...
-
-
-@wraps(stmts.Depolarize)
-def depolarize(p: float) -> Op: ...
-
-
-@wraps(stmts.Depolarize2)
-def depolarize2(p: float) -> Op: ...
-
-
-@wraps(stmts.SingleQubitPauliChannel)
-def single_qubit_pauli_channel(
-    params: ilist.IList[float, Literal[3]] | list[float],
-) -> Op: ...
-
-
-@wraps(stmts.TwoQubitPauliChannel)
-def two_qubit_pauli_channel(
-    params: ilist.IList[float, Literal[15]] | list[float],
-) -> Op: ...
-
-
-@wraps(stmts.QubitLoss)
-def qubit_loss(p: float) -> Op: ...

bloqade/squin/noise/rewrite.py

@@ -1,129 +0,0 @@
-import itertools
-
-from kirin import ir
-from kirin.passes import Pass
-from kirin.rewrite import Walk
-from kirin.dialects import py, ilist
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-from .stmts import (
-    QubitLoss,
-    Depolarize,
-    PauliError,
-    Depolarize2,
-    NoiseChannel,
-    TwoQubitPauliChannel,
-    SingleQubitPauliChannel,
-    StochasticUnitaryChannel,
-)
-from ..op.stmts import X, Y, Z, Kron, Identity
-
-
-class _RewriteNoiseStmts(RewriteRule):
-    """Rewrites squin noise statements to StochasticUnitaryChannel"""
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        if not isinstance(node, NoiseChannel) or isinstance(node, QubitLoss):
-            return RewriteResult()
-
-        return getattr(self, "rewrite_" + node.name)(node)
-
-    def rewrite_pauli_error(self, node: PauliError) -> RewriteResult:
-        (operators := ilist.New(values=(node.basis,))).insert_before(node)
-        (ps := ilist.New(values=(node.p,))).insert_before(node)
-        stochastic_channel = StochasticUnitaryChannel(
-            operators=operators.result, probabilities=ps.result
-        )
-
-        node.replace_by(stochastic_channel)
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_single_qubit_pauli_channel(
-        self, node: SingleQubitPauliChannel
-    ) -> RewriteResult:
-        paulis = (X(), Y(), Z())
-        paulis_ssa: list[ir.SSAValue] = []
-        for op in paulis:
-            op.insert_before(node)
-            paulis_ssa.append(op.result)
-
-        (pauli_ops := ilist.New(values=paulis_ssa)).insert_before(node)
-
-        stochastic_unitary = StochasticUnitaryChannel(
-            operators=pauli_ops.result, probabilities=node.params
-        )
-        node.replace_by(stochastic_unitary)
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_two_qubit_pauli_channel(
-        self, node: TwoQubitPauliChannel
-    ) -> RewriteResult:
-        operator_list = self._insert_two_qubit_paulis_before_node(node)
-        stochastic_unitary = StochasticUnitaryChannel(
-            operators=operator_list, probabilities=node.params
-        )
-
-        node.replace_by(stochastic_unitary)
-        return RewriteResult(has_done_something=True)
-
-    @staticmethod
-    def _insert_two_qubit_paulis_before_node(
-        node: TwoQubitPauliChannel | Depolarize2,
-    ) -> ir.ResultValue:
-        paulis = (Identity(sites=1), X(), Y(), Z())
-        for op in paulis:
-            op.insert_before(node)
-
-        # 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)
-            op.insert_before(node)
-            operators.append(op.result)
-
-        (operator_list := ilist.New(values=operators)).insert_before(node)
-        return operator_list.result
-
-    def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
-        paulis = (X(), Y(), Z())
-        operators: list[ir.SSAValue] = []
-        for op in paulis:
-            op.insert_before(node)
-            operators.append(op.result)
-
-        # NOTE: need to divide the probability by 3 to get the correct total error rate
-        (three := py.Constant(3)).insert_before(node)
-        (p_over_3 := py.Div(node.p, three.result)).insert_before(node)
-
-        (operator_list := ilist.New(values=operators)).insert_before(node)
-        (ps := ilist.New(values=[p_over_3.result for _ in range(3)])).insert_before(
-            node
-        )
-
-        stochastic_unitary = StochasticUnitaryChannel(
-            operators=operator_list.result, probabilities=ps.result
-        )
-        node.replace_by(stochastic_unitary)
-
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_depolarize2(self, node: Depolarize2) -> RewriteResult:
-        operator_list = self._insert_two_qubit_paulis_before_node(node)
-
-        # NOTE: need to divide the probability by 15 to get the correct total error rate
-        (fifteen := py.Constant(15)).insert_before(node)
-        (p_over_15 := py.Div(node.p, fifteen.result)).insert_before(node)
-        (probs := ilist.New(values=[p_over_15.result] * 15)).insert_before(node)
-
-        stochastic_unitary = StochasticUnitaryChannel(
-            operators=operator_list, probabilities=probs.result
-        )
-        node.replace_by(stochastic_unitary)
-
-        return RewriteResult(has_done_something=True)
-
-
-class RewriteNoiseStmts(Pass):
-    def unsafe_run(self, mt: ir.Method):
-        return Walk(_RewriteNoiseStmts()).rewrite(mt.code)