bloqade-circuit 0.4.0__tar.gz → 0.4.1__tar.gz

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.4.0
+Version: 0.4.1
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bloqade-circuit"
-version = "0.4.0"
+version = "0.4.1"
 description = "The software development toolkit for neutral atom arrays."
 readme = "README.md"
 authors = [

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/src/bloqade/pyqrack/__init__.py

@@ -16,7 +16,7 @@ from .task import PyQrackSimulatorTask as PyQrackSimulatorTask
 # NOTE: The following import is for registering the method tables
 from .noise import native as native
 from .qasm2 import uop as uop, core as core, glob as glob, parallel as parallel
-from .squin import op as op, qubit as qubit
+from .squin import op as op, noise as noise, qubit as qubit
 from .device import (
     StackMemorySimulator as StackMemorySimulator,
     DynamicMemorySimulator as DynamicMemorySimulator,

bloqade_circuit-0.4.1/src/bloqade/pyqrack/squin/noise/__init__.py

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

bloqade_circuit-0.4.1/src/bloqade/pyqrack/squin/noise/native.py

@@ -0,0 +1,72 @@
+import random
+import typing
+from dataclasses import dataclass
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.pyqrack import QubitState, PyQrackQubit, PyQrackInterpreter
+from bloqade.squin.noise.stmts import QubitLoss, StochasticUnitaryChannel
+from bloqade.squin.noise._dialect import dialect as squin_noise_dialect
+
+from ..runtime import OperatorRuntimeABC
+
+
+@dataclass(frozen=True)
+class StochasticUnitaryChannelRuntime(OperatorRuntimeABC):
+    operators: ilist.IList[OperatorRuntimeABC, typing.Any]
+    probabilities: ilist.IList[float, typing.Any]
+
+    @property
+    def n_sites(self) -> int:
+        n = self.operators[0].n_sites
+        for op in self.operators[1:]:
+            assert (
+                op.n_sites == n
+            ), "Encountered a stochastic unitary channel with operators of different size!"
+        return n
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        # NOTE: probabilities don't necessarily sum to 1; could be no noise event should occur
+        p_no_op = 1 - sum(self.probabilities)
+        if random.uniform(0.0, 1.0) < p_no_op:
+            return
+
+        selected_ops = random.choices(self.operators, weights=self.probabilities)
+        for op in selected_ops:
+            op.apply(*qubits, adjoint=adjoint)
+
+
+@dataclass(frozen=True)
+class QubitLossRuntime(OperatorRuntimeABC):
+    p: float
+
+    @property
+    def n_sites(self) -> int:
+        return 1
+
+    def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
+        if random.uniform(0.0, 1.0) < self.p:
+            qubit.state = QubitState.Lost
+
+
+@squin_noise_dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    @interp.impl(StochasticUnitaryChannel)
+    def stochastic_unitary_channel(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: StochasticUnitaryChannel,
+    ):
+        operators = frame.get(stmt.operators)
+        probabilities = frame.get(stmt.probabilities)
+
+        return (StochasticUnitaryChannelRuntime(operators, probabilities),)
+
+    @interp.impl(QubitLoss)
+    def qubit_loss(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: QubitLoss
+    ):
+        p = frame.get(stmt.p)
+        return (QubitLossRuntime(p),)

bloqade_circuit-0.4.1/src/bloqade/squin/__init__.py

@@ -0,0 +1,17 @@
+from . import (
+    op as op,
+    wire as wire,
+    noise as noise,
+    qubit as qubit,
+    lowering as lowering,
+)
+from .groups import wired as wired, kernel as kernel
+
+try:
+    # NOTE: make sure optional cirq dependency is installed
+    import cirq as cirq_package  # noqa: F401
+except ImportError:
+    pass
+else:
+    from . import cirq as cirq
+    from .cirq import load_circuit as load_circuit

bloqade_circuit-0.4.1/src/bloqade/squin/cirq/__init__.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+import cirq
+from kirin import ir, types
+from kirin.dialects import func
+
+from . import lowering as lowering
+from .. import kernel
+from .lowering import Squin
+
+
+def load_circuit(
+    circuit: cirq.Circuit,
+    kernel_name: str = "main",
+    dialects: ir.DialectGroup = kernel,
+    globals: dict[str, Any] | None = None,
+    file: str | None = None,
+    lineno_offset: int = 0,
+    col_offset: int = 0,
+    compactify: bool = True,
+):
+    """Converts a cirq.Circuit object into a squin kernel.
+
+    Args:
+        circuit (cirq.Circuit): The circuit to load.
+
+    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`.
+        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:
+
+    ```python
+    # from cirq's "hello qubit" example
+    import cirq
+    from bloqade import squin
+
+    # Pick a qubit.
+    qubit = cirq.GridQubit(0, 0)
+
+    # Create a circuit.
+    circuit = cirq.Circuit(
+        cirq.X(qubit)**0.5,  # Square root of NOT.
+        cirq.measure(qubit, key='m')  # Measurement.
+    )
+
+    # load the circuit as squin
+    main = squin.load_circuit(circuit)
+
+    # print the resulting IR
+    main.print()
+    ```
+    """
+
+    target = Squin(dialects=dialects, circuit=circuit)
+    body = target.run(
+        circuit,
+        source=str(circuit),  # TODO: proper source string
+        file=file,
+        globals=globals,
+        lineno_offset=lineno_offset,
+        col_offset=col_offset,
+        compactify=compactify,
+    )
+
+    # 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))
+
+    code = func.Function(
+        sym_name=kernel_name,
+        signature=func.Signature((), types.NoneType),
+        body=body,
+    )
+
+    return ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name=kernel_name,
+        arg_names=[],
+        dialects=dialects,
+        code=code,
+    )

bloqade_circuit-0.4.1/src/bloqade/squin/cirq/lowering.py

@@ -0,0 +1,303 @@
+import math
+from typing import Any
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir, lowering
+from kirin.rewrite import Walk, CFGCompactify
+from kirin.dialects import py, 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: qubit.New = 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
+
+        index_ssa = state.current_frame.push(py.Constant(index)).result
+        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg.result, 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_result = state.current_frame.get(qbits_stmt.name)
+
+        if qbits_result is not None:
+            return qbits_result
+
+        state.current_frame.push(qbits_stmt)
+        return qbits_stmt.result
+
+    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,
+    ) -> 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: 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))
+
+            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 lower_measurement(
+        self, state: lowering.State[CirqNode], node: cirq.GateOperation
+    ):
+        if len(node.qubits) == 1:
+            qbit = self.lower_qubit_getindex(state, node.qubits[0])
+            return state.current_frame.push(qubit.MeasureQubit(qbit))
+
+        qbits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(qubit.MeasureQubitList(qbits))
+
+    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 node.exponent == 1:
+            return state.current_frame.push(op.stmts.H())
+
+        return state.lower(node.in_su2())
+
+    def visit_XPowGate(self, state: lowering.State[CirqNode], node: cirq.XPowGate):
+        if 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 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 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
+        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,
+    ):
+        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))
+
+        # p1 =
+
+        # x = state.current_frame.push(op.stmts.X())
+        # noise_channel1 = noise.stmts.PauliError(basis=x.result, p=)

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/src/bloqade/squin/groups.py

@@ -3,12 +3,12 @@ from kirin.prelude import structural_no_opt
 from kirin.rewrite import Walk, Chain
 from kirin.dialects import ilist
 
-from . import op, wire, qubit
+from . import op, wire, noise, qubit
 from .op.rewrite import PyMultToSquinMult
 from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
 
 
-@ir.dialect_group(structural_no_opt.union([op, qubit]))
+@ir.dialect_group(structural_no_opt.union([op, qubit, noise]))
 def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
@@ -36,7 +36,7 @@ def kernel(self):
     return run_pass
 
 
-@ir.dialect_group(structural_no_opt.union([op, wire]))
+@ir.dialect_group(structural_no_opt.union([op, wire, noise]))
 def wired(self):
     py_mult_to_mult_pass = PyMultToSquinMult(self)
 

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/src/bloqade/squin/noise/__init__.py

@@ -4,5 +4,7 @@ from ._wrapper import (
     pp_error as pp_error,
     depolarize as depolarize,
     qubit_loss as qubit_loss,
-    pauli_channel as pauli_channel,
+    pauli_error as pauli_error,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

{bloqade_circuit-0.4.0 → bloqade_circuit-0.4.1}/src/bloqade/squin/noise/_wrapper.py

@@ -14,11 +14,15 @@ def pp_error(op: Op, p: float) -> Op: ...
 
 
 @wraps(stmts.Depolarize)
-def depolarize(n_qubits: int, p: float) -> Op: ...
+def depolarize(p: float) -> Op: ...
 
 
-@wraps(stmts.PauliChannel)
-def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(params: tuple[float, float, float]) -> Op: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(params: tuple[float, ...]) -> Op: ...
 
 
 @wraps(stmts.QubitLoss)

bloqade_circuit-0.4.1/src/bloqade/squin/noise/rewrite.py

@@ -0,0 +1,111 @@
+import itertools
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import Walk
+from kirin.dialects import ilist
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from .stmts import (
+    PPError,
+    QubitLoss,
+    Depolarize,
+    PauliError,
+    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:
+        paulis = (X(), Y(), Z(), Identity(sites=1))
+        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]
+        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)
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list.result, probabilities=node.params
+        )
+
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    def rewrite_p_p_error(self, node: PPError) -> RewriteResult:
+        (operators := ilist.New(values=(node.op,))).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_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)
+
+        (operator_list := ilist.New(values=operators)).insert_before(node)
+        (ps := ilist.New(values=[node.p 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)
+
+
+class RewriteNoiseStmts(Pass):
+    def unsafe_run(self, mt: ir.Method):
+        return Walk(_RewriteNoiseStmts()).rewrite(mt.code)