Add RUFF

PierreQuinton · PierreQuinton · commit 20fc69777808 · 2026-02-05T12:18:19.000+01:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -141,6 +141,7 @@ select = [
     "SIM",  # flake8-simplify
     "PERF", # Perflint
     "FURB", # refurb
+    "RUF",  # Ruff-specific rules
 ]
 
 ignore = [
diff --git a/src/torchjd/_linalg/__init__.py b/src/torchjd/_linalg/__init__.py
@@ -2,13 +2,13 @@
 from ._matrix import Matrix, PSDMatrix, PSDTensor, is_matrix, is_psd_matrix, is_psd_tensor
 
 __all__ = [
-    "compute_gramian",
-    "normalize",
-    "regularize",
     "Matrix",
     "PSDMatrix",
     "PSDTensor",
+    "compute_gramian",
     "is_matrix",
     "is_psd_matrix",
     "is_psd_tensor",
+    "normalize",
+    "regularize",
 ]
diff --git a/src/torchjd/aggregation/__init__.py b/src/torchjd/aggregation/__init__.py
@@ -81,6 +81,8 @@
 from ._weighting_bases import GeneralizedWeighting, Weighting
 
 __all__ = [
+    "IMTLG",
+    "MGDA",
     "Aggregator",
     "AlignedMTL",
     "AlignedMTLWeighting",
@@ -92,14 +94,12 @@
     "Flattening",
     "GeneralizedWeighting",
     "GradDrop",
-    "IMTLG",
     "IMTLGWeighting",
     "Krum",
     "KrumWeighting",
+    "MGDAWeighting",
     "Mean",
     "MeanWeighting",
-    "MGDA",
-    "MGDAWeighting",
     "PCGrad",
     "PCGradWeighting",
     "Random",
diff --git a/src/torchjd/aggregation/_aligned_mtl.py b/src/torchjd/aggregation/_aligned_mtl.py
@@ -68,8 +68,8 @@ def __init__(
 
     def __repr__(self) -> str:
         return (
-            f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, "
-            f"scale_mode={repr(self._scale_mode)})"
+            f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, "
+            f"scale_mode={self._scale_mode!r})"
         )
 
     def __str__(self) -> str:
diff --git a/src/torchjd/aggregation/_config.py b/src/torchjd/aggregation/_config.py
@@ -70,7 +70,7 @@ def forward(self, matrix: Matrix) -> Tensor:
         return length * unit_target_vector
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)})"
+        return f"{self.__class__.__name__}(pref_vector={self._pref_vector!r})"
 
     def __str__(self) -> str:
         return f"ConFIG{pref_vector_to_str_suffix(self._pref_vector)}"
diff --git a/src/torchjd/aggregation/_constant.py b/src/torchjd/aggregation/_constant.py
@@ -20,7 +20,7 @@ def __init__(self, weights: Tensor):
         self._weights = weights
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(weights={repr(self._weights)})"
+        return f"{self.__class__.__name__}(weights={self._weights!r})"
 
     def __str__(self) -> str:
         weights_str = vector_to_str(self._weights)
diff --git a/src/torchjd/aggregation/_dualproj.py b/src/torchjd/aggregation/_dualproj.py
@@ -48,8 +48,8 @@ def __init__(
 
     def __repr__(self) -> str:
         return (
-            f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="
-            f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={repr(self._solver)})"
+            f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, norm_eps="
+            f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={self._solver!r})"
         )
 
     def __str__(self) -> str:
diff --git a/src/torchjd/aggregation/_graddrop.py b/src/torchjd/aggregation/_graddrop.py
@@ -68,7 +68,7 @@ def _check_matrix_has_enough_rows(self, matrix: Tensor) -> None:
             )
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(f={repr(self.f)}, leak={repr(self.leak)})"
+        return f"{self.__class__.__name__}(f={self.f!r}, leak={self.leak!r})"
 
     def __str__(self) -> str:
         if self.leak is None:
diff --git a/src/torchjd/aggregation/_upgrad.py b/src/torchjd/aggregation/_upgrad.py
@@ -49,8 +49,8 @@ def __init__(
 
     def __repr__(self) -> str:
         return (
-            f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="
-            f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={repr(self._solver)})"
+            f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, norm_eps="
+            f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={self._solver!r})"
         )
 
     def __str__(self) -> str:
diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py
@@ -278,7 +278,7 @@ def compute_gramian(self, output: Tensor) -> Tensor:
             target_shape = []
 
         if has_non_batch_dim:
-            target_shape = [-1] + target_shape
+            target_shape = [-1, *target_shape]
 
         reshaped_output = ordered_output.reshape(target_shape)
         # There are four different cases for the shape of reshaped_output:
diff --git a/src/torchjd/autojac/_transform/_diagonalize.py b/src/torchjd/autojac/_transform/_diagonalize.py
@@ -64,7 +64,7 @@ def __call__(self, tensors: TensorDict, /) -> TensorDict:
         flattened_considered_values = [tensors[key].reshape([-1]) for key in self.key_order]
         diagonal_matrix = torch.cat(flattened_considered_values).diag()
         diagonalized_tensors = {
-            key: diagonal_matrix[:, begin:end].reshape((-1,) + key.shape)
+            key: diagonal_matrix[:, begin:end].reshape((-1, *key.shape))
             for (begin, end), key in zip(self.indices, self.key_order, strict=True)
         }
         return diagonalized_tensors
diff --git a/src/torchjd/autojac/_transform/_jac.py b/src/torchjd/autojac/_transform/_jac.py
@@ -62,7 +62,7 @@ def _differentiate(self, jac_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]
         if len(self.outputs) == 0:
             return tuple(
                 [
-                    torch.empty((0,) + input.shape, device=input.device, dtype=input.dtype)
+                    torch.empty((0, *input.shape), device=input.device, dtype=input.dtype)
                     for input in self.inputs
                 ]
             )
diff --git a/tests/unit/autogram/test_engine.py b/tests/unit/autogram/test_engine.py
@@ -536,11 +536,11 @@ def test_batched_non_batched_equivalence(shape: list[int], batch_dim: int):
     input = randn_([batch_size, input_size])
 
     engine1 = Engine(model1, batch_dim=batch_dim)
-    output1 = model1(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
+    output1 = model1(input).reshape([batch_size, *non_batched_shape]).movedim(0, batch_dim)
     gramian1 = engine1.compute_gramian(output1)
 
     engine2 = Engine(model2, batch_dim=None)
-    output2 = model2(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
+    output2 = model2(input).reshape([batch_size, *non_batched_shape]).movedim(0, batch_dim)
     gramian2 = engine2.compute_gramian(output2)
 
     assert_close(gramian1, gramian2)
diff --git a/tests/unit/autogram/test_gramian_utils.py b/tests/unit/autogram/test_gramian_utils.py
@@ -29,8 +29,8 @@
 def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]):
     """Tests that reshape_gramian is such that compute_gramian is equivariant to a reshape."""
 
-    original_matrix = randn_(original_shape + [2])
-    target_matrix = original_matrix.reshape(target_shape + [2])
+    original_matrix = randn_([*original_shape, 2])
+    target_matrix = original_matrix.reshape([*target_shape, 2])
 
     original_gramian = compute_gramian(original_matrix, 1)
     target_gramian = compute_gramian(target_matrix, 1)
@@ -56,7 +56,7 @@ def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]
     ],
 )
 def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
-    matrix = randn_(original_shape + [2])
+    matrix = randn_([*original_shape, 2])
     gramian = compute_gramian(matrix, 1)
     reshaped_gramian = reshape(gramian, target_shape)
     assert_is_psd_tensor(reshaped_gramian, atol=1e-04, rtol=0.0)
@@ -73,7 +73,7 @@ def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
     ],
 )
 def test_flatten_yields_matrix(shape: list[int]):
-    matrix = randn_(shape + [2])
+    matrix = randn_([*shape, 2])
     gramian = compute_gramian(matrix, 1)
     flattened_gramian = flatten(gramian)
     assert is_psd_matrix(flattened_gramian)
@@ -90,7 +90,7 @@ def test_flatten_yields_matrix(shape: list[int]):
     ],
 )
 def test_flatten_yields_psd(shape: list[int]):
-    matrix = randn_(shape + [2])
+    matrix = randn_([*shape, 2])
     gramian = compute_gramian(matrix, 1)
     flattened_gramian = flatten(gramian)
     assert_is_psd_matrix(flattened_gramian, atol=1e-04, rtol=0.0)
@@ -117,7 +117,7 @@ def test_flatten_yields_psd(shape: list[int]):
 def test_movedim_equivariance(shape: list[int], source: list[int], destination: list[int]):
     """Tests that movedim_gramian is such that compute_gramian is equivariant to a movedim."""
 
-    original_matrix = randn_(shape + [2])
+    original_matrix = randn_([*shape, 2])
     target_matrix = original_matrix.movedim(source, destination)
 
     original_gramian = compute_gramian(original_matrix, 1)
@@ -147,7 +147,7 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
     ],
 )
 def test_movedim_yields_psd(shape: list[int], source: list[int], destination: list[int]):
-    matrix = randn_(shape + [2])
+    matrix = randn_([*shape, 2])
     gramian = compute_gramian(matrix, 1)
     moveddim_gramian = movedim(gramian, source, destination)
     assert_is_psd_tensor(moveddim_gramian)
diff --git a/tests/unit/autojac/_transform/test_accumulate.py b/tests/unit/autojac/_transform/test_accumulate.py
@@ -97,7 +97,7 @@ def test_single_jac_accumulation():
 
     shapes = [[], [1], [2, 3]]
     keys = [zeros_(shape, requires_grad=True) for shape in shapes]
-    values = [ones_([4] + shape) for shape in shapes]
+    values = [ones_([4, *shape]) for shape in shapes]
     input = dict(zip(keys, values, strict=True))
 
     accumulate = AccumulateJac()
@@ -118,7 +118,7 @@ def test_multiple_jac_accumulations(iterations: int):
 
     shapes = [[], [1], [2, 3]]
     keys = [zeros_(shape, requires_grad=True) for shape in shapes]
-    values = [ones_([4] + shape) for shape in shapes]
+    values = [ones_([4, *shape]) for shape in shapes]
 
     accumulate = AccumulateJac()
 
diff --git a/tests/unit/autojac/_transform/test_interactions.py b/tests/unit/autojac/_transform/test_interactions.py
@@ -235,9 +235,7 @@ def test_equivalence_jac_grads():
     grad_2_A, grad_2_b, grad_2_c = grad_dict_2[A], grad_dict_2[b], grad_dict_2[c]
 
     n_outputs = len(outputs)
-    batched_grad_outputs = [
-        zeros_((n_outputs,) + grad_output.shape) for grad_output in grad_outputs
-    ]
+    batched_grad_outputs = [zeros_((n_outputs, *grad_output.shape)) for grad_output in grad_outputs]
     for i, grad_output in enumerate(grad_outputs):
         batched_grad_outputs[i][i] = grad_output
 
diff --git a/tests/utils/architectures.py b/tests/utils/architectures.py
@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Generic, TypeVar
+from typing import ClassVar, Generic, TypeVar
 
 import torch
 import torchvision
@@ -131,7 +131,7 @@ class SingleInputPyTreeOutput(ShapedModule):
     """Module taking a single input and returning a complex PyTree of tensors as output."""
 
     INPUT_SHAPES = (50,)
-    OUTPUT_SHAPES = {
+    OUTPUT_SHAPES: ClassVar = {
         "first": ((50,), [(60,), (70,)]),
         "second": (80,),
         "third": ([((90,),)],),
@@ -156,7 +156,7 @@ def forward(self, input: Tensor) -> PyTree:
 class PyTreeInputSingleOutput(ShapedModule):
     """Module taking a complex PyTree of tensors as input and returning a single output."""
 
-    INPUT_SHAPES = {
+    INPUT_SHAPES: ClassVar = {
         "one": [((10,), [(20,), (30,)]), (12,)],
         "two": (14,),
     }
@@ -193,12 +193,12 @@ class PyTreeInputPyTreeOutput(ShapedModule):
     output.
     """
 
-    INPUT_SHAPES = {
+    INPUT_SHAPES: ClassVar = {
         "one": [((10,), [(20,), (30,)]), (12,)],
         "two": (14,),
     }
 
-    OUTPUT_SHAPES = {
+    OUTPUT_SHAPES: ClassVar = {
         "first": ((50,), [(60,), (70,)]),
         "second": (80,),
         "third": ([((90,),)],),
diff --git a/tests/utils/forward_backwards.py b/tests/utils/forward_backwards.py
@@ -198,7 +198,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         for module in self.model.modules():
             self._restore_original_params(module)
 
-        return False  # don’t suppress exceptions
+        return False  # don't suppress exceptions
 
     def _restore_original_params(self, module: nn.Module):
         original_params = self._module_to_original_params.pop(module, {})
diff --git a/tests/utils/tensors.py b/tests/utils/tensors.py
@@ -40,4 +40,4 @@ def _make_tensors(batch_size: int, tensor_shapes: PyTree) -> PyTree:
     def is_leaf(s):
         return isinstance(s, tuple) and all(isinstance(e, int) for e in s)
 
-    return tree_map(lambda s: randn_((batch_size,) + s), tensor_shapes, is_leaf=is_leaf)
+    return tree_map(lambda s: randn_((batch_size, *s)), tensor_shapes, is_leaf=is_leaf)

Original file line number	Diff line number	Diff line change
`@@ -141,6 +141,7 @@ select = [`
`141`	`141`	`"SIM", # flake8-simplify`
`142`	`142`	`"PERF", # Perflint`
`143`	`143`	`"FURB", # refurb`
	`144`	`+ "RUF", # Ruff-specific rules`
`144`	`145`	`]`
`145`	`146`
`146`	`147`	`ignore = [`
Original file line number	Diff line number	Diff line change
`@@ -68,8 +68,8 @@ def __init__(`
`68`	`68`
`69`	`69`	`def __repr__(self) -> str:`
`70`	`70`	`return (`
`71`		`- f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, "`
`72`		`- f"scale_mode={repr(self._scale_mode)})"`
	`71`	`+ f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, "`
	`72`	`+ f"scale_mode={self._scale_mode!r})"`
`73`	`73`	`)`
`74`	`74`
`75`	`75`	`def __str__(self) -> str:`
Original file line number	Diff line number	Diff line change
`@@ -48,8 +48,8 @@ def __init__(`
`48`	`48`
`49`	`49`	`def __repr__(self) -> str:`
`50`	`50`	`return (`
`51`		`- f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="`
`52`		`- f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={repr(self._solver)})"`
	`51`	`+ f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, norm_eps="`
	`52`	`+ f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={self._solver!r})"`
`53`	`53`	`)`
`54`	`54`
`55`	`55`	`def __str__(self) -> str:`
Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ def _check_matrix_has_enough_rows(self, matrix: Tensor) -> None:`
`68`	`68`	`)`
`69`	`69`
`70`	`70`	`def __repr__(self) -> str:`
`71`		`- return f"{self.__class__.__name__}(f={repr(self.f)}, leak={repr(self.leak)})"`
	`71`	`+ return f"{self.__class__.__name__}(f={self.f!r}, leak={self.leak!r})"`
`72`	`72`
`73`	`73`	`def __str__(self) -> str:`
`74`	`74`	`if self.leak is None:`
Original file line number	Diff line number	Diff line change
`@@ -49,8 +49,8 @@ def __init__(`
`49`	`49`
`50`	`50`	`def __repr__(self) -> str:`
`51`	`51`	`return (`
`52`		`- f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="`
`53`		`- f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={repr(self._solver)})"`
	`52`	`+ f"{self.__class__.__name__}(pref_vector={self._pref_vector!r}, norm_eps="`
	`53`	`+ f"{self._norm_eps}, reg_eps={self._reg_eps}, solver={self._solver!r})"`
`54`	`54`	`)`
`55`	`55`
`56`	`56`	`def __str__(self) -> str:`