[FIX] MLPMultivariate incorrect static_exog parsing (#1170)

Nixtla · Oct 9, 2024 · 170f973 · 170f973
1 parent 8d378c6
commit 170f973
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/nbs/models.mlpmultivariate.ipynb b/nbs/models.mlpmultivariate.ipynb
@@ -210,7 +210,11 @@
     "            x = torch.cat(( x, futr_exog.reshape(batch_size, -1) ), dim=1)\n",
     "\n",
     "        if self.stat_exog_size > 0:\n",
-    "            x = torch.cat(( x, stat_exog.reshape(batch_size, -1) ), dim=1)\n",
+    "            stat_exog = stat_exog.reshape(-1)                   #   [N, S] -> [N * S]\n",
+    "            stat_exog = stat_exog.unsqueeze(0)\\\n",
+    "                                 .repeat(batch_size, \n",
+    "                                         1)                     #   [N * S] -> [B, N * S]            \n",
+    "            x = torch.cat((x, stat_exog), dim=1)\n",
     "\n",
     "        for layer in self.mlp:\n",
     "             x = torch.relu(layer(x))\n",
@@ -362,6 +366,8 @@
     "model = MLPMultivariate(h=12, \n",
     "            input_size=24,\n",
     "            n_series=2,\n",
+    "            stat_exog_list=['airline1'],\n",
+    "            futr_exog_list=['trend'],            \n",
     "            loss = MAE(),\n",
     "            scaler_type='robust',\n",
     "            learning_rate=1e-3,\n",

diff --git a/neuralforecast/models/mlpmultivariate.py b/neuralforecast/models/mlpmultivariate.py
@@ -160,7 +160,11 @@ def forward(self, windows_batch):
             x = torch.cat((x, futr_exog.reshape(batch_size, -1)), dim=1)
 
         if self.stat_exog_size > 0:
-            x = torch.cat((x, stat_exog.reshape(batch_size, -1)), dim=1)
+            stat_exog = stat_exog.reshape(-1)  #   [N, S] -> [N * S]
+            stat_exog = stat_exog.unsqueeze(0).repeat(
+                batch_size, 1
+            )  #   [N * S] -> [B, N * S]
+            x = torch.cat((x, stat_exog), dim=1)
 
         for layer in self.mlp:
             x = torch.relu(layer(x))