mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
b6abac68ec
Changes: 1. Move `polyfill.py` -> `polyfills/__init__.py`. It can be used as `polyfill.xxx` -> `polyfills.xxx`. 2. Move submodule loading from `polyfills/__init__.py` to `polyfills/loader.py`. Merge `polyfill.py` and `polyfills/` packages. Each polyfill module have its own namespace for better code organization. The ultimate goal is make `polyfills/__init__.py` empty and all polyfill functions move to its own namespace. Pull Request resolved: https://github.com/pytorch/pytorch/pull/133977 Approved by: https://github.com/jansel
29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
# 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, ...] = ()
|
|
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))
|