Change the type hint for nn.Module.__call__ to be friendly to overrides.

[ppeetteerrs]

🚀 The feature, motivation and pitch

Currently, nn.Module.__call__ has type hints defined using __call__ : Callable[..., Any] = _call_impl. However, this declared type makes it difficult for the user to override the type hint of nn.Module.__call__ using inferred types. For example:

from typing import Callable, TypeVar, cast

from torch import Tensor, nn

C = TypeVar("C", bound=Callable)


def proxy(f: C) -> C:
    return cast(C, lambda self, *x, **y: super(self.__class__, self).__call__(*x, **y))


class MyModule(nn.Module):
    def forward(self, x: Tensor, y: Tensor) -> Tensor:
        return x + y

    __call__ = proxy(forward)  # __call__ exists in nn.Module

    another_call = proxy(forward)  # another_call does not exist in nn.Module


x, y = Tensor(), Tensor()

model = MyModule()

z1 = model(x, y)  # Type hint of z1 is Any, same for model.__call__(x, y)

z2 = model.another_call(x, y)  # Type hint of z2 is Tensor

If you paste the above code in VSCode, you can see that the type for z1 is Any (wrong) while that for z2 is Tensor (correct).

Alternatives

Instead, if we remove the type annotations in nn.Module and do the following instead, the above example would work as expected.

class Module:
    __call__ = _call_impl
    def _call_impl(self, *input: Any, **kwargs: Any) -> Any:
        # actual implementation

Additional context

No response

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @ezyang @malfet @rgommers @xuzhao9 @gramster

[jonasrauber]

Actually, using mypy works just fine with your example: reveal_type(z1) and reveal_type(z2) both show Revealed type is "torch._tensor.Tensor". This might just be a limitation for VSCode's type checking.

[ppeetteerrs]

There might be some difference in the type inference mechanism? Just for reference I also discussed this issue in the pyright repo: microsoft/pyright#3249 (comment)

[mert-kurttutan]

For nn.Linear, reveal_type returns Any, instead of torch.Tensor. To replicate this, run the code below

pytorch=1.13
mypy=0.991

# main.py
import torch
from torch import nn


class LinWrapper(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.lin = nn.Linear(12,2)

    def forward(
        self, x: torch.Tensor
    ) -> torch.Tensor:
        out = self.lin(x)
        reveal_type(out)
        return out

Bash command to run

mypy main.py

Output:

main.py:15: note: Revealed type is "Any"

Note when tested with custom module (with type hints), it reveals to be torch.Tensor, just like @jonasrauber obtained.

[ezyang]

Can someone pr the proposed change so we can see how it cis

[mert-kurttutan]

One possible way to go is to copy the type signature of forward method if there is any type annotation of forward function. If there is not, continue using current default one Callable[..., Any] for type hinting.
For instance,

def get_type_signature(func):
    # if func has type annotation
    # return the annotation
    # otherwise, return Callable[..., Any]
    pass

forward_type_sign = get_type_signature(self.forward)
__call__ : forward_type_sign = _call_impl

But I am not sure how to do this in a reliable way.
Maybe one solution here: typing/issues/270
even though this does not do exactly what I described, but annotates the with type of another function.

[hchau630]

any updates to this? if not, is there any simple workaround?

[milmillin]

I have tried many different things. Here is the cleanest workaround I have:

_P = ParamSpec("_P")
_R = TypeVar("_R")

def _returns_nn_module_call(*args):
    return nn.Module.__call__

def patch_call(src_func: Callable[_P, _R]) -> Callable[..., Callable[_P, _R]]:
    return _returns_nn_module_call

class MyModule(nn.Module):
    def forward(self, x: Tensor, y: Tensor) -> Tensor:
        return x + y

    @patch_call(forward)
    def __call__(self): ...

Defining a function would override the inherited type annotation. Here, I use the decorator to replace the call signature with that of the forward function and replace __call__ with nn.Module.__call__. This won't add additional stack depth during runtime. For older Python versions without ParamSpec. I think you can use a TypeVar bound to a Callable as done by the OP. Also, the solution above assumes that you directly inherits from nn.Module. If you want to actually call super().__call__, you can do:

def patch_call(src_func: Callable[_P, _R]):
    def decorator(target_fn) -> Callable[_P, _R]:
        def inner(self, *args, **kwargs):
            return super(self.__class__, self).__call__(*args, **kwargs)
        return inner
    return decorator

with the cost of one additional stack depth at every call.