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] unpack sequence lazily for list extend/deque extendleft (#150965)

Browse Source 
Fixes https://github.com/pytorch/pytorch/issues/133063.

We were unpacking generators/iterators eagerly when we should be unpacking them one-by-one.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150965
Approved by: https://github.com/jansel
This commit is contained in:
William Wen
2025-04-09 15:21:48 -07:00
committed by PyTorch MergeBot
parent 389cd15265
commit 4161c752bb
6 changed files with 70 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
torch/_dynamo/variables/lists.py
View File

@ -373,9 +373,9 @@ class CommonListMethodsVariable(BaseListVariable):
):
assert not kwargs
(arg,) = args
seq = arg.force_unpack_var_sequence(tx)
tx.output.side_effects.mutation(self)
self.items.extend(seq)
arg.force_apply_to_var_sequence(
tx, lambda item: self.call_method(tx, "append", [item], {})
)
return ConstantVariable.create(None)
elif name == "insert" and self.is_mutable():
assert not kwargs
@ -614,9 +614,11 @@ class DequeVariable(CommonListMethodsVariable):
):
assert len(args) == 1
assert not kwargs
prefix = args[0].force_unpack_var_sequence(tx)
tx.output.side_effects.mutation(self)
self.items[:] = [*reversed(prefix), *self.items]
# NOTE this is inefficient, but the alternative is to represent self.items
# as a deque, which is a more intrusive change.
args[0].force_apply_to_var_sequence(
tx, lambda item: self.call_method(tx, "appendleft", [item], {})
)
slice_within_maxlen = slice(None, maxlen)
result = ConstantVariable.create(None)
elif name == "popleft" and self.is_mutable():