Add PYI

Author: PierreQuinton

pyproject.toml

@@ -141,6 +141,7 @@ select = [
     "SIM",  # flake8-simplify
     "ARG",  # flake8-unused-arguments
     "RET",  # flake8-return
+    "PYI",  # flake8-pyi
     "PERF", # Perflint
     "FURB", # refurb
     "RUF",  # Ruff-specific rules

src/torchjd/aggregation/_weighting_bases.py

@@ -13,7 +13,7 @@
 _FnOutputT = TypeVar("_FnOutputT", bound=Tensor)
 
 
-class Weighting(Generic[_T], nn.Module, ABC):
+class Weighting(nn.Module, ABC, Generic[_T]):
     r"""
     Abstract base class for all weighting methods. It has the role of extracting a vector of weights
     of dimension :math:`m` from some statistic of a matrix of dimension :math:`m \times n`,

src/torchjd/autojac/_transform/_init.py

@@ -1,4 +1,4 @@
-from collections.abc import Set
+from collections.abc import Set as AbstractSet
 
 import torch
 from torch import Tensor
@@ -13,7 +13,7 @@ class Init(Transform):
     :param values: Tensors for which Gradients must be returned.
     """
 
-    def __init__(self, values: Set[Tensor]):
+    def __init__(self, values: AbstractSet[Tensor]):
         self.values = values
 
     def __call__(self, _: TensorDict, /) -> TensorDict:

src/torchjd/autojac/_transform/_select.py

@@ -1,4 +1,4 @@
-from collections.abc import Set
+from collections.abc import Set as AbstractSet
 
 from torch import Tensor
 
@@ -12,7 +12,7 @@ class Select(Transform):
     :param keys: The keys that should be included in the returned subset.
     """
 
-    def __init__(self, keys: Set[Tensor]):
+    def __init__(self, keys: AbstractSet[Tensor]):
         self.keys = keys
 
     def __call__(self, tensor_dict: TensorDict, /) -> TensorDict: