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/test/FunctionHeaderWrapper.h> // Declares the operator
10 #include <executorch/kernels/test/TestUtil.h>
11 #include <executorch/kernels/test/supported_features.h>
12 #include <executorch/runtime/core/exec_aten/exec_aten.h>
13 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
14 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
15 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
16
17 #include <gtest/gtest.h>
18
19 using namespace ::testing;
20 using exec_aten::ArrayRef;
21 using exec_aten::ScalarType;
22 using exec_aten::Tensor;
23 using exec_aten::TensorList;
24 using torch::executor::testing::TensorFactory;
25 using torch::executor::testing::TensorListFactory;
26
27 class OpUnbindCopyIntOutTest : public OperatorTest {
28 protected:
op_unbind_copy_int_out(const Tensor & self,int64_t dim,TensorList out)29 void op_unbind_copy_int_out(const Tensor& self, int64_t dim, TensorList out) {
30 return torch::executor::aten::unbind_copy_outf(context_, self, dim, out);
31 }
32
33 template <ScalarType DTYPE>
make1x2x3(TensorFactory<DTYPE> & tf)34 Tensor make1x2x3(TensorFactory<DTYPE>& tf) {
35 // clang-format off
36 return tf.make(
37 /*sizes=*/{1, 2, 3},
38 /*data=*/
39 {
40 0, 1, 2, // tensor([[[ 0, 1, 2],
41 3, 4, 5, // [ 3, 4, 5]]])
42 });
43 // clang-format on
44 }
45
46 template <ScalarType DTYPE>
test_unbind_dim0()47 void test_unbind_dim0() {
48 TensorFactory<DTYPE> tf;
49 TensorListFactory<DTYPE> tlf;
50
51 // clang-format off
52 std::vector<Tensor> expected_out = {
53 tf.make(
54 /*sizes=*/{2, 3},
55 /*data=*/
56 {
57 0, 1, 2, // tensor([[ 0, 1, 2],
58 3, 4, 5, // [ 3, 4, 5]])
59 }),
60 };
61 // clang-format on
62
63 Tensor input = make1x2x3(tf);
64
65 // Output list with the same shapes/dtypes as the expected outputs.
66 TensorList out = tlf.zeros_like(expected_out);
67
68 op_unbind_copy_int_out(input, /*dim=*/0, out);
69
70 EXPECT_TENSOR_LISTS_EQ(expected_out, out);
71
72 // Also show that python negative indexing works for this case.
73 TensorList out2 = tlf.zeros_like(expected_out);
74 op_unbind_copy_int_out(input, /*dim=*/-3, out2);
75 EXPECT_TENSOR_LISTS_EQ(expected_out, out2);
76 }
77
78 template <ScalarType DTYPE>
test_unbind_dim1()79 void test_unbind_dim1() {
80 TensorFactory<DTYPE> tf;
81 TensorListFactory<DTYPE> tlf;
82
83 // clang-format off
84 std::vector<Tensor> expected_out = {
85 tf.make(
86 /*sizes=*/{1, 3},
87 /*data=*/
88 {
89 0, 1, 2, // tensor([[ 0, 1, 2]])
90 }),
91 tf.make(
92 /*sizes=*/{1, 3},
93 /*data=*/
94 {
95 3, 4, 5, // tensor([[ 3, 4, 5]])
96 }),
97 };
98 // clang-format on
99
100 Tensor input = make1x2x3(tf);
101
102 // Output list with the same shapes/dtypes as the expected outputs.
103 TensorList out = tlf.zeros_like(expected_out);
104
105 op_unbind_copy_int_out(input, /*dim=*/1, out);
106
107 EXPECT_TENSOR_LISTS_EQ(expected_out, out);
108
109 // Also show that python negative indexing works for this case.
110 TensorList out2 = tlf.zeros_like(expected_out);
111 op_unbind_copy_int_out(input, /*dim=*/-2, out2);
112 EXPECT_TENSOR_LISTS_EQ(expected_out, out2);
113 }
114
115 template <ScalarType DTYPE>
test_unbind_dim2()116 void test_unbind_dim2() {
117 TensorFactory<DTYPE> tf;
118 TensorListFactory<DTYPE> tlf;
119
120 // Splitting on dim=N with split_size=2 will produce a list of tensors where
121 // the max dim[N] is 2, and the other dims are the same as the input.
122
123 // clang-format off
124 std::vector<Tensor> expected_out = {
125 tf.make(
126 /*sizes=*/{1, 2},
127 /*data=*/
128 {
129 0, // tensor([[ 0,
130 3, // 3]]),
131 }),
132 tf.make(
133 /*sizes=*/{1, 2},
134 /*data=*/
135 {
136 1, // tensor([[ 1,
137 4, // 4]]),
138 }),
139 tf.make(
140 /*sizes=*/{1, 2},
141 /*data=*/
142 {
143 2, // tensor([[ 2,
144 5, // 5]]),
145 }),
146 };
147 // clang-format on
148
149 Tensor input = make1x2x3(tf);
150
151 // Output list with the same shapes/dtypes as the expected outputs.
152 TensorList out = tlf.zeros_like(expected_out);
153
154 op_unbind_copy_int_out(input, /*dim=*/2, out);
155
156 EXPECT_TENSOR_LISTS_EQ(expected_out, out);
157
158 // Also show that python negative indexing works for this case.
159 TensorList out2 = tlf.zeros_like(expected_out);
160 op_unbind_copy_int_out(input, /*dim=*/-1, out2);
161 EXPECT_TENSOR_LISTS_EQ(expected_out, out2);
162 }
163
164 /* %python
165 import torch
166 torch.manual_seed(0)
167 x = torch.randint(10, (2, 3, 4))
168 res = torch.unbind(x, 1)
169 op = "op_unbind_copy_int_out"
170 opt_extra_params = "1,"
171 out_args = [
172 "out_shape, dynamism",
173 "out_shape, dynamism",
174 "out_shape, dynamism"
175 ]
176 dtype = "ScalarType::Int"
177 check = "EXPECT_TENSOR_LISTS_EQ" */
178
test_dynamic_shape(const std::vector<int32_t> & out_shape,enum torch::executor::TensorShapeDynamism dynamism)179 void test_dynamic_shape(
180 const std::vector<int32_t>& out_shape,
181 enum torch::executor::TensorShapeDynamism dynamism) {
182 /* %python
183 %rewrite(unary_op_tensor_list_out) */
184
185 TensorFactory<ScalarType::Int> tf;
186
187 Tensor x = tf.make({2, 3, 4}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6,
188 6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4});
189 std::vector<Tensor> expectedv = {
190 tf.make({2, 4}, {4, 9, 3, 0, 6, 9, 8, 6}),
191 tf.make({2, 4}, {3, 9, 7, 3, 6, 8, 4, 3}),
192 tf.make({2, 4}, {7, 3, 1, 6, 6, 9, 1, 4})};
193 TensorList expected(expectedv.data(), expectedv.size());
194
195 std::vector<Tensor> outv = {
196 tf.zeros(out_shape, dynamism),
197 tf.zeros(out_shape, dynamism),
198 tf.zeros(out_shape, dynamism)};
199 TensorList out(outv.data(), outv.size());
200 op_unbind_copy_int_out(x, 1, out);
201 EXPECT_TENSOR_LISTS_EQ(out, expected);
202 }
203 };
204
205 /**
206 * Returns a 1x2x3 contiguous tensor where the underlying data counts from 0 to
207 * 26.
208 */
TEST_F(OpUnbindCopyIntOutTest,Unbind1x2x3OnDim0AllRealDtypes)209 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim0AllRealDtypes) {
210 #define TEST_ENTRY(ctype, dtype) test_unbind_dim0<ScalarType::dtype>();
211 ET_FORALL_REAL_TYPES(TEST_ENTRY);
212 #undef TEST_ENTRY
213 }
214
TEST_F(OpUnbindCopyIntOutTest,Unbind1x2x3OnDim1AllRealDTypes)215 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim1AllRealDTypes) {
216 #define TEST_ENTRY(ctype, dtype) test_unbind_dim1<ScalarType::dtype>();
217 ET_FORALL_REAL_TYPES(TEST_ENTRY);
218 #undef TEST_ENTRY
219 }
220
TEST_F(OpUnbindCopyIntOutTest,Unbind1x2x3OnDim2AllRealDTypes)221 TEST_F(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim2AllRealDTypes) {
222 #define TEST_ENTRY(ctype, dtype) test_unbind_dim2<ScalarType::dtype>();
223 ET_FORALL_REAL_TYPES(TEST_ENTRY);
224 #undef TEST_ENTRY
225 }
226
TEST_F(OpUnbindCopyIntOutTest,ZeroDimensionalInputTensorDies)227 TEST_F(OpUnbindCopyIntOutTest, ZeroDimensionalInputTensorDies) {
228 TensorFactory<ScalarType::Int> tf;
229 TensorListFactory<ScalarType::Int> tlf;
230
231 Tensor input = tf.ones(/*sizes=*/{});
232 // Arbitrary output shape since this input can't be split.
233 TensorList out = tlf.zeros_like({input});
234
235 ET_EXPECT_KERNEL_FAILURE(
236 context_, op_unbind_copy_int_out(input, /*dim=*/0, out));
237 }
238
TEST_F(OpUnbindCopyIntOutTest,UnbindWorksWithZeroSizedTensors)239 TEST_F(OpUnbindCopyIntOutTest, UnbindWorksWithZeroSizedTensors) {
240 TensorFactory<ScalarType::Int> tf;
241 TensorListFactory<ScalarType::Int> tlf;
242
243 Tensor input = tf.ones(/*sizes=*/{1, 0, 2});
244 EXPECT_EQ(input.numel(), 0);
245
246 // unbind dim 0
247 std::vector<Tensor> expected_out = {tf.ones({0, 2})};
248
249 TensorList out = tlf.zeros_like(expected_out);
250
251 op_unbind_copy_int_out(input, /*dim=*/0, out);
252 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
253
254 // unbind dim 1
255 expected_out = {};
256
257 out = tlf.zeros_like(expected_out);
258
259 op_unbind_copy_int_out(input, /*dim=*/1, out);
260 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
261
262 // unbind dim 2
263 expected_out = {tf.ones({1, 0}), tf.ones({1, 0})};
264
265 out = tlf.zeros_like(expected_out);
266
267 op_unbind_copy_int_out(input, /*dim=*/2, out);
268 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
269 }
270
TEST_F(OpUnbindCopyIntOutTest,UnbindFailsWithWronglyAllocatedOutput)271 TEST_F(OpUnbindCopyIntOutTest, UnbindFailsWithWronglyAllocatedOutput) {
272 TensorFactory<ScalarType::Int> tf;
273 TensorListFactory<ScalarType::Int> tlf;
274
275 Tensor input = tf.ones(/*sizes=*/{1, 2, 3});
276
277 // unbind dim 1
278 std::vector<Tensor> expected_out = {tf.ones({1, 3})};
279
280 TensorList out = tlf.zeros_like(expected_out);
281
282 // Die because length of the list should be 2
283 ET_EXPECT_KERNEL_FAILURE(
284 context_, op_unbind_copy_int_out(input, /*dim=*/1, out));
285
286 expected_out = {tf.ones({1, 4}), tf.ones({1, 4})};
287
288 out = tlf.zeros_like(expected_out);
289
290 // Die because output tensors in the list should be of correct sizes
291 ET_EXPECT_KERNEL_FAILURE(
292 context_, op_unbind_copy_int_out(input, /*dim=*/1, out));
293
294 expected_out = {tf.ones({1}), tf.ones({1})};
295
296 out = tlf.zeros_like(expected_out);
297
298 // Die because output tensors in the list should have correct number of dims
299 ET_EXPECT_KERNEL_FAILURE(
300 context_, op_unbind_copy_int_out(input, /*dim=*/1, out));
301 }
302
TEST_F(OpUnbindCopyIntOutTest,UnbindProduceScalarTensors)303 TEST_F(OpUnbindCopyIntOutTest, UnbindProduceScalarTensors) {
304 TensorFactory<ScalarType::Int> tf;
305 TensorListFactory<ScalarType::Int> tlf;
306
307 Tensor input = tf.make(
308 /*sizes=*/{3}, {4, 5, 6});
309
310 // unbind dim 1
311 std::vector<Tensor> expected_out = {
312 tf.make({}, {4}),
313 tf.make({}, {5}),
314 tf.make({}, {6}),
315 };
316
317 TensorList out = tlf.zeros_like(expected_out);
318
319 op_unbind_copy_int_out(input, /*dim=*/0, out);
320 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
321 }
322
TEST_F(OpUnbindCopyIntOutTest,UnbindProduceScalarLikeTensors)323 TEST_F(OpUnbindCopyIntOutTest, UnbindProduceScalarLikeTensors) {
324 TensorFactory<ScalarType::Int> tf;
325 TensorListFactory<ScalarType::Int> tlf;
326
327 Tensor input = tf.make(
328 /*sizes=*/{3, 1}, {4, 5, 6});
329
330 // unbind dim 0
331 std::vector<Tensor> expected_out = {
332 tf.make({1}, {4}),
333 tf.make({1}, {5}),
334 tf.make({1}, {6}),
335 };
336
337 TensorList out = tlf.zeros_like(expected_out);
338
339 op_unbind_copy_int_out(input, /*dim=*/0, out);
340 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
341
342 input = tf.make(
343 /*sizes=*/{1, 3}, {4, 5, 6});
344
345 // unbind dim 1
346 expected_out = {
347 tf.make({1}, {4}),
348 tf.make({1}, {5}),
349 tf.make({1}, {6}),
350 };
351
352 out = tlf.zeros_like(expected_out);
353
354 op_unbind_copy_int_out(input, /*dim=*/1, out);
355 EXPECT_TENSOR_LISTS_EQ(out, expected_out);
356 }
357
TEST_F(OpUnbindCopyIntOutTest,DynamicShapeUpperBoundSameAsExpected)358 TEST_F(OpUnbindCopyIntOutTest, DynamicShapeUpperBoundSameAsExpected) {
359 test_dynamic_shape(
360 {2, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
361 }
362
TEST_F(OpUnbindCopyIntOutTest,DynamicShapeUpperBoundLargerThanExpected)363 TEST_F(OpUnbindCopyIntOutTest, DynamicShapeUpperBoundLargerThanExpected) {
364 GTEST_SKIP() << "Dynamic shape not supported";
365 test_dynamic_shape(
366 {10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
367 }
368
TEST_F(OpUnbindCopyIntOutTest,DynamicShapeUnbound)369 TEST_F(OpUnbindCopyIntOutTest, DynamicShapeUnbound) {
370 GTEST_SKIP() << "Dynamic shape not supported";
371 test_dynamic_shape(
372 {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
373 }
374