docs: Fix monitoring example (#358)

* Rename print_similarity_with_gd to print_gd_similarity
* Change type hint of print_gd_similarity inputs to tuple[torch.Tensor, ...]
* Reorder imports in rst to match the test

Author: ValerianRey

docs/source/examples/monitoring.rst

@@ -19,8 +19,8 @@ they have a negative inner product).
 
     import torch
     from torch.nn import Linear, MSELoss, ReLU, Sequential
-    from torch.optim import SGD
     from torch.nn.functional import cosine_similarity
+    from torch.optim import SGD
 
     from torchjd import mtl_backward
     from torchjd.aggregation import UPGrad
@@ -29,7 +29,7 @@ they have a negative inner product).
         """Prints the extracted weights."""
         print(f"Weights: {weights}")
 
-    def print_similarity_with_gd(_, inputs: tuple[torch.Tensor], aggregation: torch.Tensor) -> None:
+    def print_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.Tensor) -> None:
         """Prints the cosine similarity between the aggregation and the average gradient."""
         matrix = inputs[0]
         gd_output = matrix.mean(dim=0)
@@ -50,7 +50,7 @@ they have a negative inner product).
     aggregator = UPGrad()
 
     aggregator.weighting.register_forward_hook(print_weights)
-    aggregator.register_forward_hook(print_similarity_with_gd)
+    aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
     task1_targets = torch.randn(8, 16, 1)  # 8 batches of 16 targets for the first task

tests/doc/test_rst.py

@@ -222,7 +222,7 @@ def print_weights(_, __, weights: torch.Tensor) -> None:
         """Prints the extracted weights."""
         print(f"Weights: {weights}")
 
-    def print_similarity_with_gd(_, inputs: tuple[torch.Tensor], aggregation: torch.Tensor) -> None:
+    def print_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.Tensor) -> None:
         """Prints the cosine similarity between the aggregation and the average gradient."""
         matrix = inputs[0]
         gd_output = matrix.mean(dim=0)
@@ -243,7 +243,7 @@ def print_similarity_with_gd(_, inputs: tuple[torch.Tensor], aggregation: torch.
     aggregator = UPGrad()
 
     aggregator.weighting.register_forward_hook(print_weights)
-    aggregator.register_forward_hook(print_similarity_with_gd)
+    aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
     task1_targets = torch.randn(8, 16, 1)  # 8 batches of 16 targets for the first task