From 2dc63d9ae07719cb9819a9a797b38972f09e07f3 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Mon, 9 Dec 2019 16:20:42 +0100
Subject: [PATCH] fix hidden state for LSTM and GRU

---
 pts/model/deepar/deepar_network.py | 26 +++++++++++++++-----------
 pts/modules/distribution_output.py |  1 -
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py
index 0ad7595..bfd54cf 100644
--- a/pts/model/deepar/deepar_network.py
+++ b/pts/model/deepar/deepar_network.py
@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -140,7 +140,7 @@ class DeepARNetwork(nn.Module):
                 torch.Tensor]=None,  # (batch_size, prediction_length, num_features)
             future_target: Optional[
                 torch.Tensor]=None,  # (batch_size, prediction_length, *target_shape)
-    ) -> Tuple[torch.Tensor, List, torch.Tensor, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, List], torch.Tensor, torch.Tensor]:
         
         if future_time_feat is None or future_target is None:
             time_feat = past_time_feat[:,self.history_length - self.context_length:,...]
@@ -202,7 +202,7 @@ class DeepARNetwork(nn.Module):
         outputs, state = self.rnn(inputs)
 
         # outputs: (batch_size, seq_len, num_cells)
-        # state: list of (batch_size, num_cells) tensors
+        # state: list of (num_layers, batch_size, num_cells) tensors
         # scale: (batch_size, 1, *target_shape)
         # static_feat: (batch_size, num_features + prod(target_shape))
         return outputs, state, scale, static_feat
@@ -297,7 +297,7 @@ class DeepARPredictionNetwork(DeepARNetwork):
         past_target: torch.Tensor,
         time_feat: torch.Tensor,
         scale: torch.Tensor,
-        begin_states: List,
+        begin_states: Union[torch.Tensor, List],
     ) -> torch.Tensor:
         """
         Computes sample paths by unrolling the LSTM starting with a initial
@@ -313,9 +313,9 @@ class DeepARPredictionNetwork(DeepARNetwork):
             time features. Shape: (batch_size, prediction_length, num_time_features).
         scale : Tensor
             tensor containing the scale of each element in the batch. Shape: (batch_size, 1, 1).
-        begin_states : List
-            list of initial states for the LSTM layers.
-            the shape of each tensor of the list should be (batch_size, num_cells)
+        begin_states : List or Tensor
+            list of initial states for the LSTM layers or tensor for GRU.
+            the shape of each tensor of the list should be (num_layers, batch_size, num_cells)
         Returns
         --------
         Tensor
@@ -335,10 +335,14 @@ class DeepARPredictionNetwork(DeepARNetwork):
         repeated_scale = scale.repeat_interleave(
             repeats=self.num_parallel_samples, dim=0
         )
-        repeated_states = [
-            s.repeat_interleave(repeats=self.num_parallel_samples, dim=1)
-            for s in begin_states
-        ]
+        if self.cell_type == 'LSTM':
+            repeated_states = [
+                s.repeat_interleave(repeats=self.num_parallel_samples, dim=1)
+                for s in begin_states
+            ]
+        else:
+            repeated_states = begin_states.repeat_interleave(
+                repeats=self.num_parallel_samples, dim=1)
 
         future_samples = []
 
diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py
index 3859797..3396d6c 100644
--- a/pts/modules/distribution_output.py
+++ b/pts/modules/distribution_output.py
@@ -62,7 +62,6 @@ class Output(ABC):
 
 
 class DistributionOutput(Output):
-
     distr_cls: type
 
     def distribution(