sample_binomial Class — pytorch Architecture

Source Code

aten/src/ATen/native/Distributions.h lines 217–244

template<typename scalar_t, typename accscalar_t, typename uniform_sampler_t>
C10_DEVICE scalar_t sample_binomial(scalar_t count, scalar_t prob, BaseSampler<accscalar_t, uniform_sampler_t>& standard_uniform) {
  if (count <= 0.0 || prob <= 0.0) {
    return 0;
  } else if (prob >= 1.0) {
    return count;
  } else if (prob <= 0.5) {
    if (count * prob >= 10.0) {
      // btrs
      return btrs<scalar_t, accscalar_t, uniform_sampler_t>(count, prob, standard_uniform);
    } else {
      // binomial inversion
      return binomial_inversion<scalar_t, accscalar_t, uniform_sampler_t>(count, prob, standard_uniform);
    }
  } else if (prob > 0.5) {
    scalar_t qprob = 1.0 - prob;
    if (count * qprob >= 10.0) {
      // btrs
      return count - btrs<scalar_t, accscalar_t, uniform_sampler_t>(count, qprob, standard_uniform);
    } else {
      // count - binomial inversion
      return count - binomial_inversion<scalar_t, accscalar_t, uniform_sampler_t>(count, qprob, standard_uniform);
    }
  } else {
    // prob is nan?
    return static_cast<scalar_t>(NAN);
  }
}