Adding MaybeAgg for use in places like MLPBlock

Author: marcromeyn

merlin/models/torch/transforms/agg.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Union
 
 import torch
 from torch import nn
@@ -130,3 +130,64 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor:
             raise RuntimeError("Input tensor shapes don't match for stacking.")
 
         return torch.stack(sorted_tensors, dim=self.dim)
+
+
+class MaybeAgg(nn.Module):
+    """
+    This class is designed to conditionally apply an aggregation operation
+    (e.g., Stack or Concat) on a tensor or a dictionary of tensors.
+
+    Parameters
+    ----------
+    agg : nn.Module
+        The aggregation operation to be applied.
+
+    Examples
+    --------
+    >>> stack = Stack(dim=0)
+    >>> maybe_agg = MaybeAgg(agg=stack)
+    >>> tensor1 = torch.tensor([[1, 2], [3, 4]])
+    >>> tensor2 = torch.tensor([[5, 6], [7, 8]])
+    >>> input_dict = {"tensor1": tensor1, "tensor2": tensor2}
+    >>> output = maybe_agg(input_dict)
+    >>> print(output)
+    tensor([[[1, 2],
+             [3, 4]],
+
+            [[5, 6],
+             [7, 8]]])
+
+    >>> tensor = torch.tensor([1, 2, 3])
+    >>> output = maybe_agg(tensor)
+    >>> print(output)
+    tensor([1, 2, 3])
+    """
+
+    def __init__(self, agg: nn.Module):
+        super().__init__()
+        self.agg = agg
+
+    def forward(self, inputs: Union[Dict[str, torch.Tensor], torch.Tensor]) -> torch.Tensor:
+        """
+        Conditionally applies the aggregation operation on the inputs.
+
+        Parameters
+        ----------
+        inputs : Union[Dict[str, torch.Tensor], torch.Tensor]
+            Inputs to be aggregated. If inputs is a dictionary of tensors,
+            the aggregation operation will be applied. If inputs is a single tensor,
+            it will be returned as is.
+
+        Returns
+        -------
+        torch.Tensor
+            Aggregated tensor if inputs is a dictionary, otherwise the input tensor.
+        """
+
+        if torch.jit.isinstance(inputs, Dict[str, torch.Tensor]):
+            return self.agg(inputs)
+
+        if not torch.jit.isinstance(inputs, torch.Tensor):
+            raise RuntimeError("Inputs must be either a dictionary of tensors or a single tensor.")
+
+        return inputs

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, Stack
+from merlin.models.torch.transforms.agg import Concat, MaybeAgg, Stack
 from merlin.models.torch.utils import module_utils
 
 
@@ -43,7 +43,7 @@ def test_from_registry(self):
             "a": torch.randn(2, 3),
             "b": torch.randn(2, 4),
         }
-        output = block(input_tensors)
+        output = module_utils.module_test(block, input_tensors)
         assert output.shape == (2, 7)
 
 
@@ -84,3 +84,38 @@ def test_from_registry(self):
         }
         output = block(input_tensors)
         assert output.shape == (2, 2, 3)
+
+
+class TestMaybeAgg:
+    def test_with_single_tensor(self):
+        tensor = torch.tensor([1, 2, 3])
+        stack = Stack(dim=0)
+        maybe_agg = MaybeAgg(agg=stack)
+
+        output = module_utils.module_test(maybe_agg, tensor)
+        assert torch.equal(output, tensor)
+
+    def test_with_dict(self):
+        stack = Stack(dim=0)
+        maybe_agg = MaybeAgg(agg=stack)
+
+        tensor1 = torch.tensor([[1, 2], [3, 4]])
+        tensor2 = torch.tensor([[5, 6], [7, 8]])
+        input_dict = {"tensor1": tensor1, "tensor2": tensor2}
+        expected_output = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        output = module_utils.module_test(maybe_agg, input_dict)
+
+        assert torch.equal(output, expected_output)
+
+    def test_with_incompatible_dict(self):
+        concat = Concat(dim=0)
+        maybe_agg = MaybeAgg(agg=concat)
+
+        tensor1 = torch.tensor([1, 2, 3])
+        tensor2 = torch.tensor([4, 5])
+        input_dict = {"tensor1": (tensor1, tensor2)}
+
+        with pytest.raises(
+            RuntimeError, match="Inputs must be either a dictionary of tensors or a single tensor"
+        ):
+            maybe_agg(input_dict)