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/platform/runtime.h>
16
17 #include <gtest/gtest.h>
18
19 using namespace ::testing;
20 using exec_aten::ArrayRef;
21 using exec_aten::optional;
22 using exec_aten::ScalarType;
23 using exec_aten::Tensor;
24 using torch::executor::testing::TensorFactory;
25
26 class OpAnyOutTest : public OperatorTest {
27 protected:
op_any_all_out(const Tensor & input,Tensor & out)28 Tensor& op_any_all_out(const Tensor& input, Tensor& out) {
29 return torch::executor::aten::any_outf(context_, input, out);
30 }
31
op_any_dims_out(const Tensor & input,optional<ArrayRef<int64_t>> dim,bool keepdim,Tensor & out)32 Tensor& op_any_dims_out(
33 const Tensor& input,
34 optional<ArrayRef<int64_t>> dim,
35 bool keepdim,
36 Tensor& out) {
37 return torch::executor::aten::any_outf(context_, input, dim, keepdim, out);
38 }
39
40 Tensor&
op_any_out(const Tensor & input,int64_t dim,bool keepdim,Tensor & out)41 op_any_out(const Tensor& input, int64_t dim, bool keepdim, Tensor& out) {
42 return torch::executor::aten::any_outf(context_, input, dim, keepdim, out);
43 }
44
45 template <ScalarType OUT_DTYPE>
test_any_all_out_invalid_type()46 void test_any_all_out_invalid_type() {
47 TensorFactory<ScalarType::Float> tf_float;
48 TensorFactory<OUT_DTYPE> tf_out;
49
50 Tensor in = tf_float.make(
51 {1, 4},
52 {
53 0,
54 0,
55 1,
56 0,
57 });
58 Tensor out = tf_out.zeros(/*size=*/{0});
59
60 ET_EXPECT_KERNEL_FAILURE(context_, op_any_all_out(in, out));
61 }
62
63 template <ScalarType IN_DTYPE>
test_any_all_out()64 void test_any_all_out() {
65 TensorFactory<IN_DTYPE> tf_in;
66 TensorFactory<ScalarType::Bool> tf_bool;
67 // clang-format off
68 Tensor in = tf_in.make(
69 {2, 4},
70 {
71 0, 1, 0, 1,
72 1, 0, 1, 0
73 });
74 Tensor bool_false_in = tf_bool.make(
75 {2, 4},
76 {
77 false, false, false, false,
78 false, false, false, false,
79 });
80 Tensor bool_true_in = tf_bool.make(
81 {2, 4},
82 {
83 true, true, true, true,
84 true, true, true, true,
85 });
86 // clang-format on
87
88 Tensor out = tf_bool.make({}, {false});
89
90 op_any_all_out(in, out);
91 EXPECT_TENSOR_EQ(out, tf_bool.make({}, {true}));
92
93 op_any_all_out(bool_false_in, out);
94 EXPECT_TENSOR_EQ(out, tf_bool.make({}, {false}));
95
96 op_any_all_out(bool_true_in, out);
97 EXPECT_TENSOR_EQ(out, tf_bool.make({}, {true}));
98 }
99 };
100
TEST_F(OpAnyOutTest,MismatchedDimensionsDies)101 TEST_F(OpAnyOutTest, MismatchedDimensionsDies) {
102 if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
103 GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
104 }
105 TensorFactory<ScalarType::Float> tff;
106 const std::vector<int32_t> size{2, 2};
107
108 Tensor in = tff.make(size, {0, 0, 1, 0});
109 Tensor out = tff.ones(/*size=*/{1, 1});
110
111 ET_EXPECT_KERNEL_FAILURE(context_, op_any_all_out(in, out));
112 }
113
TEST_F(OpAnyOutTest,InvalidDtypeDies)114 TEST_F(OpAnyOutTest, InvalidDtypeDies) {
115 #define TEST_ENTRY(ctype, dtype) \
116 test_any_all_out_invalid_type<ScalarType::dtype>();
117 ET_FORALL_FLOAT_TYPES(TEST_ENTRY);
118 #undef TEST_ENTRY
119 }
120
TEST_F(OpAnyOutTest,AllRealInputTypePasses)121 TEST_F(OpAnyOutTest, AllRealInputTypePasses) {
122 #define TEST_ENTRY(ctype, dtype) test_any_all_out<ScalarType::dtype>();
123 ET_FORALL_REAL_TYPES(TEST_ENTRY);
124 #undef TEST_ENTRY
125 }
126
TEST_F(OpAnyOutTest,SmokeTestDims)127 TEST_F(OpAnyOutTest, SmokeTestDims) {
128 TensorFactory<ScalarType::Bool> tfBool;
129
130 Tensor self = tfBool.make({2, 3, 1}, {true, false, true, true, false, false});
131 int64_t dims[3] = {0, 2};
132 optional<ArrayRef<int64_t>> opt_dim_list{ArrayRef<int64_t>{dims, 2}};
133 bool keepdim = true;
134 Tensor out = tfBool.zeros({1, 3, 1});
135 Tensor out_expected = tfBool.make({1, 3, 1}, {true, false, true});
136 op_any_dims_out(self, opt_dim_list, keepdim, out);
137 EXPECT_TENSOR_CLOSE(out, out_expected);
138 }
139
TEST_F(OpAnyOutTest,SmokeTest)140 TEST_F(OpAnyOutTest, SmokeTest) {
141 TensorFactory<ScalarType::Bool> tfBool;
142
143 Tensor self = tfBool.make({2, 3, 1}, {true, false, true, true, false, false});
144 int64_t dim = 0;
145 bool keepdim = false;
146 Tensor out = tfBool.zeros({3, 1});
147 Tensor out_expected = tfBool.make({3, 1}, {true, false, true});
148 op_any_out(self, dim, keepdim, out);
149 EXPECT_TENSOR_CLOSE(out, out_expected);
150 }
151