Adding ElementWiseSum

merlin/models/torch/transforms/agg.py

@@ -193,6 +193,53 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
         return torch.stack(sorted_tensors, dim=self.dim).float()
 
 
+@registry.register("element-wise-sum")
+class ElementWiseSum(AggModule):
+    """Element-wise sum of tensors.
+
+    The input dictionary will be sorted by name before concatenation.
+    The sum is computed along the first dimension (default for Stack class).
+
+    Example usage::
+        >>> ewsum = ElementWiseSum()
+        >>> feature1 = torch.tensor([[1, 2], [3, 4]])  # Shape: [batch_size, feature_dim]
+        >>> feature2 = torch.tensor([[5, 6], [7, 8]])  # Shape: [batch_size, feature_dim]
+        >>> input_dict = {"feature1": feature1, "feature2": feature2}
+        >>> output = ewsum(input_dict)
+        >>> print(output)
+        tensor([[ 6,  8],
+                [10, 12]])  # Shape: [batch_size, feature_dim]
+
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.stack = Stack(dim=0)
+
+    def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        Performs an element-wise sum of input tensors.
+
+        Parameters
+        ----------
+        inputs : Dict[str, torch.Tensor]
+            A dictionary where keys are the names of the tensors
+            and values are the tensors to be summed.
+
+        Returns
+        -------
+        torch.Tensor
+            A tensor that is the result of performing an element-wise sum
+            of the input tensors.
+
+        Raises
+        ------
+        RuntimeError
+            If the input tensor shapes don't match for stacking.
+        """
+        return self.stack(inputs).sum(dim=0)
+
+
 class MaybeAgg(BlockContainer):
     """
     This class is designed to conditionally apply an aggregation operation

tests/unit/torch/transforms/test_agg.py

@@ -2,7 +2,7 @@
 import torch
 
 from merlin.models.torch.block import Block
-from merlin.models.torch.transforms.agg import Concat, MaybeAgg, Stack
+from merlin.models.torch.transforms.agg import Concat, ElementWiseSum, MaybeAgg, Stack
 from merlin.models.torch.utils import module_utils
 from merlin.schema import Schema
 
@@ -95,6 +95,30 @@ def test_from_registry(self):
         assert output.shape == (2, 2, 3)
 
 
+class TestElementWiseSum:
+    def setup_class(self):
+        self.ewsum = ElementWiseSum()
+        self.feature1 = torch.tensor([[1, 2], [3, 4]])  # Shape: [batch_size, feature_dim]
+        self.feature2 = torch.tensor([[5, 6], [7, 8]])  # Shape: [batch_size, feature_dim]
+        self.input_dict = {"feature1": self.feature1, "feature2": self.feature2}
+
+    def test_forward(self):
+        output = self.ewsum(self.input_dict)
+        expected_output = torch.tensor([[6, 8], [10, 12]])  # Shape: [batch_size, feature_dim]
+        assert torch.equal(output, expected_output)
+
+    def test_input_tensor_shape_mismatch(self):
+        feature_mismatch = torch.tensor([1, 2, 3])  # Different shape
+        input_dict_mismatch = {"feature1": self.feature1, "feature_mismatch": feature_mismatch}
+        with pytest.raises(RuntimeError):
+            self.ewsum(input_dict_mismatch)
+
+    def test_empty_input_dict(self):
+        empty_dict = {}
+        with pytest.raises(RuntimeError):
+            self.ewsum(empty_dict)
+
+
 class TestMaybeAgg:
     def test_with_single_tensor(self):
         tensor = torch.tensor([1, 2, 3])