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::ScalarType;
20 using exec_aten::Tensor;
21 using torch::executor::testing::SupportedFeatures;
22 using torch::executor::testing::TensorFactory;
23
24 class OpMaskedScatterOutTest : public OperatorTest {
25 protected:
op_masked_scatter_out(const Tensor & in,const Tensor & mask,const Tensor & src,Tensor & out)26 Tensor& op_masked_scatter_out(
27 const Tensor& in,
28 const Tensor& mask,
29 const Tensor& src,
30 Tensor& out) {
31 return torch::executor::aten::masked_scatter_outf(
32 context_, in, mask, src, out);
33 }
34 };
35
TEST_F(OpMaskedScatterOutTest,SmokeTest)36 TEST_F(OpMaskedScatterOutTest, SmokeTest) {
37 TensorFactory<ScalarType::Int> tf;
38 TensorFactory<ScalarType::Bool> tfBool;
39
40 Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
41 Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
42 Tensor src = tf.make({3}, {10, 20, 30});
43
44 Tensor out = tf.zeros({2, 3});
45
46 op_masked_scatter_out(in, mask, src, out);
47 EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
48 }
49
TEST_F(OpMaskedScatterOutTest,BroadcastInput)50 TEST_F(OpMaskedScatterOutTest, BroadcastInput) {
51 TensorFactory<ScalarType::Int> tf;
52 TensorFactory<ScalarType::Bool> tfBool;
53
54 Tensor in = tf.make({3}, {1, 2, 3});
55 Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
56 Tensor src = tf.make({3}, {10, 20, 30});
57
58 Tensor out = tf.zeros({2, 3});
59
60 op_masked_scatter_out(in, mask, src, out);
61 EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 2, 30}));
62 }
63
TEST_F(OpMaskedScatterOutTest,BroadcastMask)64 TEST_F(OpMaskedScatterOutTest, BroadcastMask) {
65 TensorFactory<ScalarType::Int> tf;
66 TensorFactory<ScalarType::Bool> tfBool;
67
68 Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
69 Tensor mask = tfBool.make({3}, {false, true, false});
70 Tensor src = tf.make({2}, {10, 20});
71
72 Tensor out = tf.zeros({2, 3});
73
74 op_masked_scatter_out(in, mask, src, out);
75 EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {1, 10, 3, 4, 20, 6}));
76 }
77
TEST_F(OpMaskedScatterOutTest,SrcWithMoreElements)78 TEST_F(OpMaskedScatterOutTest, SrcWithMoreElements) {
79 TensorFactory<ScalarType::Int> tf;
80 TensorFactory<ScalarType::Bool> tfBool;
81
82 Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
83 Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
84 Tensor src = tf.make({4}, {10, 20, 30, 40});
85
86 Tensor out = tf.zeros({2, 3});
87
88 op_masked_scatter_out(in, mask, src, out);
89 EXPECT_TENSOR_EQ(out, tf.make({2, 3}, {10, 2, 3, 20, 5, 30}));
90 }
91
TEST_F(OpMaskedScatterOutTest,SrcWithLessElementsFails)92 TEST_F(OpMaskedScatterOutTest, SrcWithLessElementsFails) {
93 TensorFactory<ScalarType::Int> tf;
94 TensorFactory<ScalarType::Bool> tfBool;
95
96 Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
97 Tensor mask = tfBool.make({2, 3}, {true, false, false, true, false, true});
98 Tensor src = tf.make({2}, {10, 20});
99
100 Tensor out = tf.zeros({2, 3});
101
102 ET_EXPECT_KERNEL_FAILURE(context_, op_masked_scatter_out(in, mask, src, out));
103 }
104
TEST_F(OpMaskedScatterOutTest,EmptyMask)105 TEST_F(OpMaskedScatterOutTest, EmptyMask) {
106 TensorFactory<ScalarType::Int> tf;
107 TensorFactory<ScalarType::Bool> tfBool;
108
109 Tensor in = tf.make({2, 1}, {100, 200});
110 Tensor mask = tfBool.make({2, 0}, {});
111 Tensor src = tf.make({4}, {10, 20, 30, 40});
112
113 Tensor out = tf.zeros({2, 0});
114
115 op_masked_scatter_out(in, mask, src, out);
116 EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
117 }
118
TEST_F(OpMaskedScatterOutTest,EmptySrc)119 TEST_F(OpMaskedScatterOutTest, EmptySrc) {
120 TensorFactory<ScalarType::Int> tf;
121 TensorFactory<ScalarType::Bool> tfBool;
122
123 Tensor in = tf.make({2, 1}, {100, 200});
124 Tensor mask = tfBool.make({2, 1}, {false, false});
125 Tensor src = tf.make({0}, {});
126
127 Tensor out = tf.zeros({2, 1});
128
129 op_masked_scatter_out(in, mask, src, out);
130 EXPECT_TENSOR_EQ(out, tf.make({2, 1}, {100, 200}));
131 }
132
TEST_F(OpMaskedScatterOutTest,EmptyMaskAndSrc)133 TEST_F(OpMaskedScatterOutTest, EmptyMaskAndSrc) {
134 TensorFactory<ScalarType::Int> tf;
135 TensorFactory<ScalarType::Bool> tfBool;
136
137 Tensor in = tf.make({2, 1}, {100, 200});
138 Tensor mask = tfBool.make({0}, {});
139 Tensor src = tf.make({0}, {});
140
141 Tensor out = tf.zeros({2, 0});
142
143 op_masked_scatter_out(in, mask, src, out);
144 EXPECT_TENSOR_EQ(out, tf.make({2, 0}, {}));
145 }
146