Improve ConFIG implementation

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
 

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