use moduledict

pts/modules/feature.py

@@ -40,26 +40,17 @@ class FeatureEmbedder(nn.Module):
 
 
 class FeatureAssembler(nn.Module):
-    def __init__(T: int,
-                 use_static_cat: bool = False,
-                 use_static_real: bool = False,
-                 use_dynamic_cat: bool = False,
-                 use_dynamic_real: bool = False,
+    def __init__(self,
+                 T: int,
                  embed_static: Optional[FeatureEmbedder] = None,
-                 embed_dynamic: Optional[FeatureEmbedder] = None,
-                 dtype: torch.dtype = torch.float32) -> None:
+                 embed_dynamic: Optional[FeatureEmbedder] = None) -> None:
         super().__init__()
 
         self.T = T
-        self.use_static_cat = use_static_cat
-        self.use_static_real = use_static_real
-        self.use_dynamic_cat = use_dynamic_cat
-        self.use_dynamic_real = use_dynamic_real
-
-        self.embed_static: Callable[[torch.Tensor], torch.
-                                    Tensor] = embed_static or (lambda x: x)
-        self.embed_dynamic: Callable[[torch.Tensor], torch.
-                                     Tensor] = embed_dynamic or (lambda x: x)
+        self.embeddings = nn.ModuleDict({
+            'embed_static': embed_static,
+            'embed_dynamic': embed_dynamic
+        })
 
     def forward(
             self,
@@ -78,11 +69,15 @@ class FeatureAssembler(nn.Module):
         return torch.cat(processed_features, dim=-1)
 
     def process_static_cat(self, feature: torch.Tensor) -> torch.Tensor:
-        feature = self.embed_static(feature.to(self.dtype))
+        if self.embeddings['embed_static'] is not None:
+            feature = self.embeddings['embed_static'](feature)
         return feature.unsqueeze(1).expand(-1, self.T, -1)
 
     def process_dynamic_cat(self, feature: torch.Tensor) -> torch.Tensor:
-        return self.embed_dynamic(feature.to(self.dtype))
+        if self.embeddings['embed_dynamic'] is None:
+            return feature
+        else:
+            return self.embeddings['embed_dynamic'](feature)
 
     def process_static_real(self, feature: torch.Tensor) -> torch.Tensor:
         return feature.unsqueeze(1).expand(-1, self.T, -1)

test/modules/test_feature.py

@@ -59,8 +59,84 @@ def test_feature_embedder(config):
         exp_output = torch.ones(out_shape)
         
         assert act_output.shape == exp_output.shape
-        import pdb; pdb.set_trace()
         assert torch.abs(torch.sum(act_output - exp_output)) < 1e-20
 
     test_parameters_length()
     test_forward_pass()
+
+@pytest.mark.parametrize(
+    "config",
+    (
+        lambda N, T: [
+            dict(
+                N=N,
+                T=T,
+                static_cat=dict(C=2),
+                static_real=dict(C=5),
+                dynamic_cat=dict(C=3),
+                dynamic_real=dict(C=4),
+                embed_static=dict(
+                    cardinalities=[2, 4],
+                    embedding_dims=[3, 6],
+                ),
+                embed_dynamic=dict(
+                    cardinalities=[30, 30, 30],
+                    embedding_dims=[10, 20, 30],
+                ),
+            )
+        ]
+    )(10, 25),
+)
+def test_feature_assembler(config):
+    # iterate over the power-set of all possible feature types, excluding the empty set
+    feature_types = {
+        "static_cat",
+        "static_real",
+        "dynamic_cat",
+        "dynamic_real",
+    }
+    feature_combs = chain.from_iterable(
+        combinations(feature_types, r)
+        for r in range(1, len(feature_types) + 1)
+    )
+
+    # iterate over the power-set of all possible feature types, including the empty set
+    embedder_types = {"embed_static", "embed_dynamic"}
+    embedder_combs = chain.from_iterable(
+        combinations(embedder_types, r)
+        for r in range(0, len(embedder_types) + 1)
+    )
+    
+    for enabled_embedders in embedder_combs:
+        embed_static = (
+            FeatureEmbedder(**config["embed_static"])
+            if "embed_static" in enabled_embedders
+            else None
+        )
+        embed_dynamic = (
+            FeatureEmbedder(**config["embed_dynamic"])
+            if "embed_dynamic" in enabled_embedders
+            else None
+        )
+
+        for enabled_features in feature_combs:
+            assemble_feature = FeatureAssembler(
+                T=config["T"],
+                embed_static=embed_static,
+                embed_dynamic=embed_dynamic,
+            )
+            # assemble_feature.collect_params().initialize(mx.initializer.One())
+
+
+            def test_parameters_length():
+                exp_params_len = sum(
+                    [
+                        len(config[k]["embedding_dims"])
+                        for k in ["embed_static", "embed_dynamic"]
+                        if k in enabled_embedders
+                    ]
+                )
+                act_params_len = len([p for p in assemble_feature.parameters()])
+                assert exp_params_len == act_params_len
+
+            test_parameters_length()