From 2915453dc020f2687a652412d4721604b807fe91 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Mon, 26 Jan 2026 16:42:47 -0500 Subject: [PATCH 01/16] export non-Clifford gates from stim --- src/bloqade/stim/__init__.py | 5 +++++ src/bloqade/stim/_wrappers.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/bloqade/stim/__init__.py b/src/bloqade/stim/__init__.py index 32e2a20eb..c679c196e 100644 --- a/src/bloqade/stim/__init__.py +++ b/src/bloqade/stim/__init__.py @@ -40,4 +40,9 @@ pauli_channel2 as pauli_channel2, observable_include as observable_include, correlated_qubit_loss as correlated_qubit_loss, + t as t, + rotation_x as rotation_x, + rotation_y as rotation_y, + rotation_z as rotation_z, + u3 as u3, ) diff --git a/src/bloqade/stim/_wrappers.py b/src/bloqade/stim/_wrappers.py index 008e66354..0105ece69 100644 --- a/src/bloqade/stim/_wrappers.py +++ b/src/bloqade/stim/_wrappers.py @@ -72,6 +72,35 @@ def cz( def spp(targets: tuple[auxiliary.PauliString, ...], dagger=False) -> None: ... +## Non-Clifford + + +@wraps(gate.T) +def t(targets: tuple[int, ...], dagger: bool = False) -> None: ... + + +@wraps(gate.Rx) +def rotation_x( + angle: float, targets: tuple[int, ...], dagger: bool = False +) -> None: ... + + +@wraps(gate.Ry) +def rotation_y( + angle: float, targets: tuple[int, ...], dagger: bool = False +) -> None: ... + + +@wraps(gate.Rz) +def rotation_z( + angle: float, targets: tuple[int, ...], dagger: bool = False +) -> None: ... + + +@wraps(gate.U3) +def u3(theta: float, phi: float, lam: float, targets: tuple[int, ...]) -> None: ... + + # dialect:: aux @wraps(auxiliary.GetRecord) def rec(id: int) -> auxiliary.RecordResult: ... From d5d212fce786e1de00a9d1db1f133ef73b9e3726 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Mon, 26 Jan 2026 16:58:58 -0500 Subject: [PATCH 02/16] fix inheritance --- src/bloqade/stim/dialects/gate/stmts/non_clifford.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/bloqade/stim/dialects/gate/stmts/non_clifford.py b/src/bloqade/stim/dialects/gate/stmts/non_clifford.py index 936f7aa14..19476e57e 100644 --- a/src/bloqade/stim/dialects/gate/stmts/non_clifford.py +++ b/src/bloqade/stim/dialects/gate/stmts/non_clifford.py @@ -20,7 +20,7 @@ class T(SingleQubitGate): @statement(dialect=dialect) -class Rx(ir.Statement): +class Rx(SingleQubitGate): """Rx rotation gate represented as I[R_X(theta=...)] in Stim.""" name = "R_X" @@ -29,7 +29,7 @@ class Rx(ir.Statement): @statement(dialect=dialect) -class Ry(ir.Statement): +class Ry(SingleQubitGate): """Ry rotation gate represented as I[R_Y(theta=...)] in Stim.""" name = "R_Y" @@ -38,7 +38,7 @@ class Ry(ir.Statement): @statement(dialect=dialect) -class Rz(ir.Statement): +class Rz(SingleQubitGate): """Rz rotation gate represented as I[R_Z(theta=...)] in Stim.""" name = "R_Z" @@ -47,7 +47,7 @@ class Rz(ir.Statement): @statement(dialect=dialect) -class U3(ir.Statement): +class U3(SingleQubitGate): """U3 gate represented as I[U3(theta=..., phi=..., lambda=...)] in Stim.""" name = "U3" From e9bd6c3f9810ea2b778dc887c282dfcad21ea2e3 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 13:30:08 -0500 Subject: [PATCH 03/16] add tag to stim gates --- src/bloqade/stim/dialects/gate/stmts/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/bloqade/stim/dialects/gate/stmts/base.py b/src/bloqade/stim/dialects/gate/stmts/base.py index 6c1f8964e..85f077b95 100644 --- a/src/bloqade/stim/dialects/gate/stmts/base.py +++ b/src/bloqade/stim/dialects/gate/stmts/base.py @@ -10,6 +10,7 @@ class Gate(ir.Statement): traits = frozenset({lowering.FromPythonCall()}) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) dagger: bool = info.attribute(default=False) + tag: str | None = info.attribute(default=None) @statement From 3dc0ff51300073d9411a77ac73fa6bd280eff66b Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 17:48:26 -0500 Subject: [PATCH 04/16] rm redundant targets for rotations --- src/bloqade/stim/dialects/gate/stmts/non_clifford.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/bloqade/stim/dialects/gate/stmts/non_clifford.py b/src/bloqade/stim/dialects/gate/stmts/non_clifford.py index 19476e57e..070da39e7 100644 --- a/src/bloqade/stim/dialects/gate/stmts/non_clifford.py +++ b/src/bloqade/stim/dialects/gate/stmts/non_clifford.py @@ -24,7 +24,6 @@ class Rx(SingleQubitGate): """Rx rotation gate represented as I[R_X(theta=...)] in Stim.""" name = "R_X" - targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) angle: ir.SSAValue = info.argument(types.Float) @@ -33,7 +32,6 @@ class Ry(SingleQubitGate): """Ry rotation gate represented as I[R_Y(theta=...)] in Stim.""" name = "R_Y" - targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) angle: ir.SSAValue = info.argument(types.Float) @@ -42,7 +40,6 @@ class Rz(SingleQubitGate): """Rz rotation gate represented as I[R_Z(theta=...)] in Stim.""" name = "R_Z" - targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) angle: ir.SSAValue = info.argument(types.Float) From 79d4dd5bc6f4eca061e15c8f107497e59750b07a Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 18:03:23 -0500 Subject: [PATCH 05/16] add tages for gates --- src/bloqade/stim/dialects/gate/emit.py | 42 +++-- src/bloqade/stim/dialects/gate/stmts/base.py | 2 +- src/bloqade/stim/dialects/gate/stmts/pp.py | 1 + test/stim/dialects/stim/emit/test_stim_tag.py | 161 ++++++++++++++++++ 4 files changed, 190 insertions(+), 16 deletions(-) create mode 100644 test/stim/dialects/stim/emit/test_stim_tag.py diff --git a/src/bloqade/stim/dialects/gate/emit.py b/src/bloqade/stim/dialects/gate/emit.py index 599726017..e8625c50c 100644 --- a/src/bloqade/stim/dialects/gate/emit.py +++ b/src/bloqade/stim/dialects/gate/emit.py @@ -35,7 +35,9 @@ def single_qubit_gate( self, emit: EmitStimMain, frame: EmitStimFrame, stmt: SingleQubitGate ): targets: tuple[str, ...] = frame.get_values(stmt.targets) - res = f"{self.gate_1q_map[stmt.name][int(stmt.dagger)]} " + " ".join(targets) + gate_name = self.gate_1q_map[stmt.name][int(stmt.dagger)] + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -44,8 +46,9 @@ def single_qubit_gate( def t_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.T): """Emit T gate as S[T] or S_DAG[T] in Stim annotation format.""" targets: tuple[str, ...] = frame.get_values(stmt.targets) - gate_name = "S_DAG" if stmt.dagger else "S" - res = f"{gate_name}[T] " + " ".join(targets) + gate_name = "S_DAG[T]" if stmt.dagger else "S[T]" + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -55,6 +58,12 @@ def _format_angle(self, angle_str: str) -> str: pi_multiple = angle_turns * 2.0 return f"{pi_multiple}*pi" + def _format_gate_with_tag(self, gate_name: str, tag: str | None) -> str: + """Format gate name with optional tag annotation.""" + if tag: + return f"{gate_name}[{tag}]" + return gate_name + @impl(stmts.Rx) @impl(stmts.Ry) @impl(stmts.Rz) @@ -62,7 +71,9 @@ def rotation_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Rx """Emit rotation gate as I[R_X/R_Y/R_Z(theta=...)] in Stim annotation format.""" targets: tuple[str, ...] = frame.get_values(stmt.targets) angle_str: str = self._format_angle(frame.get(stmt.angle)) - res = f"I[{stmt.name}(theta={angle_str})] " + " ".join(targets) + gate_name = f"I[{stmt.name}(theta={angle_str})]" + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -73,9 +84,9 @@ def u3_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.U3): theta_str: str = self._format_angle(frame.get(stmt.theta)) phi_str: str = self._format_angle(frame.get(stmt.phi)) lam_str: str = self._format_angle(frame.get(stmt.lam)) - res = f"I[U3(theta={theta_str}, phi={phi_str}, lambda={lam_str})] " + " ".join( - targets - ) + gate_name = f"I[U3(theta={theta_str}, phi={phi_str}, lambda={lam_str})]" + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -88,9 +99,9 @@ def two_qubit_gate( self, emit: EmitStimMain, frame: EmitStimFrame, stmt: ControlledTwoQubitGate ): targets: tuple[str, ...] = frame.get_values(stmt.targets) - res = f"{self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)]} " + " ".join( - targets - ) + gate_name = self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)] + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -110,7 +121,9 @@ def ctrl_two_qubit_gate( ): controls: tuple[str, ...] = frame.get_values(stmt.controls) targets: tuple[str, ...] = frame.get_values(stmt.targets) - res = f"{self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)]} " + " ".join( + gate_name = self.gate_ctrl_2q_map[stmt.name][int(stmt.dagger)] + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join( f"{ctrl} {tgt}" for ctrl, tgt in zip(controls, targets) ) frame.write_line(res) @@ -123,10 +136,9 @@ def spp(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.SPP): targets: tuple[str, ...] = tuple( targ.upper() for targ in frame.get_values(stmt.targets) ) - if stmt.dagger: - res = "SPP_DAG " + " ".join(targets) - else: - res = "SPP " + " ".join(targets) + gate_name = "SPP_DAG" if stmt.dagger else "SPP" + gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () diff --git a/src/bloqade/stim/dialects/gate/stmts/base.py b/src/bloqade/stim/dialects/gate/stmts/base.py index 85f077b95..7052aca64 100644 --- a/src/bloqade/stim/dialects/gate/stmts/base.py +++ b/src/bloqade/stim/dialects/gate/stmts/base.py @@ -10,7 +10,7 @@ class Gate(ir.Statement): traits = frozenset({lowering.FromPythonCall()}) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) dagger: bool = info.attribute(default=False) - tag: str | None = info.attribute(default=None) + tag: str = info.attribute(types.String, default=None) @statement diff --git a/src/bloqade/stim/dialects/gate/stmts/pp.py b/src/bloqade/stim/dialects/gate/stmts/pp.py index 8085b9882..caadf04ae 100644 --- a/src/bloqade/stim/dialects/gate/stmts/pp.py +++ b/src/bloqade/stim/dialects/gate/stmts/pp.py @@ -12,4 +12,5 @@ class SPP(ir.Statement): name = "SPP" traits = frozenset({lowering.FromPythonCall()}) dagger: bool = info.attribute(types.Bool, default=False) + tag: str = info.attribute(types.String, default=None) targets: tuple[ir.SSAValue, ...] = info.argument(PauliStringType) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py new file mode 100644 index 000000000..eb806dcfd --- /dev/null +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -0,0 +1,161 @@ +import io + +from bloqade import stim +from bloqade.stim.emit import EmitStimMain +from bloqade.stim.dialects import gate + + +def test_single_qubit_gate_with_tag(): + """Test single qubit gates with tag annotation.""" + + @stim.main + def test_x_with_tag(): + stim.x(targets=(0, 1), tag="my_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_x_with_tag) + assert buf.getvalue().strip() == "X[my_tag] 0 1" + + +def test_single_qubit_gate_without_tag(): + """Test single qubit gates without tag (backward compatibility).""" + + @stim.main + def test_x_no_tag(): + stim.x(targets=(0, 1)) + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_x_no_tag) + assert buf.getvalue().strip() == "X 0 1" + + +def test_s_gate_with_tag(): + """Test S gate with tag, including dagger variant.""" + + @stim.main + def test_s_with_tag(): + stim.s(targets=(0,), tag="s_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_s_with_tag) + assert buf.getvalue().strip() == "S[s_tag] 0" + + @stim.main + def test_s_dag_with_tag(): + stim.s(targets=(0,), dagger=True, tag="s_dag_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_s_dag_with_tag) + assert buf.getvalue().strip() == "S_DAG[s_dag_tag] 0" + + +def test_controlled_gate_with_tag(): + """Test controlled two-qubit gates with tag annotation.""" + + @stim.main + def test_cx_with_tag(): + gate.CX(controls=(0,), targets=(1,), tag="cx_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_cx_with_tag) + assert buf.getvalue().strip() == "CX[cx_tag] 0 1" + + +def test_swap_gate_with_tag(): + """Test SWAP gate with tag annotation.""" + + @stim.main + def test_swap_with_tag(): + stim.swap(targets=(0, 1), tag="swap_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_swap_with_tag) + assert buf.getvalue().strip() == "SWAP[swap_tag] 0 1" + + +def test_t_gate_with_tag(): + """Test T gate with tag annotation.""" + + @stim.main + def test_t_with_tag(): + stim.t(targets=(0,), tag="t_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_t_with_tag) + assert buf.getvalue().strip() == "S[T][t_tag] 0" + + @stim.main + def test_t_dag_with_tag(): + stim.t(targets=(0,), dagger=True, tag="t_dag_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_t_dag_with_tag) + assert buf.getvalue().strip() == "S_DAG[T][t_dag_tag] 0" + + +def test_spp_with_tag(): + """Test SPP gate with tag annotation.""" + + @stim.main + def test_spp_with_tag(): + stim.spp( + targets=( + stim.pauli_string( + string=("X", "Z"), + flipped=(False, False), + targets=(0, 1), + ), + ), + tag="spp_tag", + ) + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_spp_with_tag) + assert buf.getvalue().strip() == "SPP[spp_tag] X0*Z1" + + +def test_rotation_gate_with_tag(): + """Test rotation gates (Rx, Ry, Rz) with tag annotation.""" + + @stim.main + def test_rx_with_tag(): + gate.Rx(targets=(0,), angle=0.25, tag="rx_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_rx_with_tag) + assert buf.getvalue().strip() == "I[R_X(theta=0.5*pi)][rx_tag] 0" + + @stim.main + def test_ry_with_tag(): + gate.Ry(targets=(0,), angle=0.5, tag="ry_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_ry_with_tag) + assert buf.getvalue().strip() == "I[R_Y(theta=1.0*pi)][ry_tag] 0" + + +def test_u3_gate_with_tag(): + """Test U3 gate with tag annotation.""" + + @stim.main + def test_u3_with_tag(): + gate.U3(targets=(0,), theta=0.25, phi=0.5, lam=0.125, tag="u3_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_u3_with_tag) + assert ( + buf.getvalue().strip() + == "I[U3(theta=0.5*pi, phi=1.0*pi, lambda=0.25*pi)][u3_tag] 0" + ) From 23756098980f40a9c67608f236ba3640fa192b14 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 18:14:55 -0500 Subject: [PATCH 06/16] test: add failing tests for collapse dialect tag support Add tests for tag attribute on collapse dialect statements: - test_measurement_with_tag (MZ) - test_reset_with_tag (RZ) - test_ppmeasurement_with_tag (MPP) These tests will fail until StimStatement base class is implemented. Co-Authored-By: Claude Opus 4.5 --- test/stim/dialects/stim/emit/test_stim_tag.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index eb806dcfd..e7ac44f78 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -159,3 +159,57 @@ def test_u3_with_tag(): buf.getvalue().strip() == "I[U3(theta=0.5*pi, phi=1.0*pi, lambda=0.25*pi)][u3_tag] 0" ) + + +# ============================================================================= +# Collapse dialect tag tests +# ============================================================================= + + +def test_measurement_with_tag(): + """Test measurement gates with tag annotation.""" + + @stim.main + def test_mz_with_tag(): + stim.mz(targets=(0, 1), p=0.01, tag="mz_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_mz_with_tag) + assert buf.getvalue().strip() == "MZ[mz_tag](0.01) 0 1" + + +def test_reset_with_tag(): + """Test reset gates with tag annotation.""" + + @stim.main + def test_rz_with_tag(): + stim.rz(targets=(0, 1), tag="rz_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_rz_with_tag) + assert buf.getvalue().strip() == "RZ[rz_tag] 0 1" + + +def test_ppmeasurement_with_tag(): + """Test Pauli-product measurement with tag annotation.""" + + @stim.main + def test_mpp_with_tag(): + stim.mpp( + targets=( + stim.pauli_string( + string=("X", "Z"), + flipped=(False, False), + targets=(0, 1), + ), + ), + p=0.01, + tag="mpp_tag", + ) + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_mpp_with_tag) + assert buf.getvalue().strip() == "MPP[mpp_tag](0.01) X0*Z1" From e581c0f1c7f6868f184e368330426f1a088ec4b5 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:23:16 -0500 Subject: [PATCH 07/16] feat: add tag support to collapse dialect via StimStatement base - Create StimStatement base class with tag attribute - Update Measurement, Reset, PPMeasurement to inherit from StimStatement - Update emit methods to include tag in output - Fix test expectations for float formatting Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/dialects/collapse/emit_str.py | 15 ++++++++++++--- .../stim/dialects/collapse/stmts/measure.py | 6 +++--- .../stim/dialects/collapse/stmts/pp_measure.py | 6 +++--- src/bloqade/stim/dialects/collapse/stmts/reset.py | 6 +++--- src/bloqade/stim/dialects/stim_statement.py | 11 +++++++++++ test/stim/dialects/stim/emit/test_stim_tag.py | 4 ++-- uv.lock | 14 ++++++++++++-- 7 files changed, 46 insertions(+), 16 deletions(-) create mode 100644 src/bloqade/stim/dialects/stim_statement.py diff --git a/src/bloqade/stim/dialects/collapse/emit_str.py b/src/bloqade/stim/dialects/collapse/emit_str.py index 7a99b94f7..4aeb95042 100644 --- a/src/bloqade/stim/dialects/collapse/emit_str.py +++ b/src/bloqade/stim/dialects/collapse/emit_str.py @@ -11,6 +11,12 @@ @dialect.register(key="emit.stim") class EmitStimCollapseMethods(MethodTable): + def _format_with_tag(self, name: str, tag: str | None) -> str: + """Format instruction name with optional tag annotation.""" + if tag: + return f"{name}[{tag}]" + return name + meas_map: dict[str, str] = { stmts.MX.name: "MX", stmts.MY.name: "MY", @@ -30,8 +36,9 @@ def get_measure(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: Measuremen probability: str = frame.get(stmt.p) targets: tuple[str, ...] = frame.get_values(stmt.targets) + name = self._format_with_tag(self.meas_map[stmt.name], stmt.tag) - out = f"{self.meas_map[stmt.name]}({probability}) " + " ".join(targets) + out = f"{name}({probability}) " + " ".join(targets) frame.write_line(out) return () @@ -48,8 +55,9 @@ def get_measure(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: Measuremen def get_reset(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: Reset): targets: tuple[str, ...] = frame.get_values(stmt.targets) + name = self._format_with_tag(self.reset_map[stmt.name], stmt.tag) - out = f"{self.reset_map[stmt.name]} " + " ".join(targets) + out = f"{name} " + " ".join(targets) frame.write_line(out) return () @@ -62,8 +70,9 @@ def pp_measure( targets: tuple[str, ...] = tuple( targ.upper() for targ in frame.get_values(stmt.targets) ) + name = self._format_with_tag("MPP", stmt.tag) - out = f"MPP({probability}) " + " ".join(targets) + out = f"{name}({probability}) " + " ".join(targets) frame.write_line(out) return () diff --git a/src/bloqade/stim/dialects/collapse/stmts/measure.py b/src/bloqade/stim/dialects/collapse/stmts/measure.py index bb22be79d..166595c65 100644 --- a/src/bloqade/stim/dialects/collapse/stmts/measure.py +++ b/src/bloqade/stim/dialects/collapse/stmts/measure.py @@ -1,13 +1,13 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from .._dialect import dialect +from ...stim_statement import StimStatement @statement -class Measurement(ir.Statement): +class Measurement(StimStatement): name = "measurement" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) """probability of noise introduced by measurement. For example 0.01 means 1% the measurement will be flipped""" targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) diff --git a/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py b/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py index e03b54ba8..4bf179866 100644 --- a/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py +++ b/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py @@ -1,14 +1,14 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from .._dialect import dialect from ...auxiliary.types import PauliStringType +from ...stim_statement import StimStatement @statement(dialect=dialect) -class PPMeasurement(ir.Statement): +class PPMeasurement(StimStatement): name = "MPP" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) """probability of noise introduced by measurement. For example 0.01 means 1% the measurement will be flipped""" targets: tuple[ir.SSAValue, ...] = info.argument(PauliStringType) diff --git a/src/bloqade/stim/dialects/collapse/stmts/reset.py b/src/bloqade/stim/dialects/collapse/stmts/reset.py index aa182447c..6b80cd346 100644 --- a/src/bloqade/stim/dialects/collapse/stmts/reset.py +++ b/src/bloqade/stim/dialects/collapse/stmts/reset.py @@ -1,13 +1,13 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from .._dialect import dialect +from ...stim_statement import StimStatement @statement -class Reset(ir.Statement): +class Reset(StimStatement): name = "reset" - traits = frozenset({lowering.FromPythonCall()}) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) diff --git a/src/bloqade/stim/dialects/stim_statement.py b/src/bloqade/stim/dialects/stim_statement.py new file mode 100644 index 000000000..e2ba3c3b5 --- /dev/null +++ b/src/bloqade/stim/dialects/stim_statement.py @@ -0,0 +1,11 @@ +from kirin import ir, types, lowering +from kirin.decl import info, statement + + +@statement +class StimStatement(ir.Statement): + """Base class for all stim instruction statements with tag support.""" + + name = "stim_statement" + traits = frozenset({lowering.FromPythonCall()}) + tag: str = info.attribute(types.String, default=None) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index e7ac44f78..1c9bf87aa 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -176,7 +176,7 @@ def test_mz_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_mz_with_tag) - assert buf.getvalue().strip() == "MZ[mz_tag](0.01) 0 1" + assert buf.getvalue().strip() == "MZ[mz_tag](0.01000000) 0 1" def test_reset_with_tag(): @@ -212,4 +212,4 @@ def test_mpp_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_mpp_with_tag) - assert buf.getvalue().strip() == "MPP[mpp_tag](0.01) X0*Z1" + assert buf.getvalue().strip() == "MPP[mpp_tag](0.01000000) X0*Z1" diff --git a/uv.lock b/uv.lock index 83cf5c82a..827ee0ce4 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.12' and sys_platform == 'darwin'", @@ -159,6 +159,9 @@ cirq = [ { name = "numba" }, { name = "qpsolvers", extra = ["clarabel"] }, ] +pyqrack-cuda = [ + { name = "pyqrack-cuda" }, +] pyqrack-opencl = [ { name = "pyqrack" }, ] @@ -211,6 +214,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=1.3.0,<2.11.0" }, { name = "pyqrack", marker = "extra == 'pyqrack-opencl'", specifier = ">=1.70.4" }, { name = "pyqrack-cpu", specifier = ">=1.70.4" }, + { name = "pyqrack-cuda", marker = "extra == 'pyqrack-cuda'", specifier = ">=1.70.4" }, { name = "pyqt5", marker = "sys_platform == 'darwin' and extra == 'vis'", specifier = ">=5.15.11" }, { name = "pyqt5", marker = "sys_platform == 'linux' and extra == 'vis'", specifier = ">=5.15.11" }, { name = "qbraid", marker = "extra == 'qbraid'", specifier = ">=0.9.3" }, @@ -220,7 +224,7 @@ requires-dist = [ { name = "stim", marker = "extra == 'stim'", specifier = ">=1.15.0" }, { name = "tqdm", marker = "extra == 'vis'", specifier = ">=4.66.5" }, ] -provides-extras = ["qasm2", "vis", "qbraid", "cirq", "pyqrack-opencl", "stim"] +provides-extras = ["qasm2", "vis", "qbraid", "cirq", "pyqrack-opencl", "pyqrack-cuda", "stim"] [package.metadata.requires-dev] dev = [ @@ -2387,6 +2391,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/30/18/09ec1971e1873031588396928d82983aef8a866cb275172dc0597a7323d1/pyqrack_cpu-1.71.2-py3-none-win_amd64.whl", hash = "sha256:025d9419a65430fad0019e00c4e30fdf824df7db0f22507d14a875f105367207", size = 842774, upload-time = "2025-11-13T21:47:49.368Z" }, ] +[[package]] +name = "pyqrack-cuda" +version = "1.83.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/07/e5/4d380b1fa6ebc65da45a04f51d17d0d3ff8b6dc01f4309b8856c2771740b/pyqrack_cuda-1.83.0.tar.gz", hash = "sha256:423672a3350f966ea2cf780a3e1f6b31c83d8ccfb459ad3fef79e3a3f5ed2bf9", size = 51972, upload-time = "2026-01-21T20:01:52.865Z" } + [[package]] name = "pyqt5" version = "5.15.11" From fe29a5a7c733e2e9d6765fe19c8c7ce52a6270c6 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:24:13 -0500 Subject: [PATCH 08/16] test: add failing tests for noise dialect tag support Add tests for tag attribute on noise dialect statements: - test_depolarize_with_tag (Depolarize1) - test_pauli_channel_with_tag (PauliChannel1) These tests will fail until noise statements inherit from StimStatement. Co-Authored-By: Claude Opus 4.5 --- test/stim/dialects/stim/emit/test_stim_tag.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index 1c9bf87aa..72c8681a5 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -213,3 +213,37 @@ def test_mpp_with_tag(): stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_mpp_with_tag) assert buf.getvalue().strip() == "MPP[mpp_tag](0.01000000) X0*Z1" + + +# ============================================================================= +# Noise dialect tag tests +# ============================================================================= + + +def test_depolarize_with_tag(): + """Test depolarize noise with tag annotation.""" + + @stim.main + def test_depolarize1_with_tag(): + stim.depolarize1(p=0.01, targets=(0, 1), tag="dep1_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_depolarize1_with_tag) + assert buf.getvalue().strip() == "DEPOLARIZE1[dep1_tag](0.01000000) 0 1" + + +def test_pauli_channel_with_tag(): + """Test Pauli channel noise with tag annotation.""" + + @stim.main + def test_pauli_channel1_with_tag(): + stim.pauli_channel1(px=0.01, py=0.02, pz=0.03, targets=(0,), tag="pc1_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_pauli_channel1_with_tag) + assert ( + buf.getvalue().strip() + == "PAULI_CHANNEL_1[pc1_tag](0.01000000, 0.02000000, 0.03000000) 0" + ) From 7bf569e702bbd5e4f7ac4477ea3070b0a88baec7 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:27:10 -0500 Subject: [PATCH 09/16] feat: add tag support to noise dialect Update all noise statements (Depolarize1, Depolarize2, PauliChannel1, PauliChannel2, XError, YError, ZError, TrivialError, QubitLoss, etc.) to inherit from StimStatement base class, which provides the tag attribute. Update emit methods to format instructions with optional tag annotations. Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/dialects/noise/emit.py | 25 +++++++++++++------- src/bloqade/stim/dialects/noise/stmts.py | 30 +++++++++--------------- 2 files changed, 28 insertions(+), 27 deletions(-) diff --git a/src/bloqade/stim/dialects/noise/emit.py b/src/bloqade/stim/dialects/noise/emit.py index 901e5c885..85e8a23bd 100644 --- a/src/bloqade/stim/dialects/noise/emit.py +++ b/src/bloqade/stim/dialects/noise/emit.py @@ -17,6 +17,12 @@ class EmitStimNoiseMethods(MethodTable): stmts.ZError.name: "Z_ERROR", } + def _format_with_tag(self, name: str, tag: str | None) -> str: + """Format instruction name with optional tag annotation.""" + if tag: + return f"{name}[{tag}]" + return name + @impl(stmts.XError) @impl(stmts.YError) @impl(stmts.ZError) @@ -28,7 +34,7 @@ def single_p_error( targets: tuple[str, ...] = frame.get_values(stmt.targets) p: str = frame.get(stmt.p) - name = self.single_p_error_map[stmt.name] + name = self._format_with_tag(self.single_p_error_map[stmt.name], stmt.tag) res = f"{name}({p}) " + " ".join(targets) frame.write_line(res) @@ -43,7 +49,8 @@ def pauli_channel1( px: str = frame.get(stmt.px) py: str = frame.get(stmt.py) pz: str = frame.get(stmt.pz) - res = f"PAULI_CHANNEL_1({px}, {py}, {pz}) " + " ".join(targets) + name = self._format_with_tag("PAULI_CHANNEL_1", stmt.tag) + res = f"{name}({px}, {py}, {pz}) " + " ".join(targets) frame.write_line(res) return () @@ -58,8 +65,9 @@ def pauli_channel2( :15 ] # extract the first 15 argument, which is the probabilities prob_str: str = ", ".join(prob) + name = self._format_with_tag("PAULI_CHANNEL_2", stmt.tag) - res = f"PAULI_CHANNEL_2({prob_str}) " + " ".join(targets) + res = f"{name}({prob_str}) " + " ".join(targets) frame.write_line(res) return () @@ -73,8 +81,9 @@ def non_stim_error( targets: tuple[str, ...] = frame.get_values(stmt.targets) prob: tuple[str, ...] = frame.get_values(stmt.probs) prob_str: str = ", ".join(prob) + name = self._format_with_tag(f"I_ERROR[{stmt.name}]", stmt.tag) - res = f"I_ERROR[{stmt.name}]({prob_str}) " + " ".join(targets) + res = f"{name}({prob_str}) " + " ".join(targets) frame.write_line(res) return () @@ -91,11 +100,11 @@ def non_stim_corr_error( targets: tuple[str, ...] = frame.get_values(stmt.targets) prob: tuple[str, ...] = frame.get_values(stmt.probs) prob_str: str = ", ".join(prob) - - res = ( - f"I_ERROR[{stmt.name}:{emit.correlated_error_count}]({prob_str}) " - + " ".join(targets) + name = self._format_with_tag( + f"I_ERROR[{stmt.name}:{emit.correlated_error_count}]", stmt.tag ) + + res = f"{name}({prob_str}) " + " ".join(targets) emit.correlated_error_count += 1 frame.write_line(res) diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index a63eafcd8..552b0d4c7 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/src/bloqade/stim/dialects/noise/stmts.py @@ -1,29 +1,27 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from ._dialect import dialect +from ..stim_statement import StimStatement @statement(dialect=dialect) -class Depolarize1(ir.Statement): +class Depolarize1(StimStatement): name = "Depolarize1" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement(dialect=dialect) -class Depolarize2(ir.Statement): +class Depolarize2(StimStatement): name = "Depolarize2" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement(dialect=dialect) -class PauliChannel1(ir.Statement): +class PauliChannel1(StimStatement): name = "PauliChannel1" - traits = frozenset({lowering.FromPythonCall()}) px: ir.SSAValue = info.argument(types.Float) py: ir.SSAValue = info.argument(types.Float) pz: ir.SSAValue = info.argument(types.Float) @@ -31,10 +29,9 @@ class PauliChannel1(ir.Statement): @statement(dialect=dialect) -class PauliChannel2(ir.Statement): +class PauliChannel2(StimStatement): name = "PauliChannel2" # TODO custom lowering to make sugar for this - traits = frozenset({lowering.FromPythonCall()}) pix: ir.SSAValue = info.argument(types.Float) piy: ir.SSAValue = info.argument(types.Float) piz: ir.SSAValue = info.argument(types.Float) @@ -54,41 +51,36 @@ class PauliChannel2(ir.Statement): @statement(dialect=dialect) -class XError(ir.Statement): +class XError(StimStatement): name = "X_ERROR" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement(dialect=dialect) -class YError(ir.Statement): +class YError(StimStatement): name = "Y_ERROR" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement(dialect=dialect) -class ZError(ir.Statement): +class ZError(StimStatement): name = "Z_ERROR" - traits = frozenset({lowering.FromPythonCall()}) p: ir.SSAValue = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement -class NonStimError(ir.Statement): +class NonStimError(StimStatement): name = "NonStimError" - traits = frozenset({lowering.FromPythonCall()}) probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) @statement -class NonStimCorrelatedError(ir.Statement): +class NonStimCorrelatedError(StimStatement): name = "NonStimCorrelatedError" - traits = frozenset({lowering.FromPythonCall()}) probs: tuple[ir.SSAValue, ...] = info.argument(types.Float) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) From 0164bcff731b1db2b3531120a694738bb9993428 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:28:03 -0500 Subject: [PATCH 10/16] test: add failing tests for auxiliary dialect tag support Add tests for TICK, DETECTOR, OBSERVABLE_INCLUDE, and QUBIT_COORDS instructions with tag annotations. Tests fail because the tag attribute doesn't exist on these statements yet. Co-Authored-By: Claude Opus 4.5 --- test/stim/dialects/stim/emit/test_stim_tag.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index 72c8681a5..d759d7a69 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -247,3 +247,64 @@ def test_pauli_channel1_with_tag(): buf.getvalue().strip() == "PAULI_CHANNEL_1[pc1_tag](0.01000000, 0.02000000, 0.03000000) 0" ) + + +# ============================================================================= +# Auxiliary dialect tag tests +# ============================================================================= + + +def test_tick_with_tag(): + """Test TICK instruction with tag annotation.""" + + @stim.main + def test_tick_with_tag(): + stim.tick(tag="tick_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_tick_with_tag) + assert buf.getvalue().strip() == "TICK[tick_tag]" + + +def test_detector_with_tag(): + """Test DETECTOR instruction with tag annotation.""" + + @stim.main + def test_detector_with_tag(): + stim.mz(targets=(0,)) + stim.detector(coord=(1.0, 2.0), targets=(stim.rec(-1),), tag="det_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_detector_with_tag) + output = buf.getvalue().strip().split("\n") + assert output[1] == "DETECTOR[det_tag](1.00000000, 2.00000000) rec[-1]" + + +def test_observable_include_with_tag(): + """Test OBSERVABLE_INCLUDE instruction with tag annotation.""" + + @stim.main + def test_obs_include_with_tag(): + stim.mz(targets=(0,)) + stim.observable_include(idx=0, targets=(stim.rec(-1),), tag="obs_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_obs_include_with_tag) + output = buf.getvalue().strip().split("\n") + assert output[1] == "OBSERVABLE_INCLUDE[obs_tag](0) rec[-1]" + + +def test_qubit_coordinates_with_tag(): + """Test QUBIT_COORDS instruction with tag annotation.""" + + @stim.main + def test_qubit_coords_with_tag(): + stim.qubit_coordinates(coord=(1.0, 2.0, 3.0), target=0, tag="qc_tag") + + buf = io.StringIO() + stim_emit = EmitStimMain(dialects=stim.main, io=buf) + stim_emit.run(test_qubit_coords_with_tag) + assert buf.getvalue().strip() == "QUBIT_COORDS[qc_tag](1.00000000, 2.00000000, 3.00000000) 0" From ce7d63b1394472fa8785af16713f75346bab5ae7 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:29:48 -0500 Subject: [PATCH 11/16] feat: add tag support to auxiliary dialect Update Tick, Detector, ObservableInclude, and QubitCoordinates statements to inherit from StimStatement base class, which provides the tag attribute. Update emit methods to format instructions with optional tag annotations. Fix test API usage for detector, observable_include, and qubit_coordinates. Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/dialects/auxiliary/emit.py | 21 +++++++++++++------ .../stim/dialects/auxiliary/stmts/annotate.py | 13 +++++------- test/stim/dialects/stim/emit/test_stim_tag.py | 13 +++++++----- 3 files changed, 28 insertions(+), 19 deletions(-) diff --git a/src/bloqade/stim/dialects/auxiliary/emit.py b/src/bloqade/stim/dialects/auxiliary/emit.py index 079fee6bd..30b427a01 100644 --- a/src/bloqade/stim/dialects/auxiliary/emit.py +++ b/src/bloqade/stim/dialects/auxiliary/emit.py @@ -9,6 +9,12 @@ @dialect.register(key="emit.stim") class EmitStimAuxMethods(MethodTable): + def _format_with_tag(self, name: str, tag: str | None) -> str: + """Format instruction name with optional tag annotation.""" + if tag: + return f"{name}[{tag}]" + return name + @impl(stmts.ConstInt) def const_int(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.ConstInt): @@ -57,8 +63,8 @@ def get_rec(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.GetRecor @impl(stmts.Tick) def tick(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Tick): - - frame.write_line("TICK") + name = self._format_with_tag("TICK", stmt.tag) + frame.write_line(name) return () @@ -67,13 +73,14 @@ def detector(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Detecto coords: tuple[str, ...] = frame.get_values(stmt.coord) targets: tuple[str, ...] = frame.get_values(stmt.targets) + name = self._format_with_tag("DETECTOR", stmt.tag) coord_str: str = ", ".join(coords) target_str: str = " ".join(targets) if len(coords): - frame.write_line(f"DETECTOR({coord_str}) {target_str}") + frame.write_line(f"{name}({coord_str}) {target_str}") else: - frame.write_line(f"DETECTOR {target_str}") + frame.write_line(f"{name} {target_str}") return () @impl(stmts.ObservableInclude) @@ -83,9 +90,10 @@ def obs_include( idx: str = frame.get(stmt.idx) targets: tuple[str, ...] = frame.get_values(stmt.targets) + name = self._format_with_tag("OBSERVABLE_INCLUDE", stmt.tag) target_str: str = " ".join(targets) - frame.write_line(f"OBSERVABLE_INCLUDE({idx}) {target_str}") + frame.write_line(f"{name}({idx}) {target_str}") return () @@ -111,8 +119,9 @@ def qubit_coordinates( coords: tuple[str, ...] = frame.get_values(stmt.coord) target: str = frame.get(stmt.target) + name = self._format_with_tag("QUBIT_COORDS", stmt.tag) coord_str: str = ", ".join(coords) - frame.write_line(f"QUBIT_COORDS({coord_str}) {target}") + frame.write_line(f"{name}({coord_str}) {target}") return () diff --git a/src/bloqade/stim/dialects/auxiliary/stmts/annotate.py b/src/bloqade/stim/dialects/auxiliary/stmts/annotate.py index 3c9f92832..b27bbf278 100644 --- a/src/bloqade/stim/dialects/auxiliary/stmts/annotate.py +++ b/src/bloqade/stim/dialects/auxiliary/stmts/annotate.py @@ -3,6 +3,7 @@ from ..types import RecordType, PauliStringType from .._dialect import dialect +from ...stim_statement import StimStatement PyNum = types.Union(types.Int, types.Float) @@ -16,25 +17,22 @@ class GetRecord(ir.Statement): @statement(dialect=dialect) -class Detector(ir.Statement): +class Detector(StimStatement): name = "detector" - traits = frozenset({lowering.FromPythonCall()}) coord: tuple[ir.SSAValue, ...] = info.argument(PyNum) targets: tuple[ir.SSAValue, ...] = info.argument(RecordType) @statement(dialect=dialect) -class ObservableInclude(ir.Statement): +class ObservableInclude(StimStatement): name = "obs.include" - traits = frozenset({lowering.FromPythonCall()}) idx: ir.SSAValue = info.argument(type=types.Int) targets: tuple[ir.SSAValue, ...] = info.argument(RecordType) @statement(dialect=dialect) -class Tick(ir.Statement): +class Tick(StimStatement): name = "tick" - traits = frozenset({lowering.FromPythonCall()}) @statement(dialect=dialect) @@ -48,8 +46,7 @@ class NewPauliString(ir.Statement): @statement(dialect=dialect) -class QubitCoordinates(ir.Statement): +class QubitCoordinates(StimStatement): name = "qubit_coordinates" - traits = frozenset({lowering.FromPythonCall()}) coord: tuple[ir.SSAValue, ...] = info.argument(PyNum) target: ir.SSAValue = info.argument(types.Int) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index d759d7a69..2d25c7800 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -269,11 +269,12 @@ def test_tick_with_tag(): def test_detector_with_tag(): """Test DETECTOR instruction with tag annotation.""" + from bloqade.stim.dialects import auxiliary @stim.main def test_detector_with_tag(): - stim.mz(targets=(0,)) - stim.detector(coord=(1.0, 2.0), targets=(stim.rec(-1),), tag="det_tag") + stim.mz(targets=(0,), p=0.0) + auxiliary.Detector(coord=(1.0, 2.0), targets=(stim.rec(-1),), tag="det_tag") buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) @@ -284,11 +285,12 @@ def test_detector_with_tag(): def test_observable_include_with_tag(): """Test OBSERVABLE_INCLUDE instruction with tag annotation.""" + from bloqade.stim.dialects import auxiliary @stim.main def test_obs_include_with_tag(): - stim.mz(targets=(0,)) - stim.observable_include(idx=0, targets=(stim.rec(-1),), tag="obs_tag") + stim.mz(targets=(0,), p=0.0) + auxiliary.ObservableInclude(idx=0, targets=(stim.rec(-1),), tag="obs_tag") buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) @@ -299,10 +301,11 @@ def test_obs_include_with_tag(): def test_qubit_coordinates_with_tag(): """Test QUBIT_COORDS instruction with tag annotation.""" + from bloqade.stim.dialects import auxiliary @stim.main def test_qubit_coords_with_tag(): - stim.qubit_coordinates(coord=(1.0, 2.0, 3.0), target=0, tag="qc_tag") + auxiliary.QubitCoordinates(coord=(1.0, 2.0, 3.0), target=0, tag="qc_tag") buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) From 81a89669f1e70c9773f41e1695a9bfa2fcc555fc Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:31:34 -0500 Subject: [PATCH 12/16] feat: add tag parameter to all stim wrapper functions Update all wrapper functions to expose the tag parameter for: - Gate dialect (x, y, z, h, s, t, cx, cy, cz, swap, spp, etc.) - Collapse dialect (mz, my, mx, mzz, myy, mxx, mpp, rz, ry, rx) - Noise dialect (depolarize1, depolarize2, pauli_channel1, etc.) - Auxiliary dialect (tick, detector, observable_include, qubit_coords) Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/_wrappers.py | 141 ++++++++++++++++++++++++---------- 1 file changed, 100 insertions(+), 41 deletions(-) diff --git a/src/bloqade/stim/_wrappers.py b/src/bloqade/stim/_wrappers.py index 0105ece69..d818f9372 100644 --- a/src/bloqade/stim/_wrappers.py +++ b/src/bloqade/stim/_wrappers.py @@ -8,97 +8,147 @@ # dialect:: gate ## 1q @wraps(gate.X) -def x(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def x( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.Y) -def y(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def y( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.Z) -def z(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def z( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.Identity) -def identity(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def identity( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.H) -def h(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def h( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.S) -def s(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def s( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.SqrtX) -def sqrt_x(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def sqrt_x( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.SqrtY) -def sqrt_y(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def sqrt_y( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.SqrtZ) -def sqrt_z(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def sqrt_z( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... ## clif 2q @wraps(gate.Swap) -def swap(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def swap( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... ## ctrl 2q @wraps(gate.CX) def cx( - controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False + controls: tuple[int, ...], + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... @wraps(gate.CY) def cy( - controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False + controls: tuple[int, ...], + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... @wraps(gate.CZ) def cz( - controls: tuple[int, ...], targets: tuple[int, ...], dagger: bool = False + controls: tuple[int, ...], + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... ## pp @wraps(gate.SPP) -def spp(targets: tuple[auxiliary.PauliString, ...], dagger=False) -> None: ... +def spp( + targets: tuple[auxiliary.PauliString, ...], + dagger: bool = False, + tag: str | None = None, +) -> None: ... ## Non-Clifford @wraps(gate.T) -def t(targets: tuple[int, ...], dagger: bool = False) -> None: ... +def t( + targets: tuple[int, ...], dagger: bool = False, tag: str | None = None +) -> None: ... @wraps(gate.Rx) def rotation_x( - angle: float, targets: tuple[int, ...], dagger: bool = False + angle: float, + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... @wraps(gate.Ry) def rotation_y( - angle: float, targets: tuple[int, ...], dagger: bool = False + angle: float, + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... @wraps(gate.Rz) def rotation_z( - angle: float, targets: tuple[int, ...], dagger: bool = False + angle: float, + targets: tuple[int, ...], + dagger: bool = False, + tag: str | None = None, ) -> None: ... @wraps(gate.U3) -def u3(theta: float, phi: float, lam: float, targets: tuple[int, ...]) -> None: ... +def u3( + theta: float, + phi: float, + lam: float, + targets: tuple[int, ...], + tag: str | None = None, +) -> None: ... # dialect:: aux @@ -108,18 +158,20 @@ def rec(id: int) -> auxiliary.RecordResult: ... @wraps(auxiliary.Detector) def detector( - coord: tuple[Union[int, float], ...], targets: tuple[auxiliary.RecordResult, ...] + coord: tuple[Union[int, float], ...], + targets: tuple[auxiliary.RecordResult, ...], + tag: str | None = None, ) -> None: ... @wraps(auxiliary.ObservableInclude) def observable_include( - idx: int, targets: tuple[auxiliary.RecordResult, ...] + idx: int, targets: tuple[auxiliary.RecordResult, ...], tag: str | None = None ) -> None: ... @wraps(auxiliary.Tick) -def tick() -> None: ... +def tick(tag: str | None = None) -> None: ... @wraps(auxiliary.NewPauliString) @@ -129,62 +181,66 @@ def pauli_string( @wraps(auxiliary.QubitCoordinates) -def qubit_coords(coord: tuple[Union[int, float], ...], target: int) -> None: ... +def qubit_coords( + coord: tuple[Union[int, float], ...], target: int, tag: str | None = None +) -> None: ... # dialect:: collapse @wraps(collapse.MZ) -def mz(p: float, targets: tuple[int, ...]) -> None: ... +def mz(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.MY) -def my(p: float, targets: tuple[int, ...]) -> None: ... +def my(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.MX) -def mx(p: float, targets: tuple[int, ...]) -> None: ... +def mx(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.MZZ) -def mzz(p: float, targets: tuple[int, ...]) -> None: ... +def mzz(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.MYY) -def myy(p: float, targets: tuple[int, ...]) -> None: ... +def myy(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.MXX) -def mxx(p: float, targets: tuple[int, ...]) -> None: ... +def mxx(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.PPMeasurement) -def mpp(p: float, targets: tuple[auxiliary.PauliString, ...]) -> None: ... +def mpp( + p: float, targets: tuple[auxiliary.PauliString, ...], tag: str | None = None +) -> None: ... @wraps(collapse.RZ) -def rz(targets: tuple[int, ...]) -> None: ... +def rz(targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.RY) -def ry(targets: tuple[int, ...]) -> None: ... +def ry(targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(collapse.RX) -def rx(targets: tuple[int, ...]) -> None: ... +def rx(targets: tuple[int, ...], tag: str | None = None) -> None: ... # dialect:: noise @wraps(noise.Depolarize1) -def depolarize1(p: float, targets: tuple[int, ...]) -> None: ... +def depolarize1(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(noise.Depolarize2) -def depolarize2(p: float, targets: tuple[int, ...]) -> None: ... +def depolarize2(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(noise.PauliChannel1) def pauli_channel1( - px: float, py: float, pz: float, targets: tuple[int, ...] + px: float, py: float, pz: float, targets: tuple[int, ...], tag: str | None = None ) -> None: ... @@ -206,26 +262,29 @@ def pauli_channel2( pzy: float, pzz: float, targets: tuple[int, ...], + tag: str | None = None, ) -> None: ... @wraps(noise.XError) -def x_error(p: float, targets: tuple[int, ...]) -> None: ... +def x_error(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(noise.YError) -def y_error(p: float, targets: tuple[int, ...]) -> None: ... +def y_error(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(noise.ZError) -def z_error(p: float, targets: tuple[int, ...]) -> None: ... +def z_error(p: float, targets: tuple[int, ...], tag: str | None = None) -> None: ... @wraps(noise.QubitLoss) -def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ... +def qubit_loss( + probs: tuple[float, ...], targets: tuple[int, ...], tag: str | None = None +) -> None: ... @wraps(noise.CorrelatedQubitLoss) def correlated_qubit_loss( - probs: tuple[float, ...], targets: tuple[int, ...] + probs: tuple[float, ...], targets: tuple[int, ...], tag: str | None = None ) -> None: ... From 58da116f2dccd4c3c77b1ad01b78336ed48db6df Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:34:03 -0500 Subject: [PATCH 13/16] refactor: gate dialect uses StimStatement base class Refactor Gate and SPP to inherit from StimStatement instead of ir.Statement, removing duplicate tag attribute definitions. This centralizes the tag functionality in the StimStatement base class. Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/dialects/gate/stmts/base.py | 7 +++---- src/bloqade/stim/dialects/gate/stmts/pp.py | 7 +++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/bloqade/stim/dialects/gate/stmts/base.py b/src/bloqade/stim/dialects/gate/stmts/base.py index 7052aca64..ffc2b3854 100644 --- a/src/bloqade/stim/dialects/gate/stmts/base.py +++ b/src/bloqade/stim/dialects/gate/stmts/base.py @@ -1,16 +1,15 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from bloqade.stim.dialects.auxiliary import RecordType +from bloqade.stim.dialects.stim_statement import StimStatement @statement -class Gate(ir.Statement): +class Gate(StimStatement): name = "stim_gate" - traits = frozenset({lowering.FromPythonCall()}) targets: tuple[ir.SSAValue, ...] = info.argument(types.Int) dagger: bool = info.attribute(default=False) - tag: str = info.attribute(types.String, default=None) @statement diff --git a/src/bloqade/stim/dialects/gate/stmts/pp.py b/src/bloqade/stim/dialects/gate/stmts/pp.py index caadf04ae..31f417864 100644 --- a/src/bloqade/stim/dialects/gate/stmts/pp.py +++ b/src/bloqade/stim/dialects/gate/stmts/pp.py @@ -1,16 +1,15 @@ -from kirin import ir, types, lowering +from kirin import ir, types from kirin.decl import info, statement from .._dialect import dialect +from ...stim_statement import StimStatement from ...auxiliary.types import PauliStringType # Generalized Pauli-product gates # --------------------------------------- @statement(dialect=dialect) -class SPP(ir.Statement): +class SPP(StimStatement): name = "SPP" - traits = frozenset({lowering.FromPythonCall()}) dagger: bool = info.attribute(types.Bool, default=False) - tag: str = info.attribute(types.String, default=None) targets: tuple[ir.SSAValue, ...] = info.argument(PauliStringType) From 8846eb7d7101b659911e830decf46a78fd74fc2d Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 20:48:43 -0500 Subject: [PATCH 14/16] chore: revert uv.lock to main branch version Co-Authored-By: Claude Opus 4.5 --- uv.lock | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/uv.lock b/uv.lock index 827ee0ce4..83cf5c82a 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.12' and sys_platform == 'darwin'", @@ -159,9 +159,6 @@ cirq = [ { name = "numba" }, { name = "qpsolvers", extra = ["clarabel"] }, ] -pyqrack-cuda = [ - { name = "pyqrack-cuda" }, -] pyqrack-opencl = [ { name = "pyqrack" }, ] @@ -214,7 +211,6 @@ requires-dist = [ { name = "pydantic", specifier = ">=1.3.0,<2.11.0" }, { name = "pyqrack", marker = "extra == 'pyqrack-opencl'", specifier = ">=1.70.4" }, { name = "pyqrack-cpu", specifier = ">=1.70.4" }, - { name = "pyqrack-cuda", marker = "extra == 'pyqrack-cuda'", specifier = ">=1.70.4" }, { name = "pyqt5", marker = "sys_platform == 'darwin' and extra == 'vis'", specifier = ">=5.15.11" }, { name = "pyqt5", marker = "sys_platform == 'linux' and extra == 'vis'", specifier = ">=5.15.11" }, { name = "qbraid", marker = "extra == 'qbraid'", specifier = ">=0.9.3" }, @@ -224,7 +220,7 @@ requires-dist = [ { name = "stim", marker = "extra == 'stim'", specifier = ">=1.15.0" }, { name = "tqdm", marker = "extra == 'vis'", specifier = ">=4.66.5" }, ] -provides-extras = ["qasm2", "vis", "qbraid", "cirq", "pyqrack-opencl", "pyqrack-cuda", "stim"] +provides-extras = ["qasm2", "vis", "qbraid", "cirq", "pyqrack-opencl", "stim"] [package.metadata.requires-dev] dev = [ @@ -2391,12 +2387,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/30/18/09ec1971e1873031588396928d82983aef8a866cb275172dc0597a7323d1/pyqrack_cpu-1.71.2-py3-none-win_amd64.whl", hash = "sha256:025d9419a65430fad0019e00c4e30fdf824df7db0f22507d14a875f105367207", size = 842774, upload-time = "2025-11-13T21:47:49.368Z" }, ] -[[package]] -name = "pyqrack-cuda" -version = "1.83.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/07/e5/4d380b1fa6ebc65da45a04f51d17d0d3ff8b6dc01f4309b8856c2771740b/pyqrack_cuda-1.83.0.tar.gz", hash = "sha256:423672a3350f966ea2cf780a3e1f6b31c83d8ccfb459ad3fef79e3a3f5ed2bf9", size = 51972, upload-time = "2026-01-21T20:01:52.865Z" } - [[package]] name = "pyqt5" version = "5.15.11" From aa61ed70392f76ebc2c120c1390d0187b789add4 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Tue, 27 Jan 2026 21:10:05 -0500 Subject: [PATCH 15/16] fix: resolve isort and black formatting issues - Sort imports alphabetically in src/bloqade/stim/__init__.py - Fix import order in collapse/stmts/pp_measure.py - Format long line in test_stim_tag.py per black requirements Co-Authored-By: Claude Opus 4.5 --- src/bloqade/stim/__init__.py | 10 +++++----- src/bloqade/stim/dialects/collapse/stmts/pp_measure.py | 2 +- test/stim/dialects/stim/emit/test_stim_tag.py | 5 ++++- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/bloqade/stim/__init__.py b/src/bloqade/stim/__init__.py index c679c196e..0cd23a9dd 100644 --- a/src/bloqade/stim/__init__.py +++ b/src/bloqade/stim/__init__.py @@ -3,6 +3,7 @@ from ._wrappers import ( h as h, s as s, + t as t, x as x, y as y, z as z, @@ -15,6 +16,7 @@ rx as rx, ry as ry, rz as rz, + u3 as u3, mpp as mpp, mxx as mxx, myy as myy, @@ -32,6 +34,9 @@ detector as detector, identity as identity, qubit_loss as qubit_loss, + rotation_x as rotation_x, + rotation_y as rotation_y, + rotation_z as rotation_z, depolarize1 as depolarize1, depolarize2 as depolarize2, pauli_string as pauli_string, @@ -40,9 +45,4 @@ pauli_channel2 as pauli_channel2, observable_include as observable_include, correlated_qubit_loss as correlated_qubit_loss, - t as t, - rotation_x as rotation_x, - rotation_y as rotation_y, - rotation_z as rotation_z, - u3 as u3, ) diff --git a/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py b/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py index 4bf179866..a58a623a8 100644 --- a/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py +++ b/src/bloqade/stim/dialects/collapse/stmts/pp_measure.py @@ -2,8 +2,8 @@ from kirin.decl import info, statement from .._dialect import dialect -from ...auxiliary.types import PauliStringType from ...stim_statement import StimStatement +from ...auxiliary.types import PauliStringType @statement(dialect=dialect) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index 2d25c7800..cad7935c9 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -310,4 +310,7 @@ def test_qubit_coords_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_qubit_coords_with_tag) - assert buf.getvalue().strip() == "QUBIT_COORDS[qc_tag](1.00000000, 2.00000000, 3.00000000) 0" + assert ( + buf.getvalue().strip() + == "QUBIT_COORDS[qc_tag](1.00000000, 2.00000000, 3.00000000) 0" + ) From dc49446cb1797c0f9022cd773e66c809b2c27e83 Mon Sep 17 00:00:00 2001 From: Chen Zhao Date: Wed, 28 Jan 2026 10:37:02 -0500 Subject: [PATCH 16/16] combine tags with non-Clifford tags --- src/bloqade/stim/dialects/gate/emit.py | 22 ++++++++++++------- test/stim/dialects/stim/emit/test_stim_tag.py | 10 ++++----- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/bloqade/stim/dialects/gate/emit.py b/src/bloqade/stim/dialects/gate/emit.py index e8625c50c..d3553aa17 100644 --- a/src/bloqade/stim/dialects/gate/emit.py +++ b/src/bloqade/stim/dialects/gate/emit.py @@ -9,7 +9,6 @@ @dialect.register(key="emit.stim") class EmitStimGateMethods(MethodTable): - gate_1q_map: dict[str, tuple[str, str]] = { stmts.Identity.name: ("I", "I"), stmts.X.name: ("X", "X"), @@ -46,8 +45,10 @@ def single_qubit_gate( def t_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.T): """Emit T gate as S[T] or S_DAG[T] in Stim annotation format.""" targets: tuple[str, ...] = frame.get_values(stmt.targets) - gate_name = "S_DAG[T]" if stmt.dagger else "S[T]" - gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + base_name = "S_DAG" if stmt.dagger else "S" + gate_name = ( + f"{base_name}[T" + ("(tag=" + stmt.tag + ")" if stmt.tag else "") + "]" + ) res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -71,8 +72,11 @@ def rotation_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.Rx """Emit rotation gate as I[R_X/R_Y/R_Z(theta=...)] in Stim annotation format.""" targets: tuple[str, ...] = frame.get_values(stmt.targets) angle_str: str = self._format_angle(frame.get(stmt.angle)) - gate_name = f"I[{stmt.name}(theta={angle_str})]" - gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + gate_name = ( + f"I[{stmt.name}(theta={angle_str}" + + (", tag=" + stmt.tag if stmt.tag else "") + + ")]" + ) res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -84,8 +88,11 @@ def u3_gate(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.U3): theta_str: str = self._format_angle(frame.get(stmt.theta)) phi_str: str = self._format_angle(frame.get(stmt.phi)) lam_str: str = self._format_angle(frame.get(stmt.lam)) - gate_name = f"I[U3(theta={theta_str}, phi={phi_str}, lambda={lam_str})]" - gate_name = self._format_gate_with_tag(gate_name, stmt.tag) + gate_name = ( + f"I[U3(theta={theta_str}, phi={phi_str}, lambda={lam_str}" + + (", tag=" + stmt.tag if stmt.tag else "") + + ")]" + ) res = f"{gate_name} " + " ".join(targets) frame.write_line(res) return () @@ -132,7 +139,6 @@ def ctrl_two_qubit_gate( @impl(stmts.SPP) def spp(self, emit: EmitStimMain, frame: EmitStimFrame, stmt: stmts.SPP): - targets: tuple[str, ...] = tuple( targ.upper() for targ in frame.get_values(stmt.targets) ) diff --git a/test/stim/dialects/stim/emit/test_stim_tag.py b/test/stim/dialects/stim/emit/test_stim_tag.py index cad7935c9..0179f6c73 100644 --- a/test/stim/dialects/stim/emit/test_stim_tag.py +++ b/test/stim/dialects/stim/emit/test_stim_tag.py @@ -89,7 +89,7 @@ def test_t_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_t_with_tag) - assert buf.getvalue().strip() == "S[T][t_tag] 0" + assert buf.getvalue().strip() == "S[T(tag=t_tag)] 0" @stim.main def test_t_dag_with_tag(): @@ -98,7 +98,7 @@ def test_t_dag_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_t_dag_with_tag) - assert buf.getvalue().strip() == "S_DAG[T][t_dag_tag] 0" + assert buf.getvalue().strip() == "S_DAG[T(tag=t_dag_tag)] 0" def test_spp_with_tag(): @@ -133,7 +133,7 @@ def test_rx_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_rx_with_tag) - assert buf.getvalue().strip() == "I[R_X(theta=0.5*pi)][rx_tag] 0" + assert buf.getvalue().strip() == "I[R_X(theta=0.5*pi, tag=rx_tag)] 0" @stim.main def test_ry_with_tag(): @@ -142,7 +142,7 @@ def test_ry_with_tag(): buf = io.StringIO() stim_emit = EmitStimMain(dialects=stim.main, io=buf) stim_emit.run(test_ry_with_tag) - assert buf.getvalue().strip() == "I[R_Y(theta=1.0*pi)][ry_tag] 0" + assert buf.getvalue().strip() == "I[R_Y(theta=1.0*pi, tag=ry_tag)] 0" def test_u3_gate_with_tag(): @@ -157,7 +157,7 @@ def test_u3_with_tag(): stim_emit.run(test_u3_with_tag) assert ( buf.getvalue().strip() - == "I[U3(theta=0.5*pi, phi=1.0*pi, lambda=0.25*pi)][u3_tag] 0" + == "I[U3(theta=0.5*pi, phi=1.0*pi, lambda=0.25*pi, tag=u3_tag)] 0" )