cavis/libnd4j/include/ops/declarable/helpers/cpu/segment.cpp

/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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 <ops/declarable/helpers/segment.h>
#include <helpers/ShapeUtils.h>
#include <execution/Threads.h>
#include <unordered_map>

namespace sd {
    namespace ops {
        namespace helpers {

            // segment max
            template <typename T>
            static void segmentMaxFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
                //int numClasses = output->sizeAt(0);
                // if input is a vector: (as if in doc sample)
                Nd4jLong idx = indices->e<Nd4jLong>(0);
                if (input->isVector() || input->isScalar()) {
                    T val = input->e<T>(0);

                    for (Nd4jLong e = 1; e < indices->lengthOf(); e++) {
                        if (idx == indices->e<Nd4jLong>(e)) {
                            // max
                            val = sd::math::nd4j_max<T>(val, input->t<T>(e));
                        }
                        else {
                            idx = indices->e<Nd4jLong>(e);
                            val = input->t<T>(e);
                        }
                        output->r<T>(idx) = val;
                    }
                }
                else {
                    std::vector<int> restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});
                    auto listOfTensors = input->allTensorsAlongDimension(restDims);
                    auto listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    auto numOfClasses = output->sizeAt(0); // number of classes
                    std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);
                    auto maxT = listOfOutTensors.at(idx);

                    //int pos = 0;
                    maxT->assign(listOfTensors.at(0));

                    for (Nd4jLong i = 1; i < indices->lengthOf(); i++) {
                        if (indices->e<int>(i) == idx) {

                            for (Nd4jLong e = 0; e < maxT->lengthOf(); e++) {
                                maxT->r<T>(e) = sd::math::nd4j_max(maxT->t<T>(e), listOfTensors.at(i)->t<T>(e));
                            }
                        }
                        else {
                            idx = indices->e<Nd4jLong>(i);
                            maxT = listOfOutTensors.at(idx);
                            maxT->assign(listOfTensors.at(i));
                        }

                    }
                }
            }

            // segmen min
            template <typename T>
            static void segmentMinFunctor_(NDArray* input, NDArray* indices, NDArray* output) {
                //int numClasses = output->sizeAt(0);
                // if input is a vector: (as if in doc sample)
                Nd4jLong idx = indices->e<Nd4jLong>(0);
                if (input->isVector() || input->isScalar()) {
                    T val = input->e<T>(0);

                    for (Nd4jLong e = 1; e < indices->lengthOf(); e++) {
                        if (idx == indices->e<Nd4jLong>(e)) {
                            // min
                            val = sd::math::nd4j_min<T>(val, input->t<T>(e));
                        }
                        else {
                            idx = indices->e<Nd4jLong>(e);
                            val = input->t<T>(e);
                        }
                        output->r<T>(idx) = val;
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    int numOfClasses = output->sizeAt(0); // number of classes
                    std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);
                    auto minT = listOfOutTensors.at(idx);

                    int pos = 0;
                    minT->assign(listOfTensors.at(0));

                    for (Nd4jLong i = 1; i < indices->lengthOf(); i++) {
                        if (indices->e<Nd4jLong>(i) == idx) {

                            for (Nd4jLong e = 0; e < minT->lengthOf(); e++) {
                                minT->p(e, sd::math::nd4j_min(minT->e<T>(e), listOfTensors.at(i)->e<T>(e)));
                            }
                        }
                        else {
                            idx = indices->e<Nd4jLong>(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<int>(0);
                if (input->isVector() || input->isScalar()) {
                    T val = T(0.f);
                    int count = 0;

                    for (Nd4jLong e = 0; e < indices->lengthOf(); e++) {
                        if (idx == indices->e<int>(e)) {
                            // mean
                            val += input->e<T>(e);
                            count++;
                        }
                        else {
                            output->p<T>(idx, val / count);
                            idx = indices->e<int>(e);
                            val = input->e<T>(e);
                            count = 1;
                        }
                        output->p<T>(idx, val / count);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    auto listOfTensors = input->allTensorsAlongDimension(restDims);
                    auto listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    int numOfClasses = output->sizeAt(0); // number of classes
                    std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);
                    auto meanT = listOfOutTensors.at(idx);
                    int count = 1;
                    auto meanV = meanT->dup();
                    meanV.assign(listOfTensors.at(0));

                    for (Nd4jLong i = 1; i < indices->lengthOf(); i++) {
                        if (indices->e<int>(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 {
                            //meanT->assign(meanV);
                            meanV.applyScalar(scalar::Divide, count, *meanT);
                            idx = indices->e<int>(i);
                            meanT = listOfOutTensors.at(idx);
                            meanV.assign(listOfTensors.at(i));
                            count = 1;
                        }
                        meanV.applyScalar(scalar::Divide, count, *meanT);
                    }
                }
            }

            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 (Nd4jLong e = 0; e < indices->lengthOf(); e++) {
                        if (idx == indices->e<int>(e)) {
                            // sum
                            val += input->t<T>(e);
                        }
                        else {
                            idx = indices->e<int>(e);
                            val = input->t<T>(e);
                        }
                        output->p(idx, val);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    auto listOfTensors = input->allTensorsAlongDimension(restDims);
                    auto listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    int numOfClasses = output->sizeAt(0); // number of classes
                    std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);
                    auto sumT = listOfOutTensors.at(idx);

                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        if (indices->e<int>(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));
                        }
                    }
                }
            }

            template <typename T>
            static void segmentProdFunctor_(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);
                output->assign(1.f);
                if (input->isVector() || input->isScalar()) {
                    T val = input->e<T>(0);
                    int count = 0;

                    for (Nd4jLong e = 1; e < indices->lengthOf(); e++) {
                        if (idx == indices->e<int>(e)) {
                            // sum
                            val *= input->e<T>(e);
                        }
                        else {
                            idx = indices->e<int>(e);
                            val = input->e<T>(e);
                        }
                        output->p(idx, val);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    auto listOfTensors = input->allTensorsAlongDimension(restDims);
                    auto listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    int numOfClasses = output->sizeAt(0); // number of classes
                    auto sumT = listOfOutTensors.at(idx);
                    sumT->assign(listOfTensors.at(0));
                    for (Nd4jLong i = 1; i < indices->lengthOf(); i++) {
                        if (indices->e<int>(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));
                        }
                    }
                }
            }

//    template <typename T>
//    static bool segmentIndicesValidate_(NDArray* indices, NDArray& aexpected, NDArray& anOutput) {
//      }

            void segmentMaxFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), segmentMaxFunctor_, (input, indices, output), LIBND4J_TYPES);
            }

            void segmentMinFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), segmentMinFunctor_, (input, indices, output), LIBND4J_TYPES);
            }

            void segmentMeanFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), segmentMeanFunctor_, (input, indices, output), LIBND4J_TYPES);
            }

            void segmentSumFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), segmentSumFunctor_, (input, indices, output), LIBND4J_TYPES);
            }

            void segmentProdFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), segmentProdFunctor_, (input, indices, output), LIBND4J_TYPES);
            }

            bool segmentIndicesValidate(sd::LaunchContext * context, NDArray* indices, NDArray& expected, NDArray& output) {
                auto val = indices->e(0);
                for (Nd4jLong e = 1; e < indices->lengthOf(); e++) {
                    output = indices->e(e);
                    if (val.e<Nd4jLong>(0) > output.e<Nd4jLong>(0))
                        return false;
                    val = indices->e(e);
                }

                return true;
            }

            //BUILD_SINGLE_TEMPLATE(template bool segmentIndicesValidate_, (NDArray*, NDArray&, NDArray&), LIBND4J_TYPES);
            BUILD_SINGLE_TEMPLATE(template void segmentProdFunctor_, (NDArray* input, NDArray* indices, NDArray* output), LIBND4J_TYPES);
            BUILD_SINGLE_TEMPLATE(template void segmentSumFunctor_, (NDArray* input, NDArray* indices, NDArray* output), LIBND4J_TYPES);
            BUILD_SINGLE_TEMPLATE(template void segmentMeanFunctor_, (NDArray* input, NDArray* indices, NDArray* output), LIBND4J_TYPES);
            BUILD_SINGLE_TEMPLATE(template void segmentMinFunctor_, (NDArray* input, NDArray* indices, NDArray* output), LIBND4J_TYPES);
            BUILD_SINGLE_TEMPLATE(template void segmentMaxFunctor_, (NDArray* input, NDArray* indices, NDArray* output), LIBND4J_TYPES);
            // -------------------------------------------------------------------------------------------------------------- //
            // Unsorted segment ops
            // -------------------------------------------------------------------------------------------------------------- //

            bool unsortedSegmentIndicesValidate(sd::LaunchContext * context, NDArray* indices, Nd4jLong expected, Nd4jLong& output) {
                Nd4jLong val = indices->e<Nd4jLong>(0);

                Nd4jLong maxInd = indices->argMax();
                if (indices->e<Nd4jLong>(maxInd) >= expected) {
                    output = val;
                    return false;
                }
                output = expected;
                return true;
            }

            template <typename T>
            static void unsortedSegmentMaxFunctor_(NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {

                // if input is a vector: (as if in doc sample)
                //int idx = static_cast<int>((*indices)(0.));
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());
                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(e)].push_back(e);


                //std::sort(idxs.begin(), idxs.end());

                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->e<T>(fi->second.at(0));
                        for (Nd4jLong idx = 1; idx < static_cast<Nd4jLong>(fi->second.size()); ++idx) {
                            val = sd::math::nd4j_max(val, input->e<T>(fi->second.at(idx)));
                        }
                        output->p(fi->first, val);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(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 (Nd4jLong idx = 0; idx < listOfTensors.size(); ++idx) {
                               if(idx >= fi->second.size() ||  fi->second.size() < 2 ||  fi->second.at(idx) >= listOfTensors.size()) {
                                   continue;
                            }

                            auto maxT = listOfTensors.at(fi->second.at(idx));
                            for (Nd4jLong e = 0; e < outputT->lengthOf(); ++e) {
                                T val = sd::math::nd4j_max(maxT->e<T>(e), outputT->e<T>(e));

                                outputT->p(e, val);
                            }
                        }


                    }


                }
            }
            void unsortedSegmentMaxFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentMaxFunctor_, (input, indices, numOfClasses, output), NUMERIC_TYPES);
            }
            BUILD_SINGLE_TEMPLATE(template void unsortedSegmentMaxFunctor_, (NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output), NUMERIC_TYPES);

            template <typename T>
            static void unsortedSegmentMinFunctor_(NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                // if input is a vector: (as if in doc sample)
                //int idx = static_cast<int>((*indices)(0.));
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());

                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(e)].push_back(e);

                //std::sort(idxs.begin(), idxs.end());

                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::nd4j_min(val, input->t<T>(fi->second.at(idx)));
                        }
                        output->r<T>(fi->first) = val;
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(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 (Nd4jLong e = 0; e < outputT->lengthOf(); ++e) {
                                outputT->r<T>(e) = sd::math::nd4j_min(minT->t<T>(e), outputT->t<T>(e));
                            }
                        }
                        //outputT->assign(maxT);
                    }
                }

            }
            void unsortedSegmentMinFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentMinFunctor_, (input, indices, numOfClasses, output),
                                      NUMERIC_TYPES);
            }

            BUILD_SINGLE_TEMPLATE(template void unsortedSegmentMinFunctor_, (NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output), NUMERIC_TYPES);

            void unsortedSegmentMeanFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());
                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(e)].push_back(e);

                //std::sort(idxs.begin(), idxs.end());

                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));
                        int 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 {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(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)));
                        Nd4jLong loopSize = fi->second.size();

                        for (Nd4jLong 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, Nd4jLong numOfClasses, NDArray* output) {
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());
                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(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));
                        Nd4jLong loop_size = fi->second.size();

                        // FIXME: parallelism here?
                        for (Nd4jLong idx = 1; idx < loop_size; ++idx) {
                            sumValue += input->e<double>(fi->second.at(idx));
                        }
                        output->p(fi->first, sumValue);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    for (auto fi = idxs.begin(); fi != idxs.end(); ++fi) {
                        auto outputT = listOfOutTensors.at(fi->first);
                        outputT->assign(listOfTensors.at(fi->second.at(0)));
                        Nd4jLong loop_size = fi->second.size();

                        // FIXME: parallelism here?
                        for (Nd4jLong idx = 1; idx < loop_size; ++idx) {
                            auto current = listOfTensors.at(fi->second.at(idx));
                            *(outputT) += *current;
                        }
                        //outputT->assign(maxT);
                    }
                }
            }

            template <typename T>
            void unsortedSegmentProdFunctor_(NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());
                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(e)].push_back(e);

                //std::sort(idxs.begin(), idxs.end());

                output->assign(1.f);

                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 {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(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, Nd4jLong numOfClasses, NDArray* output) {
                BUILD_SINGLE_SELECTOR(input->dataType(), unsortedSegmentProdFunctor_, (input, indices, numOfClasses, output), NUMERIC_TYPES);
            }
            BUILD_SINGLE_TEMPLATE(template void unsortedSegmentProdFunctor_, (NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output), NUMERIC_TYPES);

            void unsortedSegmentSqrtNFunctor(sd::LaunchContext * context, NDArray* input, NDArray* indices, Nd4jLong numOfClasses, NDArray* output) {
                MAP_IMPL<Nd4jLong, std::vector<Nd4jLong>> idxs;//(indices->lengthOf());
                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e)
                    idxs[indices->e<Nd4jLong>(e)].push_back(e);

                //std::sort(idxs.begin(), idxs.end());

                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::nd4j_sqrt<Nd4jLong, double>(fi->second.size()));
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(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->assign(maxT);
                        (*outputT) /= sd::math::nd4j_sqrt<size_t, double>(fi->second.size());
                    }
                }
            }

            // -------------------------------------------------------------------------------------------------------------- //
            // Backpropagate ops helpers
            // -------------------------------------------------------------------------------------------------------------- //
            // Sorted backpropagate ops
            //
            // segment max
            template <typename T>
            int segmentMaxFunctorBP_(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, NDArray* output) {
                //int numOfClasses = gradOut->sizeAt(0);
                // if input is a vector: (as if in doc sample)
                auto tempRes = gradOut->dup();
                segmentMaxFunctor_<T>(input, indices, &tempRes);
                if (input->isVector() || input->isScalar()) {
                    Nd4jLong loop_size = input->lengthOf();

                    auto func = PRAGMA_THREADS_FOR {
                        for (auto e = start; e < stop; e++) {
                            auto classNum = indices->e<Nd4jLong>(e);
                            if (sd::math::nd4j_abs(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<int> restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //int numOfClasses = tempRes.sizeAt(0); // number of classes
                    //std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);

                    auto func = PRAGMA_THREADS_FOR {
                        for (auto i = start; i < stop; i++) {
                            auto classNum = indices->e<Nd4jLong>(i);
                            auto current = listOfTensors.at(i);
                            auto currentOut = listOfOutTensors.at(i);
                            auto currentGradOut = listOfGradOuts.at(classNum);

                            for (Nd4jLong e = 0; e < current->lengthOf(); e++) {
                                if (sd::math::nd4j_abs(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());
                }

                return ND4J_STATUS_OK;
            }

            int segmentMaxFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, NDArray* output) {
                BUILD_SINGLE_SELECTOR(output->dataType(), return segmentMaxFunctorBP_, (context, input, indices, gradOut, output), NUMERIC_TYPES);
            }
            BUILD_SINGLE_TEMPLATE(template int segmentMaxFunctorBP_, (sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, NDArray* output), NUMERIC_TYPES);

            // segmen min
            int 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<Nd4jLong>(e);
                            if (sd::math::nd4j_abs(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 {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //int numOfClasses = tempRes.sizeAt(0); // number of classes
                    //std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);
                    output->assign(0.);
                    int pos = 0;

                    auto func = PRAGMA_THREADS_FOR {
                        for (auto i = start; i < stop; i++) {
                            auto classNum = indices->e<Nd4jLong>(i);
                            auto current = listOfTensors.at(i);
                            auto currentOut = listOfOutTensors.at(i);
                            auto currentGradOut = listOfGradOuts.at(classNum);

                            for (Nd4jLong e = 0; e < current->lengthOf(); e++) {
                                if (sd::math::nd4j_abs(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());
                }
                return ND4J_STATUS_OK;
            }

            // segmen mean
            int segmentMeanFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, NDArray* output) {
                int numClasses = output->sizeAt(0);
                MAP_IMPL<Nd4jLong, Nd4jLong> classCount;//(numClasses);

                for (Nd4jLong count = 0; count < numClasses; ++count) {
                    classCount[count] = 0;
                }

                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                    classCount[indices->e<Nd4jLong>(e)] ++;
                }

                // if input is a vector: (as if in doc sample)
                if (input->isVector() || input->isScalar()) {
                    for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum) / classCount[classNum]);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);
                    ;

                    int pos = 0;
                    //auto func = [&](uint64_t thread_id, uint64_t start, uint64_t stop, uint64_t increment) -> void {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto current = listOfTensors.at(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);

                        for (Nd4jLong e = 0; e < current->lengthOf(); e++) {
                            currentOut->p(e, currentGradOut->e<double>(e) / classCount.at(classNum));
                        }
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }
                return ND4J_STATUS_OK;
            }

            int 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)
                Nd4jLong idx = indices->e<Nd4jLong>(0);
                if (input->isVector() || input->isScalar()) {
                    for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum));
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto current = listOfTensors.at(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);

                        currentOut->assign(currentGradOut);
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }
                return Status::OK();
            }

            int 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 (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum) * tempRes.e<double>(classNum)/ input->e<double>(e));
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //int numOfClasses = tempRes.sizeAt(0); // number of classes
                    //std::vector<std::pair<NDArray*, int>> outputs(numOfClasses);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto current = listOfTensors.at(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);
                        auto currentFFOut = listOfBPTensors.at(classNum);

                        currentOut->assign((*currentFFOut) * (*currentGradOut) / (*current));
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }

                return ND4J_STATUS_OK;
            }

            // -------------------------------------------------------------------------------------------------------------- //
            // Unsorted backpropagate segment ops
            // -------------------------------------------------------------------------------------------------------------- //

            template <typename T>
            static int unsortedSegmentMaxFunctorBP_(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong 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 (Nd4jLong e = 0; e < input->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        if (sd::math::nd4j_abs(tempRes.e<double>(classNum) - input->e<double>(e)) < 1.e-5)
                            output->p(e, gradOut->e<T>(classNum));
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(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::nd4j_abs(listOfBPTensors.at(classNum)->e<double>(e) - current->e<double>(e)) < 1.e-5)
                                currentOut->p(e, currentGradOut->e<T>(e));
                        }
                    }
                }

                return ND4J_STATUS_OK;
            }

            int unsortedSegmentMaxFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output) {
                BUILD_SINGLE_SELECTOR(output->dataType(), return unsortedSegmentMaxFunctorBP_, (context, input, indices, gradOut, numOfClasses, output), NUMERIC_TYPES);
            }
            BUILD_SINGLE_TEMPLATE(template int unsortedSegmentMaxFunctorBP_, (sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output), NUMERIC_TYPES);

            template <typename T>
            static int unsortedSegmentMinFunctorBP_(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong 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<Nd4jLong>(e);
                            if (sd::math::nd4j_abs(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 {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto current = listOfTensors.at(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);

                        for (Nd4jLong e = 0; e < current->lengthOf(); e++) {
                            if (sd::math::nd4j_abs(listOfBPTensors.at(classNum)->t<T>(e) - current->t<T>(e)) < 1.e-6)
                                currentOut->r<T>(e) = currentGradOut->t<T>(e);
                        }
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }

                return ND4J_STATUS_OK;
            }

            int unsortedSegmentMinFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output) {
                BUILD_SINGLE_SELECTOR(output->dataType(), return unsortedSegmentMinFunctorBP_, (context, input, indices, gradOut, numOfClasses, output), NUMERIC_TYPES);
            }
            BUILD_SINGLE_TEMPLATE(template int unsortedSegmentMinFunctorBP_, (sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output), NUMERIC_TYPES);

            int unsortedSegmentMeanFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output) {

                MAP_IMPL<Nd4jLong, Nd4jLong> classCount;//(numClasses);

                for (Nd4jLong count = 0; count < numOfClasses; ++count) {
                    classCount[count] = 0;
                }

                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                    classCount[indices->e<Nd4jLong>(e)]++;
                }

                // if input is a vector: (as if in doc sample)
                if (input->isVector() || input->isScalar()) {
                    for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum) / classCount[classNum]);
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(i);
                        NDArray* current = listOfTensors.at(i);
                        NDArray* currentOut = listOfOutTensors.at(i);
                        NDArray* currentGradOut = listOfGradOuts.at(classNum);
                        currentOut->assign(*currentGradOut / double(classCount[classNum]));
                    }
                }
                return ND4J_STATUS_OK;
            }

            int unsortedSegmentSumFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output) {

                // if input is a vector: (as if in doc sample)
                Nd4jLong idx = indices->e<Nd4jLong>(0);
                if (input->isVector() || input->isScalar()) {
                    for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                        Nd4jLong classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum));
                    }
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);

                        currentOut->assign(currentGradOut);
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }
                return Status::OK();
            }

            int unsortedSegmentProdFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong 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<Nd4jLong>(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 {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfBPTensors = tempRes.allTensorsAlongDimension(restDims);
                    ResultSet listOfGradOuts = gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors = input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(i);
                        auto current = listOfTensors.at(i);
                        auto currentOut = listOfOutTensors.at(i);
                        auto currentGradOut = listOfGradOuts.at(classNum);
                        auto currentFFOut = listOfBPTensors.at(classNum);

                        currentOut->assign((*currentFFOut) * (*currentGradOut) / (*current));
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }

                return Status::OK();
            }

//    template <typename T>
            int unsortedSegmentSqrtNFunctorBP(sd::LaunchContext * context, NDArray* input, NDArray* indices, NDArray* gradOut, Nd4jLong numOfClasses, NDArray* output) {
                MAP_IMPL<Nd4jLong, Nd4jLong> classCount;//(numClasses);

                for (Nd4jLong count = 0; count < numOfClasses; ++count) {
                    classCount[count] = 0;
                }

                for (Nd4jLong e = 0; e < indices->lengthOf(); ++e) {
                    classCount[indices->e<Nd4jLong>(e)]++;
                }

                // if input is a vector: (as if in doc sample)
                if (input->isVector() || input->isScalar()) {
                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong e = 0; e < indices->lengthOf(); e++) {
                        auto classNum = indices->e<Nd4jLong>(e);
                        output->p(e, gradOut->e<double>(classNum) / sd::math::nd4j_sqrt<double, double>(classCount[classNum]));
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }
                else {
                    auto restDims = ShapeUtils::evalDimsToExclude(input->rankOf(), {0});

                    ResultSet listOfGradOuts  =gradOut->allTensorsAlongDimension(restDims);
                    ResultSet listOfTensors  =input->allTensorsAlongDimension(restDims);
                    ResultSet listOfOutTensors  =output->allTensorsAlongDimension(restDims);

                    //auto func = PRAGMA_THREADS_FOR {
                    for (Nd4jLong i = 0; i < indices->lengthOf(); i++) {
                        auto classNum = indices->e<Nd4jLong>(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::nd4j_sqrt<double, double>(classCount[classNum]));
                        }
                    }
                    //};

                    //samediff::Threads::parallel_for(func, 0, indices->lengthOf());
                }
                return Status::OK();
            }

        }
    }
}