Add ARG (flake8 unused arguments). Few fixes from there like making check_keys have a positional only list of argument and remove jacobians from _disunite.

Author: PierreQuinton

pyproject.toml

@@ -139,6 +139,7 @@ select = [
     "FIX",  # flake8-fixme
     "TID",  # flake8-tidy-imports
     "SIM",  # flake8-simplify
+    "ARG",  # flake8-unused-arguments
     "PERF", # Perflint
     "FURB", # refurb
     "RUF",  # Ruff-specific rules

src/torchjd/autogram/_module_hook_manager.py

@@ -101,7 +101,7 @@ def __init__(
 
     def __call__(
         self,
-        module: nn.Module,
+        _: nn.Module,
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
         outputs: PyTree,
@@ -157,11 +157,11 @@ class AutogramNode(torch.autograd.Function):
 
     @staticmethod
     def forward(
-        gramian_accumulation_phase: BoolRef,
-        gramian_computer: GramianComputer,
-        args: tuple[PyTree, ...],
-        kwargs: dict[str, PyTree],
-        gramian_accumulator: GramianAccumulator,
+        _: BoolRef,
+        __: GramianComputer,
+        ___: tuple[PyTree, ...],
+        ____: dict[str, PyTree],
+        _____: GramianAccumulator,
         *rg_tensors: Tensor,
     ) -> tuple[Tensor, ...]:
         return tuple(t.detach() for t in rg_tensors)

src/torchjd/autojac/_jac_to_grad.py

@@ -73,7 +73,7 @@ def jac_to_grad(
 
     jacobian_matrix = _unite_jacobians(jacobians)
     gradient_vector = aggregator(jacobian_matrix)
-    gradients = _disunite_gradient(gradient_vector, jacobians, tensors_)
+    gradients = _disunite_gradient(gradient_vector, tensors_)
     accumulate_grads(tensors_, gradients)
 
 
@@ -83,9 +83,7 @@ def _unite_jacobians(jacobians: list[Tensor]) -> Tensor:
     return jacobian_matrix
 
 
-def _disunite_gradient(
-    gradient_vector: Tensor, jacobians: list[Tensor], tensors: list[TensorWithJac]
-) -> list[Tensor]:
+def _disunite_gradient(gradient_vector: Tensor, tensors: list[TensorWithJac]) -> list[Tensor]:
     gradient_vectors = gradient_vector.split([t.numel() for t in tensors])
     gradients = [g.view(t.shape) for g, t in zip(gradient_vectors, tensors, strict=True)]
     return gradients

src/torchjd/autojac/_transform/_accumulate.py

@@ -18,7 +18,7 @@ def __call__(self, gradients: TensorDict, /) -> TensorDict:
         accumulate_grads(gradients.keys(), gradients.values())
         return {}
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, _: set[Tensor], /) -> set[Tensor]:
         return set()
 
 
@@ -35,5 +35,5 @@ def __call__(self, jacobians: TensorDict, /) -> TensorDict:
         accumulate_jacs(jacobians.keys(), jacobians.values())
         return {}
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, _: set[Tensor], /) -> set[Tensor]:
         return set()

src/torchjd/autojac/_transform/_base.py

@@ -45,7 +45,7 @@ def __call__(self, input: TensorDict, /) -> TensorDict:
         """Applies the transform to the input."""
 
     @abstractmethod
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         """
         Checks that the provided input_keys satisfy the transform's requirements and returns the
         corresponding output keys for recursion.
@@ -80,7 +80,7 @@ def __call__(self, input: TensorDict, /) -> TensorDict:
         intermediate = self.inner(input)
         return self.outer(intermediate)
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         intermediate_keys = self.inner.check_keys(input_keys)
         output_keys = self.outer.check_keys(intermediate_keys)
         return output_keys
@@ -113,7 +113,7 @@ def __call__(self, tensor_dict: TensorDict, /) -> TensorDict:
             union |= transform(tensor_dict)
         return union
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         output_keys_list = [key for t in self.transforms for key in t.check_keys(input_keys)]
         output_keys = set(output_keys_list)
 

src/torchjd/autojac/_transform/_diagonalize.py

@@ -69,7 +69,7 @@ def __call__(self, tensors: TensorDict, /) -> TensorDict:
         }
         return diagonalized_tensors
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         if not set(self.key_order) == input_keys:
             raise RequirementError(
                 f"The input_keys must match the key_order. Found input_keys {input_keys} and"

src/torchjd/autojac/_transform/_differentiate.py

@@ -55,7 +55,7 @@ def _differentiate(self, tensor_outputs: Sequence[Tensor], /) -> tuple[Tensor, .
         tensor_outputs should be.
         """
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         outputs = set(self.outputs)
         if not outputs == input_keys:
             raise RequirementError(

src/torchjd/autojac/_transform/_init.py

@@ -16,10 +16,10 @@ class Init(Transform):
     def __init__(self, values: Set[Tensor]):
         self.values = values
 
-    def __call__(self, input: TensorDict, /) -> TensorDict:
+    def __call__(self, _: TensorDict, /) -> TensorDict:
         return {value: torch.ones_like(value) for value in self.values}
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         if not input_keys == set():
             raise RequirementError(
                 f"The input_keys should be the empty set. Found input_keys {input_keys}."

src/torchjd/autojac/_transform/_select.py

@@ -19,7 +19,7 @@ def __call__(self, tensor_dict: TensorDict, /) -> TensorDict:
         output = {key: tensor_dict[key] for key in self.keys}
         return type(tensor_dict)(output)
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         keys = set(self.keys)
         if not keys.issubset(input_keys):
             raise RequirementError(

src/torchjd/autojac/_transform/_stack.py

@@ -28,7 +28,7 @@ def __call__(self, input: TensorDict, /) -> TensorDict:
         result = _stack(results)
         return result
 
-    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+    def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
         return {key for transform in self.transforms for key in transform.check_keys(input_keys)}