mean_in Class — pytorch Architecture

Source Code

aten/src/ATen/core/DistributionsHelper.h lines 171–197

template <typename T>
struct normal_distribution {

  C10_HOST_DEVICE inline normal_distribution(T mean_in, T stdv_in) : mean(mean_in), stdv(stdv_in) {
    TORCH_CHECK_IF_NOT_ON_CUDA(stdv_in >= 0, "stdv_in must be positive: ", stdv_in);
  }

  template <typename RNG>
  C10_HOST_DEVICE inline dist_acctype<T> operator()(RNG* generator) const {
    dist_acctype<T> ret;
    // return cached values if available
    if (maybe_get_next_normal_sample(generator, &ret)) {
      return transformation::normal(ret, mean, stdv);
    }

    // otherwise generate new normal values
    uniform_real_distribution<T> uniform(0.0, 1.0);
    const dist_acctype<T> u1 = uniform(generator);
    const dist_acctype<T> u2 = uniform(generator);
    const dist_acctype<T> r = ::sqrt(static_cast<T>(-2.0) * ::log1p(-u2));
    const dist_acctype<T> theta = static_cast<T>(2.0) * c10::pi<T> * u1;
    const dist_acctype<T> sample = r * ::sin(theta);
    maybe_set_next_normal_sample(generator, &sample);

    ret = r * ::cos(theta);
    return transformation::normal(ret, mean, stdv);
  }