def testReduceTreeLSTM(self):
with tf.device(self._test_device):
size = 10
tracker_size = 8
reducer = spinn.Reducer(size, tracker_size=tracker_size)
lstm_in = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
[0, -1, -2, -3, -4, -5, -6, -7, -8, -9]],
dtype=np.float32)
c1 = np.array([[0, 1], [2, 3]], dtype=np.float32)
c2 = np.array([[0, -1], [-2, -3]], dtype=np.float32)
h, c = reducer._tree_lstm(c1, c2, lstm_in)
self.assertEqual(tf.float32, h.dtype)
self.assertEqual(tf.float32, c.dtype)
self.assertEqual((2, 2), h.shape)
self.assertEqual((2, 2), c.shape)
def testReducer(self):
with tf.device(self._test_device):
batch_size = 3
size = 10
tracker_size = 8
reducer = spinn.Reducer(size, tracker_size=tracker_size)
left_in = []
right_in = []
tracking = []
for _ in range(batch_size):
left_in.append(tf.random_normal((1, size * 2)))
right_in.append(tf.random_normal((1, size * 2)))
tracking.append(tf.random_normal((1, tracker_size * 2)))
out = reducer(left_in, right_in, tracking=tracking)
self.assertEqual(batch_size, len(out))
self.assertEqual(tf.float32, out[0].dtype)
self.assertEqual((1, size * 2), out[0].shape)
