fixing tests

Author: weinbe58

src/bloqade/pyqrack/noise/native.py

@@ -2,11 +2,11 @@
 
 from kirin import interp
 
-from bloqade.noise import native
 from bloqade.pyqrack import PyQrackInterpreter, reg
+from bloqade.qasm2.dialects import noise
 
 
-@native.dialect.register(key="pyqrack")
+@noise.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
     def apply_pauli_error(
         self,
@@ -27,12 +27,12 @@ def apply_pauli_error(
 
         getattr(qarg.sim_reg, which)(qarg.addr)
 
-    @interp.impl(native.PauliChannel)
+    @interp.impl(noise.PauliChannel)
     def single_qubit_error_channel(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: native.PauliChannel,
+        stmt: noise.PauliChannel,
     ):
         qargs: List[reg.PyQrackQubit] = frame.get(stmt.qargs)
 
@@ -43,12 +43,12 @@ def single_qubit_error_channel(
 
         return ()
 
-    @interp.impl(native.CZPauliChannel)
+    @interp.impl(noise.CZPauliChannel)
     def cz_pauli_channel(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: native.CZPauliChannel,
+        stmt: noise.CZPauliChannel,
     ):
 
         qargs: List[reg.PyQrackQubit] = frame.get(stmt.qargs)
@@ -80,12 +80,12 @@ def cz_pauli_channel(
 
         return ()
 
-    @interp.impl(native.AtomLossChannel)
+    @interp.impl(noise.AtomLossChannel)
     def atom_loss_channel(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: native.AtomLossChannel,
+        stmt: noise.AtomLossChannel,
     ):
         qargs: List[reg.PyQrackQubit] = frame.get(stmt.qargs)
 

src/bloqade/qasm2/passes/noise.py

@@ -8,7 +8,7 @@
     DeadCodeElimination,
 )
 
-from bloqade.noise import native
+from bloqade.qasm2 import noise
 from bloqade.analysis import address
 from bloqade.qasm2.passes.lift_qubits import LiftQubits
 from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
@@ -25,12 +25,9 @@ class NoisePass(Pass):
 
     ```
     from bloqade import qasm2
-    from bloqade.noise import native
-    from bloqade.qasm2.passes.noise import NoisePass
+    from bloqade.qasm2.passes import NoisePass
 
-    noise_main = qasm2.extended.add(native.dialect)
-
-    @noise_main
+    @qasm2.extended
     def main():
         q = qasm2.qreg(2)
         qasm2.h(q[0])
@@ -51,9 +48,7 @@ def main():
 
     """
 
-    noise_model: native.MoveNoiseModelABC = field(
-        default_factory=native.TwoRowZoneModel
-    )
+    noise_model: noise.MoveNoiseModelABC = field(default_factory=noise.TwoRowZoneModel)
     address_analysis: address.AddressAnalysis = field(init=False)
 
     def __post_init__(self):

src/bloqade/qasm2/rewrite/heuristic_noise.py

@@ -5,9 +5,8 @@
 from kirin.rewrite import abc as rewrite_abc
 from kirin.dialects import ilist
 
-from bloqade.noise import native
 from bloqade.analysis import address
-from bloqade.qasm2.dialects import uop, glob, parallel
+from bloqade.qasm2.dialects import uop, glob, noise, parallel
 
 
 @dataclass
@@ -19,9 +18,7 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
 
     address_analysis: Dict[ir.SSAValue, address.Address]
     qubit_ssa_value: Dict[int, ir.SSAValue]
-    noise_model: native.MoveNoiseModelABC = field(
-        default_factory=native.TwoRowZoneModel
-    )
+    noise_model: noise.MoveNoiseModelABC = field(default_factory=noise.TwoRowZoneModel)
 
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         if isinstance(node, uop.SingleQubitGate):
@@ -43,10 +40,10 @@ def insert_single_qubit_noise(
         qargs: ir.SSAValue,
         probs: Tuple[float, float, float, float],
     ):
-        native.PauliChannel(qargs, px=probs[0], py=probs[1], pz=probs[2]).insert_before(
+        noise.PauliChannel(qargs, px=probs[0], py=probs[1], pz=probs[2]).insert_before(
             node
         )
-        native.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
+        noise.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
 
         return rewrite_abc.RewriteResult(has_done_something=True)
 
@@ -103,9 +100,9 @@ def move_noise_stmts(
             nodes.append(
                 qargs := ilist.New(tuple(self.qubit_ssa_value[q] for q in qubits))
             )
-            nodes.append(native.AtomLossChannel(qargs.result, prob=probs[3]))
+            nodes.append(noise.AtomLossChannel(qargs.result, prob=probs[3]))
             nodes.append(
-                native.PauliChannel(qargs.result, px=probs[0], py=probs[1], pz=probs[2])
+                noise.PauliChannel(qargs.result, px=probs[0], py=probs[1], pz=probs[2])
             )
 
         return nodes
@@ -116,7 +113,7 @@ def cz_gate_noise(
         qargs: ir.SSAValue,
     ) -> list[ir.Statement]:
         return [
-            native.CZPauliChannel(
+            noise.CZPauliChannel(
                 ctrls,
                 qargs,
                 px_ctrl=self.noise_model.cz_paired_gate_px,
@@ -127,7 +124,7 @@ def cz_gate_noise(
                 pz_qarg=self.noise_model.cz_paired_gate_pz,
                 paired=True,
             ),
-            native.CZPauliChannel(
+            noise.CZPauliChannel(
                 ctrls,
                 qargs,
                 px_ctrl=self.noise_model.cz_unpaired_gate_px,
@@ -138,8 +135,8 @@ def cz_gate_noise(
                 pz_qarg=self.noise_model.cz_unpaired_gate_pz,
                 paired=False,
             ),
-            native.AtomLossChannel(ctrls, prob=self.noise_model.cz_gate_loss_prob),
-            native.AtomLossChannel(qargs, prob=self.noise_model.cz_gate_loss_prob),
+            noise.AtomLossChannel(ctrls, prob=self.noise_model.cz_gate_loss_prob),
+            noise.AtomLossChannel(qargs, prob=self.noise_model.cz_gate_loss_prob),
         ]
 
     def rewrite_cz_gate(self, node: uop.CZ):

src/bloqade/qbraid/lowering.py

@@ -4,13 +4,13 @@
 from kirin import ir, types, passes
 from kirin.dialects import func, ilist
 
-from bloqade import noise, qasm2
+from bloqade import qasm2
 from bloqade.qbraid import schema
-from bloqade.qasm2.dialects import glob, parallel
+from bloqade.qasm2.dialects import glob, noise, parallel
 
 
 @ir.dialect_group(
-    [func, qasm2.core, qasm2.uop, parallel, glob, qasm2.expr, noise.native, ilist]
+    [func, qasm2.core, qasm2.uop, parallel, glob, qasm2.expr, noise, ilist]
 )
 def qbraid_noise(
     self,
@@ -192,7 +192,7 @@ def process_cz_pauli_error(
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
             )
             self.block_list.append(
-                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+                noise.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
             )
 
         for (p_ctrl, p_qarg), qubits in paired_layers.items():
@@ -204,7 +204,7 @@ def process_cz_pauli_error(
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
             )
             self.block_list.append(
-                noise.native.CZPauliChannel(
+                noise.CZPauliChannel(
                     paired=True,
                     px_ctrl=p_ctrl[0],
                     py_ctrl=p_ctrl[1],
@@ -226,7 +226,7 @@ def process_cz_pauli_error(
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
             )
             self.block_list.append(
-                noise.native.CZPauliChannel(
+                noise.CZPauliChannel(
                     paired=False,
                     px_ctrl=p_ctrl[0],
                     py_ctrl=p_ctrl[1],
@@ -285,7 +285,7 @@ def lower_pauli_errors(self, operator_error: schema.PauliErrorModel):
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
             )
             self.block_list.append(
-                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+                noise.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
             )
 
     def lower_measurement(self, operation: schema.Measurement):
@@ -303,7 +303,7 @@ def lower_atom_loss(self, survival_probs: Tuple[float, ...]):
         for survival_prob, qubits in layers.items():
             self.block_list.append(qargs := ilist.New(values=qubits))
             self.block_list.append(
-                noise.native.AtomLossChannel(prob=survival_prob, qargs=qargs.result)
+                noise.AtomLossChannel(prob=survival_prob, qargs=qargs.result)
             )
 
     def lower_number(self, value: float | int) -> ir.SSAValue:

test/analysis/fidelity/test_fidelity.py

@@ -3,7 +3,7 @@
 import pytest
 
 from bloqade import qasm2
-from bloqade.noise import native
+from bloqade.qasm2 import noise
 from bloqade.analysis.fidelity import FidelityAnalysis
 from bloqade.qasm2.passes.noise import NoisePass
 
@@ -15,7 +15,7 @@ def test_atom_loss_analysis():
     @qasm2.extended
     def main():
         q = qasm2.qreg(2)
-        native.atom_loss_channel([q[0]], prob=p_loss)
+        noise.atom_loss_channel([q[0]], prob=p_loss)
         return q
 
     fid_analysis = FidelityAnalysis(main.dialects)
@@ -32,7 +32,7 @@ def test_cz_noise():
     @qasm2.extended
     def main():
         q = qasm2.qreg(2)
-        native.cz_pauli_channel(
+        noise.cz_pauli_channel(
             [q[0]],
             [q[1]],
             px_ctrl=p_ch,
@@ -63,7 +63,7 @@ def test_single_qubit_noise():
     @qasm2.extended
     def main():
         q = qasm2.qreg(2)
-        native.pauli_channel([q[0]], px=p_ch, py=p_ch, pz=p_ch)
+        noise.pauli_channel([q[0]], px=p_ch, py=p_ch, pz=p_ch)
         return q
 
     main.print()
@@ -77,7 +77,7 @@ def main():
     assert math.isclose(fid_analysis.gate_fidelity, expected_fidelity)
 
 
-class NoiseTestModel(native.MoveNoiseModelABC):
+class NoiseTestModel(noise.MoveNoiseModelABC):
     def parallel_cz_errors(self, ctrls, qargs, rest):
         return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}
 

…yqrack/runtime/noise/native/test_loss.py …pyqrack/runtime/noise/qasm2/test_loss.pytest/pyqrack/runtime/noise/native/test_loss.py renamed to test/pyqrack/runtime/noise/qasm2/test_loss.py test/pyqrack/runtime/noise/native/test_loss.py renamed to test/pyqrack/runtime/noise/qasm2/test_loss.py

@@ -5,7 +5,7 @@
 from kirin.dialects import ilist
 
 from bloqade import qasm2
-from bloqade.noise import native
+from bloqade.qasm2 import noise
 from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter, reg
 from bloqade.pyqrack.base import MockMemory
 
@@ -23,8 +23,8 @@ def test_atom_loss():
     @qasm2.extended
     def test_atom_loss(c: qasm2.CReg):
         q = qasm2.qreg(2)
-        native.atom_loss_channel([q[0]], prob=0.1)
-        native.atom_loss_channel([q[1]], prob=0.05)
+        noise.atom_loss_channel([q[0]], prob=0.1)
+        noise.atom_loss_channel([q[1]], prob=0.05)
         qasm2.measure(q[0], c[0])
 
         return q

…qrack/runtime/noise/native/test_pauli.py …yqrack/runtime/noise/qasm2/test_pauli.pytest/pyqrack/runtime/noise/native/test_pauli.py renamed to test/pyqrack/runtime/noise/qasm2/test_pauli.py test/pyqrack/runtime/noise/native/test_pauli.py renamed to test/pyqrack/runtime/noise/qasm2/test_pauli.py

@@ -4,7 +4,7 @@
 from kirin import ir
 
 from bloqade import qasm2
-from bloqade.noise import native
+from bloqade.qasm2 import noise
 from bloqade.pyqrack.base import MockMemory, PyQrackInterpreter
 
 
@@ -20,13 +20,13 @@ def test_pauli_channel():
     @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
-        native.pauli_channel(
+        noise.pauli_channel(
             [q[0]],
             px=0.1,
             py=0.4,
             pz=0.3,
         )
-        native.pauli_channel(
+        noise.pauli_channel(
             [q[1]],
             px=0.1,
             py=0.4,
@@ -45,7 +45,7 @@ def test_pauli_probs_check():
     @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
-        native.pauli_channel(
+        noise.pauli_channel(
             [q[0]],
             px=0.1,
             py=0.4,
@@ -61,9 +61,9 @@ def test_cz_pauli_channel_false():
     @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
-        native.atom_loss_channel([q[0]], prob=0.4)
-        native.atom_loss_channel([q[1]], prob=0.4)
-        native.cz_pauli_channel(
+        noise.atom_loss_channel([q[0]], prob=0.4)
+        noise.atom_loss_channel([q[1]], prob=0.4)
+        noise.cz_pauli_channel(
             [q[0]],
             [q[1]],
             px_ctrl=0.1,
@@ -75,9 +75,9 @@ def test_atom_loss():
             paired=False,
         )  # no noise here
         qasm2.cz(q[0], q[1])
-        native.atom_loss_channel([q[0]], prob=0.7)
-        native.atom_loss_channel([q[1]], prob=0.4)
-        native.cz_pauli_channel(
+        noise.atom_loss_channel([q[0]], prob=0.7)
+        noise.atom_loss_channel([q[1]], prob=0.4)
+        noise.cz_pauli_channel(
             [q[0]],
             [q[1]],
             px_ctrl=0.1,
@@ -102,9 +102,9 @@ def test_cz_pauli_channel_true():
     @qasm2.extended
     def test_atom_loss():
         q = qasm2.qreg(2)
-        native.atom_loss_channel([q[0]], prob=0.4)
-        native.atom_loss_channel([q[1]], prob=0.4)
-        native.cz_pauli_channel(
+        noise.atom_loss_channel([q[0]], prob=0.4)
+        noise.atom_loss_channel([q[1]], prob=0.4)
+        noise.cz_pauli_channel(
             [q[0]],
             [q[1]],
             px_ctrl=0.1,
@@ -116,9 +116,9 @@ def test_atom_loss():
             paired=True,
         )  # no noise here
         qasm2.cz(q[0], q[1])
-        native.atom_loss_channel([q[0]], prob=0.7)
-        native.atom_loss_channel([q[1]], prob=0.4)
-        native.cz_pauli_channel(
+        noise.atom_loss_channel([q[0]], prob=0.7)
+        noise.atom_loss_channel([q[1]], prob=0.4)
+        noise.cz_pauli_channel(
             [q[0]],
             [q[1]],
             px_ctrl=0.1,