UnaryOp(); Subscript; Slice; keyword(); fmt_Tensor();

tests/test_control_flow.py

@@ -13,7 +13,6 @@ def test_constant():
         x = 2
         print(x == 2)
         """
-        <callable> : print
         True : x == 2
         None : print(x == 2)
         """
@@ -103,7 +102,6 @@ def test_for():
                 # odds.append(x)
                 """
                 [] : odds
-                <callable> : odds.append
                 1 : x
                 None : odds.append(x)
                 """
@@ -129,7 +127,6 @@ def test_for():
                 # odds.append(x)
                 """
                 [1] : odds
-                <callable> : odds.append
                 3 : x
                 None : odds.append(x)
                 """
@@ -155,7 +152,6 @@ def test_for():
                 # odds.append(x)
                 """
                 [1, 3] : odds
-                <callable> : odds.append
                 5 : x
                 None : odds.append(x)
                 """
@@ -181,7 +177,6 @@ def test_for():
                 # odds.append(x)
                 """
                 [1, 3, 5] : odds
-                <callable> : odds.append
                 7 : x
                 None : odds.append(x)
                 """
@@ -207,7 +202,6 @@ def test_for():
                 odds.append(x)
                 """
                 [1, 3, 5, 7] : odds
-                <callable> : odds.append
                 9 : x
                 None : odds.append(x)
                 """

tests/test_expressions.py

@@ -46,7 +46,6 @@ def test_call_print():
     def target():
         print('This line will be printed.')
         """
-        <callable> : print
         None : print('This line will be printed.')
         """
 ''')

tests/test_torch.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+
+from test_utils import *
+
+
+def test_torch():
+
+    @Commentor("<return>", _globals=globals())
+    def target():
+
+        x = torch.ones(4, 5)
+        for i in range(3):
+            x = x[..., None, :]
+
+        a = torch.randn(309, 110, 3)[:100]
+        f = nn.Linear(3, 128)
+        b = f(a.reshape(-1, 3)).reshape(-1, 110, 128)
+        c = torch.concat((a, b), dim=-1)
+
+        return c.flatten()
+
+    asserteq_or_print(
+        target(), '''    def target():
+        x = torch.ones(4, 5)
+        """
+        [4, 5] : torch.ones(4, 5)
+        ----------
+        [4, 5] : 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, :]
+            ----------
+            [4, 1, 1, 1, 5] : 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]
+        ----------
+        [100, 110, 3] : 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)
+        ----------
+        [100, 110, 128] : 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)
+        ----------
+        [100, 110, 131] : c
+        """
+        return c.flatten()
+        """
+        [100, 110, 131] : c
+        [1441000] : c.flatten()
+        """
+''')

tests/test_variables.py

@@ -14,7 +14,6 @@ def test_assign():
         myint = 7
         print(myint)
         """
-        <callable> : print
         7 : myint
         None : print(myint)
         """