fixing more tests, working on stim

Author: weinbe58

src/bloqade/analysis/address/impls.py

@@ -60,8 +60,8 @@ def from_literal(literal: Any) -> Address:
         match collection:
             case PartialIList(data) | PartialTuple(data):
                 return data
-            case AddressReg(data):
-                return tuple(map(AddressQubit, data))
+            case AddressReg():
+                return collection.qubits
             case ConstResult(const.Value(data)) if isinstance(data, Iterable):
                 return tuple(map(from_literal, data))
             case ConstResult(const.PartialTuple(data)):

src/bloqade/analysis/address/lattice.py

@@ -100,6 +100,10 @@ class AddressReg(Address):
     def is_subseteq(self, other: Address) -> bool:
         return isinstance(other, AddressReg) and self.data == other.data
 
+    @property
+    def qubits(self) -> tuple[AddressQubit, ...]:
+        return tuple(AddressQubit(i) for i in self.data)
+
 
 @final
 @dataclass

src/bloqade/squin/analysis/schedule.py

@@ -210,8 +210,8 @@ def _update_dag(self, stmt: ir.Statement, addr: address.Address):
                 if old_stmt is not None:
                     self.stmt_dag.add_edge(old_stmt, stmt)
                 self.use_def[idx] = stmt
-        elif isinstance(addr, address.AddressTuple):
-            for sub_addr in addr.data:
+        elif isinstance(addr, address.AddressReg):
+            for sub_addr in addr.qubits:
                 self._update_dag(stmt, sub_addr)
 
     def update_dag(self, stmt: ir.Statement, args: Sequence[ir.SSAValue]):

src/bloqade/stim/rewrite/qubit_to_stim.py

@@ -60,7 +60,6 @@ def rewrite_SingleQubitGate(
         """
 
         qubit_addr_attr = stmt.qubits.hints.get("address", None)
-
         if qubit_addr_attr is None:
             return RewriteResult()
 

src/bloqade/stim/rewrite/squin_noise.py

@@ -30,6 +30,7 @@ def rewrite_NoiseChannel(
         """Rewrite NoiseChannel statements to their stim equivalents."""
 
         rewrite_method = getattr(self, f"rewrite_{type(stmt).__name__}", None)
+
         # No rewrite method exists and the rewrite should stop
         if rewrite_method is None:
             return RewriteResult()
@@ -39,6 +40,7 @@ def rewrite_NoiseChannel(
             # support broadcasting for multi-qubit noise channels.
             # Therefore, we must expand the broadcast into individual stim statements.
             qubit_address_attr = stmt.qubits.hints.get("address", None)
+            print(qubit_address_attr)
             if not isinstance(qubit_address_attr, AddressAttribute):
                 return RewriteResult()
 

src/bloqade/stim/rewrite/util.py

@@ -21,7 +21,7 @@ def insert_qubit_idx_from_address(
     """
     address_data = address.address
     qubit_idx_ssas = []
-
+    print(address_data)
     if isinstance(address_data, AddressReg):
         for qubit_idx in address_data.data:
             create_and_insert_qubit_idx_stmt(

test/stim/passes/test_squin_noise_to_stim.py

@@ -220,12 +220,18 @@ def test():
         q = sq.qalloc(3)
         sq.correlated_qubit_loss(0.1, qubits=q[:2])
 
-    SquinToStimPass(test.dialects)(test)
+    frame, ret = AddressAnalysis(kernel).run_analysis(test)
 
+    SquinToStimPass(test.dialects)(test)
+    test.print(analysis=frame.entries)
     expected = "I_ERROR[correlated_loss:0](0.10000000) 0 1"
     assert codegen(test) == expected
 
 
+if __name__ == "__main__":
+    test_correlated_qubit_loss()
+
+
 def test_broadcast_correlated_qubit_loss():
     @kernel
     def test():