feat(aggregation): add Composition class and << operator in _aggregator_bases

src/torchjd/aggregation/_aggregator_bases.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Callable, Generic, TypeVar
 
 from torch import Tensor, nn
 
@@ -86,3 +87,32 @@
     def __init__(self, gramian_weighting: Weighting[PSDMatrix]) -> None:
         super().__init__(gramian_weighting << compute_gramian)
         self.gramian_weighting = gramian_weighting
+
+
+A = TypeVar("A", bound=nn.Module)
+F = TypeVar("F", bound=Callable)
+
+
+class Composition(nn.Module, Generic[A, F]):
+    """
+    Wraps an aggregator and a function into a composite nn.Module.
+    """
+
+    def __init__(self, outer: A, inner: F) -> None:
+        super().__init__()
+        self.outer = outer
+        self.inner = inner
+
+    def forward(self, *args, **kwargs):
+        return self.outer(self.inner(*args, **kwargs))
+
+    def __str__(self) -> str:
+        return str(self.outer) + " << " + str(self.inner)
+
+
+def compose(self: A, inner: F) -> Composition[A, F]:
+    return Composition(self, inner)
+
+
+# Bind the << operator to nn.Module
+nn.Module.__lshift__ = compose

tests/aggregation/test_composition.py

@@ -0,0 +1,47 @@
+import pytest
+import torch
+from torchjd.aggregation import Composition
+from torchjd.aggregation._aggregator import Aggregator
+
+class DummyScale(Aggregator):
+    """Scales input matrix by a factor."""
+    def __init__(self, factor: float):
+        super().__init__()
+        self.factor = factor
+
+    def __call__(self, matrix: torch.Tensor) -> torch.Tensor:
+        return matrix * self.factor
+
+class DummySum(Aggregator):
+    """Sums matrix rows."""
+    def __call__(self, matrix: torch.Tensor) -> torch.Tensor:
+        return matrix.sum(dim=0)
+
+def test_composition_lshift_operator():
+    scale = DummyScale(2.0)
+    sum_agg = DummySum()
+
+    # Test composition via operator
+    composed = sum_agg << scale
+
+    matrix = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    # Forward pass: scale matrix by 2, then sum rows
+    result = composed(matrix)
+    expected = (matrix * 2.0).sum(dim=0)
+
+    assert torch.allclose(result, expected)
+    assert str(composed) == f"{sum_agg} << {scale}"
+
+def test_composition_sequential_execution():
+    J = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    comp = Composition([DummyScale(2.0), DummySum()])
+    result = comp(J)
+    expected = torch.tensor([10.0, 14.0, 18.0])
+    assert torch.allclose(result, expected)
+
+def test_lshift_operator_syntax():
+    scale = DummyScale(2.0)
+    sum_agg = DummySum()
+    comp = sum_agg << scale
+    assert isinstance(comp, Composition)
+    assert len(comp.aggregators) == 2