/usr/local/lib64/python3.6/site-packages/torch/include/ATen/core
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/core/DistributionsHelper.h (12594B)
#pragma once
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
/**
* Distributions kernel adapted from THRandom.cpp
* The kernels try to follow std::random distributions signature
* For instance: in ATen
* auto gen = at::detail::createCPUGenerator();
* at::uniform_real_distribution uniform(0, 1);
* auto sample = uniform(gen.get());
*
* vs std::random
*
* std::mt19937 gen;
* std::uniform_real_distribution uniform(0, 1);
* auto sample = uniform(gen);
*/
namespace at {
namespace {
/**
* Samples a discrete uniform distribution in the range [base, base+range) of type T
*/
template
struct uniform_int_from_to_distribution {
C10_HOST_DEVICE inline uniform_int_from_to_distribution(uint64_t range, int64_t base) {
range_ = range;
base_ = base;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
if ((
std::is_same::value ||
std::is_same::value ||
std::is_same::value ||
std::is_same::value) && range_ >= 1ULL << 32)
{
return transformation::uniform_int_from_to(generator->random64(), range_, base_);
} else {
return transformation::uniform_int_from_to(generator->random(), range_, base_);
}
}
private:
uint64_t range_;
int64_t base_;
};
/**
* Samples a discrete uniform distribution in the range [min_value(int64_t), max_value(int64_t)]
*/
template
struct uniform_int_full_range_distribution {
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
return transformation::uniform_int_full_range(generator->random64());
}
};
/**
* Samples a discrete uniform distribution in the range [0, max_value(T)] for integral types
* and [0, 2^mantissa] for floating-point types.
*/
template
struct uniform_int_distribution {
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
if (std::is_same::value || std::is_same::value) {
return transformation::uniform_int(generator->random64());
} else {
return transformation::uniform_int(generator->random());
}
}
};
/**
* Samples a uniform distribution in the range [from, to) of type T
*/
template
struct uniform_real_distribution {
C10_HOST_DEVICE inline uniform_real_distribution(T from, T to) {
TORCH_CHECK_IF_NOT_ON_CUDA(from <= to);
TORCH_CHECK_IF_NOT_ON_CUDA(to - from <= std::numeric_limits::max());
from_ = from;
to_ = to;
}
template
C10_HOST_DEVICE inline dist_acctype operator()(RNG generator){
if(std::is_same::value) {
return transformation::uniform_real(generator->random64(), from_, to_);
} else {
return transformation::uniform_real(generator->random(), from_, to_);
}
}
private:
T from_;
T to_;
};
// The SFINAE checks introduced in #39816 looks overcomplicated and must revisited
// https://github.com/pytorch/pytorch/issues/40052
#define DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(member) \
template \
struct has_member_##member \
{ \
typedef char yes; \
typedef long no; \
template static yes test(decltype(&U::member)); \
template static no test(...); \
static constexpr bool value = sizeof(test(0)) == sizeof(yes); \
}
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(next_double_normal_sample);
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(set_next_double_normal_sample);
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(next_float_normal_sample);
DISTRIBUTION_HELPER_GENERATE_HAS_MEMBER(set_next_float_normal_sample);
#define DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(TYPE) \
\
template ::value && \
has_member_set_next_##TYPE##_normal_sample::value \
), int> = 0> \
C10_HOST_DEVICE inline bool maybe_get_next_##TYPE##_normal_sample(RNG* generator, ret_type* ret) { \
if (generator->next_##TYPE##_normal_sample()) { \
*ret = *(generator->next_##TYPE##_normal_sample()); \
generator->set_next_##TYPE##_normal_sample(c10::optional()); \
return true; \
} \
return false; \
} \
\
template ::value || \
!has_member_set_next_##TYPE##_normal_sample::value \
), int> = 0> \
C10_HOST_DEVICE inline bool maybe_get_next_##TYPE##_normal_sample(RNG* generator, ret_type* ret) { \
return false; \
} \
\
template ::value \
), int> = 0> \
C10_HOST_DEVICE inline void maybe_set_next_##TYPE##_normal_sample(RNG* generator, ret_type cache) { \
generator->set_next_##TYPE##_normal_sample(cache); \
} \
\
template ::value \
), int> = 0> \
C10_HOST_DEVICE inline void maybe_set_next_##TYPE##_normal_sample(RNG* generator, ret_type cache) { \
}
DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(double);
DISTRIBUTION_HELPER_GENERATE_NEXT_NORMAL_METHODS(float);
/**
* Samples a normal distribution using the Box-Muller method
* Takes mean and standard deviation as inputs
* Note that Box-muller method returns two samples at a time.
* Hence, we cache the "next" sample in the CPUGeneratorImpl class.
*/
template
struct normal_distribution {
C10_HOST_DEVICE inline normal_distribution(T mean_in, T stdv_in) {
TORCH_CHECK_IF_NOT_ON_CUDA(stdv_in >= 0, "stdv_in must be positive: ", stdv_in);
mean = mean_in;
stdv = stdv_in;
}
template
C10_HOST_DEVICE inline dist_acctype operator()(RNG generator){
dist_acctype ret;
// return cached values if available
if (std::is_same::value) {
if (maybe_get_next_double_normal_sample(generator, &ret)) {
return transformation::normal(ret, mean, stdv);
}
} else {
if (maybe_get_next_float_normal_sample(generator, &ret)) {
return transformation::normal(ret, mean, stdv);
}
}
// otherwise generate new normal values
uniform_real_distribution uniform(0.0, 1.0);
const dist_acctype u1 = uniform(generator);
const dist_acctype u2 = uniform(generator);
const dist_acctype r = ::sqrt(static_cast(-2.0) * ::log(static_cast(1.0)-u2));
const dist_acctype theta = static_cast(2.0) * c10::pi * u1;
if (std::is_same::value) {
maybe_set_next_double_normal_sample(generator, r * ::sin(theta));
} else {
maybe_set_next_float_normal_sample(generator, r * ::sin(theta));
}
ret = r * ::cos(theta);
return transformation::normal(ret, mean, stdv);
}
private:
T mean;
T stdv;
};
template
struct DiscreteDistributionType { using type = float; };
template <> struct DiscreteDistributionType { using type = double; };
/**
* Samples a bernoulli distribution given a probability input
*/
template
struct bernoulli_distribution {
C10_HOST_DEVICE inline bernoulli_distribution(T p_in) {
TORCH_CHECK_IF_NOT_ON_CUDA(p_in >= 0 && p_in <= 1);
p = p_in;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
uniform_real_distribution uniform(0.0, 1.0);
return transformation::bernoulli(uniform(generator), p);
}
private:
T p;
};
/**
* Samples a geometric distribution given a probability input
*/
template
struct geometric_distribution {
C10_HOST_DEVICE inline geometric_distribution(T p_in) {
TORCH_CHECK_IF_NOT_ON_CUDA(p_in > 0 && p_in < 1);
p = p_in;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
uniform_real_distribution uniform(0.0, 1.0);
return transformation::geometric(uniform(generator), p);
}
private:
T p;
};
/**
* Samples an exponential distribution given a lambda input
*/
template
struct exponential_distribution {
C10_HOST_DEVICE inline exponential_distribution(T lambda_in) {
lambda = lambda_in;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
uniform_real_distribution uniform(0.0, 1.0);
return transformation::exponential(uniform(generator), lambda);
}
private:
T lambda;
};
/**
* Samples a cauchy distribution given median and sigma as inputs
*/
template
struct cauchy_distribution {
C10_HOST_DEVICE inline cauchy_distribution(T median_in, T sigma_in) {
median = median_in;
sigma = sigma_in;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator) {
uniform_real_distribution uniform(0.0, 1.0);
return transformation::cauchy(uniform(generator), median, sigma);
}
private:
T median;
T sigma;
};
/**
* Samples a lognormal distribution
* Takes mean and standard deviation as inputs
* Outputs two samples at a time
*/
template
struct lognormal_distribution {
C10_HOST_DEVICE inline lognormal_distribution(T mean_in, T stdv_in) {
TORCH_CHECK_IF_NOT_ON_CUDA(stdv_in > 0);
mean = mean_in;
stdv = stdv_in;
}
template
C10_HOST_DEVICE inline T operator()(RNG generator){
normal_distribution normal(mean, stdv);
return transformation::log_normal(normal(generator));
}
private:
T mean;
T stdv;
};
}
} // namespace at