pytorch/torch/_dynamo/polyfills/loader.py

# 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 Tuple, TYPE_CHECKING

from .. import polyfills, trace_rules


if TYPE_CHECKING:
    from types import ModuleType


POLYFILLED_MODULE_NAMES: Tuple[str, ...] = (
    "builtins",
    "functools",
    "itertools",
    "os",
    "sys",
)
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))

        # Unregister the class object if the original function is its __new__ method
        if original_fn.__name__ == "__new__" and isinstance(
            getattr(original_fn, "__self__", None), type
        ):
            trace_rules._builtin_function_ids.remove(id(original_fn.__self__))