idx Class — pytorch Architecture

Source Code

aten/src/ATen/test/vec_test_all_types.cpp lines 1080–1105

    template<typename vec, typename VT, int64_t idx, int64_t N>
    std::enable_if_t<(!is_complex<VT>::value && idx == N), bool>
    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
    test_blendv(VT expected_val[vec::size()], VT a[vec::size()], VT b[vec::size()], VT mask[vec::size()]) {
        using bit_rep = BitType<VT>;
        // generate expected_val
        for (int64_t i = 0; i < vec::size(); i++) {
            bit_rep hex_mask = 0;
            hex_mask=c10::bit_cast<bit_rep>(mask[i]);
            expected_val[i] = (hex_mask & 0x01) ? b[i] : a[i];
        }
        // test with blendv
        auto vec_a = vec::loadu(a);
        auto vec_b = vec::loadu(b);
        auto vec_m = vec::loadu(mask);
        auto expected = vec::loadu(expected_val);
        auto actual = vec::blendv(vec_a, vec_b, vec_m);
        auto mask_str = std::string("\nblendv mask: ");
        for (int64_t i = 0; i < vec::size(); i++) {
            mask_str += std::to_string(mask[i]) + " ";
        }
        if (AssertVectorized<vec>(std::string(NAME_INFO(test_blendv)) + mask_str, expected, actual).check()) {
            return false;
        }
        return true;
    }