Merge branch 'main' into fix-diagonalize-key-order

Author: ValerianRey

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
     -   id: check-merge-conflict  # Check for files that contain merge conflict strings.
 
 -   repo: https://github.com/PyCQA/flake8
-    rev: 7.1.2
+    rev: 7.2.0
     hooks:
     -   id: flake8  # Check style and syntax. Does not modify code, issues have to be solved manually.
         args: [

CHANGELOG.md

@@ -19,6 +19,8 @@ changes that do not affect the user.
 - Refactored internal verifications in the autojac engine so that they do not run at runtime
   anymore. This should minimally improve the performance and reduce the memory usage of `backward`
   and `mtl_backward`.
+- Improved the implementation of `ConFIG` to be simpler and safer when normalizing vectors. It
+  should slightly improve the performance of `ConFIG` and minimally affect its behavior.
 
 ### Fixed
 

README.md

@@ -6,6 +6,7 @@
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/TorchJD/torchjd/main.svg)](https://results.pre-commit.ci/latest/github/TorchJD/torchjd/main)
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/torchjd)](https://pypistats.org/packages/torchjd)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torchjd)](https://pypi.org/project/torchjd/)
+[![Static Badge](https://img.shields.io/badge/Discord%20-%20community%20-%20%235865F2?logo=discord&logoColor=%23FFFFFF&label=Discord)](https://discord.gg/76KkRnb3nk)
 
 TorchJD is a library extending autograd to enable
 [Jacobian descent](https://arxiv.org/pdf/2406.16232) with PyTorch. It can be used to train neural

src/torchjd/aggregation/config.py

@@ -70,12 +70,9 @@ def forward(self, matrix: Tensor) -> Tensor:
         units = torch.nan_to_num((matrix / (matrix.norm(dim=1)).unsqueeze(1)), 0.0)
         best_direction = torch.linalg.pinv(units) @ weights
 
-        if best_direction.norm() == 0:
-            unit_target_vector = torch.zeros_like(best_direction)
-        else:
-            unit_target_vector = best_direction / best_direction.norm()
+        unit_target_vector = torch.nn.functional.normalize(best_direction, dim=0)
 
-        length = torch.sum(torch.stack([torch.dot(grad, unit_target_vector) for grad in matrix]))
+        length = torch.sum(matrix @ unit_target_vector)
 
         return length * unit_target_vector