apply_root Class — pytorch Architecture

Source Code

aten/src/ATen/native/SharedReduceOps.h lines 258–290

template <typename scalar_t, typename acc_t = scalar_t, typename out_t = acc_t, bool apply_root = true>
struct NormOps {
  acc_t norm_;

  inline C10_DEVICE acc_t reduce(acc_t acc, scalar_t data, int64_t /*idx*/) const {
    return acc + compat_pow(static_cast<acc_t>(std::abs(at::opmath_type<scalar_t>(data))), norm_);
  }

  inline C10_DEVICE acc_t combine(acc_t a, acc_t b) const {
    return a + b;
  }

  inline C10_DEVICE out_t project(acc_t a) const {
    if constexpr (apply_root) {
      return compat_pow(a, static_cast<acc_t>(1.0) / norm_);
    } else {
      return a;
    }
  }

  static C10_DEVICE acc_t translate_idx(acc_t acc, int64_t /*base_idx*/) {
    return acc;
  }

#if defined(__CUDACC__) || defined(__HIPCC__)
  inline C10_DEVICE acc_t warp_shfl_down(acc_t acc, int offset) const {
    return WARP_SHFL_DOWN(acc, offset);
  }
#endif

  NormOps(acc_t norm_): norm_(norm_) {
  }
};