typing: Add missing annotations (#584)

* Add ANN ruff rule

* Add noqa comment in test_lightning_integration

* Add annotations in static_plotter and rename a variable
* Add annotations in conftest.py
* Add annotations in run_profiler.py

* Add -> None to __init__, __init_subclass__, main, and to test functions
* Add -> str to __str__
* Add -> "MemoryFrame" to MemoryFrame.from_event

* Add annotations for args and kwargs (Any)
* Add annotations for ctx
* Add annotations for obj
* Add annotations for chunk_size

* Add annotations to _make_tensors
* Add annotations to CloneParams
* Add annotations to InterModuleParamReuse.forward
* Add annotations to time_call
* Add annotations to update_gradient_coordinate

Author: ValerianRey

docs/source/conf.py

@@ -100,7 +100,7 @@ def linkcode_resolve(domain: str, info: dict[str, str]) -> str | None:
     return link
 
 
-def _get_obj(_info: dict[str, str]):
+def _get_obj(_info: dict[str, str]) -> object:
     module_name = _info["module"]
     full_name = _info["fullname"]
     sub_module = sys.modules.get(module_name)
@@ -112,7 +112,7 @@ def _get_obj(_info: dict[str, str]):
     return obj
 
 
-def _get_file_name(obj) -> str | None:
+def _get_file_name(obj: object) -> str | None:
     try:
         file_name = inspect.getsourcefile(obj)
         file_name = os.path.relpath(file_name, start=_PATH_ROOT)
@@ -121,7 +121,7 @@ def _get_file_name(obj) -> str | None:
     return file_name
 
 
-def _get_line_str(obj) -> str:
+def _get_line_str(obj: object) -> str:
     source, start = inspect.getsourcelines(obj)
     end = start + len(source) - 1
     line_str = f"#L{start}-L{end}"

pyproject.toml

@@ -134,6 +134,7 @@ select = [
     "I",    # isort
     "UP",   # pyupgrade
     "ARG",  # flake8-unused-arguments
+    "ANN",  # flake-8-annotations
     "B",    # flake8-bugbear
     "C4",   # flake8-comprehensions
     "FIX",  # flake8-fixme
@@ -156,12 +157,16 @@ ignore = [
     "RUF012",  # Mutable default value for class attribute (a bit tedious to fix)
     "RET504",  # Unnecessary assignment return statement
     "COM812",  # Trailing comma missing (conflicts with formatter, see https://github.com/astral-sh/ruff/issues/9216)
+    "ANN401",  # Prevent annotating as Any (we rarely do that, and when we do it's hard to find an alternative)
 ]
 
 [tool.ruff.lint.per-file-ignores]
 "**/conftest.py" = ["ARG"]  # Can't change argument names in the functions pytest expects
 "tests/doc/test_rst.py" = ["ARG"]  # For the lightning example
 
+[tool.ruff.lint.flake8-annotations]
+suppress-dummy-args = true
+
 [tool.ruff.lint.isort]
 combine-as-imports = true
 

src/torchjd/aggregation/_aggregator_bases.py

@@ -13,7 +13,7 @@ class Aggregator(nn.Module, ABC):
     :math:`m \times n` into row vectors of dimension :math:`n`.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     @staticmethod
@@ -48,7 +48,7 @@ class WeightedAggregator(Aggregator):
     :param weighting: The object responsible for extracting the vector of weights from the matrix.
     """
 
-    def __init__(self, weighting: Weighting[Matrix]):
+    def __init__(self, weighting: Weighting[Matrix]) -> None:
         super().__init__()
         self.weighting = weighting
 
@@ -77,6 +77,6 @@ class GramianWeightedAggregator(WeightedAggregator):
         gramian.
     """
 
-    def __init__(self, gramian_weighting: Weighting[PSDMatrix]):
+    def __init__(self, gramian_weighting: Weighting[PSDMatrix]) -> None:
         super().__init__(gramian_weighting << compute_gramian)
         self.gramian_weighting = gramian_weighting

src/torchjd/aggregation/_aligned_mtl.py

@@ -61,7 +61,7 @@ def __init__(
         self,
         pref_vector: Tensor | None = None,
         scale_mode: SUPPORTED_SCALE_MODE = "min",
-    ):
+    ) -> None:
         self._pref_vector = pref_vector
         self._scale_mode: SUPPORTED_SCALE_MODE = scale_mode
         super().__init__(AlignedMTLWeighting(pref_vector, scale_mode=scale_mode))
@@ -92,7 +92,7 @@ def __init__(
         self,
         pref_vector: Tensor | None = None,
         scale_mode: SUPPORTED_SCALE_MODE = "min",
-    ):
+    ) -> None:
         super().__init__()
         self._pref_vector = pref_vector
         self._scale_mode: SUPPORTED_SCALE_MODE = scale_mode

src/torchjd/aggregation/_cagrad.py

@@ -34,7 +34,7 @@ class CAGrad(GramianWeightedAggregator):
         To install it, use ``pip install "torchjd[cagrad]"``.
     """
 
-    def __init__(self, c: float, norm_eps: float = 0.0001):
+    def __init__(self, c: float, norm_eps: float = 0.0001) -> None:
         super().__init__(CAGradWeighting(c=c, norm_eps=norm_eps))
         self._c = c
         self._norm_eps = norm_eps
@@ -67,7 +67,7 @@ class CAGradWeighting(Weighting[PSDMatrix]):
         function.
     """
 
-    def __init__(self, c: float, norm_eps: float = 0.0001):
+    def __init__(self, c: float, norm_eps: float = 0.0001) -> None:
         super().__init__()
 
         if c < 0.0:

src/torchjd/aggregation/_config.py

@@ -50,7 +50,7 @@ class ConFIG(Aggregator):
         <https://github.com/tum-pbs/ConFIG/tree/main/conflictfree>`_.
     """
 
-    def __init__(self, pref_vector: Tensor | None = None):
+    def __init__(self, pref_vector: Tensor | None = None) -> None:
         super().__init__()
         self.weighting = pref_vector_to_weighting(pref_vector, default=SumWeighting())
         self._pref_vector = pref_vector

src/torchjd/aggregation/_constant.py

@@ -15,7 +15,7 @@ class Constant(WeightedAggregator):
     :param weights: The weights associated to the rows of the input matrices.
     """
 
-    def __init__(self, weights: Tensor):
+    def __init__(self, weights: Tensor) -> None:
         super().__init__(weighting=ConstantWeighting(weights=weights))
         self._weights = weights
 
@@ -35,7 +35,7 @@ class ConstantWeighting(Weighting[Matrix]):
     :param weights: The weights to return at each call.
     """
 
-    def __init__(self, weights: Tensor):
+    def __init__(self, weights: Tensor) -> None:
         if weights.dim() != 1:
             raise ValueError(
                 "Parameter `weights` should be a 1-dimensional tensor. Found `weights.shape = "

src/torchjd/aggregation/_dualproj.py

@@ -33,7 +33,7 @@ def __init__(
         norm_eps: float = 0.0001,
         reg_eps: float = 0.0001,
         solver: SUPPORTED_SOLVER = "quadprog",
-    ):
+    ) -> None:
         self._pref_vector = pref_vector
         self._norm_eps = norm_eps
         self._reg_eps = reg_eps
@@ -77,7 +77,7 @@ def __init__(
         norm_eps: float = 0.0001,
         reg_eps: float = 0.0001,
         solver: SUPPORTED_SOLVER = "quadprog",
-    ):
+    ) -> None:
         super().__init__()
         self._pref_vector = pref_vector
         self.weighting = pref_vector_to_weighting(pref_vector, default=MeanWeighting())

src/torchjd/aggregation/_flattening.py

@@ -20,7 +20,7 @@ class Flattening(GeneralizedWeighting):
     :param weighting: The weighting to apply to the Gramian matrix.
     """
 
-    def __init__(self, weighting: Weighting):
+    def __init__(self, weighting: Weighting) -> None:
         super().__init__()
         self.weighting = weighting
 

src/torchjd/aggregation/_graddrop.py

@@ -26,7 +26,7 @@ class GradDrop(Aggregator):
         through. Defaults to None, which means no leak.
     """
 
-    def __init__(self, f: Callable = _identity, leak: Tensor | None = None):
+    def __init__(self, f: Callable = _identity, leak: Tensor | None = None) -> None:
         if leak is not None and leak.dim() != 1:
             raise ValueError(
                 "Parameter `leak` should be a 1-dimensional tensor. Found `leak.shape = "