RNN.call should get initial state from full input spec

keras/layers/recurrent.py

@@ -551,6 +551,14 @@ def call(self,
         # note that the .build() method of subclasses MUST define
         # self.input_spec and self.state_spec with complete input shapes.
         if isinstance(inputs, list):
+            # get initial_state from full input spec
+            # as they could be copied to multiple GPU.
+            if self._num_constants is None:
+                initial_state = inputs[1:]
+            else:
+                initial_state = inputs[1:-self._num_constants]
+            if len(initial_state) == 0:
+                initial_state = None
             inputs = inputs[0]
         if initial_state is not None:
             pass

keras/layers/wrappers.py

@@ -496,10 +496,27 @@ def call(self,
             kwargs['constants'] = constants
 
         if initial_state is not None and has_arg(self.layer.call, 'initial_state'):
-            forward_state = initial_state[:len(initial_state) // 2]
-            backward_state = initial_state[len(initial_state) // 2:]
-            y = self.forward_layer.call(inputs, initial_state=forward_state, **kwargs)
-            y_rev = self.backward_layer.call(inputs, initial_state=backward_state, **kwargs)
+            forward_inputs = [inputs[0]]
+            backward_inputs = [inputs[0]]
+            pivot = len(initial_state) // 2 + 1
+            # add forward initial state
+            forward_state = inputs[1:pivot]
+            forward_inputs += forward_state
+            if self._num_constants is None:
+                # add backward initial state
+                backward_state = inputs[pivot:]
+                backward_inputs += backward_state
+            else:
+                # add backward initial state
+                backward_state = inputs[pivot:-self._num_constants]
+                backward_inputs += backward_state
+                # add constants for forward and backward layers
+                forward_inputs += inputs[-self._num_constants:]
+                backward_inputs += inputs[-self._num_constants:]
+            y = self.forward_layer.call(forward_inputs,
+                                        initial_state=forward_state, **kwargs)
+            y_rev = self.backward_layer.call(backward_inputs,
+                                             initial_state=backward_state, **kwargs)
         else:
             y = self.forward_layer.call(inputs, **kwargs)
             y_rev = self.backward_layer.call(inputs, **kwargs)

tests/keras/utils/multi_gpu_test.py

@@ -272,5 +272,19 @@ def multi_gpu_application_folder_generator_benchmark():
         print('%d gpus training:' % i, total_time)
 
 
+@keras_test
+def test_multi_gpu_with_multi_input_layers():
+    inputs = keras.Input((4, 3))
+    init_state = keras.Input((3,))
+    outputs = keras.layers.SimpleRNN(
+        3, return_sequences=True)(inputs, initial_state=init_state)
+    x = [np.random.randn(2, 4, 3), np.random.randn(2, 3)]
+    y = np.random.randn(2, 4, 3)
+    model = keras.models.Model([inputs, init_state], outputs)
+    parallel_model = multi_gpu_model(model, 2)
+    parallel_model.compile(loss='mean_squared_error', optimizer='adam')
+    parallel_model.train_on_batch(x, y)
+
+
 if __name__ == '__main__':
     pytest.main([__file__])