chore: Add ARG rule to ruff (#577)

* Add ARG rule to ruff
* Add [tool.ruff.lint.per-file-ignores] to ignore the ARG rule in `conftest.py` and `test_rst.py`
* Add `_` in front of unused parameter names that we have to keep
* Remove unused `jacobians` parameter of `_disunite_gradient`
* Make the `input_keys` parameter of `check_keys` positional-only, so that we can rename it in subclasses

Author: ValerianRey

pyproject.toml

@@ -133,6 +133,7 @@ select = [
     "W",    # pycodestyle Warning
     "I",    # isort
     "UP",   # pyupgrade
+    "ARG",  # flake8-unused-arguments
     "B",    # flake8-bugbear
     "C4",   # flake8-comprehensions
     "FIX",  # flake8-fixme
@@ -157,6 +158,10 @@ ignore = [
     "COM812",  # Trailing comma missing (conflicts with formatter, see https://github.com/astral-sh/ruff/issues/9216)
 ]
 
+[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.isort]
 combine-as-imports = true
 

src/torchjd/autogram/_module_hook_manager.py

@@ -101,7 +101,7 @@ def __init__(
 
     def __call__(
         self,
-        module: nn.Module,
+        _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,
+        _gramian_accumulation_phase: BoolRef,
+        _gramian_computer: GramianComputer,
+        _args: tuple[PyTree, ...],
+        _kwargs: dict[str, PyTree],
+        _gramian_accumulator: GramianAccumulator,
         *rg_tensors: Tensor,
     ) -> tuple[Tensor, ...]:
         return tuple(t.detach() for t in rg_tensors)

src/torchjd/autojac/_jac_to_grad.py

@@ -75,7 +75,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)
 
 
@@ -87,7 +87,6 @@ def _unite_jacobians(jacobians: list[Tensor]) -> Tensor:
 
 def _disunite_gradient(
     gradient_vector: Tensor,
-    jacobians: list[Tensor],
     tensors: list[TensorWithJac],
 ) -> list[Tensor]:
     gradient_vectors = gradient_vector.split([t.numel() for t in tensors])

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, _input_keys: 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, _input_keys: set[Tensor], /) -> set[Tensor]:
         return set()

src/torchjd/autojac/_transform/_base.py

@@ -43,7 +43,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.
@@ -78,7 +78,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
@@ -111,7 +111,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: AbstractSet[Tensor]):
         self.values = values
 
-    def __call__(self, input: TensorDict, /) -> TensorDict:
+    def __call__(self, _input: 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)}