wehub-resource-sync
571 lines
23 KiB
C++
571 lines
23 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
#ifndef LIBND4J_SPECIAL_RANDOM_OPS_H
|
|
#define LIBND4J_SPECIAL_RANDOM_OPS_H
|
|
#include <execution/Threads.h>
|
|
#include <graph/RandomGenerator.h>
|
|
#include <helpers/shape.h>
|
|
#include <ops/random_ops.h>
|
|
#include <ops/specials_cuda.h>
|
|
|
|
namespace randomOps {
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class Choice {
|
|
public:
|
|
method_idx method_X method_XY
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
// ... (CUDA implementation remains unchanged)
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
sd::LongType yLength = shape::length(yShapeBuffer);
|
|
|
|
int elementsPerThread = zLength / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e, zRank, zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
T prob = rng->relativeT<T>(e);
|
|
T cumProb = (T)0.0f;
|
|
for (sd::LongType f = 0; f < yLength; f++) {
|
|
sd::LongType yCoords[SD_MAX_RANK];
|
|
INDEX2COORDS(f, yRank, yShape, yCoords);
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
|
T relProb = y[yOffset];
|
|
cumProb += relProb;
|
|
|
|
if (prob <= cumProb || f == yLength - 1) {
|
|
sd::LongType xCoords[SD_MAX_RANK];
|
|
INDEX2COORDS(f, xRank, xShape, xCoords);
|
|
sd::LongType xOffset;
|
|
COORDS2INDEX(xRank,xStride , xCoords, xOffset);
|
|
z[zOffset] = x[xOffset];
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, zLength, 1, _threads);
|
|
}
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class GaussianDistribution {
|
|
public:
|
|
method_XY method_X method_idx
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
// ... (CUDA implementation remains unchanged)
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
const T two_pi = static_cast<T>(2.0f) * static_cast<T>(3.14159265358979323846);
|
|
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
auto middle = zLength % 2 + zLength / 2;
|
|
|
|
int elementsPerThread = middle / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
|
|
sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
const T mean = extraArguments[0];
|
|
const T stddev = extraArguments[1];
|
|
|
|
const T epsilon = static_cast<T>(1e-5);
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e, zRank, zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
auto epm = e + middle;
|
|
|
|
// we need to get random values
|
|
T r0 = rng->relativeT<T>(e, epsilon, static_cast<T>(1.0f));
|
|
T r1 = rng->relativeT<T>(epm, epsilon, static_cast<T>(1.0f));
|
|
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(yRank, yStride, coords, yOffset);
|
|
T realMean0 = y == z ? mean : y[yOffset];
|
|
|
|
z[zOffset] = (sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_cos<T, T>(two_pi * r1)) * stddev + realMean0;
|
|
|
|
if (epm < zLength) {
|
|
INDEX2COORDS(epm, zRank, zShape, coords);
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
COORDS2INDEX(yRank, yStride, coords, yOffset);
|
|
T realMean1 = y == z ? mean : y[yOffset];
|
|
z[zOffset] = (sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_sin<T, T>(two_pi * r1)) * stddev + realMean1;
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, middle, 1, _threads);
|
|
}
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class BinomialDistribution {
|
|
public:
|
|
method_XY method_X method_idx
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
// ... (CUDA implementation remains unchanged)
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
int trials = (int)extraArguments[0];
|
|
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
|
|
int elementsPerThread = zLength / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
|
|
T prob = extraArguments[1];
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e,zRank, zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
int success = 0;
|
|
for (int t = 1; t <= trials; t++) {
|
|
T randVal = rng->relativeT<T>((e + 1) * t);
|
|
if (y != z) {
|
|
// we're using external probs
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(yRank,yStride, coords, yOffset);
|
|
prob = y[yOffset];
|
|
}
|
|
|
|
if (randVal < prob) success++;
|
|
}
|
|
|
|
// if trials is set to 0, effectively we just have successful memset
|
|
z[zOffset] = static_cast<T>(success);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, zLength, 1, _threads);
|
|
}
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class BinomialDistributionEx {
|
|
public:
|
|
method_XY method_X method_idx
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
// ... (CUDA implementation remains unchanged)
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
int trials = (int)extraArguments[0];
|
|
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
|
|
int elementsPerThread = zLength / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
T prob = extraArguments[1];
|
|
|
|
auto rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e,zRank, zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
int success = 0;
|
|
for (int t = 1; t <= trials; t++) {
|
|
T randVal = rng->relativeT<T>((e + 1) * t);
|
|
if (y != z) {
|
|
// we're using external probs
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(shape::rank(yShapeBuffer), shape::stride(yShapeBuffer), coords, yOffset);
|
|
prob = y[yOffset];
|
|
}
|
|
|
|
if (randVal < prob) success++;
|
|
}
|
|
|
|
// if trials is set to 0, effectively we just have successful memset
|
|
z[zOffset] = static_cast<T>(success);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, zLength, 1, _threads);
|
|
}
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class TruncatedNormalDistribution {
|
|
private:
|
|
static SD_INLINE SD_HOST_DEVICE T step(sd::graph::RandomGenerator *rng, T mean, T stddev, sd::LongType e,
|
|
sd::LongType middle, T &z) {
|
|
auto epm = e + middle;
|
|
const T two_pi = static_cast<T>(2.0f) * static_cast<T>(3.14159265358979323846);
|
|
const T epsilon = static_cast<T>(1.e-5f);
|
|
// we need to get random values
|
|
T r0 = rng->relativeT<T>(e, epsilon, static_cast<T>(1.0f));
|
|
T r1 = rng->relativeT<T>(epm, epsilon, static_cast<T>(1.0f));
|
|
|
|
T realMean0 = mean;
|
|
|
|
auto z0 = (sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_cos<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean0;
|
|
z = z0;
|
|
if (epm < middle) {
|
|
T realMean1 = mean;
|
|
auto z1 = (sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_sin<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean1;
|
|
z = z1;
|
|
}
|
|
return z;
|
|
}
|
|
|
|
public:
|
|
method_XY method_X method_idx
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
// ... (CUDA implementation remains unchanged)
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
GaussianDistribution<T>::specialOp(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments);
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
auto rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
T mean = extraArguments[0];
|
|
T stddev = extraArguments[1];
|
|
T ds = sd::math::sd_abs<T,T>(stddev) * (T)2.0f;
|
|
sd::LongType middle = zLength / 2 + (zLength % 2);
|
|
int elementsPerThread = middle / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e, zRank,zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
|
|
if (z[zOffset] > mean + ds || z[zOffset] < mean - ds) {
|
|
z[zOffset] = step(rng, mean, stddev, e, middle, z[zOffset]);
|
|
|
|
if (z[zOffset] > mean + ds || z[zOffset] < mean - ds) z[zOffset] = mean + sd::DataTypeUtils::min_positive<T>();
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, zLength, 1, _threads);
|
|
}
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
class LogNormalDistribution {
|
|
public:
|
|
method_XY method_X method_idx
|
|
|
|
static const bool requiresSpecial = true;
|
|
|
|
#ifdef __CUDACC__
|
|
static SD_INLINE SD_DEVICE void specialOpCuda(sd::Pointer state, T const *x, sd::LongType const *xShapeBuffer,
|
|
T const *y, sd::LongType const *yShapeBuffer, T *z,
|
|
sd::LongType const *zShapeBuffer, T *extraArguments) {
|
|
__shared__ T epsilon;
|
|
__shared__ T two_pi;
|
|
|
|
__shared__ sd::LongType zLength;
|
|
__shared__ T mean;
|
|
__shared__ T stddev;
|
|
__shared__ int step;
|
|
|
|
__shared__ T *tZ;
|
|
|
|
__shared__ sd::graph::RandomGenerator *rng;
|
|
__shared__ unsigned char *cB;
|
|
__shared__ unsigned char *dB;
|
|
__shared__ sd::graph::RandomGenerator *devRng;
|
|
__shared__ sd::LongType yRank;
|
|
__shared__ sd::LongType *yShape;
|
|
__shared__ sd::LongType *yStride;
|
|
__shared__ sd::LongType *xShape;
|
|
__shared__ sd::LongType xRank;
|
|
__shared__ sd::LongType *xStride;
|
|
__shared__ sd::LongType *zShape;
|
|
__shared__ sd::LongType *zStride;
|
|
__shared__ sd::LongType zRank;
|
|
|
|
|
|
if (threadIdx.x == 0) {
|
|
extern __shared__ unsigned char shmem[];
|
|
rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
cB = shmem;
|
|
devRng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
|
|
dB = reinterpret_cast<unsigned char *>(state);
|
|
|
|
tZ = reinterpret_cast<T *>(shmem + sizeof(sd::graph::RandomGenerator));
|
|
|
|
zLength = shape::length(zShapeBuffer);
|
|
|
|
epsilon = static_cast<T>(1e-5);
|
|
two_pi = static_cast<T>(2.0f) * static_cast<T>(3.14159265358979323846);
|
|
|
|
mean = extraArguments[0];
|
|
stddev = extraArguments[1];
|
|
|
|
step = (blockDim.x * gridDim.x);
|
|
xRank = shape::rank(xShapeBuffer);
|
|
xShape = shape::shapeOf(xShapeBuffer);
|
|
xStride = shape::stride(xShapeBuffer);
|
|
yRank = shape::rank(yShapeBuffer);
|
|
yShape = shape::shapeOf(yShapeBuffer);
|
|
yStride = shape::stride(yShapeBuffer);
|
|
zRank = shape::rank(zShapeBuffer);
|
|
zShape = shape::shapeOf(zShapeBuffer);
|
|
zStride = shape::stride(zShapeBuffer);
|
|
|
|
}
|
|
__syncthreads();
|
|
|
|
// using this loop instead of memcpy
|
|
for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x) cB[e] = dB[e];
|
|
|
|
__syncthreads();
|
|
|
|
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
int middle = zLength % 2 == 0 ? zLength / 2 : zLength / 2 + 1;
|
|
|
|
for (sd::LongType e = tid; e < middle; e += step) {
|
|
auto epm = e + middle;
|
|
|
|
// we need to get random values
|
|
T r0 = rng->relativeT<T>(e, epsilon, static_cast<T>(1.0f));
|
|
T r1 = rng->relativeT<T>(epm, epsilon, static_cast<T>(1.0f));
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
sd::LongType yCoords[SD_MAX_RANK];
|
|
INDEX2COORDS(e, yRank, yShape, yCoords);
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
|
sd::LongType zOffset;
|
|
INDEX2COORDS(e, zRank, zShape, coords);
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
T realMean = y == z ? mean : y[yOffset];
|
|
z[zOffset] =
|
|
sd::math::sd_exp<T, T>((sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_cos<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean);
|
|
|
|
if (epm < zLength) {
|
|
realMean = y == z ? mean : y[epm + yOffset];
|
|
z[epm + zOffset] =
|
|
sd::math::sd_exp<T, T>((sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_sin<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean);
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
|
|
static inline void specialOp(sd::Pointer state, const T *x, const sd::LongType *xShapeBuffer, const T *y,
|
|
const sd::LongType *yShapeBuffer, T *z, const sd::LongType *zShapeBuffer,
|
|
T *extraArguments) {
|
|
const T two_pi = static_cast<T>(2.0f) * static_cast<T>(3.14159265358979323846);
|
|
|
|
sd::LongType zLength = shape::length(zShapeBuffer);
|
|
auto middle = zLength % 2 == 0 ? zLength / 2 : zLength / 2 + 1;
|
|
|
|
int elementsPerThread = middle / TAD_THRESHOLD;
|
|
int _threads = sd::math::sd_max<int>(1, elementsPerThread);
|
|
_threads = sd::math::sd_min<int>(_threads, sd::Environment::getInstance().maxThreads());
|
|
|
|
auto rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
|
|
|
|
const T mean = extraArguments[0];
|
|
const T stddev = extraArguments[1];
|
|
const T epsilon = static_cast<T>(1e-5);
|
|
|
|
sd::LongType zRank = shape::rank(zShapeBuffer);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeBuffer);
|
|
sd::LongType *zStride = shape::stride(zShapeBuffer);
|
|
sd::LongType yRank = shape::rank(yShapeBuffer);
|
|
sd::LongType *yShape = shape::shapeOf(yShapeBuffer);
|
|
sd::LongType *yStride = shape::stride(yShapeBuffer);
|
|
sd::LongType *xShape = shape::shapeOf(xShapeBuffer);
|
|
sd::LongType xRank = shape::rank(xShapeBuffer);
|
|
sd::LongType *xStride = shape::stride(xShapeBuffer);
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(e, zRank, zShape, coords);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
auto epm = e + middle;
|
|
|
|
// we need to get random values
|
|
T r0 = rng->relativeT<T>(e, epsilon, static_cast<T>(1.0f));
|
|
T r1 = rng->relativeT<T>(epm, epsilon, static_cast<T>(1.0f));
|
|
|
|
sd::LongType yOffset;
|
|
COORDS2INDEX(yRank, yStride, coords, yOffset);
|
|
T realMean = y == z ? mean : y[yOffset];
|
|
|
|
z[zOffset] =
|
|
sd::math::sd_exp<T, T>((sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_cos<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean);
|
|
|
|
if (epm < zLength) {
|
|
INDEX2COORDS(epm,zRank, zShape, coords);
|
|
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
|
COORDS2INDEX(yRank, yStride, coords, yOffset);
|
|
realMean = y == z ? mean : y[yOffset];
|
|
z[zOffset] =
|
|
sd::math::sd_exp<T, T>((sd::math::sd_sqrt<T, T>(static_cast<T>(-2.0f) * sd::math::sd_log<T, T>(r0)) *
|
|
sd::math::sd_sin<T, T>(two_pi * r1)) *
|
|
stddev +
|
|
realMean);
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, middle, 1, _threads);
|
|
}
|
|
};
|
|
|
|
} // namespace randomOps
|
|
|
|
#endif // LIBND4J_SPECIAL_RANDOM_OPS_H
|