Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add lazy IteratorVariable implementations for map and zip (#131413)
Summary: Fixes pytorch/pytorch#130750. Repro of lazy/eager `map` discrepancy without `islice`: ```python def fn(a, b): y = 1 def f(x): nonlocal y y += 1 return x l = list(zip([a, b], map(f, [1, 2, 3, 4]))) return a + y ``` The major change is that we implement `MapVariable` and `ZipVariable` based on `IteratorVariable`. Before, `map` and `zip` were being traced by immediately unpacking the result as a `TupleVariable`, which is wrong in cases such as the example above. `MapVariable`s are not allowed to be unpacked while `ZipVariable`s can only be unpacked if all of its iterables can also be unpacked. We also add new `[has_]force_unpack_var_sequence` methods to `VariableTracker` for the case where it is safe to unpack the entire sequence lazily, e.g., when building a list from a map (i.e. `list(map(f, ...))`). X-link: pytorch/pytorch#131413 Approved by: https://github.com/anijain2305 Reviewed By: clee2000 Differential Revision: D60322948 Pulled By: williamwen42 fbshipit-source-id: 52f7763d58943696c0ed2abf8fa03fa6795d1be9
- Loading branch information