geometric_kernel Class — pytorch Architecture

Source Code

aten/src/ATen/native/cuda/DistributionTemplates.h lines 539–549

template<typename RNG>
void geometric_kernel(TensorIteratorBase& iter, double p, RNG gen) {
  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "geometric_cuda", [&] {
    using accscalar_t = at::DiscreteDistributionType<scalar_t>::type;
    // define lambda for geometric transformation
    auto geometric_func = [p] __device__ (accscalar_t rand) {
      return static_cast<scalar_t>(transformation::geometric<accscalar_t>(rand, p));
    };
    uniform_and_transform<scalar_t, accscalar_t>(iter, gen, geometric_func);
  });
}