[iter] support iter(callable, sentinel) (#156416)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156416
Approved by: https://github.com/XuehaiPan, https://github.com/zou3519
ghstack dependencies: #156371
This commit is contained in:
Guilherme Leobas
2025-07-29 16:11:37 -03:00
committed by PyTorch MergeBot
parent fcf59df2b6
commit 4e3e3dc0a7
6 changed files with 58 additions and 23 deletions

View File

@ -1810,13 +1810,9 @@ class BuiltinVariable(VariableTracker):
# (e.g. when __iter__ just returns iter(self.list)) or return a user-defined iterator.
# If the object implements a __getitem__ method, iter(...) will call obj.__getitem__()
# with an integer argument starting at 0, until __getitem__ raises IndexError
return variables.UserFunctionVariable(
polyfills.builtins.iter
).call_function(
tx,
[obj],
{},
)
ret = variables.UserFunctionVariable(
polyfills.builtins.iter_
).call_function(tx, [obj, *args], {})
return ret