mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
0d4dbfd9ed
This builds off the previous patch in the stack, and fully fixes https://github.com/huggingface/diffusers/issues/10795. Essentially, tensor subclass in the issue uses `torch.Tensor._make_subclass`, which has a pretty simple shallow-copy plus type change semantics, as far as Dynamo is concerned. So this patch adds a polyfill for it. As a result, this allows us to trace through many user-defined `__new__` in tensor subclass (it's similar to how we trace through user-defined `__new__` for `UserDefinedClassVariable`), so this patch also faithfully trace through these `__new__` methods. Differential Revision: [D71906139](https://our.internmc.facebook.com/intern/diff/D71906139) Pull Request resolved: https://github.com/pytorch/pytorch/pull/149483 Approved by: https://github.com/zou3519, https://github.com/mlazos ghstack dependencies: #149482
40 lines
1.2 KiB
Python
40 lines
1.2 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 TYPE_CHECKING
|
|
|
|
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, ...] = (
|
|
"builtins",
|
|
"functools",
|
|
"itertools",
|
|
"operator",
|
|
"os",
|
|
"pytree",
|
|
"sys",
|
|
"fx",
|
|
"tensor",
|
|
)
|
|
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))
|