From d889f4909e8020c70e4412cb4b4c3baf91d049c2 Mon Sep 17 00:00:00 2001
From: SundaramGupta <powerofaisinstudy@gmail.com>
Date: Mon, 14 Sep 2026 08:58:20 +0530
Subject: [PATCH 1/3] test(aggregation): add unit tests for Composition and
 __lshift__  Step 2: Open the Draft PR

---
 tests/aggregation/test_composition.py | 36 +++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 tests/aggregation/test_composition.py

diff --git a/tests/aggregation/test_composition.py b/tests/aggregation/test_composition.py
new file mode 100644
index 00000000..445a175a
--- /dev/null
+++ b/tests/aggregation/test_composition.py
@@ -0,0 +1,36 @@
+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_empty_raises():
+    with pytest.raises(ValueError, match="cannot be empty"):
+        Composition([])
+
+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

From 8d99b6801e85a79f6ec3233ebaff79b5e2ea8b81 Mon Sep 17 00:00:00 2001
From: SundaramGupta <powerofaisinstudy@gmail.com>
Date: Mon, 14 Sep 2026 22:46:35 +0530
Subject: [PATCH 2/3] feat(aggregation): add Composition class and << operator
 in _aggregator_bases

.
---
 src/torchjd/aggregation/_aggregator_bases.py | 30 ++++++++++++++++++++
 1 file changed, 30 insertions(+)

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

From 58908b3d15e12e00803953db2a3520d392107e0e Mon Sep 17 00:00:00 2001
From: SundaramGupta <powerofaisinstudy@gmail.com>
Date: Mon, 14 Sep 2026 23:05:19 +0530
Subject: [PATCH 3/3] test(aggregation): add unit tests for Composition and
 __lshift__ operator

---
 tests/aggregation/test_composition.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/tests/aggregation/test_composition.py b/tests/aggregation/test_composition.py
index 445a175a..ff1b4865 100644
--- a/tests/aggregation/test_composition.py
+++ b/tests/aggregation/test_composition.py
@@ -17,9 +17,20 @@ class DummySum(Aggregator):
     def __call__(self, matrix: torch.Tensor) -> torch.Tensor:
         return matrix.sum(dim=0)
 
-def test_composition_empty_raises():
-    with pytest.raises(ValueError, match="cannot be empty"):
-        Composition([])
+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]])