def _testACT(self, input_size, hidden_size, output_size, seq_len,
batch_size, core, get_state_for_halting):
threshold = 0.99
act = pondering_rnn.ACTCore(core, output_size, threshold,
get_state_for_halting)
seq_input = tf.random_uniform(shape=(seq_len, batch_size, input_size))
initial_state = core.initial_state(batch_size)
seq_output = tf.nn.dynamic_rnn(act,
seq_input,
time_major=True,
initial_state=initial_state)
for tensor in nest.flatten(seq_output):
self.assertEqual(seq_input.dtype, tensor.dtype)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(seq_output)
(final_out, (iteration, r_t)), final_cumul_state = output
self.assertEqual((seq_len, batch_size, output_size), final_out.shape)
self.assertEqual((seq_len, batch_size, 1), iteration.shape)
self.assertTrue(np.all(iteration == np.floor(iteration)))
state_shape = get_state_for_halting(
initial_state).get_shape().as_list()
self.assertEqual(tuple(state_shape),
get_state_for_halting(final_cumul_state).shape)
self.assertEqual((seq_len, batch_size, 1), r_t.shape)
self.assertTrue(np.all(r_t >= 0))
self.assertTrue(np.all(r_t <= threshold))
def testOutput2D(self):
core = Output2DCore(name="output_2d_core")
err = "Output of core should be 1D."
with self.assertRaisesRegexp(ValueError, err):
pondering_rnn.ACTCore(core, 1, 0.99, lambda state: state)
def testOutputTuple(self):
core = OutputTupleCore(name="output_tuple_core")
err = "Output of core should be single Tensor."
with self.assertRaisesRegexp(ValueError, err):
pondering_rnn.ACTCore(core, 1, 0.99, lambda state: state)
