Remove plural types and aliases for native pytorch types (#677)

* First pass fix ensure_tuple

* Remove endline

* Adjust state docstring

* Adjust docstring

* Address comments

* Ooops typo

* Add tests

* Lint

* Fix types

* Fixed type errors

* Docstring styling

* Errr docstring links?

* Attempt at removing plural types pt. 1

* Remove plural types pt. 2

* Clean up types.py

Remove re-exported types from types.py. This should help with code readability, as now the underlying type is directly imported.

* Rearranged import

* Fix circular import

* Fix doctests

* Replaced StateDict type annotation with `Dict[str, Any]`

`Dict[str, Any]` is almost as short but clearer on what it is (a dictionary with string keys).

* Removed some more plural types

* Cleanup tensor_to_device

* Fixed typo

* Allowing generic sequences

* Fix LoggerDestination issues

Co-authored-by: ravi-mosaicml <[email protected]>

Author: Landanjs

STYLE_GUIDE.md

@@ -119,13 +119,13 @@ The following rules apply to public APIs:
 
 1. Parameters that could take a sequence of elements should also allow `None` or a singleton.
     This simplifies the user API by not having to construct a list (or tuple) to hold a single element
-    (or no element). For example, `Tensors = Union[Tensor, Tuple[Tensor, ...], List[Tensor]]`.
+    (or no element). For example, use `Optional[Union[torch.Tensor, Sequence[torch.Tensor]]`.
 
     The `composer.utils.ensure_tuple` helper method can convert a singleton, list, or tuple into a tuple.
     For example
 
     ```python
-    def foo(x: Optional[Tensors]) -> Tuple[Tensor, ...]:
+    def foo(x: Optional[Union[Tensor, Sequence[Tensor]]) -> Tuple[Tensor, ...]:
         return ensure_tuple(x)  # ensures that the result is always a (potentially empty) tuple of tensors
     ```
 

composer/algorithms/alibi/alibi.py

@@ -9,12 +9,12 @@
 import math
 from operator import attrgetter
 from types import MethodType, ModuleType
-from typing import Any, Callable, Optional, Tuple, Type, Union, cast
+from typing import Any, Callable, Optional, Sequence, Tuple, Type, Union, cast
 
 import torch
+from torch.optim import Optimizer
 
 from composer.core import Algorithm, Event, State
-from composer.core.types import Optimizers
 from composer.loggers import Logger
 from composer.utils import module_surgery
 
@@ -32,7 +32,7 @@ def apply_alibi(
     attr_to_replace: str,
     alibi_attention: Callable,
     mask_replacement_function: Optional[Callable[[torch.nn.Module, int], torch.nn.Module]] = None,
-    optimizers: Optional[Optimizers] = None,
+    optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None,
 ) -> None:
     """Removes position embeddings and replaces the attention function and attention mask
     according as per :class:`~composer.algorithms.alibi.alibi.Alibi`. Note that the
@@ -83,10 +83,10 @@ def apply_alibi(
             ``max_sequence_length``. For example,
             ``composer.algorithms.alibi._gpt2_alibi.enlarge_mask``. Default: ``None``,
             which means no modification of the model's default attention mask.
-        optimizers (Optimizers, optional): Existing optimizers bound to ``model.parameters()``.
-            All optimizers that have already been constructed with
-            ``model.parameters()`` must be specified here so they will optimize
-            the correct parameters. Default: ``None``.
+        optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional):
+            Existing optimizers bound to ``model.parameters()``. All optimizers that have already been
+            constructed with ``model.parameters()`` must be specified here so
+            they will optimize the correct parameters.
 
             If the optimizer(s) are constructed *after* calling this function,
             then it is safe to omit this parameter. These optimizers will see the correct

composer/algorithms/augmix/augmix.py

@@ -15,8 +15,7 @@
 
 from composer.algorithms.utils import augmentation_sets
 from composer.algorithms.utils.augmentation_common import map_pillow_function
-from composer.core.event import Event
-from composer.core.types import Algorithm, Event, State
+from composer.core import Algorithm, Event, State
 from composer.datasets.utils import add_vision_dataset_transform
 from composer.loggers import Logger
 

composer/algorithms/blurpool/blurpool.py

@@ -4,14 +4,14 @@
 
 import functools
 import logging
-from typing import Optional
+from typing import Optional, Sequence, Union
 
 import numpy as np
 import torch
+from torch.optim import Optimizer
 
 from composer.algorithms.blurpool.blurpool_layers import BlurConv2d, BlurMaxPool2d
 from composer.core import Algorithm, Event, State
-from composer.core.types import Optimizers
 from composer.loggers import Logger
 from composer.utils import module_surgery
 
@@ -22,7 +22,7 @@ def apply_blurpool(model: torch.nn.Module,
                    replace_convs: bool = True,
                    replace_maxpools: bool = True,
                    blur_first: bool = True,
-                   optimizers: Optional[Optimizers] = None) -> torch.nn.Module:
+                   optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None) -> torch.nn.Module:
     """Add anti-aliasing filters to the strided :class:`torch.nn.Conv2d` and/or :class:`torch.nn.MaxPool2d` modules
     within `model`.
 
@@ -41,8 +41,8 @@ def apply_blurpool(model: torch.nn.Module,
             overhead (though more closely matching
             `the paper <http://proceedings.mlr.press/v97/zhang19a.html>`_).
             See :class:`.BlurConv2d` for further discussion. Default: ``True``.
-        optimizers (Optimizers, optional):  Existing optimizers bound to
-            ``model.parameters()``. All optimizers that have already been
+        optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional):
+            Existing optimizers bound to ``model.parameters()``. All optimizers that have already been
             constructed with ``model.parameters()`` must be specified here so
             they will optimize the correct parameters.
 

composer/algorithms/channels_last/channels_last.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from composer.core.types import Algorithm, Event, State
+from composer.core import Algorithm, Event, State
 from composer.loggers import Logger
 
 log = logging.getLogger(__name__)

composer/algorithms/colout/colout.py

@@ -11,11 +11,11 @@
 
 import torch
 from PIL.Image import Image as PillowImage
+from torch import Tensor
 from torchvision.datasets import VisionDataset
 
 from composer.algorithms.utils.augmentation_common import image_as_type
 from composer.core import Algorithm, Event, State
-from composer.core.types import Tensor
 from composer.datasets.utils import add_vision_dataset_transform
 from composer.loggers import Logger
 

composer/algorithms/cutmix/cutmix.py

@@ -9,9 +9,10 @@
 
 import numpy as np
 import torch
+from torch import Tensor
 from torch.nn import functional as F
 
-from composer.core.types import Algorithm, Event, State, Tensor
+from composer.core import Algorithm, Event, State
 from composer.loggers import Logger
 from composer.models.loss import _check_for_index_targets
 

composer/algorithms/cutout/cutout.py

@@ -10,9 +10,10 @@
 import numpy as np
 import torch
 from PIL.Image import Image as PillowImage
+from torch import Tensor
 
 from composer.algorithms.utils.augmentation_common import image_as_type
-from composer.core.types import Algorithm, Event, State, Tensor
+from composer.core import Algorithm, Event, State
 from composer.loggers import Logger
 
 log = logging.getLogger(__name__)

composer/algorithms/factorize/factorize.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional, Type, Union, cast
+from typing import Optional, Sequence, Type, Union, cast
 
 import torch
+from torch.optim import Optimizer
 
 from composer.algorithms.factorize.factorize_modules import (FactorizedConv2d, FactorizedLinear,
                                                              factorizing_could_speedup)
 from composer.core import Algorithm, Event, State
-from composer.core.types import Optimizers
 from composer.loggers import Logger
 from composer.utils import module_surgery
 
@@ -27,7 +27,7 @@ def apply_factorization(model: torch.nn.Module,
                         latent_channels: Union[int, float] = 0.25,
                         min_features: int = 512,
                         latent_features: Union[int, float] = 0.25,
-                        optimizers: Optional[Optimizers] = None) -> torch.nn.Module:
+                        optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None) -> torch.nn.Module:
     """Replaces :class:`~torch.nn.Linear` and :class:`~torch.nn.Conv2d` modules and with
     :class:`~composer.algorithms.factorize.FactorizedLinear` and
     :class:`~composer.algorithms.factorize.FactorizedConv2d` modules.
@@ -62,8 +62,8 @@ def apply_factorization(model: torch.nn.Module,
             ``min(in_features, out_features)`` for each :class:`~torch.nn.Linear`
             module, and is converted to the equivalent integer value, with a
             minimum of 1. Default: ``0.25``.
-        optimizers (Optimizers, optional):  Existing optimizers bound to
-            ``model.parameters()``. All optimizers that have already been
+        optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional):
+            Existing optimizers bound to ``model.parameters()``. All optimizers that have already been
             constructed with ``model.parameters()`` must be specified here so
             they will optimize the correct parameters.
 
@@ -217,7 +217,7 @@ def _python_log_surgery_result(model: torch.nn.Module, new_class: Type[torch.nn.
 def _factorize_conv2d_modules(model: torch.nn.Module,
                               min_channels: int = 512,
                               latent_channels: Union[int, float] = 0.25,
-                              optimizers: Optional[Optimizers] = None):
+                              optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None):
     """Replaces :class:`~torch.nn.Conv2d` modules in ``model`` with
     :class:`~composer.algorithms.factorize.FactorizedConv2d` modules.
 
@@ -241,7 +241,7 @@ def _maybe_replace_conv2d(module: torch.nn.Module, module_index: int) -> Optiona
 def _factorize_linear_modules(model: torch.nn.Module,
                               min_features: int = 512,
                               latent_features: Union[int, float] = 0.25,
-                              optimizers: Optional[Optimizers] = None):
+                              optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None):
     """Replaces :class:`~torch.nn.Linear` modules in ``model`` with
     :class:`~composer.algorithms.factorize.FactorizedLinear` modules.
 

composer/algorithms/ghost_batchnorm/ghost_batchnorm.py

@@ -3,13 +3,13 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional
+from typing import Optional, Sequence, Union
 
 import numpy as np
 import torch
+from torch.optim import Optimizer
 
 from composer.core import Algorithm, Event, State
-from composer.core.types import Optimizers
 from composer.loggers import Logger
 from composer.utils import module_surgery
 
@@ -20,7 +20,7 @@
 
 def apply_ghost_batchnorm(model: torch.nn.Module,
                           ghost_batch_size: int = 32,
-                          optimizers: Optional[Optimizers] = None) -> torch.nn.Module:
+                          optimizers: Optional[Union[Optimizer, Sequence[Optimizer]]] = None) -> torch.nn.Module:
     """Replace batch normalization modules with ghost batch normalization modules.
 
     Ghost batch normalization modules split their input into chunks of
@@ -30,10 +30,10 @@ def apply_ghost_batchnorm(model: torch.nn.Module,
     Args:
         model (torch.nn.Module): the model to modify in-place
         ghost_batch_size (int, optional): size of sub-batches to normalize over. Default: ``32``.
-        optimizers (Optimizers, optional):  Existing optimizers bound to ``model.parameters()``.
-            All optimizers that have already been constructed with
-            ``model.parameters()`` must be specified here so they will optimize
-            the correct parameters.
+        optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional):
+            Existing optimizers bound to ``model.parameters()``. All optimizers that have already been
+            constructed with ``model.parameters()`` must be specified here so
+            they will optimize the correct parameters.
 
             If the optimizer(s) are constructed *after* calling this function,
             then it is safe to omit this parameter. These optimizers will see the correct