Adding registry

merlin/models/torch/block.py

@@ -22,9 +22,11 @@
 
 from merlin.models.torch.batch import Batch
 from merlin.models.torch.container import BlockContainer, BlockContainerDict
+from merlin.models.torch.registry import registry
+from merlin.models.utils.registry import RegistryMixin
 
 
-class Block(BlockContainer):
+class Block(BlockContainer, RegistryMixin):
     """A base-class that calls it's modules sequentially.
 
     Parameters
@@ -35,6 +37,8 @@ class Block(BlockContainer):
         The name of the block. If None, no name is assigned.
     """
 
+    registry = registry
+
     def __init__(self, *module: nn.Module, name: Optional[str] = None):
         super().__init__(*module, name=name)
 

merlin/models/torch/registry.py

@@ -0,0 +1,3 @@
+from merlin.models.utils.registry import Registry
+
+registry: Registry = Registry.class_registry("modules")

tests/unit/torch/test_block.py

@@ -82,6 +82,20 @@ def test_repeat(self):
         with pytest.raises(ValueError, match="n must be greater than 0"):
             block.repeat(0)
 
+    def test_from_registry(self):
+        @Block.registry.register("my_block")
+        class MyBlock(Block):
+            def forward(self, inputs):
+                _inputs = inputs + 1
+
+                return super().forward(_inputs)
+
+        block = Block.parse("my_block")
+        assert block.__class__ == MyBlock
+
+        inputs = torch.randn(1, 3)
+        assert torch.equal(block(inputs), inputs + 1)
+
 
 class TestParallelBlock:
     def test_init(self):