Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,700 @@
|
||||
"""
|
||||
Python polyfills for common builtins.
|
||||
"""
|
||||
|
||||
# NOTE: 1. Please do not import any submodule in the directory here to avoid circular imports.
|
||||
# 2. While adding a new polyfill module, also add it to POLYFILLED_MODULE_NAMES in loader.py.
|
||||
# Add it in the TYPE_CHECKING block below as well.
|
||||
|
||||
import types
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
|
||||
from itertools import repeat as _repeat
|
||||
from operator import eq, ne
|
||||
from typing import Any, TYPE_CHECKING, TypeGuard, TypeVar
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
C = TypeVar("C")
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..utils import dict_keys
|
||||
|
||||
# Load by torch._dynamo.polyfills.loader
|
||||
# See also the POLYFILLED_MODULE_NAMES in torch/_dynamo/polyfills/loader.py
|
||||
# Put the submodules here to avoid circular imports
|
||||
from . import (
|
||||
_collections as _collections,
|
||||
builtins as builtins,
|
||||
functools as functools,
|
||||
itertools as itertools,
|
||||
operator as operator,
|
||||
os as os,
|
||||
pytree as pytree,
|
||||
struct as struct,
|
||||
sys as sys,
|
||||
torch_c_nn as torch_c_nn,
|
||||
traceback as traceback,
|
||||
)
|
||||
|
||||
from torch.overrides import BaseTorchFunctionMode
|
||||
|
||||
|
||||
# These classes handle support for TorchFunctionModes across
|
||||
# graph breaks
|
||||
# Today the TorchFunctionMode enter (for the classes we support)
|
||||
# simply pushes the mode onto the stack. Since after this occurs
|
||||
# the stack is mutated, and we replay these mutations, we don't need
|
||||
# any cleanup logic to be run once the graph break occurs, we simply replay
|
||||
# these mutations to ensure at the graph break the torch function mode stack is correct
|
||||
# and reconstruct the torch function mode stack normally
|
||||
# when we compile the resume function on the other side of the break.
|
||||
# However, to ensure we exit properly
|
||||
# in the resume function, we need to re-enter the contexts as we do other contexts.
|
||||
# These contexts do nothing on enter, but provide the correct exit logic to ensure
|
||||
# the stack state is correct.
|
||||
class NoEnterTorchFunctionMode(BaseTorchFunctionMode):
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# Used by WrappedUserFunctionVariable and similar to inline decorated function
|
||||
# calls with bytecode backing. Without this, the context enter/exit happens in
|
||||
# Python-level VT code, so a nested graph break inside `fn` would skip applying
|
||||
# the context in the compiled fn/resume. By inlining through this polyfill, the
|
||||
# `with` statement has real bytecode that the resume function can continue from.
|
||||
def _fn_with_ctx(ctx: Any, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
|
||||
with ctx:
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
|
||||
def index(
|
||||
iterator: Iterator[T], item: T, start: int = 0, end: int | None = None
|
||||
) -> int:
|
||||
from itertools import islice
|
||||
|
||||
for i, elem in islice(enumerate(iterator), start, end):
|
||||
if item == elem:
|
||||
return i
|
||||
# This will not run in dynamo
|
||||
raise ValueError(f"{item} is not in {type(iterator)}")
|
||||
|
||||
|
||||
def repeat(item: T, count: int) -> Iterator[T]:
|
||||
for _ in range(count):
|
||||
yield item
|
||||
|
||||
|
||||
def radians(x: float) -> float:
|
||||
import math
|
||||
|
||||
return math.pi / 180.0 * x
|
||||
|
||||
|
||||
def impl_IS_MAPPING(a: object) -> TypeIs[Mapping[Any, Any]]:
|
||||
return isinstance(a, Mapping)
|
||||
|
||||
|
||||
def impl_MATCH_SEQUENCE(a: object) -> TypeGuard[Sequence[Any]]:
|
||||
return isinstance(a, Sequence) and not isinstance(a, (str, bytes, bytearray))
|
||||
|
||||
|
||||
def _match_class_attr(obj: object, name: str, seen: set[str]) -> object:
|
||||
if name in seen:
|
||||
raise TypeError(f"{type(obj)} got multiple sub-patterns for attribute {name}")
|
||||
|
||||
attr = getattr(obj, name)
|
||||
seen.add(name)
|
||||
return attr
|
||||
|
||||
|
||||
def impl_MATCH_CLASS(
|
||||
subject: object, cls: type, nargs: int, kwargs: tuple[str, ...]
|
||||
) -> tuple[object, ...] | None:
|
||||
if not isinstance(cls, type):
|
||||
raise TypeError("called match pattern must be a class")
|
||||
|
||||
if not isinstance(subject, cls):
|
||||
return None
|
||||
|
||||
typ = type(subject)
|
||||
match_self = False
|
||||
match_args = ()
|
||||
|
||||
attrs = []
|
||||
seen = set()
|
||||
|
||||
if nargs:
|
||||
if hasattr(typ, "__match_args__"):
|
||||
match_args = typ.__match_args__
|
||||
|
||||
if not isinstance(match_args, tuple):
|
||||
raise TypeError(
|
||||
f"{typ}.__match_args__ must be a tuple, (got {type(match_args)})"
|
||||
)
|
||||
|
||||
for name in match_args[:nargs]:
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(
|
||||
f"__match_args__ elements must be strings (got {type(name)})"
|
||||
)
|
||||
attrs.append(_match_class_attr(subject, name, seen))
|
||||
else:
|
||||
# We should somehow check if the type has TPFLAGS_MATCH_SELF set
|
||||
# match_self is only true if TPFLAGS_MATCH_SELF is set, but there is
|
||||
# no way to check for it directly in Python. So we assume it is set
|
||||
# if there are no __match_args__
|
||||
match_self = True
|
||||
attrs.append(subject)
|
||||
|
||||
allowed = 1 if match_self else len(match_args)
|
||||
if allowed < nargs:
|
||||
raise TypeError(
|
||||
f"accepts {allowed} positional sub-patterns ({nargs} given)"
|
||||
)
|
||||
|
||||
for name in kwargs:
|
||||
attrs.append(_match_class_attr(subject, name, seen))
|
||||
|
||||
return tuple(attrs)
|
||||
|
||||
|
||||
def impl_MATCH_KEYS(obj: Mapping[T, U], keys: tuple[T, ...]) -> tuple[U, ...] | None:
|
||||
assert isinstance(obj, Mapping)
|
||||
if all(key in obj for key in keys):
|
||||
return tuple(obj[key] for key in keys)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def impl_CONTAINS_OP_fallback(a: T, b: Iterable[T]) -> bool:
|
||||
# performs fallback "a in b"
|
||||
if hasattr(b, "__iter__"):
|
||||
# use __iter__ if __contains__ is not available
|
||||
for x in b:
|
||||
if x == a:
|
||||
return True
|
||||
return False
|
||||
raise TypeError(f"argument of type {type(b)} is not iterable")
|
||||
|
||||
|
||||
def accumulate_grad(x: torch.Tensor, new_grad: torch.Tensor | None) -> None:
|
||||
# polyfills according to the Gradient Layout Contract
|
||||
if new_grad is None:
|
||||
return
|
||||
new_grad_strided = torch.empty_like(x)
|
||||
new_grad_strided.copy_(new_grad)
|
||||
if x.grad is None:
|
||||
x.grad = new_grad_strided
|
||||
elif torch.is_grad_enabled():
|
||||
x.grad = x.grad + new_grad_strided
|
||||
else:
|
||||
x.grad.add_(new_grad_strided)
|
||||
|
||||
|
||||
# This mirrors
|
||||
# https://github.com/python/cpython/blob/a1c52d1265c65bcf0d9edf87e143843ad54f9b8f/Objects/listobject.c#L3352-L3413
|
||||
def list_cmp(
|
||||
op: Callable[[Any, Any], bool], left: Sequence[T], right: Sequence[T]
|
||||
) -> bool:
|
||||
"""emulate `(1,2,3) > (1,2)` etc"""
|
||||
|
||||
# Optimization: For equality, short-circuit if lengths differ
|
||||
# This avoids iterating through elements and triggering guards on SymInts
|
||||
left_len = len(left)
|
||||
right_len = len(right)
|
||||
|
||||
if op is eq and left_len != right_len:
|
||||
return False
|
||||
if op is ne and left_len != right_len:
|
||||
return True
|
||||
|
||||
# Apply `op` to the first pair that differ
|
||||
for a, b in zip(left, right):
|
||||
if a != b:
|
||||
return op(a, b)
|
||||
|
||||
# No more pairs to compare, so compare sizes.
|
||||
return op(left_len, right_len)
|
||||
|
||||
|
||||
def dict___eq__(d: dict[T, U], other: dict[T, U]) -> bool:
|
||||
if (len(d) != len(other)) or (d.keys() != other.keys()):
|
||||
return False
|
||||
|
||||
if all(isinstance(a, OrderedDict) for a in (d, other)):
|
||||
return list(d.items()) == list(other.items())
|
||||
|
||||
for k, v in d.items():
|
||||
if v != other[k]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def set_symmetric_difference(
|
||||
set1: Iterable[T],
|
||||
set2: Iterable[T],
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
symmetric_difference_set: set[T] = set()
|
||||
for x in set1:
|
||||
if x not in set2:
|
||||
symmetric_difference_set.add(x)
|
||||
for x in set2:
|
||||
if x not in set1:
|
||||
symmetric_difference_set.add(x)
|
||||
return cls(symmetric_difference_set)
|
||||
|
||||
|
||||
def set_symmetric_difference_update(set1: set[T], set2: set[T]) -> None:
|
||||
result = set1.symmetric_difference(set2)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def set_isdisjoint(set1: set[T], set2: set[T]) -> bool:
|
||||
if not isinstance(set2, Iterable):
|
||||
raise TypeError(f"'{type(set2)}' object is not iterable")
|
||||
|
||||
for x in set1:
|
||||
for y in set2:
|
||||
if not isinstance(y, Hashable):
|
||||
raise TypeError(f"unhashable type: '{type(y)}'")
|
||||
if x == y:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def set_intersection(
|
||||
set1: set[T],
|
||||
*others: Iterable[T],
|
||||
# See facebook/pyrefly#1496 - leave generic
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.difference expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
# return a new set with elements common in all sets
|
||||
intersection_set = set()
|
||||
for x in set1:
|
||||
for set2 in others:
|
||||
if not any(x == y for y in set2):
|
||||
break
|
||||
else:
|
||||
intersection_set.add(x)
|
||||
return cls(intersection_set)
|
||||
|
||||
|
||||
def set_intersection_update(set1: set[T], *others: Iterable[T]) -> None:
|
||||
result = set1.intersection(*others)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def set_union(
|
||||
set1: set[T], *others: Iterable[T], cls: type[C] | None = None
|
||||
) -> C | set[T]:
|
||||
# frozenset also uses this function
|
||||
if cls is None:
|
||||
# pyrefly: ignore[bad-assignment]
|
||||
cls = type(set1)
|
||||
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.union expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
union_set = set(set1.copy())
|
||||
for set2 in others:
|
||||
set_update(union_set, set2)
|
||||
|
||||
# frozenset also uses this function
|
||||
# pyrefly: ignore [bad-argument-count, not-callable]
|
||||
return cls(union_set)
|
||||
|
||||
|
||||
# pyrefly: ignore [bad-return]
|
||||
def set_update(set1: set[T], *others: Iterable[T]) -> set[T]:
|
||||
if len(others) == 0:
|
||||
return set1
|
||||
|
||||
for set2 in others:
|
||||
for x in set2:
|
||||
if x not in set1:
|
||||
set1.add(x)
|
||||
|
||||
|
||||
def set_difference(
|
||||
set1: set[T],
|
||||
*others: Iterable[T],
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.difference expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
difference_set = set()
|
||||
for x in set1:
|
||||
for set2 in others:
|
||||
if x in set2:
|
||||
break
|
||||
else:
|
||||
difference_set.add(x)
|
||||
return cls(difference_set)
|
||||
|
||||
|
||||
def set_difference_update(set1: set[T], *others: Iterable[T]) -> None:
|
||||
result = set1.difference(*others)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def assert_dict_equal(
|
||||
self_: Any, d1: dict[T, U], d2: dict[T, U], msg: str | None = None
|
||||
) -> None:
|
||||
self_.assertTrue(d1 == d2, msg)
|
||||
|
||||
|
||||
def assert_multi_line_equal(
|
||||
self_: Any, first: T, second: T, msg: str | None = None
|
||||
) -> None:
|
||||
return self_.assertTrue(first == second, msg)
|
||||
|
||||
|
||||
# The original impl. uses difflib
|
||||
def assert_sequence_equal(
|
||||
self_: Any,
|
||||
seq1: Sequence[T],
|
||||
seq2: Sequence[T],
|
||||
msg: str | None = None,
|
||||
seq_type: type[Any] | None = None,
|
||||
) -> None:
|
||||
return self_.assertTrue(seq1 == seq2, msg)
|
||||
|
||||
|
||||
def getattr_and_trace(*args: Any, **kwargs: Any) -> Any:
|
||||
wrapper_obj = args[0]
|
||||
attr_name = args[1]
|
||||
fn = getattr(wrapper_obj, attr_name)
|
||||
return fn(*args[2:], **kwargs)
|
||||
|
||||
|
||||
def mapping_get(obj: Mapping[T, U], key: T, value: U | None = None, /) -> U | None:
|
||||
try:
|
||||
return obj.__getitem__(key)
|
||||
except KeyError:
|
||||
return value
|
||||
|
||||
|
||||
def instantiate_user_defined_class_object(
|
||||
cls: type[T], /, *args: Any, **kwargs: Any
|
||||
) -> T:
|
||||
obj = cls.__new__(cls, *args, **kwargs)
|
||||
|
||||
# Only call __init__ if the object's type is a subclass of cls.
|
||||
# CPython uses PyType_IsSubtype(Py_TYPE(obj), type) at the C level, which does NOT
|
||||
# go through metaclass __instancecheck__. Using isinstance() here would be wrong
|
||||
# for classes with custom __instancecheck__ (e.g. torch.ByteStorage).
|
||||
# Reference: https://github.com/python/cpython/blob/3.12/Objects/typeobject.c#L1670-L1673
|
||||
if issubclass(type(obj), cls):
|
||||
obj.__init__(*args, **kwargs)
|
||||
return obj
|
||||
|
||||
|
||||
def mutable_mapping_update(
|
||||
self,
|
||||
data: Mapping[T, U] | Iterable[tuple[T, U]] = (),
|
||||
/,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
if isinstance(data, Mapping):
|
||||
# Merge standard mapping with PyMapping_Items
|
||||
for key, value in data.items():
|
||||
self[key] = value
|
||||
# FIXME: Enabling the `elif`-branch below needs too many `VariableClass.call_obj_hasattr` changes.
|
||||
# >>> class Foo:
|
||||
# ... def __init__(self):
|
||||
# ... self.keys = lambda: ['a', 'b', 'c'] # not required to be a method
|
||||
# ...
|
||||
# ... def __getitem__(self, key):
|
||||
# ... return 0
|
||||
# ...
|
||||
# >>> dict(Foo())
|
||||
# {'a': 0, 'b': 0, 'c': 0}
|
||||
#
|
||||
# > This is a rare case, so we comment it out for now.
|
||||
#
|
||||
# elif hasattr(data, "keys"):
|
||||
# # Merge mapping-like object with PyMapping_Keys + PyObject_GetItem
|
||||
# for key in data.keys():
|
||||
# self[key] = data[key]
|
||||
else:
|
||||
if not isinstance(data, Iterable):
|
||||
raise TypeError(f"{type(data).__name__!r} object is not iterable")
|
||||
# Likely a sequence of pairs
|
||||
for key, value in data:
|
||||
self[key] = value
|
||||
|
||||
if kwargs:
|
||||
for key, value in kwargs.items():
|
||||
self[key] = value
|
||||
|
||||
|
||||
# Used with something like dict(obj)
|
||||
def construct_dict(
|
||||
cls: type[T],
|
||||
data: Mapping[object, object] | Iterable[tuple[object, object]] = (),
|
||||
/,
|
||||
**kwargs: Any,
|
||||
) -> T:
|
||||
self = cls.__new__(cls)
|
||||
mutable_mapping_update(self, data, **kwargs)
|
||||
return self
|
||||
|
||||
|
||||
def foreach_map_fn(*args: Any) -> Any:
|
||||
op = args[0]
|
||||
new_args: list[Any] = []
|
||||
at_least_one_list = False
|
||||
for arg in args[1:]:
|
||||
if not isinstance(arg, (list, tuple)):
|
||||
new_args.append(_repeat(arg))
|
||||
else:
|
||||
at_least_one_list = True
|
||||
new_args.append(arg)
|
||||
|
||||
# Just apply op once to args if there are no lists
|
||||
if not at_least_one_list:
|
||||
return op(*args[1:])
|
||||
|
||||
out = []
|
||||
for unpacked in zip(*new_args):
|
||||
out.append(op(*unpacked))
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def foreach_lerp_inplace(
|
||||
self,
|
||||
end: list[torch.Tensor] | tuple[torch.Tensor, ...],
|
||||
weight: float | int | torch.Tensor,
|
||||
) -> None:
|
||||
# Decompose lerp via addcmul_ for FMA. Uses the same dual-formula
|
||||
# approach as CUDA's native lerp to get bitwise identical results:
|
||||
# |w| < 0.5 (low): fma(w, diff, start)
|
||||
# |w| >= 0.5 (high): fma(-(1-w), diff, end)
|
||||
# For tensor weights (e.g. 0-dim tensor from tensor betas in Adam) the
|
||||
# low formula is always used because the native lerp_scalar lowering
|
||||
# would crash on float(weight) for symbolic expressions.
|
||||
diff = torch._foreach_sub(end, self)
|
||||
if isinstance(weight, torch.Tensor):
|
||||
# Select base and weight for the dual formula before a single addcmul:
|
||||
# low (|w| < 0.5): fma(w, diff, self)
|
||||
# high (|w| >= 0.5): fma(-(1-w), diff, end)
|
||||
mask = weight.abs() >= 0.5
|
||||
neg_omw = -(1.0 - weight)
|
||||
w = torch.where(mask, neg_omw, weight)
|
||||
bases = [torch.where(mask, e, s) for s, e in zip(self, end)]
|
||||
w_list = [w] * len(diff)
|
||||
torch._foreach_addcmul_(bases, w_list, diff)
|
||||
for s, b in zip(self, bases):
|
||||
s.copy_(b)
|
||||
else:
|
||||
abs_weight = weight if weight >= 0 else -weight
|
||||
if abs_weight >= 0.5:
|
||||
# High formula: end + (-(1-w)) * diff → fma(-(1-w), diff, end)
|
||||
# Compute 1-w in target dtype to match CUDA rounding.
|
||||
d0 = self[0]
|
||||
neg_omw = -(1.0 - torch.tensor(weight, dtype=d0.dtype, device=d0.device))
|
||||
neg_omw_list = [neg_omw] * len(diff)
|
||||
for s, e in zip(self, end):
|
||||
s.copy_(e)
|
||||
torch._foreach_addcmul_(self, neg_omw_list, diff)
|
||||
else:
|
||||
# Low formula: start + w * diff → fma(w, diff, start)
|
||||
weights = [torch.full_like(d, weight) for d in diff]
|
||||
torch._foreach_addcmul_(self, weights, diff)
|
||||
return self
|
||||
|
||||
|
||||
def foreach_pow_scalar(
|
||||
scalar: Any, exps: Sequence[bool | complex | float | int]
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
return torch._foreach_pow([scalar for _ in exps], exps)
|
||||
|
||||
|
||||
def predicate(obj: object) -> bool:
|
||||
# This will cause the rest of dynamo to handle the if statement correctly, so we don't have to rewrite it here.
|
||||
# We can't just use bool() here since we can't trace into that in general.
|
||||
if obj:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def cmp_eq(a: object, b: object) -> bool:
|
||||
# Note that the commented `is` check should ideally be removed. This is a
|
||||
# CPython optimization that skips the __eq__ checks it the obj id's are
|
||||
# same. But, these lines adds many `is` nodes in the Fx graph for
|
||||
# SymNodeVariable. For now, we can just skip this check. This is STILL
|
||||
# correct because one of the __eq__ checks will pass later, just could be
|
||||
# slow in some corner cases.
|
||||
# if a is b:
|
||||
# return True
|
||||
if isinstance(a, type):
|
||||
# Default metaclass equality is identity-based. Preserve the reflected
|
||||
# operand fallback without tracing through type.__eq__.
|
||||
if type(a).__eq__ is type.__eq__:
|
||||
result = True if a is b else NotImplemented
|
||||
else:
|
||||
result = type(a).__eq__(a, b)
|
||||
else:
|
||||
result = a.__eq__(b)
|
||||
if result is NotImplemented:
|
||||
if isinstance(b, type):
|
||||
if type(b).__eq__ is type.__eq__:
|
||||
result = True if a is b else NotImplemented
|
||||
else:
|
||||
result = type(b).__eq__(b, a)
|
||||
else:
|
||||
result = b.__eq__(a)
|
||||
return result is not NotImplemented and result
|
||||
|
||||
|
||||
def cmp_ne(a: object, b: object) -> bool:
|
||||
if isinstance(a, type):
|
||||
if type(a).__ne__ is type.__ne__:
|
||||
result = False if a is b else NotImplemented
|
||||
else:
|
||||
result = type(a).__ne__(a, b)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
elif isinstance(type(a).__ne__, types.FunctionType):
|
||||
result = a.__ne__(b)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
# Fall through to try b.__ne__(a) or cmp_eq
|
||||
if isinstance(b, type):
|
||||
if type(b).__ne__ is type.__ne__:
|
||||
result = False if a is b else NotImplemented
|
||||
else:
|
||||
result = type(b).__ne__(b, a)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
elif isinstance(type(b).__ne__, types.FunctionType):
|
||||
result = b.__ne__(a)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
return not cmp_eq(a, b)
|
||||
|
||||
|
||||
def cmp_lt(a: Any, b: Any) -> bool:
|
||||
result = a.__lt__(b)
|
||||
if result is NotImplemented:
|
||||
raise TypeError(f"{type(a)} does not support the < operator")
|
||||
return result
|
||||
|
||||
|
||||
def cmp_le(a: Any, b: Any) -> bool:
|
||||
# Check if __le__ is overridden
|
||||
if isinstance(type(a).__le__, types.FunctionType):
|
||||
return a.__le__(b)
|
||||
return cmp_eq(a, b) or cmp_lt(a, b)
|
||||
|
||||
|
||||
def cmp_gt(a: Any, b: Any) -> bool:
|
||||
# Check if __gt__ is overridden
|
||||
if isinstance(type(a).__gt__, types.FunctionType):
|
||||
return a.__gt__(b)
|
||||
# a > b is equivalent to b < a
|
||||
return cmp_lt(b, a)
|
||||
|
||||
|
||||
def cmp_ge(a: Any, b: Any) -> bool:
|
||||
# Check if __ge__ is overridden
|
||||
if isinstance(type(a).__ge__, types.FunctionType):
|
||||
return a.__ge__(b)
|
||||
return cmp_eq(a, b) or cmp_gt(a, b)
|
||||
|
||||
|
||||
def group_tensors_by_device_and_dtype(
|
||||
tensorlistlist: list[list[torch.Tensor | None]], with_indices: bool = False
|
||||
) -> dict[tuple[torch.device, torch.dtype], tuple[list[list[Any]], list[int]]]:
|
||||
"""Pure Python implementation of torch._C._group_tensors_by_device_and_dtype.
|
||||
|
||||
Groups tensors by their device and dtype. This is useful before sending
|
||||
tensors off to a foreach implementation, which requires tensors to be on
|
||||
one device and dtype.
|
||||
|
||||
Args:
|
||||
tensorlistlist: A list of lists of tensors (tensors can be None).
|
||||
with_indices: If True, track original indices in the output.
|
||||
|
||||
Returns:
|
||||
A dict mapping (device, dtype) tuples to (grouped_tensorlistlist, indices).
|
||||
"""
|
||||
# Result dict: (device, dtype) -> (list of lists, indices)
|
||||
result: dict[
|
||||
tuple[torch.device, torch.dtype], tuple[list[list[Any]], list[int]]
|
||||
] = {}
|
||||
|
||||
if not tensorlistlist or not tensorlistlist[0]:
|
||||
return result
|
||||
|
||||
num_lists = len(tensorlistlist)
|
||||
num_tensors = len(tensorlistlist[0])
|
||||
|
||||
for idx in range(num_tensors):
|
||||
# Find the first non-None tensor at this index to get device and dtype
|
||||
first_tensor = None
|
||||
for tlist in tensorlistlist:
|
||||
if tlist is not None and idx < len(tlist) and tlist[idx] is not None:
|
||||
first_tensor = tlist[idx]
|
||||
break
|
||||
|
||||
if first_tensor is None:
|
||||
# All tensors at this index are None, skip
|
||||
continue
|
||||
|
||||
key = (first_tensor.device, first_tensor.dtype)
|
||||
|
||||
if key not in result:
|
||||
# Initialize empty lists for each tensorlist
|
||||
result[key] = ([[] for _ in range(num_lists)], [])
|
||||
|
||||
grouped_lists, indices = result[key]
|
||||
|
||||
# Add tensors from each list at this index
|
||||
for list_idx, tlist in enumerate(tensorlistlist):
|
||||
if tlist is not None and idx < len(tlist):
|
||||
grouped_lists[list_idx].append(tlist[idx])
|
||||
else:
|
||||
grouped_lists[list_idx].append(None)
|
||||
|
||||
if with_indices:
|
||||
indices.append(idx)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,33 @@
|
||||
"""
|
||||
Python polyfills for builtins
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, MutableMapping
|
||||
from typing import TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
try:
|
||||
import _collections # type: ignore[import-not-found]
|
||||
|
||||
@substitute_in_graph(_collections._count_elements)
|
||||
def _count_elements(
|
||||
mapping: MutableMapping[T, int],
|
||||
iterable: Iterable[T],
|
||||
) -> None:
|
||||
"Tally elements from the iterable."
|
||||
mapping_get = mapping.get
|
||||
for elem in iterable:
|
||||
mapping[elem] = mapping_get(elem, 0) + 1
|
||||
|
||||
__all__.append("_count_elements")
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Python polyfills for builtins
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import functools
|
||||
import operator
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
__all__ = [
|
||||
"all",
|
||||
"any",
|
||||
"cast",
|
||||
"enumerate",
|
||||
"sum",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.all, can_constant_fold_through=True)
|
||||
def all(iterable: Iterable[object], /) -> bool:
|
||||
for elem in iterable:
|
||||
if not elem:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.any, can_constant_fold_through=True)
|
||||
def any(iterable: Iterable[object], /) -> bool:
|
||||
for elem in iterable:
|
||||
if elem:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.enumerate, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def enumerate(iterable: Iterable[_T], start: int = 0) -> Iterable[tuple[int, _T]]:
|
||||
if not isinstance(start, int):
|
||||
raise TypeError(
|
||||
f"{type(start).__name__!r} object cannot be interpreted as an integer"
|
||||
)
|
||||
|
||||
for x in iterable:
|
||||
yield start, x
|
||||
start += 1
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.sum, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def sum(iterable: Iterable[_T], /, start: _T = 0) -> _T: # type: ignore[assignment]
|
||||
return functools.reduce(operator.add, iterable, start)
|
||||
|
||||
|
||||
class _CallableIterator:
|
||||
def __init__(self, fn, sentinel): # type: ignore[no-untyped-def]
|
||||
self.fn = fn
|
||||
self.sentinel = sentinel
|
||||
|
||||
def __iter__(self): # type: ignore[no-untyped-def]
|
||||
return self
|
||||
|
||||
def __next__(self): # type: ignore[no-untyped-def]
|
||||
# The iterator created in this case will call object with no arguments
|
||||
# for each call to its __next__() method;
|
||||
r = self.fn()
|
||||
|
||||
# If the value returned is equal to sentinel, StopIteration will be raised
|
||||
if r == self.sentinel:
|
||||
raise StopIteration
|
||||
|
||||
# otherwise the value will be returned.
|
||||
return r
|
||||
|
||||
|
||||
_sentinel_missing = object()
|
||||
|
||||
|
||||
# TODO(guilhermeleobas): use substitute_in_graph for iter()
|
||||
def iter_(fn_or_iterable, sentinel=_sentinel_missing, /): # type: ignore[no-untyped-def]
|
||||
# Without a second argument, object must be a collection object which supports
|
||||
# the iterable (__iter__) or the sequence protocol (__getitem__ with an integer
|
||||
# starting at 0)
|
||||
if sentinel is _sentinel_missing:
|
||||
iterable = fn_or_iterable
|
||||
if hasattr(iterable, "__iter__"):
|
||||
iterator = iterable.__iter__()
|
||||
if hasattr(iterator, "__next__"):
|
||||
return iterator
|
||||
else:
|
||||
raise TypeError(f"'{type(iterator)}' object is not iterable")
|
||||
if hasattr(iterable, "__getitem__"):
|
||||
# Needs to be a new function to avoid iter becoming a generator
|
||||
def sequence_protocol(iterable): # type: ignore[no-untyped-def]
|
||||
i = 0
|
||||
while True:
|
||||
try:
|
||||
yield iterable.__getitem__(i)
|
||||
i += 1
|
||||
except IndexError:
|
||||
break
|
||||
|
||||
return sequence_protocol(iterable)
|
||||
raise TypeError(f"'{type(iterable)}' object is not iterable")
|
||||
else:
|
||||
# If the second argument, sentinel, is given, then object must be a
|
||||
# callable object.
|
||||
fn = fn_or_iterable
|
||||
|
||||
if not isinstance(fn, Callable): # type: ignore[arg-type]
|
||||
raise TypeError("iter(v, w): v must be a callable")
|
||||
|
||||
return _CallableIterator(fn, sentinel)
|
||||
|
||||
|
||||
@substitute_in_graph(typing.cast, can_constant_fold_through=True)
|
||||
def cast(typ: type, val: _T) -> _T: # type: ignore[type-var]
|
||||
return val
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Python polyfills for copy
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"reduce_ex_user_defined_object",
|
||||
]
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@substitute_in_graph(object.__reduce_ex__, skip_signature_check=True) # type: ignore[arg-type]
|
||||
def reduce_ex_user_defined_object(obj: T, protocol: int, /) -> tuple: # type: ignore[type-arg]
|
||||
"""Traceable polyfill for object.__reduce_ex__ on user-defined objects.
|
||||
|
||||
Returns the same tuple that CPython's _common_reduce produces:
|
||||
(copyreg.__newobj__, (cls,), obj.__dict__, None, None).
|
||||
copy._reconstruct then calls cls.__new__(cls) and updates __dict__.
|
||||
"""
|
||||
import copyreg
|
||||
|
||||
cls = type(obj)
|
||||
return (
|
||||
copyreg.__newobj__, # pyrefly: ignore[missing-attribute]
|
||||
(cls,),
|
||||
obj.__dict__,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
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")
|
||||
|
||||
|
||||
_initial_missing = object()
|
||||
|
||||
|
||||
# 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
|
||||
@@ -0,0 +1,41 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from torch._C import _fx_map_aggregate, _fx_map_arg
|
||||
from torch.fx.immutable_collections import immutable_dict, immutable_list
|
||||
from torch.fx.node import Node
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph(_fx_map_arg, can_constant_fold_through=True)
|
||||
def map_arg(a: Any, fn: Callable[[Node], Any]) -> Any:
|
||||
return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
|
||||
|
||||
|
||||
@substitute_in_graph(_fx_map_aggregate, can_constant_fold_through=True)
|
||||
def map_aggregate(a: Any, fn: Callable[[Any], Any]) -> Any:
|
||||
result: Any
|
||||
if isinstance(a, tuple):
|
||||
it = (map_aggregate(elem, fn) for elem in a)
|
||||
# Support NamedTuple (if it has `_fields`) by repacking into original type.
|
||||
result = type(a)(*it) if hasattr(a, "_fields") else tuple(it)
|
||||
elif isinstance(a, list):
|
||||
result = immutable_list([map_aggregate(elem, fn) for elem in a])
|
||||
elif isinstance(a, dict):
|
||||
result = immutable_dict([(k, map_aggregate(v, fn)) for k, v in a.items()])
|
||||
elif isinstance(a, slice):
|
||||
result = slice(
|
||||
map_aggregate(a.start, fn),
|
||||
map_aggregate(a.stop, fn),
|
||||
map_aggregate(a.step, fn),
|
||||
)
|
||||
else:
|
||||
result = fn(a)
|
||||
return result
|
||||
|
||||
|
||||
__all__ = [
|
||||
"map_arg",
|
||||
"map_aggregate",
|
||||
]
|
||||
@@ -0,0 +1,119 @@
|
||||
"""
|
||||
Python polyfills for heapq
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import importlib
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
# Partially copied from CPython test/support/import_helper.py
|
||||
# https://github.com/python/cpython/blob/bb8791c0b75b5970d109e5557bfcca8a578a02af/Lib/test/support/import_helper.py
|
||||
def _save_and_remove_modules(names: set[str]) -> dict[str, ModuleType]:
|
||||
orig_modules = {}
|
||||
prefixes = tuple(name + "." for name in names)
|
||||
for modname in list(sys.modules):
|
||||
if modname in names or modname.startswith(prefixes):
|
||||
orig_modules[modname] = sys.modules.pop(modname)
|
||||
return orig_modules
|
||||
|
||||
|
||||
def import_fresh_module(name: str, blocked: list[str]) -> ModuleType:
|
||||
# Keep track of modules saved for later restoration as well
|
||||
# as those which just need a blocking entry removed
|
||||
names = {name, *blocked}
|
||||
orig_modules = _save_and_remove_modules(names)
|
||||
for modname in blocked:
|
||||
sys.modules[modname] = None # type: ignore[assignment]
|
||||
|
||||
try:
|
||||
return importlib.import_module(name)
|
||||
finally:
|
||||
_save_and_remove_modules(names)
|
||||
sys.modules.update(orig_modules)
|
||||
|
||||
|
||||
# Import the pure Python heapq module, blocking the C extension
|
||||
py_heapq = import_fresh_module("heapq", blocked=["_heapq"])
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_heapify_max",
|
||||
"_heappop_max",
|
||||
"_heapreplace_max",
|
||||
"heapify",
|
||||
"heappop",
|
||||
"heappush",
|
||||
"heappushpop",
|
||||
"heapreplace",
|
||||
"merge",
|
||||
"nlargest",
|
||||
"nsmallest",
|
||||
]
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heapify_max)
|
||||
def _heapify_max(heap: list[_T], /) -> None:
|
||||
return py_heapq._heapify_max(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heappop_max) # type: ignore[attr-defined]
|
||||
def _heappop_max(heap: list[_T]) -> _T:
|
||||
return py_heapq._heappop_max(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heapreplace_max) # type: ignore[attr-defined]
|
||||
def _heapreplace_max(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq._heapreplace_max(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heapify)
|
||||
def heapify(heap: list[_T], /) -> None:
|
||||
return py_heapq.heapify(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappop)
|
||||
def heappop(heap: list[_T], /) -> _T:
|
||||
return py_heapq.heappop(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappush)
|
||||
def heappush(heap: list[_T], item: _T) -> None:
|
||||
return py_heapq.heappush(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappushpop)
|
||||
def heappushpop(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq.heappushpop(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heapreplace)
|
||||
def heapreplace(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq.heapreplace(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.merge) # type: ignore[arg-type]
|
||||
def merge(*iterables, key=None, reverse=False): # type: ignore[no-untyped-def]
|
||||
return py_heapq.merge(*iterables, key=key, reverse=reverse)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.nlargest) # type: ignore[arg-type]
|
||||
def nlargest(n, iterable, key=None): # type: ignore[no-untyped-def]
|
||||
return py_heapq.nlargest(n, iterable, key=key)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.nsmallest) # type: ignore[arg-type]
|
||||
def nsmallest(n, iterable, key=None): # type: ignore[no-untyped-def]
|
||||
return py_heapq.nsmallest(n, iterable, key=key)
|
||||
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Python polyfills for itertools
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import operator
|
||||
from collections.abc import Callable
|
||||
from typing import overload, TYPE_CHECKING, TypeAlias, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable, Iterator
|
||||
|
||||
|
||||
__all__ = [
|
||||
"accumulate",
|
||||
"chain",
|
||||
"chain_from_iterable",
|
||||
"compress",
|
||||
"cycle",
|
||||
"dropwhile",
|
||||
"filterfalse",
|
||||
"islice",
|
||||
"pairwise",
|
||||
"starmap",
|
||||
"takewhile",
|
||||
"tee",
|
||||
"zip_longest",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_U = TypeVar("_U")
|
||||
_Predicate: TypeAlias = Callable[[_T], object]
|
||||
_T1 = TypeVar("_T1")
|
||||
_T2 = TypeVar("_T2")
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.chain
|
||||
@substitute_in_graph(itertools.chain, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def chain(*iterables: Iterable[_T]) -> Iterator[_T]:
|
||||
for iterable in iterables:
|
||||
yield from iterable
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.accumulate
|
||||
@substitute_in_graph(itertools.accumulate, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def accumulate(
|
||||
iterable: Iterable[_T],
|
||||
func: Callable[[_T, _T], _T] | None = None,
|
||||
*,
|
||||
initial: _T | None = None,
|
||||
) -> Iterator[_T]:
|
||||
# call iter outside of the generator to match cypthon behavior
|
||||
iterator = iter(iterable)
|
||||
if func is None:
|
||||
func = operator.add
|
||||
|
||||
def _accumulate(iterator: Iterator[_T]) -> Iterator[_T]:
|
||||
total = initial
|
||||
if total is None:
|
||||
try:
|
||||
total = next(iterator)
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
yield total
|
||||
for element in iterator:
|
||||
total = func(total, element)
|
||||
yield total
|
||||
|
||||
return _accumulate(iterator)
|
||||
|
||||
|
||||
@substitute_in_graph(itertools.chain.from_iterable) # type: ignore[arg-type]
|
||||
def chain_from_iterable(iterable: Iterable[Iterable[_T]], /) -> Iterator[_T]:
|
||||
# previous version of this code was:
|
||||
# return itertools.chain(*iterable)
|
||||
# If iterable is an infinite generator, this will lead to infinite recursion
|
||||
for it in iterable:
|
||||
yield from it
|
||||
|
||||
|
||||
chain.from_iterable = chain_from_iterable # type: ignore[attr-defined]
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.compress
|
||||
@substitute_in_graph(itertools.compress, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def compress(data: Iterable[_T], selectors: Iterable[_U], /) -> Iterator[_T]:
|
||||
return (datum for datum, selector in zip(data, selectors) if selector)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.cycle
|
||||
@substitute_in_graph(itertools.cycle, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def cycle(iterable: Iterable[_T]) -> Iterator[_T]:
|
||||
iterator = iter(iterable)
|
||||
|
||||
def _cycle(iterator: Iterator[_T]) -> Iterator[_T]:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
saved = []
|
||||
for element in iterable:
|
||||
yield element
|
||||
saved.append(element)
|
||||
|
||||
while saved:
|
||||
for element in saved:
|
||||
yield element
|
||||
|
||||
return _cycle(iterator)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.dropwhile
|
||||
@substitute_in_graph(itertools.dropwhile, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def dropwhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
# dropwhile(lambda x: x < 5, [1, 4, 6, 3, 8]) -> 6 3 8
|
||||
if not callable(predicate):
|
||||
raise TypeError(f"'{type(predicate).__name__}' object is not callable")
|
||||
|
||||
iterator = iter(iterable)
|
||||
for x in iterator:
|
||||
if not predicate(x):
|
||||
yield x
|
||||
break
|
||||
|
||||
yield from iterator
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.takewhile
|
||||
@substitute_in_graph(itertools.takewhile, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def takewhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
# takewhile(lambda x: x<5, [1,4,6,3,8]) → 1 4
|
||||
if not callable(predicate):
|
||||
raise TypeError(f"'{type(predicate).__name__}' object is not callable")
|
||||
|
||||
for x in iterable:
|
||||
if not predicate(x):
|
||||
break
|
||||
yield x
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[], _U],
|
||||
iterable: Iterable[tuple[()]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T], _U],
|
||||
iterable: Iterable[tuple[_T]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T, _T1], _U],
|
||||
iterable: Iterable[tuple[_T, _T1]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T, _T1, _T2], _U],
|
||||
iterable: Iterable[tuple[_T, _T1, _T2]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.starmap
|
||||
@substitute_in_graph(itertools.starmap, is_embedded_type=True) # type: ignore[arg-type]
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def starmap(function: Callable[..., _T], iterable: Iterable, /) -> Iterable[_T]:
|
||||
# starmap(pow, [(2,5), (3,2), (10,3)]) → 32 9 1000
|
||||
if not callable(function):
|
||||
raise TypeError(f"'{type(function).__name__}' object is not callable")
|
||||
|
||||
for args in iterable:
|
||||
yield function(*args)
|
||||
|
||||
|
||||
@substitute_in_graph(itertools.filterfalse, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def filterfalse(function: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
it = iter(iterable)
|
||||
if function is None:
|
||||
return filter(operator.not_, it)
|
||||
else:
|
||||
return filter(lambda x: not function(x), it)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.islice
|
||||
@substitute_in_graph(itertools.islice, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def islice(iterable: Iterable[_T], /, *args: int | None) -> Iterator[_T]:
|
||||
s = slice(*args)
|
||||
start = 0 if s.start is None else s.start
|
||||
stop = s.stop
|
||||
step = 1 if s.step is None else s.step
|
||||
if start < 0 or (stop is not None and stop < 0) or step <= 0:
|
||||
raise ValueError(
|
||||
"Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.",
|
||||
)
|
||||
|
||||
if stop is None:
|
||||
# TODO: use indices = itertools.count() and merge implementation with the else branch
|
||||
# when we support infinite iterators
|
||||
next_i = start
|
||||
for i, element in enumerate(iterable):
|
||||
if i == next_i:
|
||||
yield element
|
||||
next_i += step
|
||||
else:
|
||||
indices = range(max(start, stop))
|
||||
next_i = start
|
||||
for i, element in zip(indices, iterable):
|
||||
if i == next_i:
|
||||
yield element
|
||||
next_i += step
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.pairwise
|
||||
@substitute_in_graph(itertools.pairwise, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def pairwise(iterable: Iterable[_T], /) -> Iterator[tuple[_T, _T]]:
|
||||
a = None
|
||||
first = True
|
||||
for b in iterable:
|
||||
if first:
|
||||
first = False
|
||||
else:
|
||||
yield a, b # type: ignore[misc]
|
||||
a = b
|
||||
|
||||
|
||||
# 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))
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
/,
|
||||
*,
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T1]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
iter2: Iterable[_T2],
|
||||
/,
|
||||
) -> Iterator[tuple[_T1 | None, _T2 | None]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
iter2: Iterable[_T2],
|
||||
/,
|
||||
*,
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T1 | _U, _T2 | _U]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T],
|
||||
iter2: Iterable[_T],
|
||||
iter3: Iterable[_T],
|
||||
/,
|
||||
*iterables: Iterable[_T],
|
||||
) -> Iterator[tuple[_T | None, ...]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T],
|
||||
iter2: Iterable[_T],
|
||||
iter3: Iterable[_T],
|
||||
/,
|
||||
*iterables: Iterable[_T],
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T | _U, ...]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.zip_longest
|
||||
@substitute_in_graph(itertools.zip_longest, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def zip_longest(
|
||||
*iterables: Iterable[_T],
|
||||
fillvalue: _U = None, # type: ignore[assignment]
|
||||
) -> Iterator[tuple[_T | _U, ...]]:
|
||||
# zip_longest('ABCD', 'xy', fillvalue='-') -> Ax By C- D-
|
||||
|
||||
iterators = list(map(iter, iterables))
|
||||
num_active = len(iterators)
|
||||
if not num_active:
|
||||
return
|
||||
|
||||
while True:
|
||||
values = []
|
||||
for i, iterator in enumerate(iterators):
|
||||
try:
|
||||
value = next(iterator)
|
||||
except StopIteration:
|
||||
num_active -= 1
|
||||
if not num_active:
|
||||
return
|
||||
iterators[i] = itertools.repeat(fillvalue) # type: ignore[arg-type]
|
||||
value = fillvalue # type: ignore[assignment]
|
||||
values.append(value)
|
||||
yield tuple(values)
|
||||
@@ -0,0 +1,48 @@
|
||||
# Used to load and initialize polyfill handlers when importing torch._dynamo
|
||||
# Please add a new import when adding a new polyfill module.
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch.utils._pytree as python_pytree
|
||||
|
||||
from .. import polyfills, trace_rules
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
|
||||
# See also the TYPE_CHECKING block in torch/_dynamo/polyfills/__init__.py
|
||||
POLYFILLED_MODULE_NAMES: tuple[str, ...] = (
|
||||
"_collections",
|
||||
"builtins",
|
||||
"copy",
|
||||
"functools",
|
||||
"itertools",
|
||||
"operator",
|
||||
"os",
|
||||
"struct",
|
||||
"sys",
|
||||
"fx",
|
||||
"tensor",
|
||||
"torch_c_nn",
|
||||
"traceback",
|
||||
)
|
||||
if python_pytree._cxx_pytree_dynamo_traceable:
|
||||
POLYFILLED_MODULE_NAMES += ("pytree",)
|
||||
|
||||
POLYFILLED_MODULES: tuple["ModuleType", ...] = tuple(
|
||||
importlib.import_module(f".{submodule}", package=polyfills.__name__)
|
||||
for submodule in POLYFILLED_MODULE_NAMES
|
||||
)
|
||||
|
||||
|
||||
# Unregister the builtin functions from _builtin_function_ids to let them to be
|
||||
# dispatched with the appropriate VariableTracker type. Otherwise, they will be
|
||||
# dispatched with BuiltinVariable if present in _builtin_function_ids.
|
||||
for polyfill_module in POLYFILLED_MODULES:
|
||||
for polyfill_name in polyfill_module.__all__:
|
||||
polyfill_handler = getattr(polyfill_module, polyfill_name)
|
||||
original_fn = polyfill_handler.__torch_dynamo_original__
|
||||
trace_rules._builtin_function_ids.remove(id(original_fn))
|
||||
@@ -0,0 +1,119 @@
|
||||
"""
|
||||
Python polyfills for operator
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Any, overload, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
|
||||
# Most unary and binary operators are handled by BuiltinVariable (e.g., `pos`, `add`)
|
||||
__all__ = ["attrgetter", "itemgetter", "methodcaller", "countOf"]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_T1 = TypeVar("_T1")
|
||||
_T2 = TypeVar("_T2")
|
||||
_Ts = TypeVarTuple("_Ts")
|
||||
_U = TypeVar("_U")
|
||||
_U1 = TypeVar("_U1")
|
||||
_U2 = TypeVar("_U2")
|
||||
_Us = TypeVarTuple("_Us")
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def attrgetter(attr: str, /) -> Callable[[Any], _U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def attrgetter(
|
||||
attr1: str, attr2: str, /, *attrs: str
|
||||
) -> Callable[[Any], tuple[_U1, _U2, Unpack[_Us]]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.attrgetter
|
||||
@substitute_in_graph(operator.attrgetter, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def attrgetter(*attrs: str) -> Callable[[Any], Any | tuple[Any, ...]]:
|
||||
if len(attrs) == 0:
|
||||
raise TypeError("attrgetter expected 1 argument, got 0")
|
||||
|
||||
if any(not isinstance(attr, str) for attr in attrs):
|
||||
raise TypeError("attribute name must be a string")
|
||||
|
||||
def resolve_attr(obj: Any, attr: str) -> Any:
|
||||
for name in attr.split("."):
|
||||
obj = getattr(obj, name)
|
||||
return obj
|
||||
|
||||
if len(attrs) == 1:
|
||||
attr = attrs[0]
|
||||
|
||||
def getter(obj: Any) -> Any:
|
||||
return resolve_attr(obj, attr)
|
||||
|
||||
else:
|
||||
|
||||
def getter(obj: Any) -> tuple[Any, ...]: # type: ignore[misc]
|
||||
return tuple(resolve_attr(obj, attr) for attr in attrs)
|
||||
|
||||
return getter
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def itemgetter(item: _T, /) -> Callable[[Any], _U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def itemgetter(
|
||||
item1: _T1, item2: _T2, /, *items: Unpack[_Ts]
|
||||
) -> Callable[[Any], tuple[_U1, _U2, Unpack[_Us]]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.itemgetter
|
||||
@substitute_in_graph(operator.itemgetter, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def itemgetter(*items: Any) -> Callable[[Any], Any | tuple[Any, ...]]:
|
||||
if len(items) == 0:
|
||||
raise TypeError("itemgetter expected 1 argument, got 0")
|
||||
|
||||
if len(items) == 1:
|
||||
item = items[0]
|
||||
|
||||
def getter(obj: Any) -> Any:
|
||||
return obj[item]
|
||||
|
||||
else:
|
||||
|
||||
def getter(obj: Any) -> tuple[Any, ...]: # type: ignore[misc]
|
||||
return tuple(obj[item] for item in items)
|
||||
|
||||
return getter
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.methodcaller
|
||||
@substitute_in_graph(operator.methodcaller, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def methodcaller(name: str, /, *args: Any, **kwargs: Any) -> Callable[[Any], Any]:
|
||||
if not isinstance(name, str):
|
||||
raise TypeError("method name must be a string")
|
||||
|
||||
def caller(obj: Any) -> Any:
|
||||
return getattr(obj, name)(*args, **kwargs)
|
||||
|
||||
return caller
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.countOf
|
||||
@substitute_in_graph(operator.countOf, can_constant_fold_through=True) # type: ignore[arg-type,misc]
|
||||
def countOf(a: Iterable[_T], b: _T, /) -> int:
|
||||
return sum(it is b or it == b for it in a)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
Python polyfills for os
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import AnyStr
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = ["fspath"]
|
||||
|
||||
|
||||
# Copied from os.py in the standard library
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
@substitute_in_graph(os.fspath, can_constant_fold_through=True)
|
||||
def fspath(path: AnyStr | os.PathLike[AnyStr]) -> AnyStr:
|
||||
if isinstance(path, (str, bytes)):
|
||||
return path
|
||||
|
||||
path_type = type(path)
|
||||
try:
|
||||
path_repr = path_type.__fspath__(path) # type: ignore[arg-type]
|
||||
except AttributeError:
|
||||
if hasattr(path_type, "__fspath__"):
|
||||
raise
|
||||
raise TypeError(
|
||||
f"expected str, bytes or os.PathLike object, not {path_type.__name__}",
|
||||
) from None
|
||||
if isinstance(path_repr, (str, bytes)):
|
||||
return path_repr # type: ignore[return-value]
|
||||
raise TypeError(
|
||||
f"expected {path_type.__name__}.__fspath__() to return str or bytes, "
|
||||
f"not {type(path_repr).__name__}",
|
||||
)
|
||||
@@ -0,0 +1,759 @@
|
||||
# Owner(s): ["module: pytree"]
|
||||
|
||||
"""
|
||||
Python polyfills for torch.utils.pytree
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
|
||||
import optree
|
||||
import optree._C
|
||||
import optree.utils
|
||||
from optree import (
|
||||
is_namedtuple,
|
||||
is_namedtuple_class,
|
||||
is_namedtuple_instance,
|
||||
is_structseq,
|
||||
is_structseq_class,
|
||||
is_structseq_instance,
|
||||
namedtuple_fields,
|
||||
structseq_fields,
|
||||
)
|
||||
|
||||
import torch.utils._cxx_pytree as cxx_pytree # noqa: F401
|
||||
import torch.utils._pytree as python_pytree
|
||||
from torch.utils._pytree import BUILTIN_TYPES, STANDARD_DICT_TYPES
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import builtins
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from typing_extensions import Self, TypeIs
|
||||
|
||||
from torch.utils._cxx_pytree import PyTree
|
||||
|
||||
|
||||
__all__ = [
|
||||
"is_namedtuple",
|
||||
"is_namedtuple_class",
|
||||
"is_namedtuple_instance",
|
||||
"is_structseq",
|
||||
"is_structseq_class",
|
||||
"is_structseq_instance",
|
||||
"namedtuple_fields",
|
||||
"structseq_fields",
|
||||
"treespec_leaf",
|
||||
"treespec_tuple",
|
||||
"treespec_dict",
|
||||
"tree_is_leaf",
|
||||
"tree_iter",
|
||||
"tree_leaves",
|
||||
"tree_flatten",
|
||||
"tree_flatten_with_path",
|
||||
"tree_structure",
|
||||
"tree_unflatten",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_KT = TypeVar("_KT")
|
||||
_VT = TypeVar("_VT")
|
||||
|
||||
|
||||
@substitute_in_graph(
|
||||
optree._C.is_dict_insertion_ordered,
|
||||
can_constant_fold_through=True,
|
||||
)
|
||||
def _(*args: Any, **kwargs: Any) -> bool:
|
||||
# In namespace 'torch', the dictionary is always traversed in insertion order.
|
||||
# This function returns True.
|
||||
raise ValueError(
|
||||
"Should not be called directly "
|
||||
"because the original function will be called in the constant fold path."
|
||||
)
|
||||
|
||||
|
||||
__name = ""
|
||||
for __name, __func in (
|
||||
("is_namedtuple", is_namedtuple),
|
||||
("is_namedtuple_class", is_namedtuple_class),
|
||||
("is_namedtuple_instance", is_namedtuple_instance),
|
||||
("is_structseq", is_structseq),
|
||||
("is_structseq_class", is_structseq_class),
|
||||
("is_structseq_instance", is_structseq_instance),
|
||||
("namedtuple_fields", namedtuple_fields),
|
||||
("structseq_fields", structseq_fields),
|
||||
):
|
||||
globals()[__name] = substitute_in_graph(
|
||||
__func, # type: ignore[arg-type]
|
||||
can_constant_fold_through=True,
|
||||
)(__func.__python_implementation__) # type: ignore[attr-defined]
|
||||
del __func
|
||||
del __name
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_is_leaf, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def tree_is_leaf(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> bool:
|
||||
if (tree is None and none_is_leaf) or (is_leaf is not None and is_leaf(tree)):
|
||||
return True
|
||||
if optree.register_pytree_node.get(type(tree), namespace=namespace) is None:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_iter, can_constant_fold_through=False) # type: ignore[arg-type]
|
||||
def tree_iter(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> Iterable[Any]:
|
||||
stack = [tree]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
if tree_is_leaf(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
):
|
||||
yield node
|
||||
continue
|
||||
|
||||
children, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
stack.extend(reversed(children))
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_leaves, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def tree_leaves(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> list[Any]:
|
||||
return list(
|
||||
tree_iter(
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _Asterisk(str):
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls) -> Self:
|
||||
return super().__new__(cls, "*")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "*" # no quotes
|
||||
|
||||
|
||||
_asterisk = _Asterisk()
|
||||
del _Asterisk
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PyTreeSpec:
|
||||
"""Analog for :class:`optree.PyTreeSpec` in Python."""
|
||||
|
||||
_children: tuple[PyTreeSpec, ...]
|
||||
_type: builtins.type | None
|
||||
_metadata: Any
|
||||
_entries: tuple[Any, ...]
|
||||
_unflatten_func: Callable[[Any | None, Iterable[PyTree]], PyTree] | None
|
||||
none_is_leaf: bool
|
||||
namespace: str
|
||||
|
||||
num_nodes: int = field(init=False)
|
||||
num_leaves: int = field(init=False)
|
||||
num_children: int = field(init=False)
|
||||
|
||||
def __post_init__(self, /) -> None:
|
||||
if self._type is None:
|
||||
assert len(self._children) == 0
|
||||
assert self._metadata is None
|
||||
assert self._entries == ()
|
||||
assert self._unflatten_func is None
|
||||
num_nodes = 1
|
||||
num_leaves = 1
|
||||
num_children = 0
|
||||
else:
|
||||
assert callable(self._unflatten_func)
|
||||
num_nodes = 1
|
||||
num_leaves = 0
|
||||
for child in self._children:
|
||||
num_nodes += child.num_nodes
|
||||
num_leaves += child.num_leaves
|
||||
num_children = len(self._children)
|
||||
|
||||
object.__setattr__(self, "num_nodes", num_nodes)
|
||||
object.__setattr__(self, "num_leaves", num_leaves)
|
||||
object.__setattr__(self, "num_children", num_children)
|
||||
|
||||
def __repr__(self, /) -> str:
|
||||
def helper(treespec: PyTreeSpec) -> str:
|
||||
if treespec.is_leaf():
|
||||
assert treespec.type is None
|
||||
return _asterisk
|
||||
|
||||
assert treespec.type is not None
|
||||
assert callable(treespec._unflatten_func)
|
||||
children_representations = [
|
||||
helper(subspec) for subspec in treespec._children
|
||||
]
|
||||
if (
|
||||
treespec.type in BUILTIN_TYPES
|
||||
or (treespec.type is type(None) and not self.none_is_leaf)
|
||||
or optree.is_namedtuple_class(treespec.type)
|
||||
or optree.is_structseq_class(treespec.type)
|
||||
):
|
||||
return treespec._unflatten_func(
|
||||
treespec._metadata,
|
||||
children_representations,
|
||||
)
|
||||
return (
|
||||
f"CustomTreeNode({treespec.type.__name__}[{treespec._metadata!r}], "
|
||||
f"[{', '.join(children_representations)}])"
|
||||
)
|
||||
|
||||
inner = [
|
||||
str(helper(self)),
|
||||
*(["NoneIsLeaf"] if self.none_is_leaf else []),
|
||||
f"namespace={self.namespace!r}",
|
||||
]
|
||||
return f"PyTreeSpec({', '.join(inner)})"
|
||||
|
||||
def __len__(self, /) -> int:
|
||||
return self.num_leaves
|
||||
|
||||
@property
|
||||
def type(self, /) -> builtins.type | None:
|
||||
return self._type
|
||||
|
||||
def is_leaf(self, /) -> bool:
|
||||
return self.num_nodes == 1 and self.num_leaves == 1
|
||||
|
||||
def paths(self, /) -> list[tuple[Any, ...]]:
|
||||
def helper(treespec: PyTreeSpec, path_prefix: list[Any]) -> None:
|
||||
if treespec.is_leaf():
|
||||
paths.append(path_prefix)
|
||||
return
|
||||
|
||||
for entry, subspec in zip(
|
||||
treespec._entries,
|
||||
treespec._children,
|
||||
strict=True,
|
||||
):
|
||||
helper(subspec, path_prefix + [entry])
|
||||
|
||||
paths: list[list[Any]] = []
|
||||
helper(self, [])
|
||||
return [tuple(path) for path in paths]
|
||||
|
||||
def accessors(self, /) -> list[optree.PyTreeAccessor]:
|
||||
def helper(
|
||||
treespec: PyTreeSpec,
|
||||
entry_path_prefix: list[optree.PyTreeEntry],
|
||||
) -> None:
|
||||
if treespec.is_leaf():
|
||||
entry_paths.append(entry_path_prefix)
|
||||
return
|
||||
|
||||
node_type = treespec.type
|
||||
assert node_type is not None
|
||||
handler = optree.register_pytree_node.get(
|
||||
node_type, namespace=treespec.namespace
|
||||
)
|
||||
assert handler is not None
|
||||
kind: optree.PyTreeKind = handler.kind
|
||||
path_entry_type: type[optree.PyTreeEntry] = handler.path_entry_type
|
||||
|
||||
for entry, subspec in zip(
|
||||
treespec._entries,
|
||||
treespec._children,
|
||||
strict=True,
|
||||
):
|
||||
helper(
|
||||
subspec,
|
||||
entry_path_prefix + [path_entry_type(entry, node_type, kind)],
|
||||
)
|
||||
|
||||
entry_paths: list[list[optree.PyTreeEntry]] = []
|
||||
helper(self, [])
|
||||
return [optree.PyTreeAccessor(path) for path in entry_paths]
|
||||
|
||||
def children(self, /) -> list[PyTreeSpec]:
|
||||
return list(self._children)
|
||||
|
||||
def child(self, index: int, /) -> PyTreeSpec:
|
||||
return self._children[index]
|
||||
|
||||
def entries(self, /) -> list[Any]:
|
||||
return list(self._entries)
|
||||
|
||||
def entry(self, index: int, /) -> Any:
|
||||
return self._entries[index]
|
||||
|
||||
def flatten_up_to(self, tree: PyTree, /) -> list[PyTree]:
|
||||
def helper(
|
||||
treespec: PyTreeSpec,
|
||||
node: PyTree,
|
||||
subtrees: list[PyTree],
|
||||
) -> None:
|
||||
if treespec.is_leaf():
|
||||
subtrees.append(node)
|
||||
return
|
||||
|
||||
node_type = type(node)
|
||||
if treespec.type not in BUILTIN_TYPES:
|
||||
# Always require custom node types to match exactly
|
||||
if node_type != treespec.type:
|
||||
raise ValueError(
|
||||
f"Type mismatch; "
|
||||
f"expected {treespec.type!r}, but got {node_type!r}.",
|
||||
)
|
||||
|
||||
children, metadata, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
none_is_leaf=self.none_is_leaf,
|
||||
namespace=self.namespace,
|
||||
)
|
||||
if len(children) != treespec.num_children:
|
||||
raise ValueError(
|
||||
f"Node arity mismatch; "
|
||||
f"expected {treespec.num_children}, but got {len(children)}.",
|
||||
)
|
||||
if metadata != treespec._metadata:
|
||||
raise ValueError(
|
||||
f"Node context mismatch for custom node type {treespec.type!r}.",
|
||||
)
|
||||
else:
|
||||
# For builtin dictionary types, we allow some flexibility
|
||||
# Otherwise, we require exact matches
|
||||
both_standard_dict = (
|
||||
treespec.type in STANDARD_DICT_TYPES
|
||||
and node_type in STANDARD_DICT_TYPES
|
||||
)
|
||||
if not both_standard_dict and node_type != treespec.type:
|
||||
raise ValueError(
|
||||
f"Node type mismatch; "
|
||||
f"expected {treespec.type!r}, but got {node_type!r}.",
|
||||
)
|
||||
if len(node) != treespec.num_children:
|
||||
raise ValueError(
|
||||
f"Node arity mismatch; "
|
||||
f"expected {treespec.num_children}, but got {len(node)}.",
|
||||
)
|
||||
|
||||
if both_standard_dict:
|
||||
# dictionary types are compatible with each other
|
||||
expected_keys = treespec.entries()
|
||||
got_key_set = set(node)
|
||||
expected_key_set = set(expected_keys)
|
||||
if got_key_set != expected_key_set:
|
||||
missing_keys = expected_key_set.difference(got_key_set)
|
||||
extra_keys = got_key_set.difference(expected_key_set)
|
||||
message = ""
|
||||
if missing_keys:
|
||||
message += f"; missing key(s): {missing_keys}"
|
||||
if extra_keys:
|
||||
message += f"; extra key(s): {extra_keys}"
|
||||
raise ValueError(f"Node keys mismatch{message}.")
|
||||
children = [node[key] for key in expected_keys]
|
||||
else:
|
||||
# node_type is treespec.type
|
||||
children, metadata, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
none_is_leaf=self.none_is_leaf,
|
||||
namespace=self.namespace,
|
||||
)
|
||||
if (
|
||||
node_type is not deque # ignore mismatch of `maxlen` for deque
|
||||
) and metadata != treespec._metadata:
|
||||
raise ValueError(
|
||||
f"Node metadata mismatch for node type {treespec.type!r}; "
|
||||
f"expected {treespec._metadata!r}, but got {metadata!r}.", # namedtuple type mismatch
|
||||
)
|
||||
|
||||
for subtree, subspec in zip(children, treespec._children, strict=True):
|
||||
helper(subspec, subtree, subtrees)
|
||||
|
||||
subtrees: list[PyTree] = []
|
||||
helper(self, tree, subtrees)
|
||||
return subtrees
|
||||
|
||||
def unflatten(self, leaves: Iterable[Any], /) -> PyTree:
|
||||
if not isinstance(leaves, (list, tuple)):
|
||||
leaves = list(leaves)
|
||||
if len(leaves) != self.num_leaves:
|
||||
raise ValueError(
|
||||
f"treespec.unflatten(leaves): `leaves` has length {len(leaves)} "
|
||||
f"but the spec refers to a pytree that holds {self.num_leaves} "
|
||||
f"items ({self}).",
|
||||
)
|
||||
if self.is_leaf():
|
||||
return leaves[0]
|
||||
|
||||
# Recursively unflatten the children
|
||||
start = 0
|
||||
end = 0
|
||||
subtrees = []
|
||||
for subspec in self._children:
|
||||
end += subspec.num_leaves
|
||||
subtrees.append(subspec.unflatten(leaves[start:end]))
|
||||
start = end
|
||||
|
||||
assert callable(self._unflatten_func)
|
||||
return self._unflatten_func(self._metadata, subtrees)
|
||||
|
||||
|
||||
def _is_pytreespec_instance(obj: Any, /) -> TypeIs[PyTreeSpec | python_pytree.TreeSpec]:
|
||||
return isinstance(obj, (PyTreeSpec, python_pytree.TreeSpec))
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_leaf,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_leaf(
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "", # unused
|
||||
) -> PyTreeSpec:
|
||||
return PyTreeSpec(
|
||||
(),
|
||||
None,
|
||||
None,
|
||||
(),
|
||||
None,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace="",
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_tuple,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_tuple(
|
||||
iterable: Iterable[PyTreeSpec] = (),
|
||||
/,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> PyTreeSpec:
|
||||
children = tuple(iterable)
|
||||
if any(not _is_pytreespec_instance(child) for child in children):
|
||||
raise ValueError(f"Expected a tuple of PyTreeSpecs, got: {children!r}.")
|
||||
if any(child.none_is_leaf != none_is_leaf for child in children):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `none_is_leaf` value "
|
||||
f"as the parent; expected {none_is_leaf}, got: {children!r}.",
|
||||
)
|
||||
if any(child.namespace not in (namespace, "") for child in children):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `namespace` value "
|
||||
f"as the parent; expected {namespace!r}, got: {children!r}.",
|
||||
)
|
||||
handler = optree.register_pytree_node.get(tuple, namespace=namespace)
|
||||
assert handler is not None
|
||||
return PyTreeSpec(
|
||||
tuple(children),
|
||||
tuple,
|
||||
None,
|
||||
tuple(range(len(children))),
|
||||
handler.unflatten_func,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_dict,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_dict(
|
||||
mapping: Mapping[Any, PyTreeSpec] | Iterable[tuple[Any, PyTreeSpec]] = (),
|
||||
/,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
**kwargs: PyTreeSpec,
|
||||
) -> PyTreeSpec:
|
||||
dct = dict(mapping, **kwargs)
|
||||
if any(not _is_pytreespec_instance(child) for child in dct.values()):
|
||||
raise ValueError(f"Expected a dictionary of TreeSpecs, got: {dct!r}.")
|
||||
if any(child.none_is_leaf != none_is_leaf for child in dct.values()):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `none_is_leaf` value "
|
||||
f"as the parent; expected {none_is_leaf}, got: {dct!r}.",
|
||||
)
|
||||
if any(child.namespace not in (namespace, "") for child in dct.values()):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `namespace` value "
|
||||
f"as the parent; expected {namespace!r}, got: {dct!r}.",
|
||||
)
|
||||
|
||||
(
|
||||
children,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func,
|
||||
) = optree.tree_flatten_one_level( # type: ignore[assignment,var-annotated]
|
||||
dct, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
return PyTreeSpec(
|
||||
tuple(children), # type: ignore[arg-type]
|
||||
dict,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_flatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_flatten(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[Any], PyTreeSpec]:
|
||||
def helper(node: PyTree, leaves: list[Any]) -> PyTreeSpec:
|
||||
if tree_is_leaf(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
):
|
||||
leaves.append(node)
|
||||
return PyTreeSpec(
|
||||
(),
|
||||
None,
|
||||
None,
|
||||
(),
|
||||
None,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
(
|
||||
children,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func,
|
||||
) = optree.tree_flatten_one_level(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
# Recursively flatten the children
|
||||
subspecs = tuple(helper(child, leaves) for child in children)
|
||||
return PyTreeSpec(
|
||||
subspecs,
|
||||
type(node),
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
) # type: ignore[arg-type]
|
||||
|
||||
leaves: list[Any] = []
|
||||
treespec = helper(tree, leaves)
|
||||
return leaves, treespec
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree._C.flatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def _C_flatten(
|
||||
tree: PyTree,
|
||||
/,
|
||||
leaf_predicate: Callable[[PyTree], bool] | None = None,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[Any], PyTreeSpec]:
|
||||
return tree_flatten( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=leaf_predicate,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_flatten_with_path,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_flatten_with_path(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
|
||||
leaves, treespec = tree_flatten(
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
return treespec.paths(), leaves, treespec # type: ignore[return-value]
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree._C.flatten_with_path,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def _C_flatten_with_path(
|
||||
tree: PyTree,
|
||||
/,
|
||||
leaf_predicate: Callable[[PyTree], bool] | None = None,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
|
||||
return tree_flatten_with_path( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=leaf_predicate,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_structure,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_structure(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> PyTreeSpec:
|
||||
return tree_flatten( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)[1]
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_unflatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_unflatten(treespec: PyTreeSpec, leaves: Iterable[Any]) -> PyTree:
|
||||
if not _is_pytreespec_instance(treespec):
|
||||
raise TypeError(
|
||||
f"Expected `treespec` to be an instance of "
|
||||
f"PyTreeSpec but got item of type {type(treespec)}."
|
||||
)
|
||||
return treespec.unflatten(leaves)
|
||||
|
||||
|
||||
_none_registration = optree.register_pytree_node.get(type(None))
|
||||
assert _none_registration is not None
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_none_registration.unflatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def none_unflatten(_: None, children: Iterable[_T], /) -> None:
|
||||
if len(list(children)) != 0:
|
||||
raise ValueError("Expected no children.")
|
||||
return None
|
||||
|
||||
|
||||
with optree.dict_insertion_ordered(False, namespace="torch"):
|
||||
_dict_registration = optree.register_pytree_node.get(dict)
|
||||
assert _dict_registration is not None
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_dict_registration.flatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def dict_flatten(
|
||||
dct: dict[_KT, _VT], /
|
||||
) -> tuple[list[_VT], tuple[list[_KT], list[_KT]], tuple[_KT, ...]]:
|
||||
sorted_keys = optree.utils.total_order_sorted(dct)
|
||||
values = [dct[key] for key in sorted_keys]
|
||||
original_keys = list(dct)
|
||||
return values, (original_keys, sorted_keys), tuple(sorted_keys)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_dict_registration.unflatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def dict_unflatten(
|
||||
metadata: tuple[list[_KT], list[_KT]],
|
||||
values: Iterable[_VT],
|
||||
/,
|
||||
) -> dict[_KT, _VT]:
|
||||
original_keys, sorted_keys = metadata
|
||||
d = dict.fromkeys(original_keys)
|
||||
d.update(zip(sorted_keys, values, strict=True))
|
||||
return d # type: ignore[return-value]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
Python polyfills for struct
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from typing import Any
|
||||
from typing_extensions import Buffer
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"pack",
|
||||
"unpack",
|
||||
]
|
||||
|
||||
|
||||
@substitute_in_graph(struct.pack, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def pack(fmt: bytes | str, /, *v: Any) -> bytes:
|
||||
return struct.pack(fmt, *v)
|
||||
|
||||
|
||||
@substitute_in_graph(struct.unpack, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def unpack(format: bytes | str, buffer: Buffer, /) -> tuple[Any, ...]:
|
||||
return struct.unpack(format, buffer)
|
||||
@@ -0,0 +1,39 @@
|
||||
"""
|
||||
Python polyfills for sys
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"intern",
|
||||
"getrecursionlimit",
|
||||
]
|
||||
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
@substitute_in_graph(sys.intern, can_constant_fold_through=True)
|
||||
def intern(string: str, /) -> str:
|
||||
return string
|
||||
|
||||
|
||||
@substitute_in_graph(sys.getrecursionlimit, can_constant_fold_through=True)
|
||||
def getrecursionlimit() -> int:
|
||||
return sys.getrecursionlimit()
|
||||
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
|
||||
@substitute_in_graph(sys.get_int_max_str_digits, can_constant_fold_through=True)
|
||||
def get_int_max_str_digits() -> int:
|
||||
return sys.get_int_max_str_digits()
|
||||
|
||||
@substitute_in_graph(sys.set_int_max_str_digits, can_constant_fold_through=True)
|
||||
def set_int_max_str_digits(maxdigits: int) -> None:
|
||||
sys.set_int_max_str_digits(maxdigits)
|
||||
|
||||
__all__ += ["get_int_max_str_digits", "set_int_max_str_digits"]
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
torch.Tensor._make_subclass
|
||||
)
|
||||
def make_subclass(
|
||||
cls: type[Any], data: torch.Tensor, requires_grad: bool = False, **kwargs: Any
|
||||
) -> Any:
|
||||
with torch._C.DisableTorchFunctionSubclass():
|
||||
# This is a rough approximation of `THPVariable_make_subclass`. It should
|
||||
# suffice for most of Dynamo tracing purposes.
|
||||
# https://github.com/pytorch/pytorch/blob/ccfde4dadfa3c342076a1ee387017f84dd4ad2f7/torch/csrc/autograd/python_variable.cpp#L597-L650
|
||||
assert len(kwargs) == 0, (
|
||||
"_make_subclass only supports requires_grad as keyword arg"
|
||||
)
|
||||
data = data.detach()
|
||||
|
||||
# Avoid unnecessary `requires_grad` mutation, which isn't supported in Dynamo.
|
||||
if data.requires_grad != requires_grad:
|
||||
data.requires_grad = requires_grad
|
||||
|
||||
# Dynamo can't yet handle upcasting to base tensor type via `as_subclass`.
|
||||
if cls is torch.Tensor:
|
||||
return torch.Tensor(data)
|
||||
|
||||
# Calling `as_subclass` because
|
||||
# 1. Dynamo knows how to handle it
|
||||
# 2. the C impls match at this point -- both `THPVariable_make_subclass` and
|
||||
# `THPVariable_as_subclass` calls `THPVariable_NewWithVar`.
|
||||
return data.as_subclass(cls)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"make_subclass",
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
Polyfills for torch._C._nn functions.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.overrides import _is_torch_function_mode_enabled, _pop_mode_temporarily
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph(torch._C._nn._parse_to, skip_signature_check=True)
|
||||
def _parse_to_polyfill(
|
||||
*args: Any, **kwargs: Any
|
||||
) -> tuple[torch.device, torch.dtype, bool, torch.memory_format]:
|
||||
"""
|
||||
Polyfill for torch._C._nn._parse_to that parses arguments to nn.Module.to().
|
||||
|
||||
Signature mirrors torch._C._nn._parse_to which accepts:
|
||||
- to(device) - device as string or torch.device
|
||||
- to(dtype) - dtype as torch.dtype
|
||||
- to(tensor) - extracts device and dtype from tensor
|
||||
- to(device=..., dtype=..., non_blocking=..., memory_format=...)
|
||||
|
||||
Returns:
|
||||
tuple: (device, dtype, non_blocking, memory_format)
|
||||
"""
|
||||
# Check for __torch_function__ mode and dispatch using handle_torch_function pattern
|
||||
if _is_torch_function_mode_enabled():
|
||||
with _pop_mode_temporarily() as mode:
|
||||
result = mode.__torch_function__(
|
||||
torch._C._nn._parse_to, tuple(), args, kwargs or {}
|
||||
)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
|
||||
# Default implementation
|
||||
device = None
|
||||
dtype = None
|
||||
non_blocking = False
|
||||
memory_format = None
|
||||
|
||||
# Handle positional arguments
|
||||
if len(args) == 1:
|
||||
arg = args[0]
|
||||
# Check if it's a tensor
|
||||
if isinstance(arg, torch.Tensor):
|
||||
device = arg.device
|
||||
dtype = arg.dtype
|
||||
# Check if it's a dtype
|
||||
elif isinstance(arg, torch.dtype):
|
||||
dtype = arg
|
||||
# Check if it's a device (string or torch.device)
|
||||
elif isinstance(arg, (str, torch.device)):
|
||||
device = torch.device(arg) if isinstance(arg, str) else arg
|
||||
else:
|
||||
raise TypeError(
|
||||
f"to() received an invalid combination of arguments. Got: {type(arg)}"
|
||||
)
|
||||
elif len(args) > 1:
|
||||
raise TypeError(
|
||||
f"to() received too many positional arguments. Got {len(args)}, expected at most 1"
|
||||
)
|
||||
|
||||
# Handle keyword arguments
|
||||
if "device" in kwargs:
|
||||
device_arg = kwargs["device"]
|
||||
if device_arg is not None:
|
||||
device = (
|
||||
torch.device(device_arg) if isinstance(device_arg, str) else device_arg
|
||||
)
|
||||
|
||||
if "dtype" in kwargs:
|
||||
dtype = kwargs["dtype"]
|
||||
|
||||
if "non_blocking" in kwargs:
|
||||
non_blocking = kwargs["non_blocking"]
|
||||
|
||||
if "memory_format" in kwargs:
|
||||
memory_format = kwargs["memory_format"]
|
||||
|
||||
# pyrefly: ignore[bad-return]
|
||||
return (device, dtype, non_blocking, memory_format)
|
||||
|
||||
|
||||
@substitute_in_graph(torch.__future__.get_swap_module_params_on_conversion)
|
||||
def get_swap_module_params_on_conversion_polyfill() -> bool:
|
||||
"""
|
||||
Polyfill for torch.__future__.get_swap_module_params_on_conversion.
|
||||
|
||||
Returns the actual value from the underlying global variable.
|
||||
"""
|
||||
# Access the module's global variable directly to avoid recursion
|
||||
import torch.__future__ as torch_future
|
||||
|
||||
return torch_future._swap_module_params_on_conversion
|
||||
|
||||
|
||||
@substitute_in_graph(torch._has_compatible_shallow_copy_type)
|
||||
def _has_compatible_shallow_copy_type_polyfill(
|
||||
input: torch.Tensor, from_: torch.Tensor
|
||||
) -> bool:
|
||||
"""
|
||||
Polyfill for torch._has_compatible_shallow_copy_type.
|
||||
|
||||
Checks if two tensors have compatible types for shallow copying.
|
||||
The C++ implementation checks if input's TensorImpl has compatible shallow copy type
|
||||
with from_'s key_set. We approximate this by checking if both tensors are the same type.
|
||||
"""
|
||||
# Check if both tensors are the same type (handles both regular tensors and subclasses)
|
||||
# This is more permissive than checking exact torch.Tensor type equality
|
||||
# but properly handles subclasses by allowing same-type shallow copies
|
||||
return type(input) is type(from_)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_parse_to_polyfill",
|
||||
"get_swap_module_params_on_conversion_polyfill",
|
||||
"_has_compatible_shallow_copy_type_polyfill",
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
"""
|
||||
Python polyfills for traceback
|
||||
"""
|
||||
|
||||
import traceback
|
||||
from traceback import StackSummary
|
||||
from types import TracebackType
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = ["extract_tb", "clear_frames"]
|
||||
|
||||
|
||||
@substitute_in_graph(traceback.extract_tb, can_constant_fold_through=True)
|
||||
def extract_tb(tb: TracebackType | None, limit: int | None = None) -> StackSummary:
|
||||
if tb is None:
|
||||
return traceback.StackSummary.from_list([])
|
||||
# pyrefly: ignore [implicit-any]
|
||||
frame_summary = []
|
||||
while tb is not None:
|
||||
if limit:
|
||||
if len(frame_summary) < limit:
|
||||
frame_summary.append(
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
tb.frame_summary
|
||||
)
|
||||
else:
|
||||
break
|
||||
else:
|
||||
frame_summary.append(tb.frame_summary) # pyrefly: ignore[missing-attribute]
|
||||
tb = tb.tb_next
|
||||
return traceback.StackSummary.from_list(frame_summary)
|
||||
|
||||
|
||||
@substitute_in_graph(traceback.clear_frames, can_constant_fold_through=True)
|
||||
def clear_frames(tb: TracebackType | None) -> None:
|
||||
# no-op
|
||||
return None
|
||||
Reference in New Issue
Block a user