sam/trace_commentor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Lambda(), AnnAssign(), IfExp(), Set(), Dict(), FormattedValue(), JoinedStr(); some beautification.

Browse Source
This commit is contained in:
Yuyao Huang
2024-04-23 20:46:52 +08:00
parent 1b4eb16d51
commit f448993c0d
16 changed files with 192 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_control_flow.py
+9 -117
View File
@@ -43,12 +43,15 @@ def test_if():
if x > 3: # False
x = 2 * x # skipped
y = 1 # skipped
elif x > 2: # False
x = 4 * x # skipped
y = 2 # skipped
elif x > 3: # False
x = 4 * x # skipped
y = 3 # skipped
else: # True
x = 8 * x
"""
@@ -57,7 +60,9 @@ def test_if():
----------
16 : x
"""
y = 5
''')
@@ -81,133 +86,20 @@ def test_for():
odds = []
for x in range(10):
###### !new iteration! ######
"""
0 : __REG__for_loop_iter_once
----------
0 : x
"""
# if x % 2 == 0: # True
# continue # True
# odds.append(x) # skipped
###### !new iteration! ######
"""
1 : __REG__for_loop_iter_once
----------
1 : x
"""
# if x % 2 == 0: # False
# continue # skipped
# odds.append(x)
"""
[] : odds
1 : x
None : odds.append(x)
"""
###### !new iteration! ######
"""
2 : __REG__for_loop_iter_once
----------
2 : x
"""
# if x % 2 == 0: # True
# continue # True
# odds.append(x) # skipped
###### !new iteration! ######
"""
3 : __REG__for_loop_iter_once
----------
3 : x
"""
# if x % 2 == 0: # False
# continue # skipped
# odds.append(x)
"""
[1] : odds
3 : x
None : odds.append(x)
"""
###### !new iteration! ######
"""
4 : __REG__for_loop_iter_once
----------
4 : x
"""
# if x % 2 == 0: # True
# continue # True
# odds.append(x) # skipped
###### !new iteration! ######
"""
5 : __REG__for_loop_iter_once
----------
5 : x
"""
# if x % 2 == 0: # False
# continue # skipped
# odds.append(x)
"""
[1, 3] : odds
5 : x
None : odds.append(x)
"""
###### !new iteration! ######
"""
6 : __REG__for_loop_iter_once
----------
6 : x
"""
# if x % 2 == 0: # True
# continue # True
# odds.append(x) # skipped
###### !new iteration! ######
"""
7 : __REG__for_loop_iter_once
----------
7 : x
"""
# if x % 2 == 0: # False
# continue # skipped
# odds.append(x)
"""
[1, 3, 5] : odds
7 : x
None : odds.append(x)
"""
###### !new iteration! ######
"""
8 : __REG__for_loop_iter_once
----------
8 : x
"""
# if x % 2 == 0: # True
# continue # True
# odds.append(x) # skipped
###### !new iteration! ######
"""
9 : __REG__for_loop_iter_once
----------
9 : x
"""
if x % 2 == 0: # False
continue # skipped
odds.append(x)
"""
[1, 3, 5, 7] : odds
9 : x
None : odds.append(x)
"""
return odds
"""
[1, 3, 5, 7, 9] : odds
"""''')
"""
''')
tests/test_definitions.py
+1
View File
@@ -53,6 +53,7 @@ def test_args():
5 : c
2 : k
"""
return a + k
"""
1 : a
tests/test_torch.py
+27 -54
View File
@@ -21,87 +21,60 @@ def test_torch():
return c.flatten()
asserteq_or_print(
target(), ''' def target():
target(), '''
def target():
x = torch.ones(4, 5)
"""
[4, 5] : torch.ones(4, 5)
Tensor((4, 5), f32) : torch.ones(4, 5)
----------
[4, 5] : x
Tensor((4, 5), f32) : x
"""
for i in range(3):
###### !new iteration! ######
"""
0 : __REG__for_loop_iter_once
----------
0 : i
"""
# x = x[..., None, :]
"""
[4, 5] : x
[4, 1, 5] : x[..., None, :]
----------
[4, 1, 5] : x
"""
###### !new iteration! ######
"""
1 : __REG__for_loop_iter_once
----------
1 : i
"""
# x = x[..., None, :]
"""
[4, 1, 5] : x
[4, 1, 1, 5] : x[..., None, :]
----------
[4, 1, 1, 5] : x
"""
###### !new iteration! ######
"""
2 : __REG__for_loop_iter_once
----------
2 : i
"""
x = x[..., None, :]
"""
[4, 1, 1, 5] : x
[4, 1, 1, 1, 5] : x[..., None, :]
Tensor((4, 1, 1, 5), f32) : x
Tensor((4, 1, 1, 1, 5), f32) : x[..., None, :]
----------
[4, 1, 1, 1, 5] : x
Tensor((4, 1, 1, 1, 5), f32) : x
"""
a = torch.randn(309, 110, 3)[:100]
"""
[309, 110, 3] : torch.randn(309, 110, 3)
[100, 110, 3] : torch.randn(309, 110, 3)[:100]
Tensor((309, 110, 3), f32) : torch.randn(309, 110, 3)
Tensor((100, 110, 3), f32) : torch.randn(309, 110, 3)[:100]
----------
[100, 110, 3] : a
Tensor((100, 110, 3), f32) : a
"""
f = nn.Linear(3, 128)
"""
----------
"""
b = f(a.reshape(-1, 3)).reshape(-1, 110, 128)
"""
[100, 110, 3] : a
[11000, 3] : a.reshape(-1, 3)
[11000, 128] : f(a.reshape(-1, 3))
[100, 110, 128] : f(a.reshape(-1, 3)).reshape(-1, 110, 128)
Tensor((100, 110, 3), f32) : a
Tensor((11000, 3), f32) : a.reshape(-1, 3)
Tensor((11000, 128), f32) : f(a.reshape(-1, 3))
Tensor((100, 110, 128), f32) : f(a.reshape(-1, ... pe(-1, 110, 128)
----------
[100, 110, 128] : b
Tensor((100, 110, 128), f32) : b
"""
c = torch.concat((a, b), dim=-1)
"""
[100, 110, 3] : a
[100, 110, 128] : b
[100, 110, 131] : torch.concat((a, b), dim=-1)
Tensor((100, 110, 3), f32) : a
Tensor((100, 110, 128), f32) : b
Tensor((100, 110, 131), f32) : torch.concat((a, b), dim=-1)
----------
[100, 110, 131] : c
Tensor((100, 110, 131), f32) : c
"""
return c.flatten()
"""
[100, 110, 131] : c
[1441000] : c.flatten()
Tensor((100, 110, 131), f32) : c
Tensor((1441000,), f32) : c.flatten()
"""
''')
tests/test_utils.py
+5 -2
View File
@@ -1,11 +1,14 @@
import re
from io import StringIO
from contextlib import closing
from trace_commentor import flags, Commentor
WS = re.compile(" +")
def asserteq_or_print(value, ground_truth):
if flags.DEBUG or flags.PRINT:
print(value)
else:
value = value.strip("\n").rstrip(" ")
ground_truth = ground_truth.strip("\n").rstrip(" ")
value = re.sub(WS, " ", value.strip("\n").rstrip(" ").rstrip("\n"))
ground_truth = re.sub(WS, " ", ground_truth.strip("\n").rstrip(" ").rstrip("\n"))
assert value == ground_truth, "\n".join(["\n\n<<<<<<<< VALUE", value, "========================", ground_truth, ">>>>>>>> GROUND\n"])