@@ -18,7 +18,11 @@ def test_lod_tensor_to_array_level_0(self):
1818 tensor .set_lod ([[0 , 3 , 9 , 10 ]])
1919 expect = map (lambda x : numpy .array (x ).astype ('int32' ),
2020 [[3 , 0 , 9 ], [4 , 1 ], [5 , 2 ], [6 ], [7 ], [8 ]])
21- self .main (tensor = tensor , expect_array = expect , expect_lod = [] * 6 )
21+ self .main (
22+ tensor = tensor ,
23+ expect_array = expect ,
24+ expect_lod = [] * 6 ,
25+ expect_max_len = 6 )
2226
2327 def test_lod_tensor_to_array_level_0_empty_seq (self ):
2428 tensor = core .LoDTensor ()
@@ -27,7 +31,11 @@ def test_lod_tensor_to_array_level_0_empty_seq(self):
2731 tensor .set_lod ([[0 , 3 , 9 , 9 , 10 ]])
2832 expect = map (lambda x : numpy .array (x ).astype ('int32' ),
2933 [[3 , 0 , 9 ], [4 , 1 ], [5 , 2 ], [6 ], [7 ], [8 ]])
30- self .main (tensor = tensor , expect_array = expect , expect_lod = [] * 6 )
34+ self .main (
35+ tensor = tensor ,
36+ expect_array = expect ,
37+ expect_lod = [] * 6 ,
38+ expect_max_len = 6 )
3139
3240 def test_lod_tensor_to_array_level_1 (self ):
3341 tensor = core .LoDTensor ()
@@ -44,7 +52,11 @@ def test_lod_tensor_to_array_level_1(self):
4452 ]
4553
4654 lod = [[[0 , 2 , 5 ]], [[0 , 6 , 12 ]], [[0 , 3 ]]]
47- self .main (tensor = tensor , expect_array = expect , expect_lod = lod )
55+ self .main (
56+ tensor = tensor ,
57+ expect_array = expect ,
58+ expect_lod = lod ,
59+ expect_max_len = 3 )
4860
4961 def test_lod_tensor_to_array_level_1_empty_seq (self ):
5062 tensor = core .LoDTensor ()
@@ -63,7 +75,11 @@ def test_lod_tensor_to_array_level_1_empty_seq(self):
6375 ]
6476
6577 lod = [[[0 , 5 , 8 , 8 , 15 ]], [[0 , 2 , 6 , 7 , 8 ]], [[0 , 2 , 6 ]], [[0 , 2 ]]]
66- self .main (tensor = tensor , expect_array = expect , expect_lod = lod )
78+ self .main (
79+ tensor = tensor ,
80+ expect_array = expect ,
81+ expect_lod = lod ,
82+ expect_max_len = 4 )
6783
6884 def test_lod_tensor_to_array_level_2 (self ):
6985 tensor = core .LoDTensor ()
@@ -80,22 +96,33 @@ def test_lod_tensor_to_array_level_2(self):
8096 ]
8197 lod = [[[0 , 1 , 3 , 4 ], [0 , 1 , 4 , 8 , 12 ]],
8298 [[0 , 4 , 7 ], [0 , 1 , 5 , 9 , 17 , 21 , 27 , 31 ]], [[0 , 2 ], [0 , 6 , 7 ]]]
83- self .main (tensor = tensor , expect_array = expect , expect_lod = lod )
99+ self .main (
100+ tensor = tensor ,
101+ expect_array = expect ,
102+ expect_lod = lod ,
103+ expect_max_len = 3 )
84104
85105 def test_lod_tensor_to_array_level_2_skip_level (self ):
86106 tensor = core .LoDTensor ()
87107 tensor .set (
88108 numpy .arange (50 ).reshape (50 , 1 ).astype ('int32' ), self .place ())
89109 tensor .set_lod ([[0 , 2 , 5 , 6 ], [0 , 2 , 5 , 6 , 10 , 12 , 13 ],
90110 [0 , 3 , 7 , 11 , 17 , 21 , 22 , 23 , 27 , 31 , 39 , 45 , 46 , 50 ]])
91- self .main (tensor = tensor , expect_array = None , expect_lod = None , level = 1 )
92-
93- def main (self , tensor , expect_array , expect_lod , level = 0 ):
111+ self .main (
112+ tensor = tensor ,
113+ expect_array = None ,
114+ expect_lod = None ,
115+ expect_max_len = 4 ,
116+ level = 1 )
117+
118+ def main (self , tensor , expect_array , expect_lod , expect_max_len , level = 0 ):
94119 place = self .place ()
95120 program = Program ()
96121 x = layers .data (name = 'x' , shape = [10 ], main_program = program )
97122 x .persistable = True
98123 table = layers .lod_rank_table (x , level = level , main_program = program )
124+ max_len = layers .max_sequence_len (table , main_program = program )
125+ max_len .persistable = True
99126 array = layers .lod_tensor_to_array (x , table , main_program = program )
100127 array .persistable = True
101128
@@ -110,6 +137,10 @@ def main(self, tensor, expect_array, expect_lod, level=0):
110137 self .check_array_same (array , expect_array , expect_lod )
111138 self .check_tensor_same (scope .find_var (result .name ).get_tensor (), tensor )
112139
140+ self .assertEqual (
141+ numpy .array (scope .find_var (max_len .name ).get_tensor ())[0 ],
142+ expect_max_len )
143+
113144 def check_array_same (self , array , expect_tensor , expect_lod ):
114145 self .assertEqual (len (expect_tensor ), len (array ))
115146 for i , exp in enumerate (zip (expect_tensor , expect_lod )):
0 commit comments