1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <executorch/kernels/portable/cpu/util/reduce_util.h>
10 #include <executorch/runtime/core/exec_aten/exec_aten.h>
11 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
12 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
13 #include <executorch/test/utils/DeathTest.h>
14
15 #include <gtest/gtest.h>
16
17 using namespace ::testing;
18 using exec_aten::ArrayRef;
19 using exec_aten::optional;
20 using exec_aten::ScalarType;
21 using exec_aten::Tensor;
22 using executorch::runtime::testing::TensorFactory;
23 using torch::executor::apply_over_dim;
24 using torch::executor::apply_over_dim_list;
25 using torch::executor::get_out_numel;
26
_apply_over_dim(const Tensor & in,const optional<int64_t> & dim)27 void _apply_over_dim(const Tensor& in, const optional<int64_t>& dim) {
28 int64_t* in_data = in.mutable_data_ptr<int64_t>();
29 for (size_t out_ix = 0; out_ix < get_out_numel(in, dim); ++out_ix) {
30 apply_over_dim(
31 [in_data, out_ix](size_t in_ix, size_t _) { in_data[in_ix] = out_ix; },
32 in,
33 dim,
34 out_ix);
35 }
36 }
37
_apply_over_dim_list(const Tensor & in,const optional<ArrayRef<int64_t>> & dim_list)38 void _apply_over_dim_list(
39 const Tensor& in,
40 const optional<ArrayRef<int64_t>>& dim_list) {
41 int64_t* in_data = in.mutable_data_ptr<int64_t>();
42 for (size_t out_ix = 0; out_ix < get_out_numel(in, dim_list); ++out_ix) {
43 apply_over_dim_list(
44 [in_data, out_ix](size_t in_ix) { in_data[in_ix] = out_ix; },
45 in,
46 dim_list,
47 out_ix);
48 }
49 }
50
TEST(ReduceUtilTest,ApplyOverDim)51 TEST(ReduceUtilTest, ApplyOverDim) {
52 TensorFactory<ScalarType::Long> tf;
53 optional<ArrayRef<int64_t>> dim_list;
54
55 Tensor in = tf.zeros({2, 4, 5, 3});
56 _apply_over_dim(in, 0);
57 // clang-format off
58 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
59 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
60 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
61 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
62 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
63
64 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
65 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
66 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
67 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
68 }));
69 // clang-format on
70
71 in = tf.zeros({2, 4, 5, 3});
72 _apply_over_dim(in, 1);
73 // clang-format off
74 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
75 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
76 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
77 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
78 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
79
80 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
81 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
82 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
83 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
84 }));
85 // clang-format on
86
87 in = tf.zeros({2, 4, 5, 3});
88 int64_t dim_array_2[1] = {2};
89 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_2, 1});
90 _apply_over_dim(in, 2);
91 // clang-format off
92 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
93 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
94 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
95 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8,
96 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11,
97
98 12, 13, 14, 12, 13, 14, 12, 13, 14, 12, 13, 14, 12, 13, 14,
99 15, 16, 17, 15, 16, 17, 15, 16, 17, 15, 16, 17, 15, 16, 17,
100 18, 19, 20, 18, 19, 20, 18, 19, 20, 18, 19, 20, 18, 19, 20,
101 21, 22, 23, 21, 22, 23, 21, 22, 23, 21, 22, 23, 21, 22, 23,
102 }));
103 // clang-format on
104
105 in = tf.zeros({2, 4, 5, 3});
106 int64_t dim_array_3[1] = {3};
107 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_3, 1});
108 _apply_over_dim(in, 3);
109 // clang-format off
110 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
111 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
112 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
113 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14,
114 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
115
116 20, 20, 20, 21, 21, 21, 22, 22, 22, 23, 23, 23, 24, 24, 24,
117 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 29, 29, 29,
118 30, 30, 30, 31, 31, 31, 32, 32, 32, 33, 33, 33, 34, 34, 34,
119 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39,
120 }));
121 // clang-format on
122 }
123
TEST(ReduceUtilTest,ApplyOverDimListNull)124 TEST(ReduceUtilTest, ApplyOverDimListNull) {
125 TensorFactory<ScalarType::Long> tf;
126 optional<ArrayRef<int64_t>> null_dim_list;
127
128 Tensor in = tf.ones({2, 4, 5, 3});
129 _apply_over_dim_list(in, null_dim_list);
130 EXPECT_TENSOR_EQ(in, tf.zeros({2, 4, 5, 3}));
131 }
132
TEST(ReduceUtilTest,ApplyOverZeroDimListEmpty)133 TEST(ReduceUtilTest, ApplyOverZeroDimListEmpty) {
134 TensorFactory<ScalarType::Long> tf;
135 optional<ArrayRef<int64_t>> null_dim_list;
136
137 Tensor in = tf.ones({});
138 _apply_over_dim_list(in, null_dim_list);
139 EXPECT_TENSOR_EQ(in, tf.zeros({}));
140 }
141
TEST(ReduceUtilTest,ApplyOverZeroDim)142 TEST(ReduceUtilTest, ApplyOverZeroDim) {
143 TensorFactory<ScalarType::Long> tf;
144 optional<ArrayRef<int64_t>> dim_list;
145 int64_t dim_array_0[1] = {0};
146 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_0, 1});
147
148 Tensor in = tf.ones({});
149 _apply_over_dim_list(in, dim_list);
150 EXPECT_TENSOR_EQ(in, tf.zeros({}));
151 }
152
TEST(ReduceUtilTest,ApplyOverDimListEmpty)153 TEST(ReduceUtilTest, ApplyOverDimListEmpty) {
154 TensorFactory<ScalarType::Long> tf;
155 optional<ArrayRef<int64_t>> empty_dim_list{ArrayRef<int64_t>{}};
156
157 Tensor in = tf.ones({2, 4, 5, 3});
158 _apply_over_dim_list(in, empty_dim_list);
159 EXPECT_TENSOR_EQ(in, tf.zeros({2, 4, 5, 3}));
160 }
161
TEST(ReduceUtilTest,ApplyOverDimListLength1)162 TEST(ReduceUtilTest, ApplyOverDimListLength1) {
163 TensorFactory<ScalarType::Long> tf;
164 optional<ArrayRef<int64_t>> dim_list;
165
166 Tensor in = tf.zeros({2, 4, 5, 3});
167 int64_t dim_array_0[1] = {0};
168 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_0, 1});
169 _apply_over_dim_list(in, dim_list);
170 // clang-format off
171 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
172 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
173 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
174 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
175 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
176
177 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
178 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
179 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
180 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
181 }));
182 // clang-format on
183
184 in = tf.zeros({2, 4, 5, 3});
185 int64_t dim_array_1[1] = {1};
186 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_1, 1});
187 _apply_over_dim_list(in, dim_list);
188 // clang-format off
189 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
190 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
191 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
192 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
193 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
194
195 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
196 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
197 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
198 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
199 }));
200 // clang-format on
201
202 in = tf.zeros({2, 4, 5, 3});
203 int64_t dim_array_2[1] = {2};
204 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_2, 1});
205 _apply_over_dim_list(in, dim_list);
206 // clang-format off
207 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
208 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
209 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
210 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8,
211 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11,
212
213 12, 13, 14, 12, 13, 14, 12, 13, 14, 12, 13, 14, 12, 13, 14,
214 15, 16, 17, 15, 16, 17, 15, 16, 17, 15, 16, 17, 15, 16, 17,
215 18, 19, 20, 18, 19, 20, 18, 19, 20, 18, 19, 20, 18, 19, 20,
216 21, 22, 23, 21, 22, 23, 21, 22, 23, 21, 22, 23, 21, 22, 23,
217 }));
218 // clang-format on
219
220 in = tf.zeros({2, 4, 5, 3});
221 int64_t dim_array_3[1] = {3};
222 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_3, 1});
223 _apply_over_dim_list(in, dim_list);
224 // clang-format off
225 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
226 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
227 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
228 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14,
229 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
230
231 20, 20, 20, 21, 21, 21, 22, 22, 22, 23, 23, 23, 24, 24, 24,
232 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 29, 29, 29,
233 30, 30, 30, 31, 31, 31, 32, 32, 32, 33, 33, 33, 34, 34, 34,
234 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39,
235 }));
236 // clang-format on
237 }
238
TEST(ReduceUtilTest,ApplyOverDimListLength2)239 TEST(ReduceUtilTest, ApplyOverDimListLength2) {
240 TensorFactory<ScalarType::Long> tf;
241 optional<ArrayRef<int64_t>> dim_list;
242
243 Tensor in = tf.zeros({2, 4, 5, 3});
244 int64_t dim_array_01[2] = {0, 1};
245 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_01, 2});
246 _apply_over_dim_list(in, dim_list);
247 // clang-format off
248 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
249 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
250 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
251 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
252 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
253
254 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
255 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
256 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
257 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
258 }));
259 // clang-format on
260
261 in = tf.zeros({2, 4, 5, 3});
262 int64_t dim_array_02[2] = {0, 2};
263 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_02, 2});
264 _apply_over_dim_list(in, dim_list);
265 // clang-format off
266 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
267 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
268 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
269 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8,
270 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11,
271
272 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
273 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
274 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8, 6, 7, 8,
275 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11, 9, 10, 11,
276 }));
277 // clang-format on
278
279 in = tf.zeros({2, 4, 5, 3});
280 int64_t dim_array_03[2] = {0, 3};
281 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_03, 2});
282 _apply_over_dim_list(in, dim_list);
283 // clang-format off
284 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
285 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
286 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
287 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14,
288 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
289
290 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
291 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
292 10, 10, 10, 11, 11, 11, 12, 12, 12, 13, 13, 13, 14, 14, 14,
293 15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
294 }));
295 // clang-format on
296
297 in = tf.zeros({2, 4, 5, 3});
298 int64_t dim_array_12[2] = {1, 2};
299 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_12, 2});
300 _apply_over_dim_list(in, dim_list);
301 // clang-format off
302 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
303 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
304 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
305 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
306 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
307
308 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
309 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
310 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
311 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5,
312 }));
313 // clang-format on
314
315 in = tf.zeros({2, 4, 5, 3});
316 int64_t dim_array_13[2] = {1, 3};
317 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_13, 2});
318 _apply_over_dim_list(in, dim_list);
319 // clang-format off
320 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
321 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
322 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
323 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
324 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
325
326 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
327 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
328 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
329 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9,
330 }));
331 // clang-format on
332
333 in = tf.zeros({2, 4, 5, 3});
334 int64_t dim_array_23[2] = {2, 3};
335 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_23, 2});
336 _apply_over_dim_list(in, dim_list);
337 // clang-format off
338 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
339 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
340 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
341 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
342 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
343
344 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
345 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
346 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
347 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
348 }));
349 // clang-format on
350 }
351
TEST(ReduceUtilTest,ApplyOverDimListLength3)352 TEST(ReduceUtilTest, ApplyOverDimListLength3) {
353 TensorFactory<ScalarType::Long> tf;
354 optional<ArrayRef<int64_t>> dim_list;
355
356 Tensor in = tf.zeros({2, 4, 5, 3});
357 int64_t dim_array_012[3] = {0, 1, 2};
358 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_012, 3});
359 _apply_over_dim_list(in, dim_list);
360 // clang-format off
361 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
362 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
363 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
364 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
365 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
366
367 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
368 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
369 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
370 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2,
371 }));
372 // clang-format on
373
374 in = tf.zeros({2, 4, 5, 3});
375 int64_t dim_array_013[3] = {0, 1, 3};
376 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_013, 3});
377 _apply_over_dim_list(in, dim_list);
378 // clang-format off
379 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
380 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
381 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
382 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
383 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
384
385 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
386 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
387 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
388 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4,
389 }));
390 // clang-format on
391
392 in = tf.zeros({2, 4, 5, 3});
393 int64_t dim_array_023[3] = {0, 2, 3};
394 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_023, 3});
395 _apply_over_dim_list(in, dim_list);
396 // clang-format off
397 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
398 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
399 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
400 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
401 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
402
403 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
404 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
405 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
406 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
407 }));
408 // clang-format on
409
410 in = tf.zeros({2, 4, 5, 3});
411 int64_t dim_array_123[3] = {1, 2, 3};
412 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_123, 3});
413 _apply_over_dim_list(in, dim_list);
414 // clang-format off
415 EXPECT_TENSOR_EQ(in, tf.make({2, 4, 5, 3}, {
416 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
417 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
418 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
419 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
420
421 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
422 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
423 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
424 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
425 }));
426 // clang-format on
427 }
428
TEST(ReduceUtilTest,ApplyOverDimListLength4)429 TEST(ReduceUtilTest, ApplyOverDimListLength4) {
430 TensorFactory<ScalarType::Long> tf;
431 optional<ArrayRef<int64_t>> dim_list;
432
433 Tensor in = tf.ones({2, 4, 5, 3});
434 int64_t dim_array_0123[4] = {0, 1, 2, 3};
435 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_0123, 4});
436 _apply_over_dim_list(in, dim_list);
437 EXPECT_TENSOR_EQ(in, tf.zeros({2, 4, 5, 3}));
438 }
439
TEST(ReduceUtilTest,ApplyOnZeroDimTensorOverDim)440 TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDim) {
441 TensorFactory<ScalarType::Long> tf;
442
443 Tensor in = tf.ones({});
444 _apply_over_dim(in, 0);
445 EXPECT_TENSOR_EQ(in, tf.make({}, {0}));
446 }
447
TEST(ReduceUtilTest,ApplyOnZeroDimTensorOverDimListNull)448 TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDimListNull) {
449 TensorFactory<ScalarType::Long> tf;
450 optional<ArrayRef<int64_t>> null_dim_list;
451
452 Tensor in = tf.ones({});
453 _apply_over_dim_list(in, null_dim_list);
454 EXPECT_TENSOR_EQ(in, tf.make({}, {0}));
455 }
456
TEST(ReduceUtilTest,ApplyOnZeroDimTensorOverDimListEmpty)457 TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDimListEmpty) {
458 TensorFactory<ScalarType::Long> tf;
459 optional<ArrayRef<int64_t>> empty_dim_list{ArrayRef<int64_t>{}};
460
461 Tensor in = tf.ones({});
462 _apply_over_dim_list(in, empty_dim_list);
463 EXPECT_TENSOR_EQ(in, tf.make({}, {0}));
464 }
465
TEST(ReduceUtilTest,ApplyOnZeroDimTensorOverDimListNonEmpty)466 TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDimListNonEmpty) {
467 TensorFactory<ScalarType::Long> tf;
468 int64_t dim_array_0[1] = {0};
469 optional<ArrayRef<int64_t>> dim_list =
470 optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_0, 1});
471
472 Tensor in = tf.ones({});
473 _apply_over_dim_list(in, dim_list), "";
474 EXPECT_TENSOR_EQ(in, tf.make({}, {0}));
475 }
476
TEST(ReduceUtilTest,ApplyOnEmptyTensorOverDim)477 TEST(ReduceUtilTest, ApplyOnEmptyTensorOverDim) {
478 TensorFactory<ScalarType::Long> tf;
479 optional<ArrayRef<int64_t>> dim_list;
480
481 Tensor in = tf.zeros({2, 0, 5, 3});
482 Tensor out = tf.zeros({2, 5, 3});
483
484 // dim = 1
485 int64_t dim = 1;
486 EXPECT_TRUE(in.numel() == 0);
487 EXPECT_TRUE(out.numel() == 30 && out.numel() == get_out_numel(in, dim));
488
489 int64_t* in_data = in.mutable_data_ptr<int64_t>();
490 int64_t* out_data = out.mutable_data_ptr<int64_t>();
491 for (size_t out_ix = 0; out_ix < get_out_numel(in, dim); ++out_ix) {
492 out_data[out_ix] = 1;
493 apply_over_dim(
494 [in_data, out_data, out_ix](size_t in_ix, size_t _) {
495 in_data[in_ix] = out_ix; // Should be ignored.
496 out_data[out_ix] = 2; // Should be ignored.
497 },
498 in,
499 dim,
500 out_ix);
501 }
502 EXPECT_TENSOR_EQ(out, tf.ones({2, 5, 3}));
503
504 // dim = 0
505 dim = 0;
506 EXPECT_TRUE(in.numel() == 0);
507 EXPECT_TRUE(get_out_numel(in, dim) == 0);
508 // Should die if called on empty tensor with dim that also produces
509 // empty tensor, because out_ix will be out of bounds
510 ET_EXPECT_DEATH(
511 apply_over_dim([](size_t in_ix, size_t _) { return; }, in, dim, 0), "");
512 }
513
TEST(ReduceUtilTest,ApplyOnEmptyTensorOverDimList)514 TEST(ReduceUtilTest, ApplyOnEmptyTensorOverDimList) {
515 TensorFactory<ScalarType::Long> tf;
516 optional<ArrayRef<int64_t>> dim_list;
517
518 Tensor in = tf.zeros({2, 0, 5, 3});
519 Tensor out = tf.zeros({5, 3});
520
521 // dim list = {0, 1}
522 int64_t dim_array_01[2] = {0, 1};
523 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_01, 2});
524
525 EXPECT_TRUE(in.numel() == 0);
526 EXPECT_TRUE(out.numel() == 15 && out.numel() == get_out_numel(in, dim_list));
527
528 int64_t* in_data = in.mutable_data_ptr<int64_t>();
529 int64_t* out_data = out.mutable_data_ptr<int64_t>();
530 for (size_t out_ix = 0; out_ix < get_out_numel(in, dim_list); ++out_ix) {
531 out_data[out_ix] = 1;
532 apply_over_dim_list(
533 [in_data, out_data, out_ix](size_t in_ix) {
534 in_data[in_ix] = out_ix; // Should be ignored.
535 out_data[out_ix] = 2; // Should be ignored.
536 },
537 in,
538 dim_list,
539 out_ix);
540 }
541 EXPECT_TENSOR_EQ(out, tf.ones({5, 3}));
542
543 // dim list = {0, 2}
544 int64_t dim_array_02[2] = {0, 2};
545 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_02, 2});
546
547 EXPECT_TRUE(in.numel() == 0);
548 EXPECT_TRUE(get_out_numel(in, dim_list) == 0);
549 // Should die if called on empty tensor with dim list that also produces
550 // empty tensor, because out_ix will be out of bounds
551 ET_EXPECT_DEATH(
552 apply_over_dim_list([](size_t in_ix) { return; }, in, dim_list, 0), "");
553 }
554
TEST(ReduceUtilTest,ApplyOverDimListInvalid)555 TEST(ReduceUtilTest, ApplyOverDimListInvalid) {
556 TensorFactory<ScalarType::Long> tf;
557 optional<ArrayRef<int64_t>> dim_list;
558
559 Tensor in = tf.zeros({2, 4, 5, 3});
560 int64_t dim_array_09[2] = {0, 9};
561 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_09, 2});
562
563 ET_EXPECT_DEATH(
564 apply_over_dim_list([](size_t in_ix) { return; }, in, dim_list, 0), "");
565
566 int64_t dim_array_neg[3] = {0, -5, 3};
567 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_neg, 3});
568
569 ET_EXPECT_DEATH(
570 apply_over_dim_list([](size_t in_ix) { return; }, in, dim_list, 0), "");
571
572 int64_t dim_array_011[3] = {0, 1, 1};
573 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_011, 3});
574
575 ET_EXPECT_DEATH(
576 apply_over_dim_list([](size_t in_ix) { return; }, in, dim_list, 0), "");
577
578 int64_t dim_array_1_3[2] = {1, -3};
579 dim_list = optional<ArrayRef<int64_t>>(ArrayRef<int64_t>{dim_array_1_3, 2});
580
581 ET_EXPECT_DEATH(
582 apply_over_dim_list([](size_t in_ix) { return; }, in, dim_list, 0), "");
583 }
584