calc_i0 Class — pytorch Architecture

Source Code

aten/src/ATen/native/Math.h lines 1473–1485

template <typename T>
inline typename std::enable_if_t<std::is_floating_point_v<T>, T>
calc_i0(T _x) {
  T x = std::abs(_x);

  if (x <= T{8.0}) {
    auto [A, len] = chebyshev_coefficients_i0e_A<T>();
    T y = (x / T{2.0}) - T{2.0};
    return static_cast<T>(std::exp(x) * chbevl(y, A, len));
  }
  auto [B, len] = chebyshev_coefficients_i0e_B<T>();
  return std::exp(x) * chbevl(T{32.0} / x - T{2.0}, B, len) / std::sqrt(x);
}