Added nd4j_gamma, nd4j_igamma and nd4j_igammac functions.

master
shugeo 2019-10-16 13:53:31 +03:00
parent df2448613e
commit f90e6da97e
1 changed files with 17 additions and 16 deletions

View File

@ -731,28 +731,29 @@ namespace nd4j {
template<typename T> template<typename T>
math_def inline void nd4j_swap(T &val1, T &val2) { math_def inline void nd4j_swap(T &val1, T &val2) {
T temp = val1; val1=val2; val2=temp; T temp = val1; val1=val2; val2=temp;
}; };
// template <typename X, typename Z> template <typename X, typename Z>
// math_def inline Z nd4j_gamma(X a) { math_def inline Z nd4j_gamma(X a) {
// return (Z)std::tgamma(a);
// } }
template <typename X, typename Z>
math_def inline Z nd4j_igamma(X a, X x) { template <typename X, typename Y, typename Z>
Z aim = nd4j_pow<X, X, Z>(x, a) / (nd4j_exp<X,Z>(x) * std::tgamma(a)); math_def inline Z nd4j_igamma(X a, Y x) {
Z aim = nd4j_pow<X, X, Z>(x, a) / (nd4j_exp<X, Z>(x) * nd4j_gamma<Y, Z>(a));
auto sum = Z(0.); auto sum = Z(0.);
auto denom = Z(1.); auto denom = Z(1.);
for (auto i = 0; Z(1./denom) > Z(1.0e-6); i++) { for (int i = 0; Z(1./denom) > Z(1.0e-12); i++) {
denom *= (a + i); denom *= (a + i);
sum += nd4j_pow<X, int>(x, i) / denom; sum += nd4j_pow<X, int, Z>(x, i) / denom;
} }
return aim * sum; return aim * sum;
} }
template <typename X, typename Z> template <typename X, typename Y, typename Z>
math_def inline Z nd4j_igammac(X a, X x) { math_def inline Z nd4j_igammac(X a, Y x) {
return Z(1.) - nd4j_igamma<X, Z>(a, x); return Z(1.) - nd4j_igamma<X, Y, Z>(a, x);
} }
#ifdef __CUDACC__ #ifdef __CUDACC__
namespace atomics { namespace atomics {
@ -1494,4 +1495,4 @@ inline __device__ bfloat16 nd4j_atomicDiv<bfloat16>(bfloat16* address, bfloat16
} }
#endif /* TEMPLATEMATH_H_ */ #endif /* TEMPLATEMATH_H_ */