Merge pull request #109 from Yoctol/fix-grucell-initial-state

make GRUCell.get_initial_state feedable

talos/layers/__init__.py

@@ -123,7 +123,6 @@
 from tensorflow.python.keras.layers.recurrent import RNN
 from tensorflow.python.keras.layers.recurrent import StackedRNNCells
 from tensorflow.python.keras.layers.recurrent import SimpleRNNCell
-from tensorflow.python.keras.layers.recurrent import GRUCell
 from tensorflow.python.keras.layers.recurrent import LSTMCell
 from tensorflow.python.keras.layers.recurrent import SimpleRNN
 from tensorflow.python.keras.layers.recurrent import GRU
@@ -149,6 +148,7 @@
 from .embeddings import Embedding
 from .layer_norm import LayerNormalization
 from .positional_encode import PositionalEncode
+from .recurrent import GRUCell
 
 from .cudnn_recurrent import CuDNNLSTM
 from .cudnn_recurrent import CuDNNGRU

talos/layers/recurrent.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+
+class GRUCell(tf.keras.layers.GRUCell):
+
+    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+        return [super().get_initial_state(inputs, batch_size, dtype)]

talos/layers/tests/test_recurrent.py

@@ -0,0 +1,9 @@
+import tensorflow as tf
+
+from ..recurrent import GRUCell
+
+
+def test_initial_state():
+    cell = GRUCell(10)
+    x = tf.zeros([3, 4])
+    cell(x, cell.get_initial_state(x))