mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
[BE]: Update Typeguard to TypeIs for better type inference (#133814)
Uses TypeIs instead of TypeGuard for better inference. See https://peps.python.org/pep-0742/ Pull Request resolved: https://github.com/pytorch/pytorch/pull/133814 Approved by: https://github.com/ezyang
This commit is contained in:
committed by
PyTorch MergeBot
parent
9bb327bfc6
commit
16caa8c1b3
@ -4,7 +4,7 @@ import contextlib
|
||||
import warnings
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Set, Union, Protocol, Tuple, Sequence, overload, Deque, Type
|
||||
from typing_extensions import TypeGuard
|
||||
from typing_extensions import TypeIs
|
||||
from collections import deque
|
||||
|
||||
import torch
|
||||
@ -365,7 +365,7 @@ class TensorWithFlatten(Protocol):
|
||||
|
||||
|
||||
|
||||
def is_traceable_wrapper_subclass(t: object) -> TypeGuard[TensorWithFlatten]:
|
||||
def is_traceable_wrapper_subclass(t: object) -> TypeIs[TensorWithFlatten]:
|
||||
"""
|
||||
Returns whether or not a tensor subclass that implements __torch_dispatch__
|
||||
is 'traceable' with torch.compile.
|
||||
@ -402,7 +402,7 @@ def is_traceable_wrapper_subclass(t: object) -> TypeGuard[TensorWithFlatten]:
|
||||
and hasattr(t, "__tensor_unflatten__")
|
||||
)
|
||||
|
||||
def is_traceable_wrapper_subclass_type(t: Type) -> TypeGuard[Type[TensorWithFlatten]]:
|
||||
def is_traceable_wrapper_subclass_type(t: Type) -> TypeIs[Type[TensorWithFlatten]]:
|
||||
"""Same as above, but takes a type argument instead of an instance."""
|
||||
return (issubclass(t, torch.Tensor) and t != torch.Tensor
|
||||
and hasattr(t, "__tensor_flatten__") and hasattr(t, "__tensor_unflatten__"))
|
||||
|
Reference in New Issue
Block a user