Fix tests and add some mandatory exceptions

Author: PierreQuinton

pyproject.toml

@@ -151,8 +151,9 @@ select = [
 ]
 
 ignore = [
-    "E501", # line-too-long (handled by the formatter)
-    "E402", # module-import-not-at-top-of-file
+    "E501",   # line-too-long (handled by the formatter)
+    "E402",   # module-import-not-at-top-of-file
+    "COM812",
 ]
 
 [tool.ruff.lint.isort]

src/torchjd/aggregation/_nash_mtl.py

@@ -189,7 +189,7 @@ def _init_optim_problem(self) -> None:
 
         G_alpha = self.G_param @ self.alpha_param
         constraint = [
-            -cp.log(self.a * self.normalization_factor_param) - cp.log(G_a) <= 0
+            -cp.log(a * self.normalization_factor_param) - cp.log(G_a) <= 0
             for a, G_a in zip(self.alpha_param, G_alpha, strict=True)
         ]
         obj = cp.Minimize(cp.sum(G_alpha) + self.phi_alpha / self.normalization_factor_param)

src/torchjd/autogram/_engine.py

@@ -288,7 +288,9 @@ def compute_gramian(self, output: Tensor) -> Tensor:
         self._module_hook_manager.gramian_accumulation_phase.value = True
 
         try:
-            square_gramian = self._compute_square_gramian(reshaped_output, has_non_batch_dim)
+            square_gramian = self._compute_square_gramian(
+                reshaped_output, has_non_batch_dim=has_non_batch_dim
+            )
         finally:
             # Reset everything that has a state, even if the previous call raised an exception
             self._module_hook_manager.gramian_accumulation_phase.value = False

src/torchjd/autojac/_backward.py

@@ -100,7 +100,7 @@ def _create_transform(
     diag = Diagonalize(tensors)
 
     # Transform that computes the required Jacobians.
-    jac = Jac(tensors, inputs, parallel_chunk_size, retain_graph)
+    jac = Jac(tensors, inputs, chunk_size=parallel_chunk_size, retain_graph=retain_graph)
 
     # Transform that accumulates the result in the .jac field of the inputs.
     accumulate = AccumulateJac()

src/torchjd/autojac/_jac.py

@@ -148,6 +148,6 @@ def _create_transform(
     diag = Diagonalize(outputs)
 
     # Transform that computes the required Jacobians.
-    jac = Jac(outputs, inputs, parallel_chunk_size, retain_graph)
+    jac = Jac(outputs, inputs, chunk_size=parallel_chunk_size, retain_graph=retain_graph)
 
     return jac << diag << init

src/torchjd/autojac/_mtl_backward.py

@@ -132,7 +132,7 @@ def _create_transform(
             features,
             task_params,
             OrderedSet([loss]),
-            retain_graph,
+            retain_graph=retain_graph,
         )
         for task_params, loss in zip(tasks_params, losses, strict=True)
     ]
@@ -142,7 +142,7 @@ def _create_transform(
     stack = Stack(task_transforms)
 
     # Transform that computes the Jacobians of the losses w.r.t. the shared parameters.
-    jac = Jac(features, shared_params, parallel_chunk_size, retain_graph)
+    jac = Jac(features, shared_params, chunk_size=parallel_chunk_size, retain_graph=retain_graph)
 
     # Transform that accumulates the result in the .jac field of the shared parameters.
     accumulate = AccumulateJac()
@@ -165,7 +165,7 @@ def _create_task_transform(
 
     # Transform that computes the gradients of the loss w.r.t. the task-specific parameters and
     # the features.
-    grad = Grad(loss, to_differentiate, retain_graph)
+    grad = Grad(loss, to_differentiate, retain_graph=retain_graph)
 
     # Transform that accumulates the gradients w.r.t. the task-specific parameters into their
     # .grad fields.

src/torchjd/autojac/_transform/_grad.py

@@ -33,7 +33,7 @@ def __init__(
         retain_graph: bool = False,
         create_graph: bool = False,
     ):
-        super().__init__(outputs, inputs, retain_graph, create_graph)
+        super().__init__(outputs, inputs, retain_graph=retain_graph, create_graph=create_graph)
 
     def _differentiate(self, grad_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]:
         """
@@ -54,4 +54,4 @@ def _differentiate(self, grad_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...
         if len(self.outputs) == 0:
             return tuple(torch.zeros_like(input) for input in self.inputs)
 
-        return self._get_vjp(grad_outputs, self.retain_graph)
+        return self._get_vjp(grad_outputs, retain_graph=self.retain_graph)

src/torchjd/autojac/_transform/_jac.py

@@ -40,7 +40,7 @@ def __init__(
         retain_graph: bool = False,
         create_graph: bool = False,
     ):
-        super().__init__(outputs, inputs, retain_graph, create_graph)
+        super().__init__(outputs, inputs, retain_graph=retain_graph, create_graph=create_graph)
         self.chunk_size = chunk_size
 
     def _differentiate(self, jac_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]:

tests/conftest.py

@@ -49,7 +49,7 @@ def pytest_collection_modifyitems(config, items):
             item.add_marker(xfail_cuda)
 
 
-def pytest_make_parametrize_id(_, val, __):
+def pytest_make_parametrize_id(config, val, argname):  # noqa: ARG001
     MAX_SIZE = 40
     optional_string = None  # Returning None means using pytest's way of making the string
 

tests/unit/autojac/test_backward.py

@@ -48,7 +48,7 @@ def test_jac_is_populated():
 @mark.parametrize("chunk_size", [1, 2, None])
 def test_value_is_correct(
     shape: tuple[int, int],
-    *manually_specify_inputs: bool,
+    manually_specify_inputs: bool,  # noqa: FBT001
     chunk_size: int | None,
 ):
     """