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::IntArrayRef;
21 using exec_aten::ScalarType;
22 using exec_aten::Tensor;
23 using torch::executor::testing::TensorFactory;
24
25 class OpZerosOutTest : public OperatorTest {
26 protected:
op_zeros_out(IntArrayRef size,Tensor & out)27 Tensor& op_zeros_out(IntArrayRef size, Tensor& out) {
28 return torch::executor::aten::zeros_outf(context_, size, out);
29 }
30
31 template <ScalarType DTYPE>
test_zeros_out(std::vector<int32_t> && size_int32_t)32 void test_zeros_out(std::vector<int32_t>&& size_int32_t) {
33 TensorFactory<DTYPE> tf;
34 std::vector<int64_t> sizes(size_int32_t.begin(), size_int32_t.end());
35 auto aref = exec_aten::ArrayRef<int64_t>(sizes.data(), sizes.size());
36 Tensor out = tf.ones(size_int32_t);
37
38 op_zeros_out(aref, out);
39
40 EXPECT_TENSOR_EQ(out, tf.zeros(size_int32_t));
41 }
42 };
43
44 #define GENERATE_TEST(_, DTYPE) \
45 TEST_F(OpZerosOutTest, DTYPE##Tensors) { \
46 test_zeros_out<ScalarType::DTYPE>({2, 3, 4}); \
47 test_zeros_out<ScalarType::DTYPE>({2, 0, 4}); \
48 test_zeros_out<ScalarType::DTYPE>({}); \
49 }
50
ET_FORALL_REAL_TYPES_AND(Bool,GENERATE_TEST)51 ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_TEST)
52
53 TEST_F(OpZerosOutTest, DynamicShapeUpperBoundSameAsExpected) {
54 TensorFactory<ScalarType::Float> tf;
55 Tensor expected = tf.zeros({3, 2});
56
57 int64_t sizes[2] = {3, 2};
58 auto sizes_aref = exec_aten::ArrayRef<int64_t>(sizes);
59 Tensor out =
60 tf.ones({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
61 op_zeros_out(sizes_aref, out);
62 EXPECT_TENSOR_EQ(out, expected);
63 }
64
TEST_F(OpZerosOutTest,DynamicShapeUpperBoundLargerThanExpected)65 TEST_F(OpZerosOutTest, DynamicShapeUpperBoundLargerThanExpected) {
66 TensorFactory<ScalarType::Float> tf;
67 Tensor expected = tf.zeros({3, 2});
68
69 int64_t sizes[2] = {3, 2};
70 auto sizes_aref = exec_aten::ArrayRef<int64_t>(sizes);
71 Tensor out =
72 tf.ones({10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
73 op_zeros_out(sizes_aref, out);
74 EXPECT_TENSOR_EQ(out, expected);
75 }
76
TEST_F(OpZerosOutTest,DynamicShapeUnbound)77 TEST_F(OpZerosOutTest, DynamicShapeUnbound) {
78 if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
79 GTEST_SKIP() << "Dynamic shape unbound not supported";
80 }
81 TensorFactory<ScalarType::Float> tf;
82 Tensor expected = tf.zeros({3, 2});
83
84 int64_t sizes[2] = {3, 2};
85 auto sizes_aref = exec_aten::ArrayRef<int64_t>(sizes);
86 Tensor out =
87 tf.ones({1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
88 op_zeros_out(sizes_aref, out);
89 EXPECT_TENSOR_EQ(out, expected);
90 }
91