为何mypy无法推断torch.nn.functional.linear的返回类型？

问题场景 #

我在PyTorch代码里用mypy做类型检查时遇到错误，代码如下：

from torch import FloatTensor
from torch.nn.functional import linear

def foo(x: FloatTensor, y: FloatTensor) -> FloatTensor:
    return linear(x, y)

mypy输出的错误信息：

error: Incompatible return value type (got "Tensor", expected "FloatTensor")  [return-value]

疑问与解答 #

1. 为什么torch.nn.functional.linear的返回类型被识别为Tensor？

PyTorch官方给linear函数标注的返回类型是基类Tensor，而非具体的FloatTensor。静态类型检查器（比如mypy）只能依赖函数的类型注解做判断，没办法在编译阶段追踪张量的dtype、设备等动态属性变化——哪怕你传入的是FloatTensor，mypy也没法从linear的注解里推断出返回的是同类型张量，只能识别成注解里写的Tensor。

2. 解决这个mypy错误的最佳方式是什么？

推荐以下几种方案，按实用性排序：

• 使用类型断言（cast）：明确告诉mypy返回值的具体类型，需要先导入cast：

  from typing import cast
  from torch import FloatTensor
  from torch.nn.functional import linear
  
  def foo(x: FloatTensor, y: FloatTensor) -> FloatTensor:
      return cast(FloatTensor, linear(x, y))
  

  这种方式既保留了你对返回类型的严格要求，又能让mypy通过检查。
• 放宽返回类型为Tensor：如果业务逻辑不需要严格限定返回FloatTensor，直接把函数的返回类型改成Tensor，和linear的注解保持一致：

  from torch import Tensor, FloatTensor
  from torch.nn.functional import linear
  
  def foo(x: FloatTensor, y: FloatTensor) -> Tensor:
      return linear(x, y)
  

• 使用更完善的类型存根：部分第三方维护的PyTorch类型存根会细化这类函数的返回类型注解，安装对应的存根包可以提升类型检查的准确性。

内容的提问来源于stack exchange，提问作者Grayscale