Improve type hints and docs in timm/models

CONTRIBUTING.md

@@ -49,7 +49,7 @@ This is YES:
    }
 ```
 
-When there is descrepancy in a given source file (there are many origins for various bits of code and not all have been updated to what I consider current goal), please follow the style in a given file.
+When there is discrepancy in a given source file (there are many origins for various bits of code and not all have been updated to what I consider current goal), please follow the style in a given file.
 
 In general, if you add new code, formatting it with black using the following options should result in a style that is compatible with the rest of the code base:
 

timm/layers/__init__.py

@@ -52,4 +52,5 @@
 from .std_conv import StdConv2d, StdConv2dSame, ScaledStdConv2d, ScaledStdConv2dSame
 from .test_time_pool import TestTimePoolHead, apply_test_time_pool
 from .trace_utils import _assert, _float_to_int
+from .typing import LayerType, PadType
 from .weight_init import trunc_normal_, trunc_normal_tf_, variance_scaling_, lecun_normal_

timm/layers/typing.py

@@ -0,0 +1,9 @@
+import functools
+import types
+from typing import Tuple, Union
+
+import torch.nn
+
+
+LayerType = Union[type, str, types.FunctionType, functools.partial, torch.nn.Module]
+PadType = Union[str, int, Tuple[int, int]]

timm/models/_efficientnet_builder.py

@@ -11,6 +11,7 @@
 import re
 from copy import deepcopy
 from functools import partial
+from typing import Any, Dict, List
 
 import torch.nn as nn
 
@@ -34,6 +35,8 @@
 BN_EPS_TF_DEFAULT = 1e-3
 _BN_ARGS_TF = dict(momentum=BN_MOMENTUM_TF_DEFAULT, eps=BN_EPS_TF_DEFAULT)
 
+BlockArgs = List[List[Dict[str, Any]]]
+
 
 def get_bn_args_tf():
     return _BN_ARGS_TF.copy()