Fix type annotations for torch.sparse, enable in CI (#43108)

rgommers · facebook-github-bot · commit 864f0cfb2d55 · 2020-08-17T14:40:11.000-07:00
Summary: Closes gh-42982 Pull Request resolved: #43108 Reviewed By: malfet Differential Revision: D23167560 Pulled By: ezyang fbshipit-source-id: 0d660ca686ada2347bf440c6349551d1539f99ef
diff --git a/mypy.ini b/mypy.ini
@@ -93,9 +93,6 @@ ignore_errors = True
 [mypy-torch.jit.*]
 ignore_errors = True
 
-[mypy-torch.sparse]
-ignore_errors = True
-
 [mypy-torch.tensor]
 ignore_errors = True
 
diff --git a/torch/sparse/__init__.py b/torch/sparse/__init__.py
@@ -1,17 +1,18 @@
 # The Tensor classes are added to this module by python_tensor.cpp
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List, Union
 
 import torch
 from torch import Tensor
 
 # A workaround to support both TorchScript and MyPy:
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from torch import dtype as DType
+    from torch.types import _dtype as DType
+    DimOrDims = Optional[Union[int, Tuple[int], List[int]]]
 else:
+    # The JIT doesn't understand Union, nor torch.dtype here
     DType = int
-# TODO: replace the above with
-# from torch.types import _dtype as DType
+    DimOrDims = Optional[Tuple[int]]
 
 
 __all__ = [
@@ -23,8 +24,8 @@
 ]
 
 
-def addmm(mat, mat1, mat2, beta=1, alpha=1):
-    # type: (Tensor, Tensor, Tensor, float, float) -> Tensor
+def addmm(mat: Tensor, mat1: Tensor, mat2: Tensor,
+          beta: float = 1., alpha: float = 1.) -> Tensor:
     r"""
     This function does exact same thing as :func:`torch.addmm` in the forward,
     except that it supports backward for sparse matrix :attr:`mat1`. :attr:`mat1`
@@ -41,7 +42,7 @@ def addmm(mat, mat1, mat2, beta=1, alpha=1):
     return torch._sparse_addmm(mat, mat1, mat2, beta=beta, alpha=alpha)
 
 
-def mm(mat1, mat2):
+def mm(mat1: Tensor, mat2: Tensor) -> Tensor:
     r"""
     Performs a matrix multiplication of the sparse matrix :attr:`mat1`
     and dense matrix :attr:`mat2`. Similar to :func:`torch.mm`, If :attr:`mat1` is a
@@ -83,8 +84,8 @@ def mm(mat1, mat2):
     return torch._sparse_mm(mat1, mat2)
 
 
-def sum(input, dim=None, dtype=None):
-    # type: (Tensor, Optional[Tuple[int]], Optional[int]) -> Tensor
+def sum(input: Tensor, dim: DimOrDims = None,
+        dtype: Optional[DType] = None) -> Tensor:
     r"""
     Returns the sum of each row of SparseTensor :attr:`input` in the given
     dimensions :attr:`dim`. If :attr:`dim` is a list of dimensions,
