From 75c8138b393443a4bdf9c41cde3d2761d26a071f Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Sat, 2 Nov 2019 09:41:57 +0100
Subject: [PATCH] added feature test

---
 pts/modules/__init__.py      |  1 +
 pts/modules/feature.py       |  2 ++
 test/modules/test_feature.py | 66 ++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 test/modules/test_feature.py

diff --git a/pts/modules/__init__.py b/pts/modules/__init__.py
index bc468ac..92a406c 100644
--- a/pts/modules/__init__.py
+++ b/pts/modules/__init__.py
@@ -1,2 +1,3 @@
 from .distribution_output import ArgProj, Output, DistributionOutput, StudentTOutput
 from .lambda_layer import LambdaLayer
+from .feature import FeatureEmbedder, FeatureAssembler
\ No newline at end of file
diff --git a/pts/modules/feature.py b/pts/modules/feature.py
index dae3244..5d85f96 100644
--- a/pts/modules/feature.py
+++ b/pts/modules/feature.py
@@ -1,3 +1,5 @@
+from typing import Callable, List, Optional 
+
 import torch
 import torch.nn as nn
 
diff --git a/test/modules/test_feature.py b/test/modules/test_feature.py
new file mode 100644
index 0000000..df595a3
--- /dev/null
+++ b/test/modules/test_feature.py
@@ -0,0 +1,66 @@
+import pytest
+from itertools import chain, combinations
+
+
+import torch
+import torch.nn as nn
+
+from pts.modules import FeatureEmbedder, FeatureAssembler
+
+@pytest.mark.parametrize(
+    "config",
+    (
+        lambda N, T: [
+            # single static feature
+            dict(
+                shape=(N, 1),
+                kwargs=dict(cardinalities=[50], embedding_dims=[10]),
+            ),
+            # single dynamic feature
+            dict(
+                shape=(N, T, 1),
+                kwargs=dict(cardinalities=[2], embedding_dims=[10]),
+            ),
+            # multiple static features
+            dict(
+                shape=(N, 4),
+                kwargs=dict(
+                    cardinalities=[50, 50, 50, 50],
+                    embedding_dims=[10, 20, 30, 40],
+                ),
+            ),
+            # multiple dynamic features
+            dict(
+                shape=(N, T, 3),
+                kwargs=dict(
+                    cardinalities=[30, 30, 30], embedding_dims=[10, 20, 30]
+                ),
+            ),
+        ]
+    )(10, 20),
+)
+def test_feature_embedder(config):
+    out_shape = config["shape"][:-1] + (
+        sum(config["kwargs"]["embedding_dims"]),
+    )
+    embed_feature = FeatureEmbedder(
+        **config["kwargs"]
+    )
+    for embed in embed_feature._FeatureEmbedder__embedders:
+        nn.init.constant_(embed.weight, 1.0)
+
+    def test_parameters_length():
+        exp_params_len = len([p for p in embed_feature.parameters()])
+        act_params_len = len(config["kwargs"]["embedding_dims"])
+        assert exp_params_len == act_params_len
+    
+    def test_forward_pass():
+        act_output = embed_feature(torch.ones(config["shape"]).to(torch.long))
+        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()