File size: 943 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
import functools
from typing import Callable
from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI
from torchgen.context import native_function_manager
from torchgen.utils import T
# Like tools.api.context.with_native_function, but for
# NativeFunctionWithDifferentiabilityInfo.
def with_native_function_with_differentiability_info(
func: Callable[[NFWDI], T]
) -> Callable[[NFWDI], T]:
@functools.wraps(func)
def wrapper(f: NFWDI) -> T:
with native_function_manager(f.func):
return func(f)
return wrapper
# Like the above but with an additional dispatch key string argument
def with_native_function_with_differentiability_info_and_key(
func: Callable[[NFWDI, str], T]
) -> Callable[[NFWDI, str], T]:
@functools.wraps(func)
def wrapper(f: NFWDI, key: str) -> T:
with native_function_manager(f.func):
return func(f, key)
return wrapper
|