mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
9d0b77f4cd
This is a follow-up of #165515. ruff `UP035` rules are applied to dynamo code to use Py 3.10+ typing. Pull Request resolved: https://github.com/pytorch/pytorch/pull/165709 Approved by: https://github.com/ezyang
48 lines
966 B
Python
48 lines
966 B
Python
"""
|
|
Python polyfills for functools
|
|
"""
|
|
|
|
import functools
|
|
from collections.abc import Callable, Iterable
|
|
from typing import TypeVar
|
|
|
|
from ..decorators import substitute_in_graph
|
|
|
|
|
|
__all__ = ["reduce"]
|
|
|
|
|
|
_T = TypeVar("_T")
|
|
_U = TypeVar("_U")
|
|
|
|
|
|
class _INITIAL_MISSING:
|
|
pass
|
|
|
|
|
|
# Reference: https://docs.python.org/3/library/functools.html#functools.reduce
|
|
@substitute_in_graph(functools.reduce)
|
|
def reduce(
|
|
function: Callable[[_U, _T], _U],
|
|
iterable: Iterable[_T],
|
|
initial: _U = _INITIAL_MISSING, # type: ignore[assignment]
|
|
/,
|
|
) -> _U:
|
|
it = iter(iterable)
|
|
|
|
value: _U
|
|
if initial is _INITIAL_MISSING:
|
|
try:
|
|
value = next(it) # type: ignore[assignment]
|
|
except StopIteration:
|
|
raise TypeError(
|
|
"reduce() of empty iterable with no initial value",
|
|
) from None
|
|
else:
|
|
value = initial
|
|
|
|
for element in it:
|
|
value = function(value, element)
|
|
|
|
return value
|