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 #include <executorch/test/utils/DeathTest.h>
17 #include <gtest/gtest.h>
18 #include <cmath>
19
20 using namespace ::testing;
21 using exec_aten::ArrayRef;
22 using exec_aten::ScalarType;
23 using exec_aten::Tensor;
24 using torch::executor::testing::TensorFactory;
25
26 class OpAmaxOutTest : public OperatorTest {
27 protected:
op_amax_out(const Tensor & in,ArrayRef<int64_t> dim,bool keepdim,Tensor & out)28 Tensor& op_amax_out(
29 const Tensor& in,
30 ArrayRef<int64_t> dim,
31 bool keepdim,
32 Tensor& out) {
33 return torch::executor::aten::amax_outf(context_, in, dim, keepdim, out);
34 }
35
36 template <ScalarType DTYPE>
test_amax_out_invalid_dimensions()37 void test_amax_out_invalid_dimensions() {
38 TensorFactory<DTYPE> tf;
39
40 // clang-format off
41 Tensor in = tf.make(
42 {2, 3, 4},
43 {
44 0, 1, 2, 4,
45 4, 2, 1, 0,
46 1, 0, 4, 2,
47
48 4, 2, 1, 0,
49 0, 1, 2, 4,
50 1, 0, 4, 2,
51 });
52 // clang-format on
53 Tensor out = tf.zeros({2, 3, 1});
54
55 // out-of-bound dim in dim list
56 int64_t dims_1[1] = {3};
57 ArrayRef<int64_t> dim_list{ArrayRef<int64_t>{dims_1, 1}};
58 ET_EXPECT_KERNEL_FAILURE(
59 context_, op_amax_out(in, dim_list, /*keepdim=*/true, out));
60
61 // the same dim appears multiple times in list of dims
62 int64_t dims_2[2] = {2, 2};
63 dim_list = ArrayRef<int64_t>{dims_2, 2};
64 ET_EXPECT_KERNEL_FAILURE(
65 context_, op_amax_out(in, dim_list, /*keepdim=*/true, out));
66 }
67
68 template <ScalarType DTYPE>
test_amax_out_invalid_shape()69 void test_amax_out_invalid_shape() {
70 TensorFactory<DTYPE> tf;
71
72 // clang-format off
73 Tensor in = tf.make(
74 {2, 3, 4},
75 {
76 0, 1, 2, 4,
77 4, 2, 1, 0,
78 1, 0, 4, 2,
79
80 4, 2, 1, 0,
81 0, 1, 2, 4,
82 1, 0, 4, 2,
83 });
84 // clang-format on
85
86 // dimension size mismatch when keepdim is true
87 Tensor out = tf.zeros({2, 4});
88
89 int64_t dims_1[1] = {1};
90 ArrayRef<int64_t> dim_list{ArrayRef<int64_t>{dims_1, 1}};
91 ET_EXPECT_KERNEL_FAILURE(
92 context_, op_amax_out(in, dim_list, /*keepdim=*/true, out));
93
94 // dimension size mismatch when keepdim is false
95 out = tf.zeros({2, 1, 4});
96 ET_EXPECT_KERNEL_FAILURE(
97 context_, op_amax_out(in, dim_list, /*keepdim=*/false, out));
98 }
99
100 template <ScalarType DTYPE>
test_amax_out_dtype()101 void test_amax_out_dtype() {
102 TensorFactory<DTYPE> tf;
103 // clang-format off
104 Tensor in = tf.make(
105 {2, 3, 4},
106 {
107 0, 1, 2, 4,
108 4, 2, 1, 0,
109 1, 5, 4, 2,
110
111 4, 2, 1, 0,
112 5, 1, 2, 4,
113 7, 5, 4, 2,
114 });
115 // clang-format on
116
117 // keepdim=true should work
118 Tensor out = tf.zeros({2, 3, 1});
119 int64_t dims_1[1] = {2};
120 ArrayRef<int64_t> dim_list{ArrayRef<int64_t>{dims_1, 1}};
121
122 op_amax_out(in, dim_list, /*keepdim=*/true, out);
123 // clang-format off
124 EXPECT_TENSOR_CLOSE(out, tf.make(
125 {2, 3, 1},
126 {4, 4, 5, 4, 5, 7}));
127 // clang-format on
128
129 // keepdim=false should work
130 out = tf.zeros({2, 3});
131 op_amax_out(in, dim_list, /*keepdim=*/false, out);
132 // clang-format off
133 EXPECT_TENSOR_CLOSE(out, tf.make(
134 {2, 3},
135 {4, 4, 5, 4, 5, 7}));
136 // clang-format on
137
138 // dim list with multiple dimensions should work
139 out = tf.zeros({1, 1, 4});
140 int64_t dims_2[2] = {0, 1};
141 dim_list = ArrayRef<int64_t>{dims_2, 2};
142 op_amax_out(in, dim_list, /*keepdim=*/true, out);
143 EXPECT_TENSOR_CLOSE(out, tf.make({1, 1, 4}, {7, 5, 4, 4}));
144
145 out = tf.zeros({4});
146 op_amax_out(in, dim_list, /*keepdim=*/false, out);
147 EXPECT_TENSOR_CLOSE(out, tf.make({4}, {7, 5, 4, 4}));
148
149 // dim list with negative dimensions should work
150 out = tf.zeros({2, 1, 4});
151 int64_t dims_3[1] = {-2};
152 dim_list = ArrayRef<int64_t>{dims_3, 1};
153 op_amax_out(in, dim_list, /*keepdim=*/true, out);
154 // clang-format off
155 EXPECT_TENSOR_CLOSE(out, tf.make(
156 {2, 1, 4},
157 {
158 4, 5, 4, 4,
159
160 7, 5, 4, 4,
161 }));
162 // clang-format on
163
164 // empty/null dim list should work
165 // clang-format off
166 in = tf.make(
167 {2, 2, 4},
168 {
169 8, 7, 5, 4,
170 4, 3, 7, 9,
171
172 4, 2, 6, 8,
173 8, 7, 3, 4,
174 });
175 // clang-format on
176 out = tf.zeros({1, 1, 1});
177 ArrayRef<int64_t> null_dim_list;
178 op_amax_out(in, null_dim_list, /*keepdim=*/true, out);
179 EXPECT_TENSOR_CLOSE(out, tf.make({1, 1, 1}, {9}));
180
181 ArrayRef<int64_t> empty_dim_list{ArrayRef<int64_t>{}};
182 op_amax_out(in, empty_dim_list, /*keepdim=*/true, out);
183 EXPECT_TENSOR_CLOSE(out, tf.make({1, 1, 1}, {9}));
184
185 out = tf.zeros({});
186 op_amax_out(in, null_dim_list, /*keepdim=*/false, out);
187 EXPECT_TENSOR_CLOSE(out, tf.make({}, {9}));
188
189 op_amax_out(in, empty_dim_list, /*keepdim=*/false, out);
190 EXPECT_TENSOR_CLOSE(out, tf.make({}, {9}));
191 }
192 };
193
194 template <>
test_amax_out_dtype()195 void OpAmaxOutTest::test_amax_out_dtype<ScalarType::Bool>() {
196 TensorFactory<ScalarType::Bool> tf_bool;
197 // clang-format off
198 Tensor in = tf_bool.make(
199 {2, 3, 4},
200 {
201 true, false, true, false,
202 false, false, false, false,
203 false, true, true, false,
204
205 false, false, true, false,
206 false, false, false, true,
207 true, true, true, true,
208 });
209 // clang-format on
210
211 Tensor out = tf_bool.zeros({2, 3, 1});
212
213 // +/-inf and nan should work
214 op_amax_out(in, /*dim=*/-1, /*keepdim=*/true, out);
215 // clang-format off
216 EXPECT_TENSOR_CLOSE(
217 out, tf_bool.make(
218 {2, 3, 1},
219 {
220 true,
221 false,
222 true,
223
224 true,
225 true,
226 true
227 }));
228 // clang-format on
229 }
230
TEST_F(OpAmaxOutTest,InvalidDimensionListDies)231 TEST_F(OpAmaxOutTest, InvalidDimensionListDies) {
232 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
233 GTEST_SKIP() << "ATen kernel test fails";
234 }
235 #define TEST_ENTRY(ctype, dtype) \
236 test_amax_out_invalid_dimensions<ScalarType::dtype>();
237 ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
238 #undef TEST_ENTRY
239 }
240
TEST_F(OpAmaxOutTest,InvalidShapeDies)241 TEST_F(OpAmaxOutTest, InvalidShapeDies) {
242 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
243 GTEST_SKIP() << "ATen kernel test fails";
244 }
245 #define TEST_ENTRY(ctype, dtype) \
246 test_amax_out_invalid_shape<ScalarType::dtype>();
247 ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
248 #undef TEST_ENTRY
249 }
250
TEST_F(OpAmaxOutTest,MismatchedDTypesDies)251 TEST_F(OpAmaxOutTest, MismatchedDTypesDies) {
252 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
253 GTEST_SKIP() << "ATen kernel test fails";
254 }
255 TensorFactory<ScalarType::Float> tf_float;
256 TensorFactory<ScalarType::Int> tf_int;
257
258 // clang-format off
259 Tensor in = tf_int.make(
260 {2, 3, 4},
261 {
262 0, 1, 2, 4,
263 4, 2, 1, 0,
264 1, 0, 4, 2,
265
266 4, 2, 1, 0,
267 0, 1, 2, 4,
268 1, 0, 4, 2,
269 });
270 // clang-format on
271
272 Tensor out = tf_float.zeros({2, 3, 1});
273 int64_t dims_1[1] = {2};
274 ArrayRef<int64_t> dim_list{ArrayRef<int64_t>{dims_1, 1}};
275
276 // out tensor should be of the same dtype with dtype when dtype is specified
277 ET_EXPECT_KERNEL_FAILURE(
278 context_, op_amax_out(in, dim_list, /*keepdim=*/true, out));
279 }
280
TEST_F(OpAmaxOutTest,AllRealInputOutputPasses)281 TEST_F(OpAmaxOutTest, AllRealInputOutputPasses) {
282 #define TEST_ENTRY(ctype, dtype) test_amax_out_dtype<ScalarType::dtype>();
283 ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
284 #undef TEST_ENTRY
285 }
286
TEST_F(OpAmaxOutTest,InfinityAndNANTest)287 TEST_F(OpAmaxOutTest, InfinityAndNANTest) {
288 TensorFactory<ScalarType::Float> tf_float;
289 // clang-format off
290 Tensor in = tf_float.make(
291 {2, 3, 4},
292 {
293 0, 1, 2, INFINITY,
294 INFINITY, -INFINITY, 1, 0,
295 NAN, INFINITY, -INFINITY, 2,
296
297 NAN, NAN, 1, 0,
298 0, INFINITY, NAN, 4,
299 1, NAN, 3.14, 2,
300 });
301 // clang-format on
302
303 Tensor out = tf_float.zeros({2, 3, 1});
304 int64_t dims[1] = {-1};
305 ArrayRef<int64_t> dim_list{ArrayRef<int64_t>{dims, 1}};
306 op_amax_out(in, dim_list, /*keepdim=*/true, out);
307 // clang-format off
308 EXPECT_TENSOR_CLOSE(
309 out, tf_float.make({2, 3, 1}, {INFINITY, INFINITY, NAN, NAN, NAN, NAN}));
310 // clang-format on
311 }
312