8000 [Dynamo][pytree] handle `isinstance(...)` check for polyfilled class by XuehaiPan · Pull Request #146921 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[Dynamo][pytree] handle isinstance(...) check for polyfilled class #146921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed

[Dynamo][pytree] handle isinstance(...) check for polyfilled class #146921

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d5d05d0
Update
XuehaiPan Feb 11, 2025
9e78423
Update
XuehaiPan Feb 11, 2025
468d321
Update
XuehaiPan Feb 11, 2025
eba607c
Update
XuehaiPan Feb 11, 2025
837db58
Update
XuehaiPan Feb 12, 2025
f417298
Update
XuehaiPan Feb 12, 2025
ac62c83
Update
XuehaiPan Feb 12, 2025
72d44f2
Update
XuehaiPan Feb 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update 
[ghstack-poisoned]
  • Loading branch information
@XuehaiPan
XuehaiPan committed Feb 11, 2025
commit 468d321cc2b18e7d560b003a2e973f505341c879
8 changes: 4 additions & 4 deletions torch/_dynamo/polyfills/pytree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Callable, Literal, TYPE_CHECKING
from typing_extensions import TypeIs

import torch.utils._pytree as python_pytree
from torch.utils._pytree import BUILTIN_TYPES, STANDARD_DICT_TYPES
Expand Down Expand Up @@ -318,9 +317,10 @@ def unflatten(self, leaves: Iterable[Any]) -> PyTree:
assert callable(self._unflatten_func)
return self._unflatten_func(self._metadata, subtrees)

cxx_pytree._pytreespec_types = (PyTreeSpec, cxx_pytree.PyTreeSpec)
polyfill_class_mapping[PyTreeSpec] = (cxx_pytree.PyTreeSpec,)
polyfill_class_mapping[cxx_pytree.PyTreeSpec] = (PyTreeSpec,)
_pytreespec_types = (PyTreeSpec, cxx_pytree.PyTreeSpec)
cxx_pytree._pytreespec_types = _pytreespec_types
polyfill_class_mapping[PyTreeSpec] = _pytreespec_types
polyfill_class_mapping[cxx_pytree.PyTreeSpec] = _pytreespec_types
_is_pytreespec_instance = cxx_pytree._is_pytreespec_instance

_LEAF_SPEC = PyTreeSpec((), None, None, (), None)
Expand Down
28 changes: 19 additions & 9 deletions torch/_dynamo/variables/builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1667,15 +1667,22 @@ def check_type(ty):
isinstance_type.__class__.__instancecheck__(isinstance_type, arg.value)
)

if isinstance(isinstance_type, type):
isinstance_type: tuple[type, ...] = (isinstance_type,)
isinstance_type_tuple: tuple[type, ...]
if isinstance(isinstance_type, (type,)) or callable(
getattr(isinstance_type, "__instancecheck__", None)
):
isinstance_type_tuple = (isinstance_type,)
elif sys.version_info >= (3, 10) and isinstance(
isinstance_type, types.UnionType
):
isinstance_type: tuple[type, ...] = isinstance_type.__args__
isinstance_type_tuple = isinstance_type.__args__
elif not (
isinstance(isinstance_type, tuple)
and all(isinstance(tp, type) for tp in isinstance_type)
and all(
isinstance(tp, (type,))
or callable(getattr(tp, "__instancecheck__", None))
for tp in isinstance_type
)
):
raise_observed_exception(
TypeError,
Expand All @@ -1684,20 +1691,23 @@ def check_type(ty):
"isinstance() arg 2 must be a type, a tuple of types, or a union"
],
)
else:
isinstance_type_tuple = isinstance_type

if any(tp in polyfill_class_mapping for tp in isinstance_type):
isinstance_type = tuple(
if any(tp in polyfill_class_mapping for tp in isinstance_type_tuple):
isinstance_type_tuple = tuple(
dict.fromkeys(
itertools.chain.from_iterable(
polyfill_class_mapping.get(tp, (tp,)) for tp in isinstance_type
polyfill_class_mapping.get(tp, (tp,))
for tp in isinstance_type_tuple
)
)
)

try:
val = issubclass(arg_type, isinstance_type)
val = issubclass(arg_type, isinstance_type_tuple)
except TypeError:
val = arg_type in isinstance_type
val = arg_type in isinstance_type_tuple
return variables.ConstantVariable.create(val)

def call_issubclass(self, tx: "InstructionTranslator", left_ty, right_ty):
Expand Down
3 changes: 2 additions & 1 deletion torch/utils/_cxx_pytree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,8 @@ def _private_register_pytree_node(
)


_pytreespec_types = (PyTreeSpec,) # will be updated in torch._dynamo.polyfilles.pytree
# Will be updated in torch._dynamo.polyfilles.pytree
_pytreespec_types: tuple[type, ...] = (PyTreeSpec,)


def _is_pytreespec_instance(obj: Any, /) -> TypeIs[TreeSpec]:
Expand Down
Loading
0