def test_gather_tree_from_array_complex_trajectory(self): # Max. time = 7, batch = 1, beam = 5. array = np.expand_dims( np.array([[[25, 12, 114, 89, 97]], [[9, 91, 64, 11, 162]], [[34, 34, 34, 34, 34]], [[2, 4, 2, 2, 4]], [[2, 3, 6, 2, 2]], [[2, 2, 2, 3, 2]], [[2, 2, 2, 2, 2]]]), -1) parent_ids = np.array([[[0, 0, 0, 0, 0]], [[0, 0, 0, 0, 0]], [[0, 1, 2, 3, 4]], [[0, 0, 1, 2, 1]], [[0, 1, 1, 2, 3]], [[0, 1, 3, 1, 2]], [[0, 1, 2, 3, 4]]]) expected_array = np.expand_dims( np.array([[[25, 25, 25, 25, 25]], [[9, 9, 91, 9, 9]], [[34, 34, 34, 34, 34]], [[2, 4, 2, 4, 4]], [[2, 3, 6, 3, 6]], [[2, 2, 2, 3, 2]], [[2, 2, 2, 2, 2]]]), -1) sequence_length = [[4, 6, 4, 7, 6]] array = tf.convert_to_tensor(array, dtype=tf.float32) parent_ids = tf.convert_to_tensor(parent_ids, dtype=tf.int32) expected_array = tf.convert_to_tensor(expected_array, dtype=tf.float32) sorted_array = beam_search_decoder.gather_tree_from_array( array, parent_ids, sequence_length) with self.cached_session() as sess: sorted_array, expected_array = sess.run( [sorted_array, expected_array]) self.assertAllEqual(expected_array, sorted_array)
def test_gather_tree_from_array_complex_trajectory(): # Max. time = 7, batch = 1, beam = 5. array = np.expand_dims( np.array( [ [[25, 12, 114, 89, 97]], [[9, 91, 64, 11, 162]], [[34, 34, 34, 34, 34]], [[2, 4, 2, 2, 4]], [[2, 3, 6, 2, 2]], [[2, 2, 2, 3, 2]], [[2, 2, 2, 2, 2]], ] ), -1, ) parent_ids = np.array( [ [[0, 0, 0, 0, 0]], [[0, 0, 0, 0, 0]], [[0, 1, 2, 3, 4]], [[0, 0, 1, 2, 1]], [[0, 1, 1, 2, 3]], [[0, 1, 3, 1, 2]], [[0, 1, 2, 3, 4]], ] ) expected_array = np.expand_dims( np.array( [ [[25, 25, 25, 25, 25]], [[9, 9, 91, 9, 9]], [[34, 34, 34, 34, 34]], [[2, 4, 2, 4, 4]], [[2, 3, 6, 3, 6]], [[2, 2, 2, 3, 2]], [[2, 2, 2, 2, 2]], ] ), -1, ) sequence_length = [[4, 6, 4, 7, 6]] array = tf.convert_to_tensor(array, dtype=tf.float32) parent_ids = tf.convert_to_tensor(parent_ids, dtype=tf.int32) expected_array = tf.convert_to_tensor(expected_array, dtype=tf.float32) sorted_array = beam_search_decoder.gather_tree_from_array( array, parent_ids, sequence_length ) np.testing.assert_equal(expected_array.numpy(), sorted_array.numpy())
def _test_gather_tree_from_array(depth_ndims=0, merged_batch_beam=False): array = np.array( [ [[1, 2, 3], [4, 5, 6], [7, 8, 9], [0, 0, 0]], [[2, 3, 4], [5, 6, 7], [8, 9, 10], [11, 12, 0]], ] ).transpose([1, 0, 2]) parent_ids = np.array( [ [[0, 0, 0], [0, 1, 1], [2, 1, 2], [-1, -1, -1]], [[0, 0, 0], [1, 1, 0], [2, 0, 1], [0, 1, 0]], ] ).transpose([1, 0, 2]) expected_array = np.array( [ [[2, 2, 2], [6, 5, 6], [7, 8, 9], [0, 0, 0]], [[2, 3, 2], [7, 5, 7], [8, 9, 8], [11, 12, 0]], ] ).transpose([1, 0, 2]) sequence_length = [[3, 3, 3], [4, 4, 3]] array = tf.convert_to_tensor(array, dtype=tf.float32) parent_ids = tf.convert_to_tensor(parent_ids, dtype=tf.int32) expected_array = tf.convert_to_tensor(expected_array, dtype=tf.float32) max_time = tf.shape(array)[0] batch_size = tf.shape(array)[1] beam_width = tf.shape(array)[2] def _tile_in_depth(tensor): # Generate higher rank tensors by concatenating tensor and # tensor + 1. for _ in range(depth_ndims): tensor = tf.stack([tensor, tensor + 1], -1) return tensor if merged_batch_beam: array = tf.reshape(array, [max_time, batch_size * beam_width]) expected_array = tf.reshape(expected_array, [max_time, batch_size * beam_width]) if depth_ndims > 0: array = _tile_in_depth(array) expected_array = _tile_in_depth(expected_array) sorted_array = beam_search_decoder.gather_tree_from_array( array, parent_ids, sequence_length ) np.testing.assert_equal(expected_array.numpy(), sorted_array.numpy())