def test_specify_initial_state_keras_tensor(self, layer_class):
input_size = 10
timesteps = 6
units = 2
num_samples = 32
num_states = 2 if layer_class is keras.layers.CuDNNLSTM else 1
inputs = keras.Input((timesteps, input_size))
initial_state = [keras.Input((units, )) for _ in range(num_states)]
layer = layer_class(units)
if len(initial_state) == 1:
output = layer(inputs, initial_state=initial_state[0])
else:
output = layer(inputs, initial_state=initial_state)
self.assertTrue(
any(initial_state[0] is t
for t in layer._inbound_nodes[0].input_tensors))
model = keras.models.Model([inputs] + initial_state, output)
model.compile(loss='categorical_crossentropy',
optimizer=RMSprop(learning_rate=0.001),
run_eagerly=testing_utils.should_run_eagerly())
inputs = np.random.random((num_samples, timesteps, input_size))
initial_state = [
np.random.random((num_samples, units)) for _ in range(num_states)
]
targets = np.random.random((num_samples, units))
model.fit([inputs] + initial_state, targets)
def test_time_major_input(self, layer_class):
input_size = 10
timesteps = 6
units = 2
num_samples = 32
model = keras.models.Sequential()
model.add(keras.layers.Lambda(lambda t: tf.transpose(t, [1, 0, 2])))
layer = layer_class(units, time_major=True, return_sequences=True)
model.add(layer)
model.add(keras.layers.Lambda(lambda t: tf.transpose(t, [1, 0, 2])))
model.compile(loss='categorical_crossentropy',
optimizer=RMSprop(learning_rate=0.001))
model.fit(np.ones((num_samples, timesteps, input_size)),
np.ones((num_samples, timesteps, units)))
out = model.predict(np.ones((num_samples, timesteps, input_size)))
self.assertEqual(out.shape, (num_samples, timesteps, units))