| /* |
| * Copyright (c) Meta Platforms, Inc. and affiliates. |
| * All rights reserved. |
| * |
| * This source code is licensed under the BSD-style license found in the |
| * LICENSE file in the root directory of this source tree. |
| */ |
| |
| #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator |
| #include <executorch/kernels/test/TestUtil.h> |
| #include <executorch/kernels/test/supported_features.h> |
| #include <executorch/runtime/core/exec_aten/exec_aten.h> |
| #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h> |
| #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h> |
| |
| #include <gtest/gtest.h> |
| |
| using namespace ::testing; |
| using exec_aten::Scalar; |
| using exec_aten::ScalarType; |
| using exec_aten::Tensor; |
| using torch::executor::testing::TensorFactory; |
| |
| class OpMaskedFillTest : public OperatorTest { |
| protected: |
| Tensor& op_masked_fill_scalar_out( |
| const Tensor& self, |
| const Tensor& mask, |
| const Scalar& value, |
| Tensor& out) { |
| return torch::executor::aten::masked_fill_outf( |
| context_, self, mask, value, out); |
| } |
| |
| // Common testing for masked fill of integer Tensor. |
| template <ScalarType DTYPE> |
| void test_integer_masked_fill_scalar_out() { |
| TensorFactory<DTYPE> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Destination for the masked_fill. |
| Tensor out = tf.zeros(sizes); |
| |
| // Masked fill half of the tensor. |
| op_masked_fill_scalar_out( |
| tf.make(sizes, /*data=*/{23, 29, 31, 37}), |
| tf_bool.make(sizes, /*data=*/{false, true, true, false}), |
| /*value=*/71, |
| out); |
| |
| // Check that it matches the expected output. |
| EXPECT_TENSOR_EQ(out, tf.make(sizes, /*data=*/{23, 71, 71, 37})); |
| } |
| |
| // Common testing for masked fill of floating point Tensor. |
| template <ScalarType DTYPE> |
| void test_floating_point_masked_fill_scalar_out() { |
| TensorFactory<DTYPE> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Destination for the masked_fill. |
| Tensor out = tf.zeros(sizes); |
| |
| // Masked fill half of the tensor. |
| op_masked_fill_scalar_out( |
| tf.make(sizes, /*data=*/{1.1, 2.2, 4.4, 8.8}), |
| tf_bool.make(sizes, /*data=*/{true, false, false, true}), |
| /*value=*/3.3, |
| out); |
| |
| // Check that it matches the expected output. |
| EXPECT_TENSOR_CLOSE(out, tf.make(sizes, /*data=*/{3.3, 2.2, 4.4, 3.3})); |
| } |
| }; |
| |
| TEST_F(OpMaskedFillTest, ByteTensors) { |
| test_integer_masked_fill_scalar_out<ScalarType::Byte>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, CharTensors) { |
| test_integer_masked_fill_scalar_out<ScalarType::Char>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, ShortTensors) { |
| test_integer_masked_fill_scalar_out<ScalarType::Short>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, IntTensors) { |
| test_integer_masked_fill_scalar_out<ScalarType::Int>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, LongTensors) { |
| test_integer_masked_fill_scalar_out<ScalarType::Long>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, IntTensorFloatAlphaDies) { |
| // add_out() doesn't handle floating alpha for intergal inputs |
| TensorFactory<ScalarType::Int> tf; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Destination for the op. |
| Tensor out = tf.zeros(sizes); |
| |
| // Elementwise add operation on two integral tensor with floating alpha |
| // should cause an assertion and kill the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, |
| op_masked_fill_scalar_out( |
| tf.ones(sizes), tf.ones(sizes), /*alpha=*/.7, out)); |
| } |
| |
| TEST_F(OpMaskedFillTest, FloatTensors) { |
| test_floating_point_masked_fill_scalar_out<ScalarType::Float>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, DoubleTensors) { |
| test_floating_point_masked_fill_scalar_out<ScalarType::Double>(); |
| } |
| |
| TEST_F(OpMaskedFillTest, BoolTensors) { |
| TensorFactory<ScalarType::Bool> tf; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Input and mask |
| Tensor self = tf.make(sizes, /*data=*/{false, true, false, true}); |
| Tensor mask = tf.make(sizes, /*data=*/{true, false, true, false}); |
| |
| // Destination for the masked_fill. |
| Tensor out = tf.zeros(sizes); |
| |
| op_masked_fill_scalar_out(self, mask, /*value=*/true, out); |
| // Check that it matches the expected output. |
| EXPECT_TENSOR_CLOSE(out, tf.ones(sizes)); |
| } |
| |
| // The input tensor and value may not have different dtypes. |
| TEST_F(OpMaskedFillTest, MismatchedInputAndValueDtypesDies) { |
| TensorFactory<ScalarType::Byte> tf_byte; |
| TensorFactory<ScalarType::Char> tf_char; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Dummy input and mask value |
| Tensor self = tf_byte.ones(sizes); |
| Tensor mask = tf_char.ones(sizes); |
| |
| // Destination for the fill; matches the type of the input. |
| Tensor out = tf_byte.zeros(sizes); |
| |
| // Filling tensor with mismatched scalar should cause an assertion and kill |
| // the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_masked_fill_scalar_out(self, mask, /*value=*/1.3, out)); |
| } |
| |
| // The output tensor may not have a dtype different from the inputs even if it |
| // has the same shape. |
| TEST_F(OpMaskedFillTest, MismatchedOutputDtypeDies) { |
| // Two different dtypes. This test uses two types with the same size to |
| // demonstrate that the ScalarType itself matters, not the size of the |
| // tensor elements. |
| TensorFactory<ScalarType::Bool> tf_bool; |
| TensorFactory<ScalarType::Byte> tf_byte; |
| TensorFactory<ScalarType::Char> tf_char; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Input and mask |
| Tensor self = tf_byte.ones(sizes); |
| Tensor mask = tf_bool.ones(sizes); |
| |
| // Destination with a dtype then input. |
| Tensor out = tf_char.zeros(sizes); |
| |
| // Filling the tensor into a mismatched output should cause an assertion and |
| // kill the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_masked_fill_scalar_out(self, mask, /*fill=*/0, out)); |
| } |
| // The mask tensor type must be bool, even if shapes are the same |
| TEST_F(OpMaskedFillTest, MismatchedMaskDtypeDies) { |
| TensorFactory<ScalarType::Int> tf; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Input and destination |
| Tensor self = tf.ones(sizes); |
| Tensor out = tf.zeros(sizes); |
| |
| // Mask tensor with non bool dtype |
| Tensor mask = tf.ones(sizes); |
| |
| // Filling the tensor using non boolean mask should cause an assertion and |
| // kill the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_masked_fill_scalar_out(self, mask, /*fill=*/0, out)); |
| } |
| |
| // Mismatched shape tests. |
| TEST_F(OpMaskedFillTest, MismatchedInputShapesDies) { |
| TensorFactory<ScalarType::Int> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| // Input and mask of different shapes that cannot be broadcasted. |
| Tensor self = tf.ones(/*sizes=*/{4}); |
| Tensor mask = tf_bool.ones(/*sizes=*/{2}); |
| |
| // Destination for the sum; matches the shape of one of the inputs. |
| Tensor out = tf.zeros(/*sizes=*/{4}); |
| |
| // Masked fill with mismatch input and mask shapes should cause an assertion |
| // and kill the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_masked_fill_scalar_out(self, mask, /*value=*/0, out)); |
| } |
| |
| TEST_F(OpMaskedFillTest, BroadcastTest) { |
| TensorFactory<ScalarType::Int> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| // Input and mask of different shapes |
| Tensor self = tf.make({2, 2}, /*data=*/{1, 2, 4, 8}); |
| Tensor mask = tf_bool.make({2}, /*data=*/{true, false}); |
| |
| // Destination for the masked_fill. |
| Tensor out = tf.zeros({2, 2}); |
| |
| // Masked fill half of the tensor. |
| op_masked_fill_scalar_out( |
| self, |
| mask, |
| /*value=*/3, |
| out); |
| |
| // Check that it matches the expected output. |
| EXPECT_TENSOR_CLOSE(out, tf.make({2, 2}, /*data=*/{3, 2, 3, 8})); |
| } |
| |
| TEST_F(OpMaskedFillTest, MismatchedOutputShapesDies) { |
| if (torch::executor::testing::SupportedFeatures::get()->is_aten) { |
| GTEST_SKIP() << "ATen kernel can handle mismatched output shape"; |
| } |
| TensorFactory<ScalarType::Int> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| const std::vector<int32_t> sizes = {2, 2}; |
| |
| // Input and mask of different shapes |
| Tensor a = tf.ones(sizes); |
| Tensor b = tf_bool.ones(sizes); |
| |
| // Destination with a different shape. |
| Tensor out = tf.zeros(/*sizes=*/{4}); |
| |
| // Mask filling the tensor into a mismatched output should cause an assertion |
| // and kill the test process. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_masked_fill_scalar_out(a, b, /*value=*/0, out)); |
| } |
| |
| TEST_F(OpMaskedFillTest, BroadcastDimSizeIsOneAB) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> bool_tf; |
| |
| Tensor x = tf.make( |
| {3, 2}, |
| {0.9701170325279236, |
| 0.4185227155685425, |
| 0.39851099252700806, |
| 0.8725584745407104, |
| 0.714692234992981, |
| 0.3167606592178345}); |
| Tensor y = bool_tf.make({1, 2}, {false, false}); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.9701170325279236, |
| 0.4185227155685425, |
| 0.39851099252700806, |
| 0.8725584745407104, |
| 0.714692234992981, |
| 0.3167606592178345}); |
| |
| Tensor out = tf.zeros({3, 2}); |
| Tensor ret = op_masked_fill_scalar_out(x, y, Scalar(3.0), out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, BroadcastDimSizeMissingAB) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> bool_tf; |
| |
| Tensor x = tf.make( |
| {3, 2}, |
| {0.9701170325279236, |
| 0.4185227155685425, |
| 0.39851099252700806, |
| 0.8725584745407104, |
| 0.714692234992981, |
| 0.3167606592178345}); |
| Tensor y = bool_tf.make({2}, {false, false}); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.9701170325279236, |
| 0.4185227155685425, |
| 0.39851099252700806, |
| 0.8725584745407104, |
| 0.714692234992981, |
| 0.3167606592178345}); |
| |
| Tensor out = tf.zeros({3, 2}); |
| Tensor ret = op_masked_fill_scalar_out(x, y, Scalar(3.0), out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, DynamicShapeUpperBoundSameAsExpected) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> bool_tf; |
| |
| Tensor x = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 0.49675875902175903}); |
| Tensor y = bool_tf.make({3, 2}, {false, false, false, false, false, true}); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 3.0}); |
| |
| Tensor out = |
| tf.zeros({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); |
| Tensor ret = op_masked_fill_scalar_out(x, y, Scalar(3.0), out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, DynamicShapeUpperBoundLargerThanExpected) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> bool_tf; |
| |
| Tensor x = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 0.49675875902175903}); |
| Tensor y = bool_tf.make({3, 2}, {false, false, false, false, false, true}); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 3.0}); |
| |
| Tensor out = |
| tf.zeros({6, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); |
| Tensor ret = op_masked_fill_scalar_out(x, y, Scalar(3.0), out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, DynamicShapeUnbound) { |
| GTEST_SKIP() << "Dynamic shape unbound not supported"; |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> bool_tf; |
| |
| Tensor x = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 0.49675875902175903}); |
| Tensor y = bool_tf.make({3, 2}, {false, false, false, false, false, true}); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.974706768989563, |
| 0.46383917331695557, |
| 0.050839245319366455, |
| 0.26296138763427734, |
| 0.8404526114463806, |
| 3.0}); |
| |
| Tensor out = |
| tf.zeros({1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); |
| Tensor ret = op_masked_fill_scalar_out(x, y, Scalar(3.0), out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, BroadcastDimSizeIsOneBA) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| auto x = tf.make( |
| {3, 2}, |
| {0.38566190004348755, |
| 0.47776442766189575, |
| 0.1954779028892517, |
| 0.6691004633903503, |
| 0.6580829620361328, |
| 0.48968571424484253}); |
| auto y = tf_bool.make({2}, {false, false}); |
| auto z = Scalar(3.0); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.38566190004348755, |
| 0.47776442766189575, |
| 0.1954779028892517, |
| 0.6691004633903503, |
| 0.6580829620361328, |
| 0.48968571424484253}); |
| |
| Tensor out = tf.zeros({3, 2}); |
| Tensor ret = op_masked_fill_scalar_out(x, y, z, out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
| |
| TEST_F(OpMaskedFillTest, BroadcastDimSizeMissingBA) { |
| TensorFactory<ScalarType::Float> tf; |
| TensorFactory<ScalarType::Bool> tf_bool; |
| |
| auto x = tf.make( |
| {3, 2}, |
| {0.38566190004348755, |
| 0.47776442766189575, |
| 0.1954779028892517, |
| 0.6691004633903503, |
| 0.6580829620361328, |
| 0.48968571424484253}); |
| auto y = tf_bool.make({2}, {false, false}); |
| auto z = Scalar(3.0); |
| Tensor expected_result = tf.make( |
| {3, 2}, |
| {0.38566190004348755, |
| 0.47776442766189575, |
| 0.1954779028892517, |
| 0.6691004633903503, |
| 0.6580829620361328, |
| 0.48968571424484253}); |
| |
| Tensor out = tf.zeros({3, 2}); |
| Tensor ret = op_masked_fill_scalar_out(x, y, z, out); |
| EXPECT_TENSOR_CLOSE(out, expected_result); |
| } |
