Support constants in StackedRNNCells (#9089)

* Support constants in StackedRNNCells

* Add test

* Add missing test wrappers

Author: myutwo150

keras/layers/recurrent.py

@@ -73,7 +73,7 @@ def state_size(self):
                 state_size.append(cell.state_size)
         return tuple(state_size)
 
-    def call(self, inputs, states, **kwargs):
+    def call(self, inputs, states, constants=None, **kwargs):
         # Recover per-cell states.
         nested_states = []
         for cell in self.cells[::-1]:
@@ -88,7 +88,12 @@ def call(self, inputs, states, **kwargs):
         # Call the cells in order and store the returned states.
         new_nested_states = []
         for cell, states in zip(self.cells, nested_states):
-            inputs, states = cell.call(inputs, states, **kwargs)
+            if has_arg(cell.call, 'constants'):
+                inputs, states = cell.call(inputs, states,
+                                           constants=constants,
+                                           **kwargs)
+            else:
+                inputs, states = cell.call(inputs, states, **kwargs)
             new_nested_states.append(states)
 
         # Format the new states as a flat list
@@ -99,9 +104,15 @@ def call(self, inputs, states, **kwargs):
         return inputs, states
 
     def build(self, input_shape):
+        if isinstance(input_shape, list):
+            constants_shape = input_shape[1:]
+            input_shape = input_shape[0]
         for cell in self.cells:
             if isinstance(cell, Layer):
-                cell.build(input_shape)
+                if has_arg(cell.call, 'constants'):
+                    cell.build([input_shape] + constants_shape)
+                else:
+                    cell.build(input_shape)
             if hasattr(cell.state_size, '__len__'):
                 output_dim = cell.state_size[0]
             else:

tests/keras/layers/recurrent_test.py

@@ -698,6 +698,7 @@ def test_batch_size_equal_one(layer_class):
     model.train_on_batch(x, y)
 
 
+@keras_test
 def test_rnn_cell_with_constants_layer():
 
     class RNNCellWithConstants(keras.layers.Layer):
@@ -778,7 +779,35 @@ def get_config(self):
     y_np_3 = model.predict([x_np, c_np])
     assert_allclose(y_np, y_np_3, atol=1e-4)
 
+    # Test stacking.
+    cells = [recurrent.GRUCell(8),
+             RNNCellWithConstants(12),
+             RNNCellWithConstants(32)]
+    layer = recurrent.RNN(cells)
+    y = layer(x, constants=c)
+    model = keras.models.Model([x, c], y)
+    model.compile(optimizer='rmsprop', loss='mse')
+    model.train_on_batch(
+        [np.zeros((6, 5, 5)), np.zeros((6, 3))],
+        np.zeros((6, 32))
+    )
+
+    # Test stacked RNN serialization.
+    x_np = np.random.random((6, 5, 5))
+    c_np = np.random.random((6, 3))
+    y_np = model.predict([x_np, c_np])
+    weights = model.get_weights()
+    config = layer.get_config()
+    with keras.utils.CustomObjectScope(custom_objects):
+        layer = recurrent.RNN.from_config(config.copy())
+    y = layer(x, constants=c)
+    model = keras.models.Model([x, c], y)
+    model.set_weights(weights)
+    y_np_2 = model.predict([x_np, c_np])
+    assert_allclose(y_np, y_np_2, atol=1e-4)
+
 
+@keras_test
 def test_rnn_cell_with_constants_layer_passing_initial_state():
 
     class RNNCellWithConstants(keras.layers.Layer):