test_set Class — pytorch Architecture

Source Code

aten/src/ATen/test/vec_test_all_types.cpp lines 1161–1180

    template<typename vec, typename VT>
    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
    void test_set(VT expected_val[vec::size()], VT a[vec::size()], VT b[vec::size()], int64_t count){
        if (count < 0) return;
        //generate expected_val
        for (int64_t i = 0; i < vec::size(); i++) {
            expected_val[i] = (i < count) ? b[i] : a[i];
        }
        // test with set
        auto vec_a = vec::loadu(a);
        auto vec_b = vec::loadu(b);
        auto expected = vec::loadu(expected_val);
        auto actual = vec::set(vec_a, vec_b, count);

        auto count_str = std::string("\ncount: ") + std::to_string(count);
        if (AssertVectorized<vec>(std::string(NAME_INFO(test_set)) + count_str, expected, actual).check()) {
          return;
        }
        test_set<vec, VT>(expected_val, a, b, (count == 0 ? -1 : count / 2));
    }