Just a temporary solution, which is just a little bit better than copying hints each time, inspired by @InSync and the decorator in related question:
from typing import TypeVar, Callable
from typing_extensions import ParamSpec # for 3.7 <= python < 3.10, import from typing for versions later
T = TypeVar('T')
P = ParamSpec('P')
def take_annotation_from(this: Callable[P, T]) -> Callable[[Callable], Callable[P, T]]:
def decorator(real_function: Callable) -> Callable[P, T]:
def new_function(*args: P.args, **kwargs: P.kwargs) -> T:
return real_function(*args, **kwargs)
new_function.__doc__ = this.__doc__
return new_function
return decorator
And use it as
from torch.nn import Module
class MyModule(Module):
def __init__(self, k: float):
...
def forward(self, ...) -> ...: # with type hints
"""docstring"""
...
@take_annotation_from(forward)
def __call__(self, *args, **kwds):
return Module.__call__(self, *args, **kwds)
And this solution may be proved if last three lines of the code above can be packed as something like macro, because it remains unchanged among different implementations of sub-nn.Module
s.