add type annotations to torch.nn.modules.normalization

mypy.ini

@@ -80,9 +80,6 @@ ignore_errors = True
 [mypy-torch.nn.modules.module]
 ignore_errors = True
 
-[mypy-torch.nn.modules.normalization]
-ignore_errors = True
-
 [mypy-torch.nn.modules.padding]
 ignore_errors = True
 

torch/nn/functional.pyi.in

@@ -205,7 +205,7 @@ def instance_norm(input: Tensor, running_mean: Optional[Tensor] = ..., running_v
                   momentum: float = ..., eps: float = ...) -> Tensor: ...
 
 
-def layer_norm(input: Tensor, normalized_shape: List[int], weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ...,
+def layer_norm(input: Tensor, normalized_shape: Sequence[int], weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ...,
                eps: float = ...) -> Tensor: ...
 
 

torch/nn/modules/normalization.py

@@ -7,7 +7,7 @@
 from .. import init
 
 from torch import Tensor, Size
-from typing import Union, List
+from typing import Union, List, Tuple
 
 
 class LocalResponseNorm(Module):
@@ -141,20 +141,21 @@ class LayerNorm(Module):
         >>> output = m(input)
     """
     __constants__ = ['normalized_shape', 'eps', 'elementwise_affine']
-    normalized_shape: _shape_t
+    normalized_shape: Tuple[int, ...] 
     eps: float
     elementwise_affine: bool
 
     def __init__(self, normalized_shape: _shape_t, eps: float = 1e-5, elementwise_affine: bool = True) -> None:
         super(LayerNorm, self).__init__()
         if isinstance(normalized_shape, numbers.Integral):
-            normalized_shape = (normalized_shape,)
-        self.normalized_shape = tuple(normalized_shape)
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
         self.eps = eps
         self.elementwise_affine = elementwise_affine
         if self.elementwise_affine:
-            self.weight = Parameter(torch.Tensor(*normalized_shape))
-            self.bias = Parameter(torch.Tensor(*normalized_shape))
+            self.weight = Parameter(torch.Tensor(*self.normalized_shape))
+            self.bias = Parameter(torch.Tensor(*self.normalized_shape))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
@@ -169,7 +170,7 @@ def forward(self, input: Tensor) -> Tensor:
         return F.layer_norm(
             input, self.normalized_shape, self.weight, self.bias, self.eps)
 
-    def extra_repr(self) -> Tensor:
+    def extra_repr(self) -> str:
         return '{normalized_shape}, eps={eps}, ' \
             'elementwise_affine={elementwise_affine}'.format(**self.__dict__)