/*
 * 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 <executorch/runtime/core/exec_aten/util/scalar_type_util.h>

#include <gtest/gtest.h>

using namespace ::testing;
using exec_aten::ArrayRef;
using exec_aten::optional;
using exec_aten::ScalarType;
using exec_aten::Tensor;
using torch::executor::testing::TensorFactory;

class OpSliceScatterTensorOutTest : public OperatorTest {
 protected:
  Tensor& op_slice_scatter_out(
      const Tensor& self,
      const Tensor& src,
      int64_t dim,
      optional<int64_t> start,
      optional<int64_t> end,
      int64_t step,
      Tensor& out) {
    return torch::executor::aten::slice_scatter_outf(
        context_, self, src, dim, start, end, step, out);
  }

  template <class CTYPE, exec_aten::ScalarType DTYPE>
  void test_dtype() {
    TensorFactory<DTYPE> tf;

    // clang-format off
    Tensor input = tf.make(
      /*sizes=*/{3, 4},
      /*data=*/{
        1,   2,   3,   4, // [0, :]
        5,   6,   7,   8, // [1, :]
        9,  10,  11,  12, // [2, :]
      });

    // op_slice_scatter_out(input, src, /*dim=*/0, /*start=*/0, /*end=*/2, /*step=*/1, out),
    // src shape should equal to input[0:2:1, :]
    Tensor src = tf.make(
      /*sizes=*/{2, 4},
      /*data=*/{
        5,   6,   7,   8, // [0, :]
        1,   2,   3,   4, // [1, :]
      });
    Tensor expect_ret = tf.make(
      /*sizes=*/{3, 4},
      /*data=*/{
        5,   6,   7,   8, // [0, :]
        1,   2,   3,   4, // [1, :]
        9,  10,  11,  12, // [2, :]
      });
    // clang-format on

    Tensor out = tf.zeros({3, 4});
    Tensor ret = op_slice_scatter_out(
        input, src, /*dim=*/0, /*start=*/0, /*end=*/2, /*step=*/1, out);

    EXPECT_TENSOR_EQ(out, ret);
    EXPECT_TENSOR_EQ(ret, expect_ret);
  }
};

TEST_F(OpSliceScatterTensorOutTest, LegalDimSupported) {
  TensorFactory<ScalarType::Double> tf;

  // clang-format off
  Tensor input = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // clang-format on

  // clang-format off
  // The size of the src tensor should follow these rules:
  // - src.size(i) shall equal input.size(i) if i != dim,
  // - src.size(i) shall equal num_values if i == dim
  //   The definition of num_values could be found at https://fburl.com/code/mnnxkowm

  // op_slice_scatter_out(input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/1, out),
  // src shape should equal to input[0:1:1，:, :]
  Tensor src_dim_0 = tf.make(
    /*sizes=*/{1, 3, 4},
    /*data=*/{
      8.,   7.,   6.,   5., // [1, :]
      4.,   3.,   2.,   1., // [0, :]
      1.,  14.,  18.,  19., // [2, :]
    });
  Tensor expected_dim_0 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       8.,   7.,   6.,   5., // [0, 1, :]
       4.,   3.,   2.,   1., // [0, 0, :]
       1.,  14.,  18.,  19., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/1, /*step=*/1, out),
  // src shape should equal to input[:，0:1:1, :]
  Tensor src_dim_1 = tf.make(
    /*sizes=*/{2, 1, 4},
    /*data=*/{
       4.,   3.,   2.,   1., // [0, :, :]
      -4.,  -3.,  -2.,  -1., // [1, :, :]
    });
  Tensor expected_dim_1 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       4.,   3.,   2.,   1., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -4.,  -3.,  -2.,  -1., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/2, /*start=*/0, /*end=*/1, /*step=*/1, out),
  // src shape should equal to input[:，:, 0:1:1]
  Tensor src_dim_2 = tf.make(
    /*sizes=*/{2, 3, 1},
    /*data=*/{
       7.,   1.,   6., // [0, :, :]
      -5.,  -9.,  -2., // [1, :, :]
    });
  Tensor expected_dim_2 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       7.,   2.,   3.,   4., // [0, 0, :]
       1.,   6.,   7.,   8., // [0, 1, :]
       6.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -5.,  -2.,  -3.,  -4., // [1, 0, :]
      -9.,  -6.,  -7.,  -8., // [1, 1, :]
      -2., -10., -11., -12., // [1, 2, :]
    });
  // clang-format on
  std::vector<Tensor> src_tensors = {
      // Source tensor for dim=-3
      src_dim_0,
      // Source tensor for dim=-2
      src_dim_1,
      // Source tensor for dim=-1
      src_dim_2,
      // Source tensor for dim=0
      src_dim_0,
      // Source tensor for dim=1
      src_dim_1,
      // Source tensor for dim=2
      src_dim_2,
  };
  std::vector<Tensor> expected_rets = {
      // Ground truth for dim=-3
      expected_dim_0,
      // Ground truth for dim=-2
      expected_dim_1,
      // Ground truth for dim=-1
      expected_dim_2,
      // Ground truth for dim=0
      expected_dim_0,
      // Ground truth for dim=1
      expected_dim_1,
      // Ground truth for dim=2
      expected_dim_2,
  };

  for (int64_t dim = -3; dim < 3; dim++) {
    int64_t testcase_idx = dim + 3;
    auto src = src_tensors[testcase_idx];
    auto expected_ret = expected_rets[testcase_idx];

    Tensor out = tf.zeros_like(expected_ret);

    // Slice input on dim with start=0, end = 0 and step = 1
    // Should always return the provided out Tensor.
    // The ret shall meet the expectation.
    Tensor ret = op_slice_scatter_out(
        input, src, dim, /*start=*/0, /*end=*/1, /*step=*/1, out);
    EXPECT_TENSOR_EQ(out, ret);
    EXPECT_TENSOR_EQ(ret, expected_rets[testcase_idx]);
  }
}

TEST_F(OpSliceScatterTensorOutTest, AllStartValsSupported) {
  TensorFactory<ScalarType::Double> tf;

  // clang-format off
  Tensor input = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // clang-format on

  // clang-format off
  // Set the end large enough to hold any start

  // The size of the src tensor should follow these rules:
  // - src.size(i) shall equal input.size(i) if i != dim,
  // - src.size(i) shall equal num_values if i == dim
  //   The definition of num_values could be found at https://fburl.com/code/mnnxkowm

  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/ <= 0, /*end=*/10, /*step=*/1, out),
  // src shape shall equal to input[:，0:3:1, :]
  Tensor src_start_0_or_below = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  Tensor expected_start_0_or_below = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/1, /*end=*/10, /*step=*/1, out),
  // src shape shall equal to input[:，1:3:1, :]
  Tensor src_start_1 = tf.make(
    /*sizes=*/{2, 2, 4},
    /*data=*/{
      // [0, :, :]
      -9., -10., -11., -12., // [0, 1, :]
      -5.,  -6.,  -7.,  -8., // [0, 0, :]

      // [1, :, :]
       9.,  10.,  11.,  12., // [1, 1, :]
       5.,   6.,   7.,   8., // [1, 0, :]
    });
  Tensor expected_start_1 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
      -9., -10., -11., -12., // [0, 1, :]
      -5.,  -6.,  -7.,  -8., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
       9.,  10.,  11.,  12., // [1, 1, :]
       5.,   6.,   7.,   8., // [1, 0, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/2, /*end=*/10, /*step=*/1, out),
  // src shape shall equal to input[:，2:3:1, :] = input
  Tensor src_start_2 = tf.make(
    /*sizes=*/{2, 1, 4},
    /*data=*/{
       1.,  19.,  18.,  17., // [0, 0, :]
      -1., -19., -18., -17., // [1, 0, :]
    });
  Tensor expected_start_2 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       1.,  19.,  18.,  17., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -1., -19., -18., -17., // [1, 2, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/ > input.size(1) = 2, /*end=*/10, /*step=*/1, out),
  // src_shape shall equal to input[:, 3:3:1, :], which is an empty tensor
  Tensor src_start_3_or_above = tf.make({2, 0, 4}, {});
  Tensor expected_start_3_or_above = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      1.,   2.,   3.,   4., // [0, 0, :]
      5.,   6.,   7.,   8., // [0, 1, :]
      9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // clang-format on

  std::vector<Tensor> src_tensors = {// start = -3
                                     src_start_0_or_below,
                                     // start = -2
                                     src_start_1,
                                     // start = -1
                                     src_start_2,
                                     // start = 0
                                     src_start_0_or_below,
                                     // start = 1
                                     src_start_1,
                                     // start = 2
                                     src_start_2,
                                     // start = 3
                                     src_start_3_or_above};
  std::vector<Tensor> expected_rets = {// start = -3
                                       expected_start_0_or_below,
                                       // start = -2
                                       expected_start_1,
                                       // start = -1
                                       expected_start_2,
                                       // start = 0
                                       expected_start_0_or_below,
                                       // start = 1
                                       expected_start_1,
                                       // start = 2
                                       expected_start_2,
                                       // start = 3
                                       expected_start_3_or_above};

  // In this test, we maintain dim and step as 1 and 1, also set the end
  // large enough to hold any start
  int64_t dim = 1;
  int64_t end = 10;
  int64_t step = 1;
  for (int64_t start = -3; start < 4; start++) {
    int64_t testcase_idx = start + 3;
    auto src = src_tensors[testcase_idx];
    auto expected_ret = expected_rets[testcase_idx];
    Tensor out = tf.zeros_like(expected_ret);

    // Should always return the provided out Tensor.
    // The ret shall meet the expectation.
    Tensor ret = op_slice_scatter_out(input, src, dim, start, end, step, out);
    EXPECT_TENSOR_EQ(out, ret);
    EXPECT_TENSOR_EQ(ret, expected_ret);
  }
}

TEST_F(OpSliceScatterTensorOutTest, AllEndValsSupported) {
  TensorFactory<ScalarType::Double> tf;

  // clang-format off
  Tensor input = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });

  // The size of expected output tensor should follow these rules:
  // - output.size(i) shall equal input.size(i) if i != dim,
  // - output.size(i) shall equal num_values if i == dim
  //   The definition of num_values could be found at https://fburl.com/code/mnnxkowm

  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/ <= 0, /*step=*/1, out),
  // src shape should equal input[:，0:0:1, :], which should be an empty tensor
  Tensor src_end_0_or_below = tf.make({2, 0, 4}, {});
  Tensor expected_end_0_or_below = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });

  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/1, /*step=*/1, out),
  // src shape should equal to input[:，0:1:1, :]
  Tensor src_end_1 = tf.make(
    /*sizes=*/{2, 1, 4},
    /*data=*/{
      -4.,  -3.,  -2.,  -1., // [0, :, :]
       4.,   3.,   2.,   1., // [1, :, :]
    });
  Tensor expected_end_1 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -4.,  -3.,  -2.,  -1., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
       4.,   3.,   2.,   1., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });

  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/2, /*step=*/1, out),
  // src shape should equal input[:，0:2:1, :]
  Tensor src_end_2 = tf.make(
    /*sizes=*/{2, 2, 4},
    /*data=*/{
      // [0, :, :]
      -8.,  -7.,  -6.,  -5., // [0, 0, :]
      -4.,  -3.,  -2.,  -1., // [0, :, :]

      // [1, :, :]
       8.,   7.,   6.,   5., // [1, 0, :]
       4.,   3.,   2.,   1., // [1, 1, :]
    });
  Tensor expected_end_2 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -8.,  -7.,  -6.,  -5., // [0, 0, :]
      -4.,  -3.,  -2.,  -1., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
       8.,   7.,   6.,   5., // [1, 0, :]
       4.,   3.,   2.,   1., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/ >= 3, /*step=*/1, out),
  // src shape should equal input[:，0:3:1, :] = input for any end >= 3
  Tensor src_end_3_or_above = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  Tensor expected_end_3_or_above = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  // clang-format on

  std::vector<Tensor> src_tensors = {// end = -3
                                     src_end_0_or_below,
                                     // end = -2
                                     src_end_1,
                                     // end = -1
                                     src_end_2,
                                     // end = 0
                                     src_end_0_or_below,
                                     // end = 1
                                     src_end_1,
                                     // end = 2
                                     src_end_2,
                                     // end = 3
                                     src_end_3_or_above};

  std::vector<Tensor> expected_rets = {// end = -3
                                       expected_end_0_or_below,
                                       // end = -2
                                       expected_end_1,
                                       // end = -1
                                       expected_end_2,
                                       // end = 0
                                       expected_end_0_or_below,
                                       // end = 1
                                       expected_end_1,
                                       // end = 2
                                       expected_end_2,
                                       // end = 3
                                       expected_end_3_or_above};

  int64_t dim = 1;
  int64_t start = 0;
  int64_t step = 1;
  for (int64_t end = -3; end < 4; end++) {
    int64_t testcase_idx = end + 3;

    auto src = src_tensors[testcase_idx];
    auto expected_ret = expected_rets[testcase_idx];
    Tensor out = tf.zeros_like(expected_ret);

    // Should always return the provided out Tensor.
    // The ret shall meet the expectation.
    Tensor ret = op_slice_scatter_out(input, src, dim, start, end, step, out);
    EXPECT_TENSOR_EQ(out, ret);
    EXPECT_TENSOR_EQ(ret, expected_ret);
  }
}

TEST_F(OpSliceScatterTensorOutTest, LegalStepsSupported) {
  TensorFactory<ScalarType::Double> tf;

  // clang-format off
  Tensor input = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
       1.,   2.,   3.,   4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
      -1.,  -2.,  -3.,  -4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });

  // Set the end large enough to hold any step

  // Expected ret for op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/10, /*step=*/1, out),
  // src shape should equal to input[:，0:3:1, :]
  Tensor src_0 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  Tensor expected_0 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -5.,  -6.,  -7.,  -8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       5.,   6.,   7.,   8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  // Expected ret for op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/10, /*step=*/2, out),
  // src shape should equal to input[:，0:3:2, :]
  Tensor src_1 = tf.make(
    /*sizes=*/{2, 2, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
      -9., -10., -11., -12., // [0, 1, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
       9.,  10.,  11.,  12., // [1, 1, :]
    });
  Tensor expected_1 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
      -9., -10., -11., -12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
       9.,  10.,  11.,  12., // [1, 2, :]
    });
  // Expected ret for op_slice_scatter_out(input, src, /*dim=*/1, /*start=*/0, /*end=*/10, /*step=*/3, out),
  // src shape should equal to input[:，0:3:3, :] = input
  Tensor src_2 = tf.make(
    /*sizes=*/{2, 1, 4},
    /*data=*/{
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
       1.,   2.,   3.,   4., // [1, 0, :]
    });
  Tensor expected_2 = tf.make(
    /*sizes=*/{2, 3, 4},
    /*data=*/{
      // [0, :, :]
      -1.,  -2.,  -3.,  -4., // [0, 0, :]
       5.,   6.,   7.,   8., // [0, 1, :]
       9.,  10.,  11.,  12., // [0, 2, :]

      // [1, :, :]
       1.,   2.,   3.,   4., // [1, 0, :]
      -5.,  -6.,  -7.,  -8., // [1, 1, :]
      -9., -10., -11., -12., // [1, 2, :]
    });
  // clang-format on

  std::vector<Tensor> src_tensors = {src_0, src_1, src_2};
  std::vector<Tensor> expected_rets = {expected_0, expected_1, expected_2};

  // In this test, we maintain start and dim as 0 and 1, also set the
  // end large enough to hold any step
  int64_t start = 0;
  int64_t dim = 1;
  int64_t end = 10;
  for (int64_t step = 1; step < 4; step++) {
    int64_t testcase_idx = step - 1;

    auto src = src_tensors[testcase_idx];
    auto expected_ret = expected_rets[testcase_idx];
    Tensor out = tf.zeros_like(expected_ret);

    // Should always return the provided out Tensor.
    // The ret shall meet the expectation.
    Tensor ret = op_slice_scatter_out(input, src, dim, start, end, step, out);
    EXPECT_TENSOR_EQ(out, ret);
    EXPECT_TENSOR_EQ(ret, expected_ret);
  }
}

/// A generic smoke test that works for any dtype that supports ones() and
/// zeros().
TEST_F(OpSliceScatterTensorOutTest, AllRealDtypesSupported) {
#define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>();
  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
#undef TEST_ENTRY
  // TODO: Also add tests for half, complex, quantized, and other types. Easiest
  // way to do that would be to make TensorFactory support zeros() and ones()
  // for those types.
}

TEST_F(OpSliceScatterTensorOutTest, EmptyInputSupported) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.ones({1, 0, 1});
  Tensor src = tf.zeros({1, 0, 1});
  Tensor out = tf.zeros({1, 0, 1});

  Tensor expect = tf.ones({1, 0, 1});

  // Some invalid dim values.
  for (int64_t dim = 0; dim > input.dim(); dim++) {
    Tensor ret = op_slice_scatter_out(
        input, src, dim, /*start=*/0, /*end=*/1, /*step=*/1, out);
    EXPECT_TENSOR_EQ(ret, out);

    // All operations in this test share same ground truth
    EXPECT_TENSOR_EQ(ret, expect);
  }
}

TEST_F(OpSliceScatterTensorOutTest, EmptySizeInputDies) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.ones({});
  Tensor src = tf.ones({});
  Tensor out = tf.ones({});

  // The operation shall die whatever the end is.
  ET_EXPECT_KERNEL_FAILURE(
      context_,
      op_slice_scatter_out(
          input, src, /*dim=*/0, /*start=*/0, /*end=*/0, /*step=*/1, out));
  ET_EXPECT_KERNEL_FAILURE(
      context_,
      op_slice_scatter_out(
          input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/1, out));
}

TEST_F(OpSliceScatterTensorOutTest, NonPostiveStepsDies) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.ones({1, 1, 1});
  Tensor src = tf.zeros({1, 1, 1});
  Tensor out = tf.zeros({1, 1, 1});

  // Some invalid step values.
  const std::vector<int64_t> invalid_steps = {-2, -1, 0};
  for (int64_t step : invalid_steps) {
    ET_EXPECT_KERNEL_FAILURE(
        context_,
        op_slice_scatter_out(
            input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/step, out));
  }
}

TEST_F(OpSliceScatterTensorOutTest, DimOutOfBoundDies) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.ones({1, 1, 1});
  Tensor src = tf.zeros({1, 1, 1});
  Tensor out = tf.zeros({1, 1, 1});

  // Some invalid dim values.
  const std::vector<int64_t> invalid_dims = {3, 4, 5, -4, -5, -6};
  for (int64_t dim : invalid_dims) {
    ET_EXPECT_KERNEL_FAILURE(
        context_,
        op_slice_scatter_out(
            input, src, dim, /*start=*/0, /*end=*/1, /*step=*/1, out));
  }
}

TEST_F(OpSliceScatterTensorOutTest, MismatchedOutDtypesDies) {
  TensorFactory<ScalarType::Int> tf_int;
  TensorFactory<ScalarType::Float> tf_float;
  Tensor input = tf_int.zeros({1, 2, 2});
  Tensor src = tf_int.zeros({1, 2, 2});

  // Size is compatible to the output, but a mismatched dtype.
  Tensor out = tf_float.ones({1, 2, 2});

  ET_EXPECT_KERNEL_FAILURE(
      context_,
      op_slice_scatter_out(
          input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/1, out));
}

TEST_F(OpSliceScatterTensorOutTest, OutSizeMismatchDimDies) {
  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
  }
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({2, 4, 7, 5});
  Tensor src = tf.zeros({2, 4, 7, 5});

  // Should be {2, 4, 7, 5}
  Tensor out = tf.zeros({2, 4, 7});

  ET_EXPECT_KERNEL_FAILURE(
      context_,
      op_slice_scatter_out(
          input, src, /*dim=*/0, /*start=*/0, /*end=*/2, /*step=*/1, out));
}

TEST_F(OpSliceScatterTensorOutTest, SrcSizeMismatchDimDies) {
  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
  }
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({2, 4, 7, 5});
  Tensor src = tf.zeros({2, 4, 7});

  // Should be {2, 4, 7, 5}
  Tensor out = tf.zeros({2, 4, 7, 5});

  ET_EXPECT_KERNEL_FAILURE(
      context_,
      op_slice_scatter_out(
          input, src, /*dim=*/0, /*start=*/0, /*end=*/2, /*step=*/1, out));
}

TEST_F(OpSliceScatterTensorOutTest, DefaultStartValSupported) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({2, 4, 7, 5});
  Tensor src = tf.ones({2, 4, 7, 5});

  Tensor out = tf.zeros({2, 4, 7, 5});
  Tensor expected = tf.ones({2, 4, 7, 5});

  Tensor ret_default_start = op_slice_scatter_out(
      input,
      src,
      /*dim=*/0,
      /*start=*/exec_aten::nullopt,
      /*end=*/2,
      /*step=*/1,
      out);
  EXPECT_TENSOR_EQ(ret_default_start, out);
  EXPECT_TENSOR_EQ(ret_default_start, expected);
}

TEST_F(OpSliceScatterTensorOutTest, DefaultEndValSupported) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({2, 4, 7, 5});
  Tensor src = tf.ones({2, 4, 7, 5});

  Tensor out = tf.zeros({2, 4, 7, 5});
  Tensor expected = tf.ones({2, 4, 7, 5});

  Tensor ret_default_end = op_slice_scatter_out(
      input,
      src,
      /*dim=*/0,
      /*start=*/0,
      /*end=*/exec_aten::nullopt,
      /*step=*/1,
      out);
  EXPECT_TENSOR_EQ(ret_default_end, out);
  EXPECT_TENSOR_EQ(ret_default_end, expected);
}

TEST_F(OpSliceScatterTensorOutTest, DynamicShapeTest) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({1, 4, 4});
  Tensor src = tf.ones({1, 4, 4});

  Tensor out =
      tf.zeros({1, 2, 8}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
  Tensor expected = tf.ones({1, 4, 4});

  Tensor ret_default_end = op_slice_scatter_out(
      input,
      src,
      /*dim=*/0,
      /*start=*/0,
      /*end=*/exec_aten::nullopt,
      /*step=*/1,
      out);
  EXPECT_TENSOR_EQ(ret_default_end, out);
  EXPECT_TENSOR_EQ(ret_default_end, expected);
}

TEST_F(OpSliceScatterTensorOutTest, LargeEndValue) {
  TensorFactory<ScalarType::Int> tf;

  Tensor input = tf.zeros({1, 1, 2, 5, 3, 3});
  Tensor src = tf.ones({1, 1, 2, 5, 3, 3});

  Tensor out = tf.zeros({1, 1, 2, 5, 3, 3});
  Tensor expected = tf.ones({1, 1, 2, 5, 3, 3});

  Tensor ret = op_slice_scatter_out(
      input,
      src,
      /*dim=*/1,
      /*start=*/0,
      /*end=*/9223372036854775807,
      /*step=*/1,
      out);
  EXPECT_TENSOR_EQ(ret, out);
  EXPECT_TENSOR_EQ(ret, expected);
}