exponential_kernel Class — pytorch Architecture

Source Code

aten/src/ATen/native/cuda/DistributionTemplates.h lines 560–572

template<typename RNG>
void exponential_kernel(TensorIteratorBase& iter, double lambda_, RNG gen) {
  TORCH_CHECK(isFloatingType(iter.dtype()), "Exponential distribution is a continuous probability distribution. dtype must be a floating point but you specified ", iter.dtype());
  AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "exponential_cuda", [&] {
    using accscalar_t = at::acc_type<scalar_t, true>;
    auto lambda = static_cast<accscalar_t>(lambda_);
    // define lambda for exponential transformation
    auto exponential_func = [lambda] __device__ (accscalar_t rand) {
      return static_cast<scalar_t>(transformation::exponential<accscalar_t>(rand, lambda));
    };
    uniform_and_transform<scalar_t, accscalar_t>(iter, gen, exponential_func);
   });
}