Skip to content

Commit 34a587e

Browse files
committed
Make non-strict zips strict in tests/scan
1 parent ff3a987 commit 34a587e

File tree

1 file changed

+21
-14
lines changed

1 file changed

+21
-14
lines changed

tests/scan/test_printing.py

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,10 @@ def test_debugprint_sitsot():
6262
Scan{scan_fn, while_loop=False, inplace=none} [id C]
6363
← Mul [id W] (inner_out_sit_sot-0)
6464
├─ *0-<Vector(float64, shape=(?,))> [id X] -> [id E] (inner_in_sit_sot-0)
65-
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M] (inner_in_non_seqs-0)"""
65+
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M] (inner_in_non_seqs-0)
66+
"""
6667

67-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
68+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
6869
assert truth.strip() == out.strip()
6970

7071

@@ -120,9 +121,10 @@ def test_debugprint_sitsot_no_extra_info():
120121
Scan{scan_fn, while_loop=False, inplace=none} [id C]
121122
← Mul [id W]
122123
├─ *0-<Vector(float64, shape=(?,))> [id X] -> [id E]
123-
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M]"""
124+
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M]
125+
"""
124126

125-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
127+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
126128
assert truth.strip() == out.strip()
127129

128130

@@ -188,9 +190,10 @@ def test_debugprint_nitsot():
188190
├─ *0-<Scalar(float64, shape=())> [id Y] -> [id S] (inner_in_seqs-0)
189191
└─ Pow [id Z]
190192
├─ *2-<Scalar(float64, shape=())> [id BA] -> [id W] (inner_in_non_seqs-0)
191-
└─ *1-<Scalar(int64, shape=())> [id BB] -> [id U] (inner_in_seqs-1)"""
193+
└─ *1-<Scalar(int64, shape=())> [id BB] -> [id U] (inner_in_seqs-1)
194+
"""
192195

193-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
196+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
194197
assert truth.strip() == out.strip()
195198

196199

@@ -303,9 +306,10 @@ def compute_A_k(A, k):
303306
Scan{scan_fn, while_loop=False, inplace=none} [id BE]
304307
← Mul [id CA] (inner_out_sit_sot-0)
305308
├─ *0-<Vector(float64, shape=(?,))> [id CB] -> [id BG] (inner_in_sit_sot-0)
306-
└─ *1-<Vector(float64, shape=(?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)"""
309+
└─ *1-<Vector(float64, shape=(?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)
310+
"""
307311

308-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
312+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
309313
assert truth.strip() == out.strip()
310314

311315
fg = FunctionGraph([c, k, A], [final_result])
@@ -402,9 +406,10 @@ def compute_A_k(A, k):
402406
→ *1-<Vector(float64, shape=(?,))> [id CB] -> [id BA] (inner_in_non_seqs-0)
403407
← Mul [id CC] (inner_out_sit_sot-0)
404408
├─ *0-<Vector(float64, shape=(?,))> [id CA] (inner_in_sit_sot-0)
405-
└─ *1-<Vector(float64, shape=(?,))> [id CB] (inner_in_non_seqs-0)"""
409+
└─ *1-<Vector(float64, shape=(?,))> [id CB] (inner_in_non_seqs-0)
410+
"""
406411

407-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
412+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
408413
assert truth.strip() == out.strip()
409414

410415

@@ -477,9 +482,10 @@ def fn(a_m2, a_m1, b_m2, b_m1):
477482
└─ *0-<Scalar(int64, shape=())> [id BD] -> [id E] (inner_in_mit_sot-0-0)
478483
← Add [id BE] (inner_out_mit_sot-1)
479484
├─ *3-<Scalar(int64, shape=())> [id BF] -> [id O] (inner_in_mit_sot-1-1)
480-
└─ *2-<Scalar(int64, shape=())> [id BG] -> [id O] (inner_in_mit_sot-1-0)"""
485+
└─ *2-<Scalar(int64, shape=())> [id BG] -> [id O] (inner_in_mit_sot-1-0)
486+
"""
481487

482-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
488+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
483489
assert truth.strip() == out.strip()
484490

485491

@@ -613,9 +619,10 @@ def test_debugprint_mitmot():
613619
Scan{scan_fn, while_loop=False, inplace=none} [id F]
614620
← Mul [id CV] (inner_out_sit_sot-0)
615621
├─ *0-<Vector(float64, shape=(?,))> [id CT] -> [id H] (inner_in_sit_sot-0)
616-
└─ *1-<Vector(float64, shape=(?,))> [id CW] -> [id P] (inner_in_non_seqs-0)"""
622+
└─ *1-<Vector(float64, shape=(?,))> [id CW] -> [id P] (inner_in_non_seqs-0)
623+
"""
617624

618-
for truth, out in zip(expected_output.split("\n"), lines, strict=False):
625+
for truth, out in zip(expected_output.split("\n"), lines, strict=True):
619626
assert truth.strip() == out.strip()
620627

621628

0 commit comments

Comments
 (0)