diff --git a/src/torchjd/aggregation/_aggregator_bases.py b/src/torchjd/aggregation/_aggregator_bases.py index 9dd014ca..85be8a44 100644 --- a/src/torchjd/aggregation/_aggregator_bases.py +++ b/src/torchjd/aggregation/_aggregator_bases.py @@ -44,6 +44,15 @@ def __repr__(self) -> str: 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): """ diff --git a/torchjd/aggregation/composition.py b/torchjd/aggregation/composition.py new file mode 100644 index 00000000..dec608a3 --- /dev/null +++ b/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)})"