wehub-resource-sync
1072 lines
44 KiB
C++
1072 lines
44 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
|
|
* *****************************************************************************
|
|
*/
|
|
|
|
//
|
|
// @author GS <sgazeos@gmail.com>
|
|
//
|
|
#include <execution/Threads.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <ops/declarable/helpers/segment.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
|
|
#include <unordered_map>
|
|
#if NOT_EXCLUDED(OP_segment)
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
// segment max
|
|
template <typename T>
|
|
static void segmentMaxFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
sd::LongType idx = indices->e<sd::LongType>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
T val = input->e<T>(0);
|
|
|
|
for (sd::LongType e = 1; e < indices->lengthOf(); e++) {
|
|
if (idx == indices->e<sd::LongType>(e)) {
|
|
// max
|
|
val = sd::math::sd_max<T>(val, input->t<T>(e));
|
|
} else {
|
|
idx = indices->e<sd::LongType>(e);
|
|
val = input->t<T>(e);
|
|
}
|
|
output->r<T>(idx) = val;
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
auto numOfClasses = output->sizeAt(0); // number of classes
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(numOfClasses);
|
|
auto maxT = listOfOutTensors.at(idx);
|
|
|
|
// int pos = 0;
|
|
maxT->assign(listOfTensors.at(0));
|
|
|
|
for (sd::LongType i = 1; i < indices->lengthOf(); i++) {
|
|
if (indices->e<int>(i) == idx) {
|
|
for (sd::LongType e = 0; e < maxT->lengthOf(); e++) {
|
|
maxT->r<T>(e) = sd::math::sd_max(maxT->t<T>(e), listOfTensors.at(i)->t<T>(e));
|
|
}
|
|
} else {
|
|
idx = indices->e<sd::LongType>(i);
|
|
maxT = listOfOutTensors.at(idx);
|
|
maxT->assign(listOfTensors.at(i));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// segmen min
|
|
template <typename T>
|
|
static void segmentMinFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
sd::LongType idx = indices->e<sd::LongType>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
T val = input->e<T>(0);
|
|
|
|
for (sd::LongType e = 1; e < indices->lengthOf(); e++) {
|
|
if (idx == indices->e<sd::LongType>(e)) {
|
|
// min
|
|
val = sd::math::sd_min<T>(val, input->t<T>(e));
|
|
} else {
|
|
idx = indices->e<sd::LongType>(e);
|
|
val = input->t<T>(e);
|
|
}
|
|
output->r<T>(idx) = val;
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
|
|
int numOfClasses = output->sizeAt(0); // number of classes
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(numOfClasses);
|
|
auto minT = listOfOutTensors.at(idx);
|
|
|
|
int pos = 0;
|
|
minT->assign(listOfTensors.at(0));
|
|
|
|
for (sd::LongType i = 1; i < indices->lengthOf(); i++) {
|
|
if (indices->e<sd::LongType>(i) == idx) {
|
|
for (sd::LongType e = 0; e < minT->lengthOf(); e++) {
|
|
minT->p(e, sd::math::sd_min(minT->e<T>(e), listOfTensors.at(i)->e<T>(e)));
|
|
}
|
|
} else {
|
|
idx = indices->e<sd::LongType>(i);
|
|
minT = listOfOutTensors.at(idx);
|
|
minT->assign(listOfTensors.at(i));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// segmen mean
|
|
template <typename T>
|
|
static void segmentMeanFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
|
|
int numClasses = output->sizeAt(0);
|
|
// if input is a vector: (as if in doc sample)
|
|
int idx = indices->e<sd::LongType>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
T val = T(0.f);
|
|
int count = 0;
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); e++) {
|
|
if (idx == indices->e<sd::LongType>(e)) {
|
|
// mean
|
|
val += input->e<T>(e);
|
|
count++;
|
|
} else {
|
|
output->p<T>(idx, val / count);
|
|
idx = indices->e<sd::LongType>(e);
|
|
val = input->e<T>(e);
|
|
count = 1;
|
|
}
|
|
output->p<T>(idx, val / count);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
int numOfClasses = output->sizeAt(0); // number of classes
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(numOfClasses);
|
|
auto meanT = listOfOutTensors.at(idx);
|
|
int count = 1;
|
|
auto meanV = meanT->dup();
|
|
meanV->assign(listOfTensors.at(0));
|
|
|
|
for (sd::LongType i = 1; i < indices->lengthOf(); i++) {
|
|
if (indices->e<sd::LongType>(i) == idx) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
meanV->p<T>(e, meanV->e<T>(e) + listOfTensors.at(i)->e<T>(e));
|
|
}
|
|
};
|
|
samediff::Threads::parallel_for(func, 0, meanT->lengthOf());
|
|
|
|
count++;
|
|
} else {
|
|
meanV->applyScalar(scalar::Divide, count, meanT);
|
|
idx = indices->e<sd::LongType>(i);
|
|
meanT = listOfOutTensors.at(idx);
|
|
meanV->assign(listOfTensors.at(i));
|
|
count = 1;
|
|
}
|
|
meanV->applyScalar(scalar::Divide, count, meanT);
|
|
}
|
|
delete meanV; // Clean up duped array
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void segmentSumFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
|
|
int numClasses = output->sizeAt(0);
|
|
// if input is a vector: (as if in doc sample)
|
|
int idx = indices->e<int>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
T val = T(0.f);
|
|
int count = 0;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); e++) {
|
|
if (idx == indices->e<sd::LongType>(e)) {
|
|
// sum
|
|
val += input->t<T>(e);
|
|
} else {
|
|
idx = indices->e<sd::LongType>(e);
|
|
val = input->t<T>(e);
|
|
}
|
|
output->p(idx, val);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
int numOfClasses = output->sizeAt(0); // number of classes
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(numOfClasses);
|
|
auto sumT = listOfOutTensors.at(idx);
|
|
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
if (indices->e<sd::LongType>(i) == idx) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sumT->p(e, sumT->e<T>(e) + listOfTensors.at(i)->e<T>(e));
|
|
}
|
|
};
|
|
samediff::Threads::parallel_for(func, 0, sumT->lengthOf());
|
|
} else {
|
|
idx = indices->e<sd::LongType>(i);
|
|
sumT = listOfOutTensors.at(idx);
|
|
sumT->assign(listOfTensors.at(i));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void segmentProdFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
|
|
// int numClasses = output->sizeAt(0);
|
|
int idx = indices->e<sd::LongType>(0);
|
|
float one = 1.f;
|
|
output->assign(one);
|
|
if (input->isVector() || input->isScalar()) {
|
|
T val = input->e<T>(0);
|
|
int count = 0;
|
|
|
|
for (sd::LongType e = 1; e < indices->lengthOf(); e++) {
|
|
if (idx == indices->e<sd::LongType>(e)) {
|
|
// sum
|
|
val *= input->e<T>(e);
|
|
} else {
|
|
idx = indices->e<sd::LongType>(e);
|
|
val = input->e<T>(e);
|
|
}
|
|
output->p(idx, val);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
int numOfClasses = output->sizeAt(0); // number of classes
|
|
auto sumT = listOfOutTensors.at(idx);
|
|
sumT->assign(listOfTensors.at(0));
|
|
for (sd::LongType i = 1; i < indices->lengthOf(); i++) {
|
|
if (indices->e<sd::LongType>(i) == idx) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
sumT->p(e, sumT->e<T>(e) * listOfTensors.at(i)->e<T>(e));
|
|
}
|
|
};
|
|
samediff::Threads::parallel_for(func, 0, sumT->lengthOf());
|
|
} else {
|
|
idx = indices->e<int>(i);
|
|
sumT = listOfOutTensors.at(idx);
|
|
sumT->assign(listOfTensors.at(i));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
void segmentMaxFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), segmentMaxFunctor_, (input, indices, output), SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
void segmentMinFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), segmentMinFunctor_, (input, indices, output), SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
void segmentMeanFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), segmentMeanFunctor_, (input, indices, output), SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
void segmentSumFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), segmentSumFunctor_, (input, indices, output), SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
void segmentProdFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), segmentProdFunctor_, (input, indices, output), SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
bool segmentIndicesValidate(sd::LaunchContext* context, NDArray* indices, NDArray& expected, NDArray& output) {
|
|
auto val = indices->e(0);
|
|
for (sd::LongType e = 1; e < indices->lengthOf(); e++) {
|
|
output = indices->e(e);
|
|
if (val.e<sd::LongType>(0) > output.e<sd::LongType>(0)) return false;
|
|
val = indices->e(e);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void segmentProdFunctor_, (NDArray * input, NDArray* indices, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( void segmentSumFunctor_, (NDArray * input, NDArray* indices, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( void segmentMeanFunctor_, (NDArray * input, NDArray* indices, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( void segmentMinFunctor_, (NDArray * input, NDArray* indices, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( void segmentMaxFunctor_, (NDArray * input, NDArray* indices, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// Unsorted segment ops
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
|
|
bool unsortedSegmentIndicesValidate(sd::LaunchContext* context, NDArray* indices, sd::LongType expected,
|
|
sd::LongType& output) {
|
|
sd::LongType val = indices->e<sd::LongType>(0);
|
|
|
|
sd::LongType maxInd = indices->argMax();
|
|
if (indices->e<sd::LongType>(maxInd) >= expected) {
|
|
output = val;
|
|
return false;
|
|
}
|
|
output = expected;
|
|
return true;
|
|
}
|
|
|
|
template <typename T>
|
|
static void unsortedSegmentMaxFunctor_(NDArray* input, NDArray* indices, sd::LongType numOfClasses, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
T maxVal = DataTypeUtils::max<T>();
|
|
T negMaxVal = -maxVal;
|
|
output->assign(negMaxVal);
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
T val = input->e<T>(fi->second.at(0));
|
|
for (sd::LongType idx = 1; idx < static_cast<sd::LongType>(fi->second.size()); ++idx) {
|
|
val = sd::math::sd_max(val, input->e<T>(fi->second.at(idx)));
|
|
}
|
|
output->p(fi->first, val);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
T maxVal = DataTypeUtils::max<T>();
|
|
output->assign(maxVal);
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
for (sd::LongType idx = 0; idx < listOfTensors.size(); ++idx) {
|
|
if (static_cast<size_t>(idx) >= fi->second.size() || fi->second.size() < 2 || fi->second.at(idx) >= listOfTensors.size()) {
|
|
continue;
|
|
}
|
|
|
|
auto maxT = listOfTensors.at(fi->second.at(idx));
|
|
for (sd::LongType e = 0; e < outputT->lengthOf(); ++e) {
|
|
T val = sd::math::sd_max(maxT->e<T>(e), outputT->e<T>(e));
|
|
|
|
outputT->p(e, val);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
void unsortedSegmentMaxFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, sd::LongType numOfClasses,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentMaxFunctor_, (input, indices, numOfClasses, output),
|
|
SD_NUMERIC_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( void unsortedSegmentMaxFunctor_,
|
|
(NDArray * input, NDArray* indices, sd::LongType numOfClasses, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
template <typename T>
|
|
static void unsortedSegmentMinFunctor_(NDArray* input, NDArray* indices, sd::LongType numOfClasses, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs;
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
T maxVal = DataTypeUtils::max<T>();
|
|
output->assign(maxVal);
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
T val = input->t<T>(fi->second.at(0));
|
|
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
val = sd::math::sd_min(val, input->t<T>(fi->second.at(idx)));
|
|
}
|
|
output->r<T>(fi->first) = val;
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
T maxVal = DataTypeUtils::max<T>();
|
|
output->assign(maxVal);
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
auto minT = listOfTensors.at(fi->second.at(idx));
|
|
|
|
for (sd::LongType e = 0; e < outputT->lengthOf(); ++e) {
|
|
outputT->r<T>(e) = sd::math::sd_min(minT->t<T>(e), outputT->t<T>(e));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
void unsortedSegmentMinFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, sd::LongType numOfClasses,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentMinFunctor_, (input, indices, numOfClasses, output),
|
|
SD_NUMERIC_TYPES);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void unsortedSegmentMinFunctor_,
|
|
(NDArray * input, NDArray* indices, sd::LongType numOfClasses, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
void unsortedSegmentMeanFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, sd::LongType numOfClasses,
|
|
NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
double sumValue = input->e<double>(fi->second.at(0));
|
|
size_t loop_size = fi->second.size();
|
|
|
|
// FIXME: parallelism here?
|
|
for (size_t idx = 1; idx < loop_size; ++idx) {
|
|
sumValue += input->e<double>(fi->second.at(idx));
|
|
}
|
|
|
|
output->p(fi->first, sumValue / fi->second.size());
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
// FIXME: parallelism here?
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
sd::LongType loopSize = fi->second.size();
|
|
|
|
for (sd::LongType idx = 1; idx < loopSize; ++idx) {
|
|
auto current = listOfTensors.at(fi->second.at(idx));
|
|
*outputT += *current;
|
|
}
|
|
(*outputT) /= double(fi->second.size());
|
|
}
|
|
}
|
|
}
|
|
|
|
void unsortedSegmentSumFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, sd::LongType numOfClasses,
|
|
NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs; //(indices->lengthOf());
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
double sumValue = input->e<double>(fi->second.at(0));
|
|
sd::LongType loop_size = fi->second.size();
|
|
|
|
// FIXME: parallelism here?
|
|
for (sd::LongType idx = 1; idx < loop_size; ++idx) {
|
|
sumValue += input->e<double>(fi->second.at(idx));
|
|
}
|
|
output->p(fi->first, sumValue);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
sd::LongType loop_size = fi->second.size();
|
|
|
|
// FIXME: parallelism here?
|
|
for (sd::LongType idx = 1; idx < loop_size; ++idx) {
|
|
auto current = listOfTensors.at(fi->second.at(idx));
|
|
*(outputT) += *current;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void unsortedSegmentProdFunctor_(NDArray* input, NDArray* indices, sd::LongType numOfClasses, NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs; //(indices->lengthOf());
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
float one = 1.f;
|
|
output->assign(one);
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
T prodValue = input->e<T>(fi->second.at(0));
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
prodValue *= input->e<T>(fi->second.at(idx));
|
|
}
|
|
output->p(fi->first, prodValue);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
auto current = listOfTensors.at(fi->second.at(idx));
|
|
|
|
*outputT *= *current;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void unsortedSegmentProdFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices, sd::LongType numOfClasses,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentProdFunctor_, (input, indices, numOfClasses, output),
|
|
SD_NUMERIC_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( void unsortedSegmentProdFunctor_,
|
|
(NDArray * input, NDArray* indices, sd::LongType numOfClasses, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
void unsortedSegmentSqrtNFunctor(sd::LaunchContext* context, NDArray* input, NDArray* indices,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, std::vector<sd::LongType>> idxs;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) idxs[indices->e<sd::LongType>(e)].push_back(e);
|
|
|
|
|
|
if (input->isVector() || input->isScalar()) { // 1D case
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
double sumValue = input->e<double>(fi->second.at(0));
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
sumValue += input->e<double>(fi->second.at(idx));
|
|
}
|
|
output->p(fi->first, sumValue / sd::math::sd_sqrt<sd::LongType, double>(fi->second.size()));
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
auto listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
auto listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
|
|
auto outputT = listOfOutTensors.at(fi->first);
|
|
outputT->assign(listOfTensors.at(fi->second.at(0)));
|
|
for (size_t idx = 1; idx < fi->second.size(); ++idx) {
|
|
auto current = listOfTensors.at(fi->second.at(idx));
|
|
*outputT += *current;
|
|
}
|
|
(*outputT) /= sd::math::sd_sqrt<size_t, double>(fi->second.size());
|
|
}
|
|
}
|
|
}
|
|
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// Backpropagate ops helpers
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// Sorted backpropagate ops
|
|
//
|
|
// segment max
|
|
template <typename T>
|
|
sd::Status segmentMaxFunctorBP_(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
auto tempRes = gradOut->dup();
|
|
segmentMaxFunctor_<T>(input, indices, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
sd::LongType loop_size = input->lengthOf();
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
auto classNum = indices->e<sd::LongType>(e);
|
|
if (sd::math::sd_abs<T,T>(tempRes->e<T>(classNum) - input->e<T>(e)) <= T(1.e-6))
|
|
output->p(e, gradOut->e<T>(classNum));
|
|
}
|
|
};
|
|
samediff::Threads::parallel_for(func, 0, loop_size);
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
for (sd::LongType e = 0; e < current->lengthOf(); e++) {
|
|
if (sd::math::sd_abs<T,T>(listOfBPTensors.at(classNum)->e<T>(e) - current->e<T>(e)) <= T(1.e-6))
|
|
currentOut->p(e, currentGradOut->e<T>(e));
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, indices->lengthOf());
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status segmentMaxFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(output->dataType(), return segmentMaxFunctorBP_, (context, input, indices, gradOut, output),
|
|
SD_NUMERIC_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( sd::Status segmentMaxFunctorBP_,
|
|
(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
// segmen min
|
|
sd::Status segmentMinFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
NDArray *tempRes = gradOut->dup();
|
|
segmentMinFunctor(context, input, indices, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
auto classNum = indices->e<sd::LongType>(e);
|
|
if (sd::math::sd_abs<double,double>(tempRes->e<double>(classNum) - input->e<double>(e)) < 1.e-5)
|
|
output->p(e, gradOut->e<double>(classNum));
|
|
}
|
|
};
|
|
samediff::Threads::parallel_for(func, 0, input->lengthOf());
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
double zero = 0.;
|
|
output->assign(zero);
|
|
int pos = 0;
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
for (sd::LongType e = 0; e < current->lengthOf(); e++) {
|
|
if (sd::math::sd_abs<double,double>(listOfBPTensors.at(classNum)->e<double>(e) - current->e<double>(e)) < 1.e-5)
|
|
currentOut->p(e, currentGradOut->e<double>(e));
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, indices->lengthOf());
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
// segmen mean
|
|
sd::Status segmentMeanFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
int numClasses = output->sizeAt(0);
|
|
SD_MAP_IMPL<sd::LongType, sd::LongType> classCount; //(numClasses);
|
|
|
|
for (sd::LongType count = 0; count < numClasses; ++count) {
|
|
classCount[count] = 0;
|
|
}
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
classCount[indices->e<sd::LongType>(e)]++;
|
|
}
|
|
|
|
// if input is a vector: (as if in doc sample)
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum) / classCount[classNum]);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
int pos = 0;
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
for (sd::LongType e = 0; e < current->lengthOf(); e++) {
|
|
currentOut->p(e, currentGradOut->e<double>(e) / classCount.at(classNum));
|
|
}
|
|
}
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status segmentSumFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
// int numClasses = output->sizeAt(0);
|
|
// if input is a vector: (as if in doc sample)
|
|
sd::LongType idx = indices->e<sd::LongType>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum));
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
currentOut->assign(currentGradOut);
|
|
}
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status segmentProdFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
auto tempRes = gradOut->dup();
|
|
segmentProdFunctor(context, input, indices, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum) * tempRes->e<double>(classNum) / input->e<double>(e));
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
|
|
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
auto currentFFOut = listOfBPTensors.at(classNum);
|
|
auto mul = (*currentFFOut) * (*currentGradOut);
|
|
auto assign = (*mul) / (*current);
|
|
currentOut->assign(assign);
|
|
delete mul;
|
|
delete assign;
|
|
}
|
|
delete restDims; // Clean up allocated vector
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// Unsorted backpropagate segment ops
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
|
|
template <typename T>
|
|
static sd::Status unsortedSegmentMaxFunctorBP_(sd::LaunchContext* context, NDArray* input, NDArray* indices,
|
|
NDArray* gradOut, sd::LongType numOfClasses, NDArray* output) {
|
|
// int numOfClasses = gradOut->sizeAt(0);
|
|
// if input is a vector: (as if in doc sample)
|
|
auto tempRes = gradOut->dup();
|
|
unsortedSegmentMaxFunctor(context, input, indices, numOfClasses, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < input->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
if (sd::math::sd_abs<double,double>(tempRes->e<double>(classNum) - input->e<double>(e)) < 1.e-5)
|
|
output->p(e, gradOut->e<T>(classNum));
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(i);
|
|
NDArray* current = listOfTensors.at(i);
|
|
NDArray* currentOut = listOfOutTensors.at(i);
|
|
NDArray* currentGradOut = listOfGradOuts.at(classNum);
|
|
for (int e = 0; e < current->lengthOf(); e++) {
|
|
if (sd::math::sd_abs<double,double>(listOfBPTensors.at(classNum)->e<double>(e) - current->e<double>(e)) < 1.e-5)
|
|
currentOut->p(e, currentGradOut->e<T>(e));
|
|
}
|
|
}
|
|
|
|
delete restDims;
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status unsortedSegmentMaxFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(output->dataType(), return unsortedSegmentMaxFunctorBP_,
|
|
(context, input, indices, gradOut, numOfClasses, output), SD_NUMERIC_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( sd::Status unsortedSegmentMaxFunctorBP_,
|
|
(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
template <typename T>
|
|
static sd::Status unsortedSegmentMinFunctorBP_(sd::LaunchContext* context, NDArray* input, NDArray* indices,
|
|
NDArray* gradOut, sd::LongType numOfClasses, NDArray* output) {
|
|
auto tempRes = gradOut->dup();
|
|
unsortedSegmentMinFunctor(context, input, indices, numOfClasses, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
auto classNum = indices->e<sd::LongType>(e);
|
|
if (sd::math::sd_abs<T,T>(tempRes->t<T>(classNum) - input->t<T>(e)) < 1.e-6)
|
|
output->r<T>(e) = gradOut->t<T>(classNum);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, input->lengthOf());
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
for (sd::LongType e = 0; e < current->lengthOf(); e++) {
|
|
if (sd::math::sd_abs<T,T>(listOfBPTensors.at(classNum)->t<T>(e) - current->t<T>(e)) < 1.e-6)
|
|
currentOut->r<T>(e) = currentGradOut->t<T>(e);
|
|
}
|
|
}
|
|
delete restDims; // Clean up allocated vector
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status unsortedSegmentMinFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(output->dataType(), return unsortedSegmentMinFunctorBP_,
|
|
(context, input, indices, gradOut, numOfClasses, output), SD_NUMERIC_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( sd::Status unsortedSegmentMinFunctorBP_,
|
|
(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output),
|
|
SD_NUMERIC_TYPES);
|
|
|
|
sd::Status unsortedSegmentMeanFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, sd::LongType> classCount; //(numClasses);
|
|
|
|
for (sd::LongType count = 0; count < numOfClasses; ++count) {
|
|
classCount[count] = 0;
|
|
}
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
classCount[indices->e<sd::LongType>(e)]++;
|
|
}
|
|
|
|
// if input is a vector: (as if in doc sample)
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum) / classCount[classNum]);
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(i);
|
|
NDArray* current = listOfTensors.at(i);
|
|
NDArray* currentOut = listOfOutTensors.at(i);
|
|
NDArray* currentGradOut = listOfGradOuts.at(classNum);
|
|
auto assign = *currentGradOut / double(classCount[classNum]);
|
|
currentOut->assign(assign);
|
|
delete assign;
|
|
}
|
|
|
|
delete restDims;
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status unsortedSegmentSumFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
sd::LongType idx = indices->e<sd::LongType>(0);
|
|
if (input->isVector() || input->isScalar()) {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
sd::LongType classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum));
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
currentOut->assign(currentGradOut);
|
|
}
|
|
|
|
delete restDims;
|
|
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
sd::Status unsortedSegmentProdFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
auto tempRes = gradOut->dup();
|
|
unsortedSegmentProdFunctor(context, input, indices, numOfClasses, tempRes);
|
|
if (input->isVector() || input->isScalar()) {
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto e = start; e < stop; e++) {
|
|
auto classNum = indices->e<sd::LongType>(e);
|
|
output->p<double>(e, gradOut->e<double>(classNum) * tempRes->e<double>(classNum) / input->e<double>(e));
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, indices->lengthOf());
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfBPTensors = tempRes->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
auto currentFFOut = listOfBPTensors.at(classNum);
|
|
auto div = (*currentGradOut) / (*current);
|
|
auto assign = (*currentFFOut) * *div;
|
|
currentOut->assign(assign);
|
|
delete div;
|
|
delete assign;
|
|
}
|
|
delete restDims; // Clean up allocated vector
|
|
}
|
|
delete tempRes; // Clean up duped array
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
// template <typename T>
|
|
sd::Status unsortedSegmentSqrtNFunctorBP(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
sd::LongType numOfClasses, NDArray* output) {
|
|
SD_MAP_IMPL<sd::LongType, sd::LongType> classCount; //(numClasses);
|
|
|
|
for (sd::LongType count = 0; count < numOfClasses; ++count) {
|
|
classCount[count] = 0;
|
|
}
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
|
|
classCount[indices->e<sd::LongType>(e)]++;
|
|
}
|
|
|
|
// if input is a vector: (as if in doc sample)
|
|
if (input->isVector() || input->isScalar()) {
|
|
// auto func = PRAGMA_THREADS_FOR {
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); e++) {
|
|
auto classNum = indices->e<sd::LongType>(e);
|
|
output->p(e, gradOut->e<double>(classNum) / sd::math::sd_sqrt<double, double>(classCount[classNum]));
|
|
}
|
|
|
|
} else {
|
|
std::vector<sd::LongType> zeroVec = {0};
|
|
std::vector<sd::LongType> *restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,zeroVec.data());
|
|
ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(*restDims);
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(*restDims);
|
|
delete restDims;
|
|
for (sd::LongType i = 0; i < indices->lengthOf(); i++) {
|
|
auto classNum = indices->e<sd::LongType>(i);
|
|
auto current = listOfTensors.at(i);
|
|
auto currentOut = listOfOutTensors.at(i);
|
|
auto currentGradOut = listOfGradOuts.at(classNum);
|
|
|
|
for (int e = 0; e < current->lengthOf(); e++) {
|
|
currentOut->p<double>(e,
|
|
currentGradOut->e<double>(e) / sd::math::sd_sqrt<double, double>(classCount[classNum]));
|
|
}
|
|
}
|
|
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|
|
#endif |