feat(aggregation): Add Composition class and << operator

src/torchjd/aggregation/_aggregator_bases.py

@@ -44,6 +44,15 @@
     def __str__(self) -> str:
         return f"{self.__class__.__name__}"
 
+    def __lshift__(self, other: "Aggregator") -> "Composition":
+        from torchjd.aggregation.composition import Composition
+
+        if isinstance(other, Composition):
+            return Composition([*other.aggregators, self])
+        elif isinstance(self, Composition):
+            return Composition([*self.aggregators, other])
+        return Composition([other, self])
+
 
 class WeightedAggregator(Aggregator):
     """

torchjd/aggregation/composition.py

@@ -0,0 +1,25 @@
+import torch
+from torch import Tensor
+from typing import Sequence
+from torchjd.aggregation.abstract import Aggregator
+
+class Composition(Aggregator):
+    """
+    Composes a sequence of aggregators/transforms into a single aggregator.
+    Executes each aggregator sequentially on the Jacobian matrix.
+    """
+    def __init__(self, aggregators: Sequence[Aggregator]):
+        super().__init__()
+        if not aggregators:
+            raise ValueError("Aggregators sequence cannot be empty.")
+        self.aggregators = list(aggregators)
+
+    def __call__(self, matrix: Tensor) -> Tensor:
+        out = matrix
+        for aggregator in self.aggregators:
+            out = aggregator(out)
+        return out
+
+    def __repr__(self) -> str:
+        names = [agg.__class__.__name__ for agg in self.aggregators]
+        return f"Composition({', '.join(names)})"