frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo][itertools] refactor itertools.chain and itertools.chain.from_iterable to use polyfills (#133864)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133864
Approved by: https://github.com/jansel
ghstack dependencies: #133769, #133778, #133779
This commit is contained in:
Xuehai Pan
2024-08-30 00:47:04 +08:00
committed by PyTorch MergeBot
parent eaa449fbf0
commit 1b70366957
8 changed files with 91 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
torch/_dynamo/variables/iter.py
View File

@ -52,14 +52,6 @@ class ItertoolsVariable(VariableTracker):
for item in itertools.product(*seqs):
items.append(variables.TupleVariable(list(item)))
return variables.ListIteratorVariable(items, mutable_local=MutableLocal())
elif (
self.value is itertools.chain
and not kwargs
and all(arg.has_unpack_var_sequence(tx) for arg in args)
):
seqs = [arg.unpack_var_sequence(tx) for arg in args]
items = list(itertools.chain.from_iterable(seqs))
return variables.ListIteratorVariable(items, mutable_local=MutableLocal())
elif self.value is itertools.accumulate:
from .builtin import BuiltinVariable