mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vec256_complex_double.h lines 52–69

  template <int64_t mask>
  static Vectorized<c10::complex<double>> blend(
      const Vectorized<c10::complex<double>>& a,
      const Vectorized<c10::complex<double>>& b) {
    // convert c10::complex<V> index mask to V index mask: xy -> xxyy
    static_assert(mask > -1 && mask < 4, "Unexpected mask value");
    switch (mask) {
      case 0:
        return a;
      case 1:
        return _mm256_blend_pd(a.values, b.values, 0x03);
      case 2:
        return _mm256_blend_pd(a.values, b.values, 0x0c);
      case 3:
        break;
    }
    return b;
  }