Skip to content

Commit 339e60d

Browse files
authored
Add PauliString impl to NsitesAnalysis (#442)
Yet another bug found in #436
1 parent 14fdf7a commit 339e60d

File tree

2 files changed

+23
-0
lines changed

2 files changed

+23
-0
lines changed

src/bloqade/squin/analysis/nsites/impls.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,13 @@ def scale(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Scal
8181
op_sites = frame.get(stmt.op)
8282
return (op_sites,)
8383

84+
@interp.impl(op.stmts.PauliString)
85+
def pauli_string(
86+
self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.PauliString
87+
):
88+
s = stmt.string
89+
return (NumberSites(sites=len(s)),)
90+
8491

8592
@scf.dialect.register(key="op.nsites")
8693
class ScfSquinOp(ScfTypeInfer):

test/squin/analysis/test_nsites_analysis.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,3 +248,19 @@ def test():
248248
assert has_n_sites[0].sites == 1
249249
for n_site in has_n_sites[1:]:
250250
assert n_site.sites == 2
251+
252+
253+
def test_pauli_string():
254+
@squin.kernel(fold=False)
255+
def main():
256+
squin.op.pauli_string(string="XYZ")
257+
258+
nsites_frame, _ = nsites.NSitesAnalysis(main.dialects).run_analysis(main)
259+
260+
has_n_sites = []
261+
for nsites_type in nsites_frame.entries.values():
262+
if isinstance(nsites_type, nsites.NumberSites):
263+
has_n_sites.append(nsites_type)
264+
265+
assert len(has_n_sites) == 1
266+
assert has_n_sites[0].sites == 3

0 commit comments

Comments
 (0)