1# Used to load and initialize polyfill handlers when importing torch._dynamo 2# Please add a new import when adding a new polyfill module. 3 4import importlib 5from typing import Tuple, TYPE_CHECKING 6 7from .. import polyfills, trace_rules 8 9 10if TYPE_CHECKING: 11 from types import ModuleType 12 13 14# See also the TYPE_CHECKING block in torch/_dynamo/polyfills/__init__.py 15POLYFILLED_MODULE_NAMES: Tuple[str, ...] = ( 16 "builtins", 17 "functools", 18 "itertools", 19 "os", 20 "sys", 21) 22POLYFILLED_MODULES: Tuple["ModuleType", ...] = tuple( 23 importlib.import_module(f".{submodule}", package=polyfills.__name__) 24 for submodule in POLYFILLED_MODULE_NAMES 25) 26 27 28# Unregister the builtin functions from _builtin_function_ids to let them to be 29# dispatched with the appropriate VariableTracker type. Otherwise, they will be 30# dispatched with BuiltinVariable if present in _builtin_function_ids. 31for polyfill_module in POLYFILLED_MODULES: 32 for polyfill_name in polyfill_module.__all__: 33 polyfill_handler = getattr(polyfill_module, polyfill_name) 34 original_fn = polyfill_handler.__torch_dynamo_original__ 35 trace_rules._builtin_function_ids.remove(id(original_fn)) 36