mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 13:44:15 +08:00
598e7e5f19
Changes Python Mode name to Torch Dispatch Mode because there is now a Torch Function Mode, so Torch Dispatch Mode and Torch Function Mode are consistent with each other Pull Request resolved: https://github.com/pytorch/pytorch/pull/76562 Approved by: https://github.com/zou3519, https://github.com/albanD
35 lines
1.7 KiB
Python
35 lines
1.7 KiB
Python
import torch
|
|
import contextlib
|
|
from typing import Iterator
|
|
|
|
# Context manager that causes all pytorch operators to dispatch to the passed-in
|
|
# type's __torch_dispatch__ function.
|
|
# operation that accepts no tensors but returns a tensor.
|
|
#
|
|
# enable_torch_dispatch_mode is affected by torch._C._DisableTorchDispatch.
|
|
#
|
|
# NB: Calling an operator inside __torch_dispatch__ does go through
|
|
# __torch_dispatch__ again. Please use _DisableTorchDispatch inside
|
|
# __torch_dispatch__ to prevent infinite recursion.
|
|
#
|
|
# TODO: Limitations and things about enable_torch_dispatch_mode we should fix before exposing it:
|
|
# - it currently cannot be nested. This should be simple to implement; we need a
|
|
# stack of TorchDispatchTypeObjects and the next bullet point.
|
|
# - We need a better user-facing api for torch._C._DisableTorchDispatch that
|
|
# is able to selectively disable __torch_dispatch__ of a particular class.
|
|
# - It doesn't work with the tensor constructors (torch.tensor, torch.Tensor)
|
|
# - Better name (see https://github.com/pytorch/pytorch/pull/63496#discussion_r694091694)
|
|
@contextlib.contextmanager
|
|
def enable_torch_dispatch_mode(cls) -> Iterator[None]:
|
|
if cls.__torch_dispatch__ is torch.Tensor.__torch_dispatch__:
|
|
raise ValueError('The class passed to enable_torch_dispatch_mode '
|
|
'must have a non-default __torch_dispatch__ classmethod')
|
|
if not isinstance(cls, type) or not issubclass(cls, (torch.Tensor,)):
|
|
raise ValueError('The argument passed to enable_torch_dispatch_mode '
|
|
'must be the type of a Tensor subclass')
|
|
torch._C._enter_torch_dispatch_mode(cls)
|
|
try:
|
|
yield
|
|
finally:
|
|
torch._C._exit_torch_dispatch_mode()
|