Type Hints for Lightning Core (#946)

* first pass for LightningModule typehints

* fix return types

* add missing types

* add type annotations to grads.py

* add type annotations to hooks.py

* add type annotation to memory.py

* proper docstring quotation marks

* add type annotations to saving.py

* fix cyclic import problem

* fix cyclic import problem

* add missing whitespace

* finish type hints for load_from_ methods

* docs: prepare_data does not return anything

* fix auto types in docs

* revert typehint for trainer in hook

* remove unnecessary return docs

* some fixes for memory docs

* revert typing for args kwargs

* added all missing None return types

* remove unused import

* add more details to dict/list return types

* fix line too long

* optimize imports

* linted

* Revert "linted"

This reverts commit 85559611e84e312bce64f4e73b638d4999a8439e.

* remove whitespace

* update

* update

* update

* update

* update

* changelog

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: William Falcon <waf2107@columbia.edu>

pytorch_lightning/core/saving.py

@@ -2,11 +2,12 @@ import csv
 import logging as log
 import os
 from argparse import Namespace
+from typing import Union, Dict, Any
 
 
 class ModelIO(object):
 
-    def on_load_checkpoint(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         """
         Do something with the checkpoint
         Gives model a chance to load something before state_dict is restored
@@ -14,7 +15,7 @@ class ModelIO(object):
         :return:
         """
 
-    def on_save_checkpoint(self, checkpoint):
+    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         """
         Give the model a chance to add something to the checkpoint.
         state_dict is already there
@@ -23,20 +24,18 @@ class ModelIO(object):
     # -------------------------
     # OPTIONAL HOOKS
     # -------------------------
-    def on_hpc_save(self, checkpoint):
+    def on_hpc_save(self, checkpoint: Dict[str, Any]) -> None:
         """
         Hook to do whatever you need right before Slurm manager saves the model
-        :return:
         """
 
-    def on_hpc_load(self, checkpoint):
+    def on_hpc_load(self, checkpoint: Dict[str, Any]) -> None:
         """
         Hook to do whatever you need right before Slurm manager loads the model
-        :return:
         """
 
 
-def load_hparams_from_tags_csv(tags_csv) -> Namespace:
+def load_hparams_from_tags_csv(tags_csv: str) -> Namespace:
     if not os.path.isfile(tags_csv):
         log.warning(f'Missing Tags: {tags_csv}.')
         return Namespace()
@@ -48,7 +47,7 @@ def load_hparams_from_tags_csv(tags_csv) -> Namespace:
     return ns
 
 
-def convert(val):
+def convert(val: str) -> Union[int, float, bool, str]:
     constructors = [int, float, str]
 
     if isinstance(val, str):