Fix slicing list returning wrong result (#222)

* Correct ListVariable source

* Fix lint

Author: yanboliang

tests/test_misc.py

@@ -1157,3 +1157,15 @@ def fn():
         inst = dis.get_instructions(fn)
         result = bytecode_transformation.assemble(inst, fn.__code__.co_firstlineno)
         self.assertTrue(result[1] == fn.__code__.co_lnotab)
+
+    def test_python_slice(self):
+        def fn(input):
+            y = 0
+            for i, x in enumerate(input[2:], 1):
+                y = y + x
+            return y
+
+        cnts = torchdynamo.testing.CompileCounter()
+        with torchdynamo.optimize(cnts):
+            z = fn([1, 2, 3, 5])
+        self.assertEqual(z, 8)

torchdynamo/utils.py

@@ -319,7 +319,7 @@ def rot_n_helper(n):
 def is_safe_constant(v):
     if istype(v, (tuple, frozenset)):
         return all(map(is_safe_constant, v))
-    return istype(v, (types.CodeType, int, float, bool, str, bytes, type(None)))
+    return istype(v, (types.CodeType, int, float, bool, str, bytes, type(None), slice))
 
 
 def check_constant_args(args, kwargs):

torchdynamo/variables/lists.py

@@ -6,6 +6,7 @@
 from .. import variables
 from ..bytecode_transformation import create_instruction
 from ..exc import unimplemented
+from ..source import GetItemSource
 from ..utils import namedtuple_fields
 from .base import MutableLocal
 from .base import VariableTracker
@@ -40,9 +41,11 @@ def as_proxy(self):
     def getitem_const(self, arg: VariableTracker):
         index = arg.as_python_constant()
         if isinstance(index, slice):
-            return self.clone(items=self.items[index], mutable_local=None).add_options(
-                arg, self
-            )
+            return self.clone(
+                items=self.items[index],
+                source=GetItemSource(self.source, index),
+                mutable_local=None,
+            ).add_options(arg, self)
         else:
             assert isinstance(index, int)
             return self.items[index].add_options(arg, self)