diff --git a/src/torchjd/aggregation/_aggregator_bases.py b/src/torchjd/aggregation/_aggregator_bases.py index 9dd014ca..d8c15370 100644 --- a/src/torchjd/aggregation/_aggregator_bases.py +++ b/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 @@ class GramianWeightedAggregator(WeightedAggregator): 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 diff --git a/tests/aggregation/test_composition.py b/tests/aggregation/test_composition.py new file mode 100644 index 00000000..ff1b4865 --- /dev/null +++ b/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