mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
70853b792a
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133771 Approved by: https://github.com/jansel ghstack dependencies: #133801
37 lines
875 B
Python
37 lines
875 B
Python
"""
|
|
Python polyfills for itertools
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import itertools
|
|
from typing import Iterable, Iterator, 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))
|