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
16 #include <gtest/gtest.h>
17
18 using namespace ::testing;
19 using exec_aten::ArrayRef;
20 using exec_aten::ScalarType;
21 using exec_aten::Tensor;
22 using torch::executor::testing::TensorFactory;
23
24 class OpTransposeIntCopyTest : public OperatorTest {
25 protected:
op_transpose_copy_int_out(const Tensor & self,int64_t dim0,int64_t dim1,Tensor & out)26 Tensor& op_transpose_copy_int_out(
27 const Tensor& self,
28 int64_t dim0,
29 int64_t dim1,
30 Tensor& out) {
31 return torch::executor::aten::transpose_copy_outf(
32 context_, self, dim0, dim1, out);
33 }
34 };
35
TEST_F(OpTransposeIntCopyTest,TwoDTranspose)36 TEST_F(OpTransposeIntCopyTest, TwoDTranspose) {
37 TensorFactory<ScalarType::Int> tf;
38
39 // clang-format off
40 Tensor t_int = tf.make({2, 3}, {
41 // 2x3 data block
42 0, 1, 2,
43 3, 4, 5
44 });
45 // clang-format on
46
47 const std::vector<int32_t> new_sizes = {3, 2};
48 Tensor out = tf.zeros(new_sizes);
49
50 op_transpose_copy_int_out(t_int, 1, 0, out);
51 // clang-format off
52 EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
53 // 3x2 data block
54 0, 3,
55 1, 4,
56 2, 5
57 }));
58 // clang-format on
59 }
60
TEST_F(OpTransposeIntCopyTest,TwoDNegativeIndices)61 TEST_F(OpTransposeIntCopyTest, TwoDNegativeIndices) {
62 TensorFactory<ScalarType::Int> tf;
63
64 // clang-format off
65 Tensor t_int = tf.make({2, 3}, {
66 // 2x3 data block
67 0, 1, 2,
68 3, 4, 5
69 });
70 // clang-format on
71
72 const std::vector<int32_t> new_sizes = {3, 2};
73 Tensor out = tf.zeros(new_sizes);
74
75 op_transpose_copy_int_out(t_int, -1, -2, out);
76 // clang-format off
77 EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
78 // 3x2 data block
79 0, 3,
80 1, 4,
81 2, 5
82 }));
83 // clang-format on
84 }
85
TEST_F(OpTransposeIntCopyTest,TransposeNoDatachange)86 TEST_F(OpTransposeIntCopyTest, TransposeNoDatachange) {
87 TensorFactory<ScalarType::Int> tf;
88
89 // clang-format off
90 Tensor t_int = tf.make({2, 1, 3}, {
91 // 2 1x3 data blocks
92 0, 1, 2,
93
94 3, 4, 5
95 });
96 // clang-format on
97
98 const std::vector<int32_t> new_sizes = {2, 3, 1};
99 Tensor out = tf.zeros(new_sizes);
100
101 op_transpose_copy_int_out(t_int, 1, 2, out);
102 // clang-format off
103 EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
104 // 2 3x1 data blocks
105 0,
106 1,
107 2,
108
109 3,
110 4,
111 5,
112 }));
113 // clang-format on
114 }
115
TEST_F(OpTransposeIntCopyTest,ThreeDTranspose)116 TEST_F(OpTransposeIntCopyTest, ThreeDTranspose) {
117 TensorFactory<ScalarType::Int> tf;
118
119 // clang-format off
120 Tensor t_int = tf.make({2, 2, 3}, {
121 // 2 2x3 data blocks
122 0, 1, 2,
123 3, 4, 5,
124
125 6, 7, 8,
126 9, 10, 11
127 });
128 // clang-format on
129
130 const std::vector<int32_t> new_sizes = {3, 2, 2};
131 Tensor out = tf.zeros(new_sizes);
132
133 op_transpose_copy_int_out(t_int, 0, 2, out);
134 // clang-format off
135 EXPECT_TENSOR_EQ(out, tf.make(new_sizes, {
136 // 3 2x2 data blocks
137 0, 6,
138 3, 9,
139
140 1, 7,
141 4, 10,
142
143 2, 8,
144 5, 11
145 }));
146 // clang-format on
147 }
148
149 // transpose an out of bounds dim
TEST_F(OpTransposeIntCopyTest,OutOfBoundDimDies)150 TEST_F(OpTransposeIntCopyTest, OutOfBoundDimDies) {
151 TensorFactory<ScalarType::Float> tf;
152
153 Tensor a = tf.ones(/*sizes=*/{2, 3});
154 Tensor out = tf.ones(/*sizes=*/{3, 2});
155
156 ET_EXPECT_KERNEL_FAILURE(context_, op_transpose_copy_int_out(a, 0, -3, out));
157 }
158
159 // transpose a 3d tensor into a 2d one
TEST_F(OpTransposeIntCopyTest,MismatchedDimDies)160 TEST_F(OpTransposeIntCopyTest, MismatchedDimDies) {
161 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
162 GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
163 }
164 TensorFactory<ScalarType::Float> tf;
165
166 Tensor a = tf.ones(/*sizes=*/{4, 2, 3});
167 Tensor out = tf.ones(/*sizes=*/{2, 2});
168
169 ET_EXPECT_KERNEL_FAILURE(context_, op_transpose_copy_int_out(a, 0, 1, out));
170 }
171
172 /* %python
173 import torch
174 torch.manual_seed(0)
175 x = torch.randint(10, (2, 2, 3))
176 res = torch.transpose(x, 0, 2)
177 op = "op_transpose_copy_int_out"
178 opt_extra_params = "0, 2,"
179 dtype = "ScalarType::Int"
180 check = "EXPECT_TENSOR_EQ" */
181
TEST_F(OpTransposeIntCopyTest,DynamicShapeUpperBoundSameAsExpected)182 TEST_F(OpTransposeIntCopyTest, DynamicShapeUpperBoundSameAsExpected) {
183 /* %python
184 out_args = "{3, 2, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
185 %rewrite(unary_op) */
186
187 TensorFactory<ScalarType::Int> tf;
188
189 Tensor x = tf.make({2, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6});
190 Tensor expected = tf.make({3, 2, 2}, {4, 7, 0, 3, 9, 3, 3, 1, 3, 7, 9, 6});
191
192 Tensor out =
193 tf.zeros({3, 2, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
194 op_transpose_copy_int_out(x, 0, 2, out);
195 EXPECT_TENSOR_EQ(out, expected);
196 }
197
TEST_F(OpTransposeIntCopyTest,DynamicShapeUpperBoundLargerThanExpected)198 TEST_F(OpTransposeIntCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
199 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
200 GTEST_SKIP() << "Dynamic shape not supported";
201 }
202 /* %python
203 out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
204 %rewrite(unary_op) */
205
206 TensorFactory<ScalarType::Int> tf;
207
208 Tensor x = tf.make({2, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6});
209 Tensor expected = tf.make({3, 2, 2}, {4, 7, 0, 3, 9, 3, 3, 1, 3, 7, 9, 6});
210
211 Tensor out =
212 tf.zeros({5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
213 op_transpose_copy_int_out(x, 0, 2, out);
214 EXPECT_TENSOR_EQ(out, expected);
215 }
216
TEST_F(OpTransposeIntCopyTest,DynamicShapeUnbound)217 TEST_F(OpTransposeIntCopyTest, DynamicShapeUnbound) {
218 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
219 GTEST_SKIP() << "Dynamic shape not supported";
220 }
221 /* %python
222 out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
223 %rewrite(unary_op) */
224
225 TensorFactory<ScalarType::Int> tf;
226
227 Tensor x = tf.make({2, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6});
228 Tensor expected = tf.make({3, 2, 2}, {4, 7, 0, 3, 9, 3, 3, 1, 3, 7, 9, 6});
229
230 Tensor out = tf.zeros(
231 {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
232 op_transpose_copy_int_out(x, 0, 2, out);
233 EXPECT_TENSOR_EQ(out, expected);
234 }
235