From 9352266c91eb1d24921680b31c51d6e09bb41a49 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 6 Oct 2025 16:36:14 -0700 Subject: [PATCH 01/14] Add support for correlated loss --- src/bloqade/pyqrack/squin/noise/native.py | 11 ++++++++ src/bloqade/squin/__init__.py | 1 + src/bloqade/squin/noise/_interface.py | 4 +++ src/bloqade/squin/noise/stmts.py | 10 +++++++ .../squin/stdlib/broadcast/__init__.py | 1 + src/bloqade/squin/stdlib/broadcast/noise.py | 14 ++++++++++ src/bloqade/squin/stdlib/simple/__init__.py | 1 + src/bloqade/squin/stdlib/simple/noise.py | 16 ++++++++++- src/bloqade/stim/__init__.py | 1 + src/bloqade/stim/_wrappers.py | 6 ++++ src/bloqade/stim/dialects/noise/emit.py | 1 + src/bloqade/stim/dialects/noise/stmts.py | 9 ++++-- src/bloqade/stim/rewrite/qubit_to_stim.py | 4 +++ src/bloqade/stim/rewrite/squin_noise.py | 6 ++-- src/bloqade/stim/rewrite/util.py | 28 +++++++++++++++++++ src/bloqade/stim/rewrite/wire_to_stim.py | 4 +++ 16 files changed, 112 insertions(+), 5 deletions(-) diff --git a/src/bloqade/pyqrack/squin/noise/native.py b/src/bloqade/pyqrack/squin/noise/native.py index b955e096..d503a87a 100644 --- a/src/bloqade/pyqrack/squin/noise/native.py +++ b/src/bloqade/pyqrack/squin/noise/native.py @@ -3,6 +3,7 @@ from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter from bloqade.squin.noise.stmts import ( QubitLoss, + CorrelatedQubitLoss, Depolarize, Depolarize2, TwoQubitPauliChannel, @@ -88,6 +89,16 @@ def qubit_loss( if interp.rng_state.uniform(0.0, 1.0) <= p: qbit.drop() + @interp.impl(CorrelatedQubitLoss) + def correlated_qubit_loss( + self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: CorrelatedQubitLoss + ): + p = frame.get(stmt.p) + qubits: list[PyQrackQubit] = frame.get(stmt.qubits) + if interp.rng_state.uniform(0.0, 1.0) <= p: + for qbit in qubits: + qbit.drop() + def apply_single_qubit_pauli_error( self, interp: PyQrackInterpreter, diff --git a/src/bloqade/squin/__init__.py b/src/bloqade/squin/__init__.py index 3fb5a1dd..8875ff50 100644 --- a/src/bloqade/squin/__init__.py +++ b/src/bloqade/squin/__init__.py @@ -32,6 +32,7 @@ bit_flip as bit_flip, depolarize as depolarize, qubit_loss as qubit_loss, + correlated_qubit_loss as correlated_qubit_loss, sqrt_x_adj as sqrt_x_adj, sqrt_y_adj as sqrt_y_adj, sqrt_z_adj as sqrt_z_adj, diff --git a/src/bloqade/squin/noise/_interface.py b/src/bloqade/squin/noise/_interface.py index 4bd00cf3..2d784c73 100644 --- a/src/bloqade/squin/noise/_interface.py +++ b/src/bloqade/squin/noise/_interface.py @@ -37,3 +37,7 @@ def two_qubit_pauli_channel( @wraps(stmts.QubitLoss) def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ... + + +@wraps(stmts.CorrelatedQubitLoss) +def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ... diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index eddb7942..61ba2d8e 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -97,3 +97,13 @@ class QubitLoss(SingleQubitNoiseChannel): # NOTE: qubit loss error (not supported by Stim) p: ir.SSAValue = info.argument(types.Float) qubits: ir.SSAValue = info.argument(ilist.IListType) + + +@statement(dialect=dialect) +class CorrelatedQubitLoss(SingleQubitNoiseChannel): + """ + Apply a correlated atom loss channel. + """ + + p: ir.SSAValue = info.argument(types.Float) + qubits: ir.SSAValue = info.argument(ilist.IListType) diff --git a/src/bloqade/squin/stdlib/broadcast/__init__.py b/src/bloqade/squin/stdlib/broadcast/__init__.py index 81eba5f9..99a6c4e7 100644 --- a/src/bloqade/squin/stdlib/broadcast/__init__.py +++ b/src/bloqade/squin/stdlib/broadcast/__init__.py @@ -2,6 +2,7 @@ bit_flip as bit_flip, depolarize as depolarize, qubit_loss as qubit_loss, + correlated_qubit_loss as correlated_qubit_loss, depolarize2 as depolarize2, two_qubit_pauli_channel as two_qubit_pauli_channel, single_qubit_pauli_channel as single_qubit_pauli_channel, diff --git a/src/bloqade/squin/stdlib/broadcast/noise.py b/src/bloqade/squin/stdlib/broadcast/noise.py index 10686007..9287d1c9 100644 --- a/src/bloqade/squin/stdlib/broadcast/noise.py +++ b/src/bloqade/squin/stdlib/broadcast/noise.py @@ -103,6 +103,20 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: noise.qubit_loss(p, qubits) +@kernel +def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: + """ + Apply a correlated qubit loss channel to the given qubits. + + All qubits are lost together with a probability `p`. + + Args: + p (float): Probability of the qubits being lost. + qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied. + """ + noise.correlated_qubit_loss(p, qubits) + + # NOTE: actual stdlib that doesn't wrap statements starts here diff --git a/src/bloqade/squin/stdlib/simple/__init__.py b/src/bloqade/squin/stdlib/simple/__init__.py index 81eba5f9..99a6c4e7 100644 --- a/src/bloqade/squin/stdlib/simple/__init__.py +++ b/src/bloqade/squin/stdlib/simple/__init__.py @@ -2,6 +2,7 @@ bit_flip as bit_flip, depolarize as depolarize, qubit_loss as qubit_loss, + correlated_qubit_loss as correlated_qubit_loss, depolarize2 as depolarize2, two_qubit_pauli_channel as two_qubit_pauli_channel, single_qubit_pauli_channel as single_qubit_pauli_channel, diff --git a/src/bloqade/squin/stdlib/simple/noise.py b/src/bloqade/squin/stdlib/simple/noise.py index 4bf0ccf8..5a33de01 100644 --- a/src/bloqade/squin/stdlib/simple/noise.py +++ b/src/bloqade/squin/stdlib/simple/noise.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeVar +from typing import Any, Literal, TypeVar from kirin.dialects import ilist @@ -97,6 +97,20 @@ def qubit_loss(p: float, qubit: Qubit) -> None: broadcast.qubit_loss(p, ilist.IList([qubit])) +@kernel +def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: + """ + Apply a correlated qubit loss channel to the given qubits. + + All qubits are lost together with a probability `p`. + + Args: + p (float): Probability of the qubits being lost. + qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied. + """ + broadcast.correlated_qubit_loss(p, qubits) + + # NOTE: actual stdlib that doesn't wrap statements starts here diff --git a/src/bloqade/stim/__init__.py b/src/bloqade/stim/__init__.py index eb8644d5..52c0c3c2 100644 --- a/src/bloqade/stim/__init__.py +++ b/src/bloqade/stim/__init__.py @@ -32,6 +32,7 @@ detector as detector, identity as identity, qubit_loss as qubit_loss, + correlated_qubit_loss as correlated_qubit_loss, depolarize1 as depolarize1, depolarize2 as depolarize2, pauli_string as pauli_string, diff --git a/src/bloqade/stim/_wrappers.py b/src/bloqade/stim/_wrappers.py index 41ef5bb2..008e6635 100644 --- a/src/bloqade/stim/_wrappers.py +++ b/src/bloqade/stim/_wrappers.py @@ -194,3 +194,9 @@ def z_error(p: float, targets: tuple[int, ...]) -> None: ... @wraps(noise.QubitLoss) def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ... + + +@wraps(noise.CorrelatedQubitLoss) +def correlated_qubit_loss( + probs: tuple[float, ...], targets: tuple[int, ...] +) -> None: ... diff --git a/src/bloqade/stim/dialects/noise/emit.py b/src/bloqade/stim/dialects/noise/emit.py index 73242a2f..8f3ebefa 100644 --- a/src/bloqade/stim/dialects/noise/emit.py +++ b/src/bloqade/stim/dialects/noise/emit.py @@ -81,6 +81,7 @@ def non_stim_error( return () @impl(stmts.TrivialCorrelatedError) + @impl(stmts.CorrelatedQubitLoss) def non_stim_corr_error( self, emit: EmitStimMain, diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index bdfdfbbb..7594d127 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/src/bloqade/stim/dialects/noise/stmts.py @@ -89,8 +89,8 @@ class NonStimError(ir.Statement): class NonStimCorrelatedError(ir.Statement): name = "NonStimCorrelatedError" traits = frozenset({lowering.FromPythonCall()}) - nonce: int = ( - info.attribute() + nonce: int = info.attribute( + default_factory=lambda: __import__("random").getrandbits(32) ) # Must be a unique value, otherwise stim might merge two correlated errors with equal probabilities probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @@ -109,3 +109,8 @@ class TrivialError(NonStimError): @statement(dialect=dialect) class QubitLoss(NonStimError): name = "loss" + + +@statement(dialect=dialect) +class CorrelatedQubitLoss(NonStimCorrelatedError): + name = "correlated_loss" diff --git a/src/bloqade/stim/rewrite/qubit_to_stim.py b/src/bloqade/stim/rewrite/qubit_to_stim.py index 1c4302c8..f7c47888 100644 --- a/src/bloqade/stim/rewrite/qubit_to_stim.py +++ b/src/bloqade/stim/rewrite/qubit_to_stim.py @@ -8,6 +8,7 @@ SQUIN_STIM_OP_MAPPING, rewrite_Control, rewrite_QubitLoss, + rewrite_CorrelatedQubitLoss, insert_qubit_idx_from_address, ) @@ -38,6 +39,9 @@ def rewrite_Apply_and_Broadcast( if isinstance(applied_op, noise.stmts.QubitLoss): return rewrite_QubitLoss(stmt) + if isinstance(applied_op, noise.stmts.CorrelatedQubitLoss): + return rewrite_CorrelatedQubitLoss(stmt) + assert isinstance(applied_op, op.stmts.Operator) if isinstance(applied_op, op.stmts.Control): diff --git a/src/bloqade/stim/rewrite/squin_noise.py b/src/bloqade/stim/rewrite/squin_noise.py index 8952792a..10c3cef6 100644 --- a/src/bloqade/stim/rewrite/squin_noise.py +++ b/src/bloqade/stim/rewrite/squin_noise.py @@ -31,8 +31,10 @@ def rewrite_Apply_and_Broadcast( # this is an SSAValue, need it to be the actual operator applied_op = stmt.operator.owner - - if isinstance(applied_op, squin_noise.stmts.QubitLoss): + if isinstance( + applied_op, + (squin_noise.stmts.QubitLoss, squin_noise.stmts.CorrelatedQubitLoss), + ): return RewriteResult() if isinstance(applied_op, squin_noise.stmts.NoiseChannel): diff --git a/src/bloqade/stim/rewrite/util.py b/src/bloqade/stim/rewrite/util.py index 1efab606..e3b8100f 100644 --- a/src/bloqade/stim/rewrite/util.py +++ b/src/bloqade/stim/rewrite/util.py @@ -21,6 +21,7 @@ op.stmts.Identity: gate.Identity, op.stmts.Reset: collapse.RZ, squin_noise.stmts.QubitLoss: stim_noise.QubitLoss, + squin_noise.stmts.CorrelatedQubitLoss: stim_noise.CorrelatedQubitLoss, } # Squin allows creation of control gates where the gate can be any operator, @@ -201,6 +202,33 @@ def rewrite_QubitLoss( return RewriteResult(has_done_something=True) +def rewrite_CorrelatedQubitLoss( + stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply, +) -> RewriteResult: + """ + Rewrite CorrelatedQubitLoss statements to Stim's TrivialCorrelatedError. + """ + + squin_loss_op = stmt.operator.owner + assert isinstance(squin_loss_op, squin_noise.stmts.CorrelatedQubitLoss) + + qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt) + if qubit_idx_ssas is None: + return RewriteResult() + + stim_loss_stmt = stim_noise.CorrelatedQubitLoss( + targets=qubit_idx_ssas, + probs=(squin_loss_op.p,), + ) + + if isinstance(stmt, (wire.Apply, wire.Broadcast)): + create_wire_passthrough(stmt) + + stmt.replace_by(stim_loss_stmt) + + return RewriteResult(has_done_something=True) + + def create_wire_passthrough(stmt: wire.Apply | wire.Broadcast) -> None: for input_wire, output_wire in zip(stmt.inputs, stmt.results): diff --git a/src/bloqade/stim/rewrite/wire_to_stim.py b/src/bloqade/stim/rewrite/wire_to_stim.py index 94640ebd..05d3ea30 100644 --- a/src/bloqade/stim/rewrite/wire_to_stim.py +++ b/src/bloqade/stim/rewrite/wire_to_stim.py @@ -5,6 +5,7 @@ from bloqade.stim.rewrite.util import ( SQUIN_STIM_OP_MAPPING, rewrite_Control, + rewrite_CorrelatedQubitLoss, rewrite_QubitLoss, insert_qubit_idx_from_wire_ssa, ) @@ -29,6 +30,9 @@ def rewrite_Apply_and_Broadcast( if isinstance(applied_op, noise.stmts.QubitLoss): return rewrite_QubitLoss(stmt) + if isinstance(applied_op, noise.stmts.CorrelatedQubitLoss): + return rewrite_CorrelatedQubitLoss(stmt) + assert isinstance(applied_op, op.stmts.Operator) if isinstance(applied_op, op.stmts.Control): From 6a07247a7f99cd51b5b7213def19cc711cfba7cc Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 6 Oct 2025 17:28:40 -0700 Subject: [PATCH 02/14] Add tests for correlated qubit loss in noise simulations --- test/pyqrack/squin/test_noise.py | 16 ++++++++++ test/squin/noise/test_stdlib_noise.py | 30 +++++++++++++++++++ test/stim/dialects/stim/test_stim_circuits.py | 14 +++++++++ 3 files changed, 60 insertions(+) diff --git a/test/pyqrack/squin/test_noise.py b/test/pyqrack/squin/test_noise.py index 12020069..f746812f 100644 --- a/test/pyqrack/squin/test_noise.py +++ b/test/pyqrack/squin/test_noise.py @@ -16,6 +16,22 @@ def main(): assert not qubit.is_active() +def test_correlated_loss(): + @squin.kernel + def main(): + q = squin.qubit.new(5) + squin.correlated_qubit_loss(0.5, q[0:4]) + return q + + target = PyQrack(5) + for _ in range(10): + qubits = target.run(main) + qubits_active = [q.is_active() for q in qubits[:4]] + assert all(qubits_active) or not any(qubits_active) + + assert qubits[4].is_active() + + def test_pauli_channel(): @squin.kernel def single_qubit(): diff --git a/test/squin/noise/test_stdlib_noise.py b/test/squin/noise/test_stdlib_noise.py index 2dd90345..613dc995 100644 --- a/test/squin/noise/test_stdlib_noise.py +++ b/test/squin/noise/test_stdlib_noise.py @@ -1,3 +1,5 @@ +import pytest +import numpy as np from bloqade import squin from bloqade.pyqrack import PyQrackQubit, StackMemorySimulator @@ -19,6 +21,34 @@ def main(): assert not qubit.is_active() +@pytest.mark.parametrize( + "seed, expected_loss_triggered", + [ + (0, False), # Seed 0: no loss + (2, True), # Seed 2: qubits 0-3 are lost + ], +) +def test_correlated_loss(seed, expected_loss_triggered): + + @squin.kernel + def main(): + q = squin.qubit.new(5) + squin.correlated_qubit_loss(0.5, q[0:4]) + return q + + rng = np.random.default_rng(seed=seed) + sim = StackMemorySimulator(min_qubits=5, rng_state=rng) + qubits = sim.run(main) + + for q in qubits: + assert isinstance(q, PyQrackQubit) + + for q in qubits[:4]: + assert not q.is_active() if expected_loss_triggered else q.is_active() + + assert qubits[4].is_active() + + def test_bit_flip(): @squin.kernel diff --git a/test/stim/dialects/stim/test_stim_circuits.py b/test/stim/dialects/stim/test_stim_circuits.py index 7b7c1779..7b531b0c 100644 --- a/test/stim/dialects/stim/test_stim_circuits.py +++ b/test/stim/dialects/stim/test_stim_circuits.py @@ -1,5 +1,6 @@ from bloqade import stim from bloqade.stim.emit import EmitStimMain +import re interp = EmitStimMain(stim.main) @@ -94,6 +95,19 @@ def test_qubit_loss(): assert interp.get_output() == "\nI_ERROR[loss](0.10000000, 0.20000000) 0 1 2" +def test_correlated_qubit_loss(): + + @stim.main + def test_correlated_qubit_loss(): + stim.correlated_qubit_loss(probs=(0.1,), targets=(0, 3, 1)) + + interp.run(test_correlated_qubit_loss, args=()) + + assert re.match( + r"\nI_ERROR\[correlated_loss:\d+\]\(0\.10000000\) 0 3 1", interp.get_output() + ) + + def test_collapse(): @stim.main def test_measure(): From 2c90cacad6d4b418b424dee975f583bc3120d98e Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 6 Oct 2025 21:36:53 -0700 Subject: [PATCH 03/14] Refactor CorrelatedQubitLoss to inherit from NoiseChannel instead of SingleQubitNoiseChannel for readability --- src/bloqade/squin/noise/stmts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index 61ba2d8e..873c3bcd 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -100,7 +100,7 @@ class QubitLoss(SingleQubitNoiseChannel): @statement(dialect=dialect) -class CorrelatedQubitLoss(SingleQubitNoiseChannel): +class CorrelatedQubitLoss(NoiseChannel): """ Apply a correlated atom loss channel. """ From cd69a37ef8a05e651a6ff2af74899c81614d03f6 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 6 Oct 2025 21:43:05 -0700 Subject: [PATCH 04/14] Run precommit --- src/bloqade/pyqrack/device.py | 2 +- src/bloqade/pyqrack/squin/noise/native.py | 2 +- src/bloqade/squin/__init__.py | 2 +- src/bloqade/squin/stdlib/broadcast/__init__.py | 2 +- src/bloqade/squin/stdlib/simple/__init__.py | 2 +- src/bloqade/stim/__init__.py | 2 +- src/bloqade/stim/rewrite/wire_to_stim.py | 2 +- test/squin/noise/test_stdlib_noise.py | 3 ++- test/stim/dialects/stim/test_stim_circuits.py | 3 ++- 9 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/bloqade/pyqrack/device.py b/src/bloqade/pyqrack/device.py index ba5fe274..2b58e3af 100644 --- a/src/bloqade/pyqrack/device.py +++ b/src/bloqade/pyqrack/device.py @@ -30,7 +30,7 @@ class QuantumState(NamedTuple): A representation of a quantum state as a density matrix, where the density matrix is rho = sum_i eigenvalues[i] |eigenvectors[:,i]> Date: Tue, 7 Oct 2025 08:54:43 -0700 Subject: [PATCH 05/14] Add filterwarnings for cirq FutureWarning in pyproject.toml For python 3.10, cirq 1.5.0 is used since later versions only support 3.11. Here, the test suite raises 10k warning which all originate from within cirq. --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 8886e76c..ec421f78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,3 +118,6 @@ include = ["src/bloqade/*"] [tool.pytest.ini_options] testpaths = "test/" +filterwarnings = [ + "ignore:In cirq 1.6 the default value of `use_repetition_ids`:FutureWarning:cirq.circuits.circuit_operation", +] From 2bba610346a0c471a576d5e5df1c9179298e80dd Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Fri, 10 Oct 2025 09:46:36 -0700 Subject: [PATCH 06/14] Make type of squin.noise.CorrelatedQubitLoss.qubits list[list[Qubit]] instead of list[Qubit] --- src/bloqade/pyqrack/squin/noise/native.py | 9 ++++--- src/bloqade/squin/noise/_interface.py | 4 ++- src/bloqade/squin/noise/stmts.py | 6 +++-- src/bloqade/squin/stdlib/broadcast/noise.py | 22 ++++++++++++---- src/bloqade/squin/stdlib/simple/noise.py | 2 +- test/squin/noise/test_stdlib_noise.py | 28 +++++++++++++++++++++ 6 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/bloqade/pyqrack/squin/noise/native.py b/src/bloqade/pyqrack/squin/noise/native.py index 9342accf..6695e9f0 100644 --- a/src/bloqade/pyqrack/squin/noise/native.py +++ b/src/bloqade/pyqrack/squin/noise/native.py @@ -94,10 +94,11 @@ def correlated_qubit_loss( self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: CorrelatedQubitLoss ): p = frame.get(stmt.p) - qubits: list[PyQrackQubit] = frame.get(stmt.qubits) - if interp.rng_state.uniform(0.0, 1.0) <= p: - for qbit in qubits: - qbit.drop() + qubits: list[list[PyQrackQubit]] = frame.get(stmt.qubits) + for qubit_group in qubits: + if interp.rng_state.uniform(0.0, 1.0) <= p: + for qbit in qubit_group: + qbit.drop() def apply_single_qubit_pauli_error( self, diff --git a/src/bloqade/squin/noise/_interface.py b/src/bloqade/squin/noise/_interface.py index 2d784c73..d917b06d 100644 --- a/src/bloqade/squin/noise/_interface.py +++ b/src/bloqade/squin/noise/_interface.py @@ -40,4 +40,6 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ... @wraps(stmts.CorrelatedQubitLoss) -def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ... +def correlated_qubit_loss( + p: float, qubits: ilist.IList[ilist.IList[Qubit, Any], Any] +) -> None: ... diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index 873c3bcd..21332659 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -96,7 +96,7 @@ class QubitLoss(SingleQubitNoiseChannel): # NOTE: qubit loss error (not supported by Stim) p: ir.SSAValue = info.argument(types.Float) - qubits: ir.SSAValue = info.argument(ilist.IListType) + qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any]) @statement(dialect=dialect) @@ -106,4 +106,6 @@ class CorrelatedQubitLoss(NoiseChannel): """ p: ir.SSAValue = info.argument(types.Float) - qubits: ir.SSAValue = info.argument(ilist.IListType) + qubits: ir.SSAValue = info.argument( + ilist.IListType[ilist.IListType[QubitType, types.Any], types.Any] + ) diff --git a/src/bloqade/squin/stdlib/broadcast/noise.py b/src/bloqade/squin/stdlib/broadcast/noise.py index 9287d1c9..9f54b426 100644 --- a/src/bloqade/squin/stdlib/broadcast/noise.py +++ b/src/bloqade/squin/stdlib/broadcast/noise.py @@ -104,15 +104,27 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: @kernel -def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: +def correlated_qubit_loss( + p: float, qubits: ilist.IList[ilist.IList[Qubit, Any], Any] +) -> None: """ - Apply a correlated qubit loss channel to the given qubits. + Apply correlated qubit loss channels to groups of qubits. - All qubits are lost together with a probability `p`. + For each group of qubits, applies a correlated loss channel where all qubits + within the group are lost together with probability `p`. Loss events are independent + between different groups. Args: - p (float): Probability of the qubits being lost. - qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied. + p (float): Loss probability for each group. + qubits (IList[IList[Qubit, Any], Any]): List of qubit groups. Each sublist + represents a group of qubits to which a correlated loss channel is applied. + + Example: + >>> q1 = squin.qubit.new(3) # First group: qubits 0, 1, 2 + >>> q2 = squin.qubit.new(3) # Second group: qubits 3, 4, 5 + >>> squin.broadcast.correlated_qubit_loss(0.5, [q1, q2]) + # Each group has 50% chance: either all qubits lost or none lost. + # Group 1 and Group 2 outcomes are independent. """ noise.correlated_qubit_loss(p, qubits) diff --git a/src/bloqade/squin/stdlib/simple/noise.py b/src/bloqade/squin/stdlib/simple/noise.py index 5a33de01..8f1fb341 100644 --- a/src/bloqade/squin/stdlib/simple/noise.py +++ b/src/bloqade/squin/stdlib/simple/noise.py @@ -108,7 +108,7 @@ def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: p (float): Probability of the qubits being lost. qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied. """ - broadcast.correlated_qubit_loss(p, qubits) + broadcast.correlated_qubit_loss(p, ilist.IList([qubits])) # NOTE: actual stdlib that doesn't wrap statements starts here diff --git a/test/squin/noise/test_stdlib_noise.py b/test/squin/noise/test_stdlib_noise.py index 29b3edfa..c65b6eb9 100644 --- a/test/squin/noise/test_stdlib_noise.py +++ b/test/squin/noise/test_stdlib_noise.py @@ -50,6 +50,34 @@ def main(): assert qubits[4].is_active() +@pytest.mark.parametrize( + "seed, expected_loss_triggered", + [(0, (False, True)), (1, (False, False)), (2, (True, True)), (8, (True, False))], +) +def test_correlated_loss_broadcast(seed, expected_loss_triggered): + + @squin.kernel + def main(): + q = squin.qubit.new(6) + q1 = q[:3] + q2 = q[3:] + squin.broadcast.correlated_qubit_loss(0.5, [q1, q2]) + return q + + rng = np.random.default_rng(seed=seed) + sim = StackMemorySimulator(min_qubits=5, rng_state=rng) + qubits = sim.run(main) + + for q in qubits: + assert isinstance(q, PyQrackQubit) + + for q in qubits[:3]: + assert not q.is_active() if expected_loss_triggered[0] else q.is_active() + + for q in qubits[3:]: + assert not q.is_active() if expected_loss_triggered[1] else q.is_active() + + def test_bit_flip(): @squin.kernel From bcb3e761604a02b171fc8c6e7338287ca94b2725 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Fri, 10 Oct 2025 09:55:46 -0700 Subject: [PATCH 07/14] Adjust typing to ensure correlated errors are only broadcast over rectangular qubit matrices --- src/bloqade/squin/noise/_interface.py | 2 +- src/bloqade/squin/noise/stmts.py | 2 +- src/bloqade/squin/stdlib/broadcast/noise.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/bloqade/squin/noise/_interface.py b/src/bloqade/squin/noise/_interface.py index d917b06d..0096f148 100644 --- a/src/bloqade/squin/noise/_interface.py +++ b/src/bloqade/squin/noise/_interface.py @@ -41,5 +41,5 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ... @wraps(stmts.CorrelatedQubitLoss) def correlated_qubit_loss( - p: float, qubits: ilist.IList[ilist.IList[Qubit, Any], Any] + p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any] ) -> None: ... diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index 21332659..7ab0042e 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -107,5 +107,5 @@ class CorrelatedQubitLoss(NoiseChannel): p: ir.SSAValue = info.argument(types.Float) qubits: ir.SSAValue = info.argument( - ilist.IListType[ilist.IListType[QubitType, types.Any], types.Any] + ilist.IListType[ilist.IListType[QubitType, N], types.Any] ) diff --git a/src/bloqade/squin/stdlib/broadcast/noise.py b/src/bloqade/squin/stdlib/broadcast/noise.py index 9f54b426..ebb7d8ba 100644 --- a/src/bloqade/squin/stdlib/broadcast/noise.py +++ b/src/bloqade/squin/stdlib/broadcast/noise.py @@ -105,7 +105,7 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: @kernel def correlated_qubit_loss( - p: float, qubits: ilist.IList[ilist.IList[Qubit, Any], Any] + p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any] ) -> None: """ Apply correlated qubit loss channels to groups of qubits. @@ -116,7 +116,7 @@ def correlated_qubit_loss( Args: p (float): Loss probability for each group. - qubits (IList[IList[Qubit, Any], Any]): List of qubit groups. Each sublist + qubits (IList[IList[Qubit, N], Any]): List of qubit groups. Each sublist represents a group of qubits to which a correlated loss channel is applied. Example: From fdafe6da43ec479f47c2f1c2c1d13deb8308f7e8 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 14 Oct 2025 16:50:36 -0700 Subject: [PATCH 08/14] Add CorrelatedQubitNoise to SquinNoiseToStim rewrite --- src/bloqade/squin/noise/stmts.py | 7 ++++- src/bloqade/stim/rewrite/squin_noise.py | 40 +++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index 7ab0042e..cf810d8a 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -24,6 +24,11 @@ class TwoQubitNoiseChannel(NoiseChannel): pass +@statement +class MultiQubitNoiseChannel(NoiseChannel): + pass + + @statement(dialect=dialect) class SingleQubitPauliChannel(SingleQubitNoiseChannel): """ @@ -100,7 +105,7 @@ class QubitLoss(SingleQubitNoiseChannel): @statement(dialect=dialect) -class CorrelatedQubitLoss(NoiseChannel): +class CorrelatedQubitLoss(MultiQubitNoiseChannel): """ Apply a correlated atom loss channel. """ diff --git a/src/bloqade/stim/rewrite/squin_noise.py b/src/bloqade/stim/rewrite/squin_noise.py index c0bb9c50..a5c8097a 100644 --- a/src/bloqade/stim/rewrite/squin_noise.py +++ b/src/bloqade/stim/rewrite/squin_noise.py @@ -9,10 +9,13 @@ from bloqade.squin import noise as squin_noise from bloqade.stim.dialects import noise as stim_noise from bloqade.stim.rewrite.util import insert_qubit_idx_from_address +from bloqade.analysis.address.lattice import AddressTuple +from bloqade.squin.rewrite.wrap_analysis import AddressAttribute @dataclass class SquinNoiseToStim(RewriteRule): + _correlated_loss_counter: int = 0 def rewrite_Statement(self, node: Statement) -> RewriteResult: match node: @@ -31,6 +34,28 @@ def rewrite_NoiseChannel( if rewrite_method is None: return RewriteResult() + if isinstance(stmt, squin_noise.stmts.MultiQubitNoiseChannel): + # MultiQubitNoiseChannel represents a broadcast operation, but Stim does not + # 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) + if not isinstance(qubit_address_attr, AddressAttribute): + return RewriteResult() + + address_tuple = qubit_address_attr.address + + if not isinstance(address_tuple, AddressTuple): + return RewriteResult() + + for address in address_tuple.data: + qubit_idx_ssas = insert_qubit_idx_from_address( + AddressAttribute(address=address), stmt + ) + stim_stmt = rewrite_method(stmt, tuple(qubit_idx_ssas)) + stim_stmt.insert_before(stmt) + stmt.delete() + return RewriteResult(has_done_something=True) + if isinstance(stmt, squin_noise.stmts.SingleQubitNoiseChannel): qubit_address_attr = stmt.qubits.hints.get("address", None) if qubit_address_attr is None: @@ -96,6 +121,21 @@ def rewrite_QubitLoss( return stim_stmt + def rewrite_CorrelatedQubitLoss( + self, + stmt: squin_noise.stmts.CorrelatedQubitLoss, + qubit_idx_ssas: Tuple[SSAValue], + ) -> Statement: + """Rewrite squin.noise.CorrelatedQubitLoss to stim.CorrelatedQubitLoss.""" + stim_stmt = stim_noise.CorrelatedQubitLoss( + targets=qubit_idx_ssas, + probs=(stmt.p,), + nonce=self._correlated_loss_counter, + ) + self._correlated_loss_counter += 1 + + return stim_stmt + def rewrite_Depolarize( self, stmt: squin_noise.stmts.Depolarize, From 4c3f85215e5c33f1be982223ea304c866aa97585 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 14 Oct 2025 16:50:53 -0700 Subject: [PATCH 09/14] Add unit tests --- .../dialects/stim/emit/test_stim_noise.py | 18 ++++++++++++ test/stim/passes/test_squin_noise_to_stim.py | 28 +++++++++++++++++++ test/stim/wrapper/test_wrapper.py | 16 ++++++++++- 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/test/stim/dialects/stim/emit/test_stim_noise.py b/test/stim/dialects/stim/emit/test_stim_noise.py index 7171736b..4f1f19f4 100644 --- a/test/stim/dialects/stim/emit/test_stim_noise.py +++ b/test/stim/dialects/stim/emit/test_stim_noise.py @@ -31,3 +31,21 @@ def test_pauli2(): out.strip() == "PAULI_CHANNEL_2(0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000, 0.10000000) 0 3 4 5" ) + + +def test_qubit_loss(): + @stim.main + def test_qubit_loss(): + stim.qubit_loss(probs=(0.1,), targets=(0, 1, 2)) + + out = codegen(test_qubit_loss) + assert out.strip() == "I_ERROR[loss](0.10000000) 0 1 2" + + +def test_correlated_qubit_loss(): + @stim.main + def test_correlated_qubit_loss(): + stim.correlated_qubit_loss(probs=(0.1,), targets=(0, 1, 2), nonce=3) + + out = codegen(test_correlated_qubit_loss) + assert out.strip() == "I_ERROR[correlated_loss:3](0.10000000) 0 1 2" diff --git a/test/stim/passes/test_squin_noise_to_stim.py b/test/stim/passes/test_squin_noise_to_stim.py index ad425845..22602121 100644 --- a/test/stim/passes/test_squin_noise_to_stim.py +++ b/test/stim/passes/test_squin_noise_to_stim.py @@ -214,6 +214,34 @@ def test(): assert codegen(test) == expected_stim_program +def test_correlated_qubit_loss(): + @kernel + def test(): + q = qubit.new(3) + sq.correlated_qubit_loss(0.1, qubits=q[:2]) + + SquinToStimPass(test.dialects)(test) + + expected = "I_ERROR[correlated_loss:0](0.10000000) 0 1" + assert codegen(test) == expected + + +def test_broadcast_correlated_qubit_loss(): + @kernel + def test(): + q1 = qubit.new(3) + q2 = qubit.new(3) + sq.broadcast.correlated_qubit_loss(0.1, qubits=[q1, q2]) + + SquinToStimPass(test.dialects)(test) + + expected = ( + "I_ERROR[correlated_loss:0](0.10000000) 0 1 2\n" + "I_ERROR[correlated_loss:1](0.10000000) 3 4 5" + ) + assert codegen(test) == expected + + def get_stmt_at_idx(method: ir.Method, idx: int) -> ir.Statement: return method.callable_region.blocks[0].stmts.at(idx) diff --git a/test/stim/wrapper/test_wrapper.py b/test/stim/wrapper/test_wrapper.py index 1a8fdff8..53744a57 100644 --- a/test/stim/wrapper/test_wrapper.py +++ b/test/stim/wrapper/test_wrapper.py @@ -1,5 +1,5 @@ from bloqade import stim -from bloqade.stim.dialects import gate, collapse, auxiliary +from bloqade.stim.dialects import gate, noise, collapse, auxiliary def test_wrapper_x(): @@ -461,3 +461,17 @@ def main_ry_wrap(): stim.ry(targets=(0, 1, 2)) assert main_ry.callable_region.is_structurally_equal(main_ry_wrap.callable_region) + + +def test_wrap_correlated_qubit_loss(): + @stim.main + def main_correlated_qubit_loss(): + noise.CorrelatedQubitLoss(probs=(0.1,), targets=(0, 1, 2), nonce=3) + + @stim.main + def main_correlated_qubit_loss_wrap(): + stim.correlated_qubit_loss(probs=(0.1,), targets=(0, 1, 2), nonce=3) + + assert main_correlated_qubit_loss.callable_region.is_structurally_equal( + main_correlated_qubit_loss_wrap.callable_region + ) From 86a499fb13e5e80595d18c28d0bc2733ea705eff Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 15 Oct 2025 10:19:28 -0700 Subject: [PATCH 10/14] Remove MultiQubitNoiseChannel --- src/bloqade/squin/noise/stmts.py | 7 +------ src/bloqade/stim/rewrite/squin_noise.py | 4 ++-- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index cf810d8a..7ab0042e 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/src/bloqade/squin/noise/stmts.py @@ -24,11 +24,6 @@ class TwoQubitNoiseChannel(NoiseChannel): pass -@statement -class MultiQubitNoiseChannel(NoiseChannel): - pass - - @statement(dialect=dialect) class SingleQubitPauliChannel(SingleQubitNoiseChannel): """ @@ -105,7 +100,7 @@ class QubitLoss(SingleQubitNoiseChannel): @statement(dialect=dialect) -class CorrelatedQubitLoss(MultiQubitNoiseChannel): +class CorrelatedQubitLoss(NoiseChannel): """ Apply a correlated atom loss channel. """ diff --git a/src/bloqade/stim/rewrite/squin_noise.py b/src/bloqade/stim/rewrite/squin_noise.py index a5c8097a..b40a5d07 100644 --- a/src/bloqade/stim/rewrite/squin_noise.py +++ b/src/bloqade/stim/rewrite/squin_noise.py @@ -34,8 +34,8 @@ def rewrite_NoiseChannel( if rewrite_method is None: return RewriteResult() - if isinstance(stmt, squin_noise.stmts.MultiQubitNoiseChannel): - # MultiQubitNoiseChannel represents a broadcast operation, but Stim does not + if isinstance(stmt, squin_noise.stmts.CorrelatedQubitLoss): + # CorrelatedQubitLoss represents a broadcast operation, but Stim does not # 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) From 23fb34deb290781d88588e2efd5c2170022f6a4f Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 15 Oct 2025 10:23:18 -0700 Subject: [PATCH 11/14] Handle None values of insert_qubit_idx_from_address --- src/bloqade/stim/rewrite/squin_noise.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/bloqade/stim/rewrite/squin_noise.py b/src/bloqade/stim/rewrite/squin_noise.py index b40a5d07..ed27804d 100644 --- a/src/bloqade/stim/rewrite/squin_noise.py +++ b/src/bloqade/stim/rewrite/squin_noise.py @@ -47,13 +47,18 @@ def rewrite_NoiseChannel( if not isinstance(address_tuple, AddressTuple): return RewriteResult() - for address in address_tuple.data: - qubit_idx_ssas = insert_qubit_idx_from_address( - AddressAttribute(address=address), stmt - ) + qubit_idx_ssas_list = [ + insert_qubit_idx_from_address(AddressAttribute(address=address), stmt) + for address in address_tuple.data + ] + if None in qubit_idx_ssas_list: + return RewriteResult() + + for qubit_idx_ssas in qubit_idx_ssas_list: stim_stmt = rewrite_method(stmt, tuple(qubit_idx_ssas)) stim_stmt.insert_before(stmt) stmt.delete() + return RewriteResult(has_done_something=True) if isinstance(stmt, squin_noise.stmts.SingleQubitNoiseChannel): From 03a335812e9f8973b9491217d3fcfc5ba3df8814 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 15 Oct 2025 10:34:06 -0700 Subject: [PATCH 12/14] Remove dynamic import --- src/bloqade/stim/dialects/noise/stmts.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index 7594d127..b336f901 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/src/bloqade/stim/dialects/noise/stmts.py @@ -1,3 +1,5 @@ +import random + from kirin import ir, types, lowering from kirin.decl import info, statement @@ -89,9 +91,8 @@ class NonStimError(ir.Statement): class NonStimCorrelatedError(ir.Statement): name = "NonStimCorrelatedError" traits = frozenset({lowering.FromPythonCall()}) - nonce: int = info.attribute( - default_factory=lambda: __import__("random").getrandbits(32) - ) # Must be a unique value, otherwise stim might merge two correlated errors with equal probabilities + # nonce must be a unique value, otherwise stim might merge two correlated errors + nonce: int = info.attribute(default_factory=lambda: random.getrandbits(32)) probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) From d130835bcd563fc88e1bd2482a39a86e05650ad7 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 15 Oct 2025 13:12:26 -0700 Subject: [PATCH 13/14] Make NonStimCorrelatedError.nonce=0 instead of random value --- src/bloqade/stim/dialects/noise/stmts.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index b336f901..5eae8b11 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/src/bloqade/stim/dialects/noise/stmts.py @@ -1,5 +1,3 @@ -import random - from kirin import ir, types, lowering from kirin.decl import info, statement @@ -91,8 +89,7 @@ class NonStimError(ir.Statement): class NonStimCorrelatedError(ir.Statement): name = "NonStimCorrelatedError" traits = frozenset({lowering.FromPythonCall()}) - # nonce must be a unique value, otherwise stim might merge two correlated errors - nonce: int = info.attribute(default_factory=lambda: random.getrandbits(32)) + nonce: int = info.attribute(default=0) probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) From 20fe041548feaa8804ebfbfd807f24eca5bed59f Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 15 Oct 2025 13:28:55 -0700 Subject: [PATCH 14/14] Revert "Make NonStimCorrelatedError.nonce=0 instead of random value" This reverts commit d130835bcd563fc88e1bd2482a39a86e05650ad7. --- src/bloqade/stim/dialects/noise/stmts.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index 5eae8b11..b336f901 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/src/bloqade/stim/dialects/noise/stmts.py @@ -1,3 +1,5 @@ +import random + from kirin import ir, types, lowering from kirin.decl import info, statement @@ -89,7 +91,8 @@ class NonStimError(ir.Statement): class NonStimCorrelatedError(ir.Statement): name = "NonStimCorrelatedError" traits = frozenset({lowering.FromPythonCall()}) - nonce: int = info.attribute(default=0) + # nonce must be a unique value, otherwise stim might merge two correlated errors + nonce: int = info.attribute(default_factory=lambda: random.getrandbits(32)) probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)