Add RET (flake8 returns)

Author: PierreQuinton

docs/source/conf.py

@@ -96,8 +96,7 @@ def linkcode_resolve(domain: str, info: dict[str, str]) -> str | None:
     line_str = _get_line_str(obj)
     version_str = _get_version_str()
 
-    link = f"https://github.com/TorchJD/torchjd/blob/{version_str}/{file_name}{line_str}"
-    return link
+    return f"https://github.com/TorchJD/torchjd/blob/{version_str}/{file_name}{line_str}"
 
 
 def _get_obj(_info: dict[str, str]):
@@ -108,8 +107,7 @@ def _get_obj(_info: dict[str, str]):
     for part in full_name.split("."):
         obj = getattr(obj, part)
     # strip decorators, which would resolve to the source of the decorator
-    obj = inspect.unwrap(obj)
-    return obj
+    return inspect.unwrap(obj)
 
 
 def _get_file_name(obj) -> str | None:
@@ -124,8 +122,7 @@ def _get_file_name(obj) -> str | None:
 def _get_line_str(obj) -> str:
     source, start = inspect.getsourcelines(obj)
     end = start + len(source) - 1
-    line_str = f"#L{start}-L{end}"
-    return line_str
+    return f"#L{start}-L{end}"
 
 
 def _get_version_str() -> str:

pyproject.toml

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

src/torchjd/aggregation/_aggregator_bases.py

@@ -59,13 +59,11 @@ def combine(matrix: Matrix, weights: Tensor) -> Tensor:
         weights.
         """
 
-        vector = weights @ matrix
-        return vector
+        return weights @ matrix
 
     def forward(self, matrix: Matrix) -> Tensor:
         weights = self.weighting(matrix)
-        vector = self.combine(matrix, weights)
-        return vector
+        return self.combine(matrix, weights)
 
 
 class GramianWeightedAggregator(WeightedAggregator):

src/torchjd/aggregation/_aligned_mtl.py

@@ -101,9 +101,7 @@ def __init__(
     def forward(self, gramian: PSDMatrix, /) -> Tensor:
         w = self.weighting(gramian)
         B = self._compute_balance_transformation(gramian, self._scale_mode)
-        alpha = B @ w
-
-        return alpha
+        return B @ w
 
     @staticmethod
     def _compute_balance_transformation(
@@ -114,8 +112,7 @@ def _compute_balance_transformation(
         rank = sum(lambda_ > tol)
 
         if rank == 0:
-            identity = torch.eye(len(M), dtype=M.dtype, device=M.device)
-            return identity
+            return torch.eye(len(M), dtype=M.dtype, device=M.device)
 
         order = torch.argsort(lambda_, dim=-1, descending=True)
         lambda_, V = lambda_[order][:rank], V[:, order][:, :rank]
@@ -133,5 +130,4 @@ def _compute_balance_transformation(
                 f"Invalid scale_mode={scale_mode!r}. Expected 'min', 'median', or 'rmse'."
             )
 
-        B = scale.sqrt() * V @ sigma_inv @ V.T
-        return B
+        return scale.sqrt() * V @ sigma_inv @ V.T

src/torchjd/aggregation/_cagrad.py

@@ -101,6 +101,4 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor:
             # We are approximately on the pareto front
             weight_array = np.zeros(dimension)
 
-        weights = torch.from_numpy(weight_array).to(device=gramian.device, dtype=gramian.dtype)
-
-        return weights
+        return torch.from_numpy(weight_array).to(device=gramian.device, dtype=gramian.dtype)

src/torchjd/aggregation/_dualproj.py

@@ -88,5 +88,4 @@ def __init__(
     def forward(self, gramian: PSDMatrix, /) -> Tensor:
         u = self.weighting(gramian)
         G = regularize(normalize(gramian, self.norm_eps), self.reg_eps)
-        w = project_weights(u, G, self.solver)
-        return w
+        return project_weights(u, G, self.solver)

src/torchjd/aggregation/_flattening.py

@@ -29,5 +29,4 @@ def forward(self, generalized_gramian: PSDTensor) -> Tensor:
         shape = generalized_gramian.shape[:k]
         square_gramian = flatten(generalized_gramian)
         weights_vector = self.weighting(square_gramian)
-        weights = weights_vector.reshape(shape)
-        return weights
+        return weights_vector.reshape(shape)

src/torchjd/aggregation/_imtl_g.py

@@ -34,6 +34,4 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor:
         v = torch.linalg.pinv(gramian) @ d
         v_sum = v.sum()
 
-        weights = torch.zeros_like(v) if v_sum.abs() < 1e-12 else v / v_sum
-
-        return weights
+        return torch.zeros_like(v) if v_sum.abs() < 1e-12 else v / v_sum

src/torchjd/aggregation/_krum.py

@@ -76,9 +76,7 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor:
 
         _, selected_indices = torch.topk(scores, k=self.n_selected, largest=False)
         one_hot_selected_indices = F.one_hot(selected_indices, num_classes=gramian.shape[0])
-        weights = one_hot_selected_indices.sum(dim=0).to(dtype=gramian.dtype) / self.n_selected
-
-        return weights
+        return one_hot_selected_indices.sum(dim=0).to(dtype=gramian.dtype) / self.n_selected
 
     def _check_matrix_shape(self, gramian: PSDMatrix) -> None:
         min_rows = self.n_byzantine + 3

src/torchjd/aggregation/_mean.py

@@ -28,5 +28,4 @@ def forward(self, matrix: Tensor, /) -> Tensor:
         device = matrix.device
         dtype = matrix.dtype
         m = matrix.shape[0]
-        weights = torch.full(size=[m], fill_value=1 / m, device=device, dtype=dtype)
-        return weights
+        return torch.full(size=[m], fill_value=1 / m, device=device, dtype=dtype)