Fix numerical embeddings + add tests (#92)

* fix dimension mismatch for numerical embeddings

* add unit tests

inst/python/ResNet.py

@@ -130,9 +130,9 @@ def __init__(self,
                 nn.init.kaiming_uniform_(parameter, a=math.sqrt(5))
 
     def forward(self, input):
-        x = self.weight.unsqueeze(0) * input.unsqueeze(-1)
+        x = self.weight[None] * input[..., None]
         if self.bias is not None:
-            x = x + self.bias.unsqueeze(-1)
+            x = x + self.bias[None]
         return x
 
 

tests/testthat/test-Transformer.R

@@ -126,3 +126,32 @@ test_that("dimHidden ratio works as expected", {
                                         dimHiddenRatio = 4/3))
 
 })
+
+test_that("numerical embedding works as expected", {
+  embeddings <- 32L # size of embeddings
+  features <- 2L # number of numerical features
+  patients <- 9L 
+
+  numTensor <- torch$randn(c(patients, features))
+
+  numericalEmbeddingClass <- reticulate::import_from_path("ResNet", path=path)$NumericalEmbedding
+  numericalEmbedding <- numericalEmbeddingClass(num_embeddings = features,
+                                                embedding_dim = embeddings,
+                                                bias = TRUE)
+  out <- numericalEmbedding(numTensor)
+
+  # should be patients x features x embedding size
+  expect_equal(out$shape[[0]], patients)
+  expect_equal(out$shape[[1]], features)
+  expect_equal(out$shape[[2]], embeddings)
+
+  numericalEmbedding <- numericalEmbeddingClass(num_embeddings = features,
+                                                embedding_dim = embeddings,
+                                                bias = FALSE)
+
+  out <- numericalEmbedding(numTensor)
+  expect_equal(out$shape[[0]], patients)
+  expect_equal(out$shape[[1]], features)
+  expect_equal(out$shape[[2]], embeddings)
+
+  })