FM / utils

/usr/local/lib64/python3.6/site-packages/torch/utils/

/usr/local/lib64/python3.6/site-packages/torch/utils

Name	Size	Mode	Actions
backcompat/	-	0755	rm
benchmark/	-	0755	rm
bottleneck/	-	0755	rm
data/	-	0755	rm
ffi/	-	0755	rm
hipify/	-	0755	rm
model_dump/	-	0755	rm
tensorboard/	-	0755	rm
__pycache__/	-	0755	rm
bundled_inputs.py	22282	0644	edit dl rm
checkpoint.py	12165	0644	edit dl rm
collect_env.py	16723	0644	edit dl rm
cpp_extension.py	87739	0644	edit dl rm
dlpack.py	2045	0644	edit dl rm
file_baton.py	1396	0644	edit dl rm
hooks.py	7373	0644	edit dl rm
mkldnn.py	7219	0644	edit dl rm
mobile_optimizer.py	6264	0644	edit dl rm
model_zoo.py	117	0644	edit dl rm
show_pickle.py	5169	0644	edit dl rm
throughput_benchmark.py	6387	0644	edit dl rm
_cpp_extension_versioner.py	1991	0644	edit dl rm
_crash_handler.py	673	0644	edit dl rm
_python_dispatch.py	1636	0644	edit dl rm
_pytree.py	8174	0644	edit dl rm
__init__.py	635	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/utils/_python_dispatch.py (1636B)

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_python_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_python_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_python_mode(cls) -> Iterator[None]:
    if not hasattr(cls, '__torch_dispatch__'):
        raise ValueError('The class passed to enable_python_mode '
                         'must have a __torch_dispatch__ classmethod')
    if not isinstance(cls, type) or not issubclass(cls, (torch.Tensor,)):
        raise ValueError('The argument passed to enable_python_mode '
                         'must be the type of a Tensor subclass')
    torch._C._enter_python_mode(cls)
    try:
        yield
    finally:
        torch._C._exit_python_mode()

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_python_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_python_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_python_mode(cls) -> Iterator[None]: if not hasattr(cls, '__torch_dispatch__'): raise ValueError('The class passed to enable_python_mode ' 'must have a __torch_dispatch__ classmethod') if not isinstance(cls, type) or not issubclass(cls, (torch.Tensor,)): raise ValueError('The argument passed to enable_python_mode ' 'must be the type of a Tensor subclass') torch._C._enter_python_mode(cls) try: yield finally: torch._C._exit_python_mode()