binomial_inversion Class — pytorch Architecture

Source Code

aten/src/ATen/native/Distributions.h lines 144–162

template<typename scalar_t, typename accscalar_t, typename uniform_sampler_t>
C10_DEVICE scalar_t binomial_inversion(scalar_t count, scalar_t prob, BaseSampler<accscalar_t, uniform_sampler_t>& standard_uniform) {
  accscalar_t U;
  accscalar_t geom_sum = 0;
  scalar_t num_geom = 0;

  accscalar_t logprob = compat_log1p(-prob);

  while (true) {
    U = standard_uniform.sample();
    accscalar_t geom = compat_ceil(compat_log(U) / logprob);
    geom_sum += geom;
    if (geom_sum > count) {
      break;
    }
    num_geom = num_geom + 1;
  }
  return num_geom;
}