Improve tests for TimeDistributed

Author: fchollet

tests/keras/layers/test_wrappers.py

@@ -12,10 +12,10 @@ def test_TimeDistributed():
     model = Sequential()
     model.add(wrappers.TimeDistributed(core.Dense(2), input_shape=(3, 4)))
     model.add(core.Activation('relu'))
-
     model.compile(optimizer='rmsprop', loss='mse')
     model.fit(np.random.random((10, 3, 4)), np.random.random((10, 3, 2)), nb_epoch=1, batch_size=10)
 
+    # test config
     model.get_config()
 
     # compare to TimeDistributedDense
@@ -26,7 +26,15 @@ def test_TimeDistributed():
     reference = Sequential()
     reference.add(core.TimeDistributedDense(2, input_shape=(3, 4), weights=weights))
     reference.add(core.Activation('relu'))
+    reference.compile(optimizer='rmsprop', loss='mse')
+
+    reference_output = reference.predict(test_input)
+    assert_allclose(test_output, reference_output, atol=1e-05)
 
+    # test when specifying a batch_input_shape
+    reference = Sequential()
+    reference.add(core.TimeDistributedDense(2, batch_input_shape=(1, 3, 4), weights=weights))
+    reference.add(core.Activation('relu'))
     reference.compile(optimizer='rmsprop', loss='mse')
 
     reference_output = reference.predict(test_input)
@@ -36,13 +44,22 @@ def test_TimeDistributed():
     model = Sequential()
     model.add(wrappers.TimeDistributed(convolutional.Convolution2D(5, 2, 2, border_mode='same'), input_shape=(2, 3, 4, 4)))
     model.add(core.Activation('relu'))
-
     model.compile(optimizer='rmsprop', loss='mse')
     model.train_on_batch(np.random.random((1, 2, 3, 4, 4)), np.random.random((1, 2, 5, 4, 4)))
 
     model = model_from_json(model.to_json())
     model.summary()
 
+    # test stacked layers
+    model = Sequential()
+    model.add(wrappers.TimeDistributed(core.Dense(2), input_shape=(3, 4)))
+    model.add(wrappers.TimeDistributed(core.Dense(3)))
+    model.add(core.Activation('relu'))
+    model.compile(optimizer='rmsprop', loss='mse')
+
+    model.fit(np.random.random((10, 3, 4)), np.random.random((10, 3, 3)), nb_epoch=1, batch_size=10)
+
+
 
 if __name__ == '__main__':
     pytest.main([__file__])