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)