Fix failined nested router test

Author: marcromeyn

merlin/models/torch/batch.py

@@ -17,7 +17,6 @@
 from typing import Dict, Optional, Union
 
 import torch
-import torch.nn
 
 from merlin.dataloader.torch import Loader
 from merlin.io import Dataset
@@ -321,8 +320,6 @@ def sample_batch(
         dictionary of target tensors.
     """
 
-<<<<<<< HEAD
-<<<<<<< HEAD
     if isinstance(data, Dataset):
         if not batch_size:
             raise ValueError("Either use 'Loader' or specify 'batch_size'")
@@ -331,37 +328,16 @@ def sample_batch(
         loader = data
     else:
         raise ValueError(f"Expected Dataset or Loader instance, got: {data}")
-=======
-=======
-    if not isinstance(dataset_or_loader, (Dataset, Loader)):
-        raise ValueError(f"Expected Dataset or Loader instance, got {dataset_or_loader}")
-
->>>>>>> 519159b5 (Increasing test-coverage)
-    if isinstance(dataset_or_loader, Dataset):
-        if not batch_size:
-            raise ValueError("Either use 'Loader' or specify 'batch_size'")
-        loader = Loader(dataset_or_loader, batch_size=batch_size, shuffle=shuffle)
-    else:
-        loader = dataset_or_loader
->>>>>>> d00e67a4 (Adding sample_batch & sample_features)
 
     batch = loader.peek()
     # batch could be of type Prediction, so we can't unpack directly
     inputs, targets = batch[0], batch[1]
 
-<<<<<<< HEAD
     return Batch(inputs, targets)
 
 
 def sample_features(
     data: Union[Dataset, Loader],
-=======
-    return inputs, targets
-
-
-def sample_features(
-    dataset_or_loader: Union[Dataset, Loader],
->>>>>>> d00e67a4 (Adding sample_batch & sample_features)
     batch_size: Optional[int] = None,
     shuffle: Optional[bool] = False,
 ) -> Dict[str, torch.Tensor]:
@@ -382,8 +358,4 @@ def sample_features(
         dictionary of feature tensors.
     """
 
-<<<<<<< HEAD
     return sample_batch(data, batch_size, shuffle).features
-=======
-    return sample_batch(dataset_or_loader, batch_size, shuffle)[0]
->>>>>>> d00e67a4 (Adding sample_batch & sample_features)

merlin/models/torch/container.py

@@ -14,19 +14,9 @@
 # limitations under the License.
 #
 
-<<<<<<< HEAD
-<<<<<<< HEAD
 from copy import deepcopy
 from functools import reduce
 from typing import Dict, Iterator, Optional, Union
-=======
-from typing import Iterator, Optional, Union
->>>>>>> 2dfe2782 (Adding torch github-action + add copyright)
-=======
-from copy import deepcopy
-from functools import reduce
-from typing import Dict, Iterator, Optional, Union
->>>>>>> cbf7e956 (Adding BlockContainerDict)
 
 from torch import nn
 from torch._jit_internal import _copy_to_script_wrapper
@@ -57,33 +47,25 @@ def __init__(self, *inputs: nn.Module, name: Optional[str] = None):
 
         self._name: str = name
 
-<<<<<<< HEAD
     def append(self, module: nn.Module, link: Optional[Link] = None):
-=======
-    def append(self, module: nn.Module):
->>>>>>> 793f47ef (Adding improved doc-strings)
         """Appends a given module to the end of the list.
 
         Parameters
         ----------
         module : nn.Module
             The PyTorch module to be appended.
-<<<<<<< HEAD
         link : Optional[LinkType]
             The link to use for the module. If None, no link is used.
             This can either be a Module or a string, options are:
                 - "residual": Adds a residual connection to the module.
                 - "shortcut": Adds a shortcut connection to the module.
                 - "shortcut-concat": Adds a shortcut connection by concatenating
                 the input and output.
-=======
->>>>>>> 793f47ef (Adding improved doc-strings)
 
         Returns
         -------
         self
         """
-<<<<<<< HEAD
         _module = self._check_link(module, link=link)
         self.values.append(self.wrap_module(_module))
 
@@ -111,27 +93,6 @@ def prepend(self, module: nn.Module, link: Optional[Link] = None):
         return self.insert(0, module, link=link)
 
     def insert(self, index: int, module: nn.Module, link: Optional[Link] = None):
-=======
-        self.values.append(self.wrap_module(module))
-
-        return self
-
-    def prepend(self, module: nn.Module):
-        """Prepends a given module to the beginning of the list.
-
-        Parameters
-        ----------
-        module : nn.Module
-            The PyTorch module to be prepended.
-
-        Returns
-        -------
-        self
-        """
-        return self.insert(0, module)
-
-    def insert(self, index: int, module: nn.Module):
->>>>>>> 793f47ef (Adding improved doc-strings)
         """Inserts a given module at the specified index.
 
         Parameters
@@ -140,28 +101,20 @@ def insert(self, index: int, module: nn.Module):
             The index at which the module is to be inserted.
         module : nn.Module
             The PyTorch module to be inserted.
-<<<<<<< HEAD
         link : Optional[LinkType]
             The link to use for the module. If None, no link is used.
             This can either be a Module or a string, options are:
                 - "residual": Adds a residual connection to the module.
                 - "shortcut": Adds a shortcut connection to the module.
                 - "shortcut-concat": Adds a shortcut connection by concatenating
                 the input and output.
-=======
->>>>>>> 793f47ef (Adding improved doc-strings)
 
         Returns
         -------
         self
         """
-<<<<<<< HEAD
         _module = self._check_link(module, link=link)
         self.values.insert(index, self.wrap_module(_module))
-=======
-
-        self.values.insert(index, self.wrap_module(module))
->>>>>>> 793f47ef (Adding improved doc-strings)
 
         return self
 
@@ -222,7 +175,6 @@ def __repr__(self) -> str:
     def _get_name(self) -> str:
         return super()._get_name() if self._name is None else self._name
 
-<<<<<<< HEAD
     def _check_link(self, module: nn.Module, link: Optional[LinkType] = None) -> nn.Module:
         if link:
             linked_module: Link = Link.parse(link)
@@ -246,31 +198,18 @@ class BlockContainerDict(nn.ModuleDict):
         An optional name for the BlockContainer.
     """
 
-=======
-
-class BlockContainerDict(nn.ModuleDict):
->>>>>>> cbf7e956 (Adding BlockContainerDict)
     def __init__(
         self, *inputs: Union[nn.Module, Dict[str, nn.Module]], name: Optional[str] = None
     ) -> None:
         if not inputs:
             inputs = [{}]
 
-<<<<<<< HEAD
-<<<<<<< HEAD
-=======
-        if isinstance(inputs, tuple) and len(inputs) == 1 and isinstance(inputs[0], (list, tuple)):
-            modules = inputs[0]
->>>>>>> cbf7e956 (Adding BlockContainerDict)
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)
         if all(isinstance(x, dict) for x in inputs):
             modules = reduce(lambda a, b: dict(a, **b), inputs)  # type: ignore
 
         super().__init__(modules)
         self._name: str = name
 
-<<<<<<< HEAD
     def append_to(
         self, name: str, module: nn.Module, link: Optional[LinkType] = None
     ) -> "BlockContainerDict":
@@ -390,37 +329,6 @@ def prepend_for_each(
 
         return self
 
-=======
-    def append_to(self, name: str, module: nn.Module) -> "BlockContainerDict":
-        self._modules[name].append(module)
-
-        return self
-
-    def prepend_to(self, name: str, module: nn.Module) -> "BlockContainerDict":
-        self._modules[name].prepend(module)
-
-        return self
-
-    # Append to all branches, optionally copying
-    def append_for_each(self, module: nn.Module, shared=False) -> "BlockContainerDict":
-        for branch in self.values():
-            _module = module if shared else deepcopy(module)
-            branch.append(_module)
-
-        return self
-
-    def prepend_for_each(self, module: nn.Module, shared=False) -> "BlockContainerDict":
-        for branch in self.values():
-            _module = module if shared else deepcopy(module)
-            branch.prepend(_module)
-
-        return self
-
-    # This assumes same branches, we append it's content to each branch
-    # def append_parallel(self, module: "BlockContainerDict") -> "BlockContainerDict":
-    #     ...
-
->>>>>>> cbf7e956 (Adding BlockContainerDict)
     def add_module(self, name: str, module: Optional[nn.Module]) -> None:
         if module and not isinstance(module, BlockContainer):
             module = BlockContainer(module, name=name[0].upper() + name[1:])

merlin/models/torch/router.py

@@ -159,6 +159,10 @@ def nested_router(self) -> "RouterBlock":
             A new router block with the current block as its selectable.
         """
 
+        if hasattr(self, "_forward_called"):
+            # We don't need to track the schema since we will be using the nested router
+            self._handle.remove()
+
         return RouterBlock(self)
 
     def select(self, selection: Selection) -> "RouterBlock":

tests/unit/torch/test_block.py

@@ -13,19 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-<<<<<<< HEAD
-<<<<<<< HEAD
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)
 from typing import Dict, Tuple
 
 import pytest
-<<<<<<< HEAD
-=======
-
->>>>>>> 2dfe2782 (Adding torch github-action + add copyright)
-=======
->>>>>>> f7011343 (Increase test-coverage)
 import torch
 from torch import nn
 
@@ -104,8 +94,6 @@ def test_repeat(self):
 
         with pytest.raises(ValueError, match="n must be greater than 0"):
             block.repeat(0)
-<<<<<<< HEAD
-<<<<<<< HEAD
 
     def test_repeat_with_link(self):
         block = Block(PlusOne())
@@ -133,8 +121,6 @@ def forward(self, inputs):
 
         inputs = torch.randn(1, 3)
         assert torch.equal(block(inputs), inputs + 1)
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)
 
 
 class TestParallelBlock:
@@ -171,7 +157,6 @@ def test_forward_dict_duplicate(self):
         with pytest.raises(RuntimeError):
             pb(inputs)
 
-<<<<<<< HEAD
     def test_schema_tracking(self):
         pb = ParallelBlock({"a": PlusOne(), "b": PlusOne()})
 
@@ -186,8 +171,6 @@ def test_schema_tracking(self):
 
         assert len(schema.select_by_tag(Tags.EMBEDDING)) == 2
 
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)
     def test_forward_tuple(self):
         inputs = torch.randn(1, 3)
         pb = ParallelBlock({"test": PlusOneTuple()})
@@ -260,8 +243,3 @@ def test_getitem(self):
 
         with pytest.raises(IndexError):
             pb["invalid_key"]
-<<<<<<< HEAD
-=======
->>>>>>> f7011343 (Increase test-coverage)
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)

tests/unit/torch/test_container.py

@@ -178,17 +178,6 @@ def test_empty(self):
         container = BlockContainerDict()
         assert len(container) == 0
 
-<<<<<<< HEAD
-<<<<<<< HEAD
-=======
-    def test_list_of_dict(self):
-        container = BlockContainerDict(({"test": self.module}))
-        assert len(container) == 1
-        assert "test" in container
-
->>>>>>> cbf7e956 (Adding BlockContainerDict)
-=======
->>>>>>> 77ca69b4 (Adding ParallelBlock)
     def test_not_module(self):
         with pytest.raises(ValueError):
             BlockContainerDict({"test": "not a module"})
@@ -197,22 +186,16 @@ def test_append_to(self):
         self.container.append_to("test", self.module)
         assert "test" in self.container._modules
 
-<<<<<<< HEAD
         self.container.append_to("test", self.module, link="residual")
         assert isinstance(self.container["test"][-1], link.Residual)
 
-=======
->>>>>>> cbf7e956 (Adding BlockContainerDict)
     def test_prepend_to(self):
         self.container.prepend_to("test", self.module)
         assert "test" in self.container._modules
 
-<<<<<<< HEAD
         self.container.prepend_to("test", self.module, link="residual")
         assert isinstance(self.container["test"][0], link.Residual)
 
-=======
->>>>>>> cbf7e956 (Adding BlockContainerDict)
     def test_append_for_each(self):
         container = BlockContainerDict({"a": nn.Module(), "b": nn.Module()})
 
@@ -227,13 +210,10 @@ def test_append_for_each(self):
         assert len(container["b"]) == 3
         assert container["a"][-1] == container["b"][-1]
 
-<<<<<<< HEAD
         container.append_for_each(to_add, link="residual")
         assert isinstance(container["a"][-1], link.Residual)
         assert isinstance(container["b"][-1], link.Residual)
 
-=======
->>>>>>> cbf7e956 (Adding BlockContainerDict)
     def test_prepend_for_each(self):
         container = BlockContainerDict({"a": nn.Module(), "b": nn.Module()})
 
@@ -247,10 +227,7 @@ def test_prepend_for_each(self):
         assert len(container["a"]) == 3
         assert len(container["b"]) == 3
         assert container["a"][0] == container["b"][0]
-<<<<<<< HEAD
 
         container.prepend_for_each(to_add, link="residual")
         assert isinstance(container["a"][0], link.Residual)
         assert isinstance(container["b"][0], link.Residual)
-=======
->>>>>>> cbf7e956 (Adding BlockContainerDict)