diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py index dc9e99f7701..9689e2c12e0 100644 --- a/test/dynamo/test_misc.py +++ b/test/dynamo/test_misc.py @@ -10045,22 +10045,6 @@ def ___make_guard_fn(): self.assertEqual(eager, compiled) self.assertEqual(len(counters["graph_break"]), 0) - def test_itertools_tee(self): - counters.clear() - - def fn(l): - a, b = itertools.tee(l) - return list(a), list(b) - - l = [1, 2, 2, 3, 4, 4, 4, 1, 2] - eager = fn(l) - - compiled_fn = torch._dynamo.optimize(backend="eager", nopython=True)(fn) - compiled = compiled_fn(l) - - self.assertEqual(eager, compiled) - self.assertEqual(len(counters["graph_break"]), 0) - def test_list_iterator_contains(self): def fn(x): it = iter(["my_weight", "not_my_weight"]) diff --git a/torch/_dynamo/polyfills/itertools.py b/torch/_dynamo/polyfills/itertools.py deleted file mode 100644 index 7207acd6504..00000000000 --- a/torch/_dynamo/polyfills/itertools.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Python polyfills for itertools -""" - -import itertools -from typing import Iterable, Iterator, Tuple, TypeVar - -from ..decorators import substitute_in_graph - - -__all__ = ["tee"] - - -_T = TypeVar("_T") - - -# Reference: https://docs.python.org/3/library/itertools.html#itertools.tee -@substitute_in_graph(itertools.tee) -def tee(iterable: Iterable[_T], n: int = 2, /) -> Tuple[Iterator[_T], ...]: - iterator = iter(iterable) - shared_link = [None, None] - - def _tee(link) -> Iterator[_T]: # type: ignore[no-untyped-def] - try: - while True: - if link[1] is None: - link[0] = next(iterator) - link[1] = [None, None] - value, link = link - yield value - except StopIteration: - return - - return tuple(_tee(shared_link) for _ in range(n)) diff --git a/torch/_dynamo/polyfills/loader.py b/torch/_dynamo/polyfills/loader.py index a1c97ff0cc3..516ce1b9c94 100644 --- a/torch/_dynamo/polyfills/loader.py +++ b/torch/_dynamo/polyfills/loader.py @@ -14,7 +14,6 @@ if TYPE_CHECKING: POLYFILLED_MODULE_NAMES: Tuple[str, ...] = ( "builtins", "functools", - "itertools", ) POLYFILLED_MODULES: Tuple["ModuleType", ...] = tuple( importlib.import_module(f".{submodule}", package=polyfills.__name__)