refactor(aggregation): Add gramian_weighting field

docs/source/examples/monitoring.rst

@@ -49,7 +49,7 @@ they have a negative inner product).
     optimizer = SGD(params, lr=0.1)
     aggregator = UPGrad()
 
-    aggregator.weighting.weighting.register_forward_hook(print_weights)
+    aggregator.gramian_weighting.register_forward_hook(print_weights)
     aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10

src/torchjd/aggregation/_aggregator_bases.py

@@ -73,8 +73,10 @@ class GramianWeightedAggregator(WeightedAggregator):
     WeightedAggregator that computes the gramian of the input jacobian matrix before applying a
     Weighting to it.
 
-    :param weighting: The object responsible for extracting the vector of weights from the gramian.
+    :param gramian_weighting: The object responsible for extracting the vector of weights from the
+        gramian.
     """
 
-    def __init__(self, weighting: Weighting[PSDMatrix]):
-        super().__init__(weighting << compute_gramian)
+    def __init__(self, gramian_weighting: Weighting[PSDMatrix]):
+        super().__init__(gramian_weighting << compute_gramian)
+        self.gramian_weighting = gramian_weighting

tests/doc/test_rst.py

@@ -308,7 +308,7 @@ def print_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.
     optimizer = SGD(params, lr=0.1)
     aggregator = UPGrad()
 
-    aggregator.weighting.weighting.register_forward_hook(print_weights)
+    aggregator.gramian_weighting.register_forward_hook(print_weights)
     aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10