Disallow incomplete defs in optimizers module (#5928)

Part of #5884.

### Description

Fully type annotate any functions with at least one type annotation in
module `optimizers`.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.

Signed-off-by: Felix Schnabel <f.schnabel@tum.de>

Author: f-schnabel

monai/optimizers/lr_finder.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import pickle
+import types
 import warnings
 from functools import partial
 from typing import TYPE_CHECKING, Any, Callable
@@ -187,7 +188,7 @@ def __init__(
         memory_cache: bool = True,
         cache_dir: str | None = None,
         amp: bool = False,
-        pickle_module=pickle,
+        pickle_module: types.ModuleType = pickle,
         pickle_protocol: int = DEFAULT_PROTOCOL,
         verbose: bool = True,
     ) -> None:
@@ -389,7 +390,9 @@ def _check_for_scheduler(self):
             if "initial_lr" in param_group:
                 raise RuntimeError("Optimizer already has a scheduler attached to it")
 
-    def _train_batch(self, train_iter, accumulation_steps: int, non_blocking_transfer: bool = True) -> float:
+    def _train_batch(
+        self, train_iter: TrainDataLoaderIter, accumulation_steps: int, non_blocking_transfer: bool = True
+    ) -> float:
         self.model.train()
         total_loss = 0
 
@@ -478,7 +481,14 @@ def get_steepest_gradient(self, skip_start: int = 0, skip_end: int = 0) -> tuple
             print("Failed to compute the gradients, there might not be enough points.")
             return None, None
 
-    def plot(self, skip_start: int = 0, skip_end: int = 0, log_lr: bool = True, ax=None, steepest_lr: bool = True):
+    def plot(
+        self,
+        skip_start: int = 0,
+        skip_end: int = 0,
+        log_lr: bool = True,
+        ax: Any | None = None,
+        steepest_lr: bool = True,
+    ) -> Any | None:
         """Plots the learning rate range test.
 
         Args:

monai/optimizers/novograd.py

@@ -11,11 +11,14 @@
 
 from __future__ import annotations
 
-from typing import Callable, Iterable
+from collections.abc import Callable, Iterable
+from typing import TypeVar
 
 import torch
 from torch.optim import Optimizer
 
+T = TypeVar("T")
+
 
 class Novograd(Optimizer):
     """
@@ -67,7 +70,7 @@ def __setstate__(self, state):
         for group in self.param_groups:
             group.setdefault("amsgrad", False)
 
-    def step(self, closure: Callable | None = None):
+    def step(self, closure: Callable[[], T] | None = None) -> T | None:
         """Performs a single optimization step.
 
         Arguments:

monai/optimizers/utils.py

@@ -26,7 +26,7 @@ def generate_param_groups(
     match_types: Sequence[str],
     lr_values: Sequence[float],
     include_others: bool = True,
-):
+) -> list[dict]:
     """
     Utility function to generate parameter groups with different LR values for optimizer.
     The output parameter groups have the same order as `layer_match` functions.

setup.cfg

@@ -222,7 +222,7 @@ check_untyped_defs = True
 # Warns about usage of untyped decorators.
 disallow_untyped_decorators = True
 
-[mypy-monai.visualize.*,monai.utils.*]
+[mypy-monai.visualize.*,monai.utils.*,monai.optimizers.*]
 disallow_incomplete_defs = True
 
 [coverage:run]