chore: import upstream snapshot with attribution

libnd4j/include/ops/impl/BroadcastBoolOpsTuple.cpp

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com
+//
+#include <ops/BroadcastBoolOpsTuple.h>
+
+namespace sd {
+BroadcastBoolOpsTuple BroadcastBoolOpsTuple::custom(sd::scalar::BoolOps scalar, sd::pairwise::BoolOps pairwise,
+                                                    sd::broadcast::BoolOps broadcast) {
+  BroadcastBoolOpsTuple t(scalar, pairwise, broadcast);
+  return t;
+}
+}  // namespace sd

libnd4j/include/ops/impl/BroadcastIntOpsTuple.cpp

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com
+//
+#include <ops/BroadcastIntOpsTuple.h>
+
+namespace sd {
+BroadcastIntOpsTuple BroadcastIntOpsTuple::custom(scalar::IntOps scalar, pairwise::IntOps pairwise,
+                                                  broadcast::IntOps broadcast) {
+  BroadcastIntOpsTuple t(scalar, pairwise, broadcast);
+  return t;
+}
+}  // namespace sd

libnd4j/include/ops/impl/BroadcastOpsTuple.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com
+//
+#include <ops/BroadcastOpsTuple.h>
+
+namespace sd {
+BroadcastOpsTuple BroadcastOpsTuple::custom(sd::scalar::Ops scalar, sd::pairwise::Ops pairwise,
+                                            sd::broadcast::Ops broadcast) {
+  BroadcastOpsTuple t(scalar, pairwise, broadcast);
+  return t;
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::Add() { return custom(sd::scalar::Add, sd::pairwise::Add, sd::broadcast::Add); }
+
+BroadcastOpsTuple BroadcastOpsTuple::Assign() {
+  return custom(sd::scalar::CopyPws, sd::pairwise::CopyPws, sd::broadcast::CopyPws);
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::Divide() {
+  return custom(sd::scalar::Divide, sd::pairwise::Divide, sd::broadcast::Divide);
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::DivideNoNan() {
+  return custom(sd::scalar::DivideNoNan, sd::pairwise::DivideNoNan, sd::broadcast::DivideNoNan);
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::Multiply() {
+  return custom(sd::scalar::Multiply, sd::pairwise::Multiply, sd::broadcast::Multiply);
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::Subtract() {
+  return custom(sd::scalar::Subtract, sd::pairwise::Subtract, sd::broadcast::Subtract);
+}
+BroadcastOpsTuple BroadcastOpsTuple::IGamma() {
+  return custom(sd::scalar::IGamma, sd::pairwise::IGamma, sd::broadcast::IGamma);
+}
+BroadcastOpsTuple BroadcastOpsTuple::IGammac() {
+  return custom(sd::scalar::IGammac, sd::pairwise::IGammac, sd::broadcast::IGammac);
+}
+
+BroadcastOpsTuple BroadcastOpsTuple::Pow() { return custom(sd::scalar::Pow, sd::pairwise::Pow, sd::broadcast::Pow); }
+BroadcastOpsTuple BroadcastOpsTuple::PowDerivative() {
+  return custom(sd::scalar::PowDerivative, sd::pairwise::PowDerivative, sd::broadcast::PowDerivative);
+}
+
+}  // namespace sd

libnd4j/include/ops/impl/compilation_units/specials_double.cpp.in

@@ -0,0 +1,53 @@
+/*
+* ******************************************************************************
+* *
+* *
+* * 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 copyrig
+ * ht 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 raver119@gmail.com
+// @date Dec 26, 2024 // Adjusted date format/value to match example style
+//
+
+#include <ops/impl/specials_double.hpp> // Original include
+#include <system/type_boilerplate.h> // For type lists like SD_NUMERIC_TYPES_PART_X
+#include <loops/pairwise_instantiations.h> // For ITERATE_COMBINATIONS macro
+// Note: Instantiations are generated to prevent compiler memory issues
+
+namespace sd {
+
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_CLASSCOMBINATION,
+     sd::DoubleMethods,
+    ()
+)
+
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    sd::SpecialTypeConverter::convertGeneric,
+    (sd::Pointer * extras, void *dx, sd::LongType N, void *dz);
+)
+
+
+} // namespace sd

libnd4j/include/ops/impl/compilation_units/specials_single.cpp.in

@@ -0,0 +1,40 @@
+/*
+* ******************************************************************************
+* *
+* *
+* * 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 raver119@gmail.com
+// @date Dec 26, 2024 // Adjusted date format/value to match example style
+//
+
+#include <ops/impl/specials_single.hpp> // Original include
+#include <system/type_boilerplate.h> // For type lists like SD_NUMERIC_TYPES_PART_X
+#include <loops/pairwise_instantiations.h> // For ITERATE_COMBINATIONS macro
+// Note: Instantiations are generated to prevent compiler memory issues
+
+namespace sd {
+
+#define SPECIAL_METHODS_DECL(T) template  class SpecialMethods<GET_SECOND(T)>;
+ITERATE_LIST(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SPECIAL_METHODS_DECL
+
+)
+
+} // namespace sd

libnd4j/include/ops/impl/gemm.cpp

@@ -0,0 +1,135 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Created by raver119 on 07.10.2017.
+// Modified by GS <sgazeos@gmail.com> on 3/9/2018
+//
+#include <execution/Threads.h>
+#include <ops/gemm.h>
+#include <system/Environment.h>
+#include <types/types.h>
+
+namespace sd {
+namespace blas {
+
+template <typename T>
+void *transpose(int orderSource, int orderTarget, int rows, int cols, void *vsource) {
+  auto ret = new T[rows * cols];
+  auto source = reinterpret_cast<T *>(vsource);
+
+  // handle transpose in parallel
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto r = start; r < stop; r++) {
+      for (int c = 0; c < cols; c++) {
+        int zIdx = orderTarget == CblasRowMajor ? linearIndexC(rows, cols, r, c) : linearIndexF(rows, cols, r, c);
+        int xIdx = orderSource == CblasColMajor ? linearIndexF(rows, cols, r, c) : linearIndexC(rows, cols, r, c);
+
+        ret[zIdx] = source[xIdx];
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, rows);
+
+  return ret;
+}
+
+template <typename X, typename Y, typename Z>
+void GEMM<X, Y, Z>::op(int Order, int TransA, int TransB, int M, int N, int K, double alpha, void *vA, int lda,
+                       void *vB, int ldb, double beta, void *vC, int ldc) {
+  auto A = reinterpret_cast<X *>(vA);
+  auto B = reinterpret_cast<Y *>(vB);
+  auto C = reinterpret_cast<Z *>(vC);
+
+  bool transAFlag = TransA == CblasTrans;
+  bool transBFlag = TransB == CblasTrans;
+
+  if (beta == 0.0) {
+    Z z = 0.f;
+    int length = M * N;
+    if (length <= Environment::getInstance().elementwiseThreshold()) {
+      for (int r = 0; r < length; r++) C[r] = z;
+    } else {
+      auto func = PRAGMA_THREADS_FOR {
+        for (auto r = start; r < stop; r++) C[r] = z;
+      };
+      samediff::Threads::parallel_for(func, 0, length);
+    }
+  }
+
+  auto func = PRAGMA_THREADS_FOR_2D {
+    for (auto r = start_x; r < stop_x; r += inc_x) {
+      for (auto c = start_y; c < stop_y; c += inc_y) {
+        int zIdx = linearIndexF(M, N, r, c);
+
+        Z dot = static_cast<Z>(0.0f);
+
+        if (alpha != 0.0) {
+          int bIdx;  // = linearIndexF(K, N, 0, c);
+          int aIdx;
+
+          for (int k = 0; k < K; k++) {
+            aIdx = (transAFlag ? linearIndexC(M, K, r, k) : linearIndexF(M, K, r, k));
+            bIdx = (transBFlag ? linearIndexC(K, N, k, c) : linearIndexF(K, N, k, c));
+            dot += static_cast<Z>(alpha) * static_cast<Z>(A[aIdx]) *
+                   static_cast<Z>(B[bIdx]);  // A[aIdx]sd::math::sd_dot<T>(aX, bX, K) * alpha;
+          }
+        }
+
+        if (beta != 0.0) {
+          C[zIdx] = static_cast<Z>(dot + static_cast<Z>(beta) * C[zIdx]);
+        } else {
+          C[zIdx] = static_cast<Z>(dot);
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, M, 1, 0, N, 1);
+}
+
+template <typename X, typename Y, typename Z>
+void GEMV<X, Y, Z>::op(int TRANS, int M, int N, double alpha, void *vX, int lda, void *vY, int incx, double beta,
+                       void *vZ, int incy) {
+  auto x = reinterpret_cast<X *>(vX);
+  auto y = reinterpret_cast<Y *>(vY);
+  auto z = reinterpret_cast<Z *>(vZ);
+
+  auto aT = TRANS == CblasTrans ? reinterpret_cast<X *>(sd::blas::transpose<X>(CblasColMajor, CblasRowMajor, M, N,
+                                                                               reinterpret_cast<void *>(x)))
+                                : x;
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto r = start; r < stop; r++) {
+      int aIdx = linearIndexC(M, N, r, 0);
+      auto aX = aT + aIdx;
+
+      auto dot = sd::math::sd_dot<X, Y, Z>(aX, y, lda) * static_cast<Z>(alpha);
+      z[r] = beta == 0.0f ? dot : dot + static_cast<Z>(beta) * z[r];
+    }
+  };
+  samediff::Threads::parallel_for(func, 0, M);
+
+  if (TRANS == CblasTrans) delete[] aT;
+}
+
+// BUILD_TRIPLE_TEMPLATE(template class  GEMV, , SD_COMMON_TYPES, SD_FLOAT_TYPES, SD_FLOAT_TYPES);
+// BUILD_TRIPLE_TEMPLATE(template class  GEMM, , SD_COMMON_TYPES, SD_FLOAT_TYPES, SD_FLOAT_TYPES);
+}  // namespace blas
+}  // namespace sd

libnd4j/include/ops/impl/specials_double.hpp

@@ -0,0 +1,336 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com, created on 07.10.2017.
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <array/NDArray.h>
+#include <helpers/Loops.h>
+
+#include <helpers/shape.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/specials.h>
+#include <types/types.h>
+#include <loops/pairwise_instantiations.h>
+
+namespace sd {
+
+template <typename S, typename T>
+void SpecialTypeConverter::convertGeneric(sd::Pointer *extras, void *dx, sd::LongType N, void *dz) {
+  auto x = reinterpret_cast<S *>(dx);
+  auto z = reinterpret_cast<T *>(dz);
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto i = start; i < stop; i++) {
+      z[i] = static_cast<T>(x[i]);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, N);
+};
+
+template <typename X, typename Y>
+void quickSort_parallel_internal_key(X *key, sd::LongType const *xShapeInfo, Y *values, sd::LongType const *yShapeInfo,
+                                     LongType left, LongType right, LongType cutoff, bool descending) {
+  sd::LongType i = left, j = right;
+  X ktmp;
+  LongType pivotCoords[] = {(left + right) / 2};
+  LongType pivotIndex;
+  COORDS2INDEX(1, shape::stride(xShapeInfo), pivotCoords, pivotIndex);
+  X pivot = key[pivotIndex];
+
+  Y vtmp;
+
+  {
+    /* PARTITION PART */
+    while (i <= j) {
+      if (descending) {
+        LongType iIndex, jIndex;
+        LongType iCoords[] = {i};
+        LongType jCoords[] = {j};
+        COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iIndex);
+        COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jIndex);
+        while (key[iIndex] > pivot) {
+          i++;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iIndex);
+        }
+        while (key[jIndex] < pivot) {
+          j--;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jIndex);
+        }
+        if (i <= j) {
+          ktmp = key[iIndex];
+          key[iIndex] = key[jIndex];
+          key[jIndex] = ktmp;
+
+          LongType iValueIndex, jValueIndex;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iValueIndex);
+          COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jValueIndex);
+          vtmp = values[iValueIndex];
+          values[iValueIndex] = values[jValueIndex];
+          values[jValueIndex] = vtmp;
+
+          i++;
+          j--;
+        }
+      } else {
+        LongType iIndex, jIndex;
+        LongType iCoords[] = {i};
+        LongType jCoords[] = {j};
+        COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iIndex);
+        COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jIndex);
+        while (key[iIndex] < pivot) {
+          i++;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iIndex);
+        }
+        while (key[jIndex] > pivot) {
+          j--;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jIndex);
+        }
+        if (i <= j) {
+          ktmp = key[iIndex];
+          key[iIndex] = key[jIndex];
+          key[jIndex] = ktmp;
+
+          LongType iValueIndex, jValueIndex;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iValueIndex);
+          COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jValueIndex);
+          vtmp = values[iValueIndex];
+          values[iValueIndex] = values[jValueIndex];
+          values[jValueIndex] = vtmp;
+
+          i++;
+          j--;
+        }
+      }
+    }
+  }
+
+  if (((right - left) < cutoff)) {
+    if (left < j) {
+      quickSort_parallel_internal_key(key, xShapeInfo, values, yShapeInfo, left, j, cutoff, descending);
+    }
+    if (i < right) {
+      quickSort_parallel_internal_key(key, xShapeInfo, values, yShapeInfo, i, right, cutoff, descending);
+    }
+  } else {
+    PRAGMA_OMP_TASK {
+      quickSort_parallel_internal_key(key, xShapeInfo, values, yShapeInfo, left, j, cutoff, descending);
+    }
+    PRAGMA_OMP_TASK {
+      quickSort_parallel_internal_key(key, xShapeInfo, values, yShapeInfo, i, right, cutoff, descending);
+    }
+  }
+}
+template <typename X, typename Y>
+void quickSort_parallel_internal_value(X *key, sd::LongType const *xShapeInfo, Y *value, sd::LongType const *yShapeInfo,
+                                       LongType left, LongType right, LongType cutoff, bool descending) {
+  sd::LongType i = left, j = right;
+  X ktmp;
+  LongType pivotCoords[] = {(left + right) / 2};
+  LongType pivotIndex;
+  COORDS2INDEX(1, shape::stride(yShapeInfo), pivotCoords, pivotIndex);
+  Y pivot = value[pivotIndex];
+
+  Y vtmp;
+
+  {
+    /* PARTITION PART */
+    while (i <= j) {
+      if (descending) {
+        LongType iIndex, jIndex;
+        LongType iCoords[] = {i};
+        LongType jCoords[] = {j};
+        COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iIndex);
+        COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jIndex);
+        while (value[iIndex] > pivot) {
+          i++;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iIndex);
+        }
+        while (value[jIndex] < pivot) {
+          j--;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jIndex);
+        }
+        if (i <= j) {
+          LongType iKeyIndex, jKeyIndex;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iKeyIndex);
+          COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jKeyIndex);
+          ktmp = key[iKeyIndex];
+          key[iKeyIndex] = key[jKeyIndex];
+          key[jKeyIndex] = ktmp;
+
+          vtmp = value[iIndex];
+          value[iIndex] = value[jIndex];
+          value[jIndex] = vtmp;
+
+          i++;
+          j--;
+        }
+      } else {
+        LongType iIndex, jIndex;
+        LongType iCoords[] = {i};
+        LongType jCoords[] = {j};
+        COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iIndex);
+        COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jIndex);
+        while (value[iIndex] < pivot) {
+          i++;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), iCoords, iIndex);
+        }
+        while (value[jIndex] > pivot) {
+          j--;
+          COORDS2INDEX(1, shape::stride(yShapeInfo), jCoords, jIndex);
+        }
+        if (i <= j) {
+          LongType iKeyIndex, jKeyIndex;
+          COORDS2INDEX(1, shape::stride(xShapeInfo), iCoords, iKeyIndex);
+          COORDS2INDEX(1, shape::stride(xShapeInfo), jCoords, jKeyIndex);
+          ktmp = key[iKeyIndex];
+          key[iKeyIndex] = key[jKeyIndex];
+          key[jKeyIndex] = ktmp;
+
+          vtmp = value[iIndex];
+          value[iIndex] = value[jIndex];
+          value[jIndex] = vtmp;
+
+          i++;
+          j--;
+        }
+      }
+    }
+  }
+
+  if (((right - left) < cutoff)) {
+    if (left < j) {
+      quickSort_parallel_internal_value(key, xShapeInfo, value, yShapeInfo, left, j, cutoff, descending);
+    }
+    if (i < right) {
+      quickSort_parallel_internal_value(key, xShapeInfo, value, yShapeInfo, i, right, cutoff, descending);
+    }
+  } else {
+    PRAGMA_OMP_TASK {
+      quickSort_parallel_internal_value(key, xShapeInfo, value, yShapeInfo, left, j, cutoff, descending);
+    }
+    PRAGMA_OMP_TASK {
+      quickSort_parallel_internal_value(key, xShapeInfo, value, yShapeInfo, i, right, cutoff, descending);
+    }
+  }
+}
+template <typename X, typename Y>
+static void quickSort_parallel_key(NDArray *x, NDArray *y, sd::LongType lenArray, int numThreads,
+                                   bool descending) {
+  auto array = reinterpret_cast<X *>(x->bufferAsT<X>());
+  auto values = reinterpret_cast<Y *>(y->bufferAsT<Y>());
+  int cutoff = 1000;
+
+  PRAGMA_OMP_PARALLEL_THREADS(numThreads) {
+    PRAGMA_OMP_SINGLE_ARGS(nowait) {
+      quickSort_parallel_internal_key(array, x->shapeInfo(), values, y->shapeInfo(), 0, lenArray - 1, cutoff, descending);
+    }
+  }
+}
+
+template <typename X, typename Y>
+static void quickSort_parallel_value(NDArray *x, NDArray *y, sd::LongType lenArray, int numThreads,
+                                     bool descending) {
+  auto array = reinterpret_cast<X *>(x->bufferAsT<X>());
+  auto values = reinterpret_cast<Y *>(y->bufferAsT<Y>());
+  int cutoff = 1000;
+
+  PRAGMA_OMP_PARALLEL_THREADS(numThreads) {
+    PRAGMA_OMP_SINGLE_ARGS(nowait) {
+      quickSort_parallel_internal_value(array, x->shapeInfo(), values,y->shapeInfo(), 0, lenArray - 1, cutoff, descending);
+    }
+  }
+}
+
+template <typename X, typename Y>
+void DoubleMethods<X, Y>::sortByKey(NDArray *x,NDArray *y,
+                                    bool descending) {
+  quickSort_parallel_key<X, Y>(x,y, x->lengthOf(),Environment::getInstance().maxMasterThreads(),
+                               descending);
+}
+
+template <typename X, typename Y>
+void DoubleMethods<X, Y>::sortByValue(NDArray *x,NDArray *y,
+                                      bool descending) {
+  quickSort_parallel_value<X, Y>(x,y,x->lengthOf(),Environment::getInstance().maxMasterThreads(),
+                                 descending);
+}
+
+template <typename X, typename Y>
+void DoubleMethods<X, Y>::sortTadByKey(NDArray *xArr,NDArray *yArr,
+                                       NDArray *dimension, bool descending) {
+  auto x = xArr->bufferAsT<X>();
+  auto y = yArr->bufferAsT<Y>();
+  auto dimensionData = dimension->bufferAsT<sd::LongType>();
+  auto dimensionLength = dimension->lengthOf();
+  auto packX = ConstantTadHelper::getInstance().tadForDimensions(xArr->shapeInfo(), dimensionData, dimensionLength);
+  auto packY = ConstantTadHelper::getInstance().tadForDimensions(yArr->shapeInfo(), dimensionData, dimensionLength);
+
+  auto xLength = xArr->lengthOf();
+  auto xTadLength = shape::length(packX->primaryShapeInfo());
+  auto numTads = packX->numberOfTads();
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto r = start; r < stop; r++) {
+      NDArray *xView = packX->extractTadView(xArr,r);
+      NDArray *yView = packY->extractTadView(yArr,r);
+      quickSort_parallel_key<X, Y>(xView,
+                                   yView, xTadLength, 1,
+                                   descending);
+      delete xView;
+      delete yView;
+    }
+  };
+
+  samediff::Threads::parallel_tad(func, 0, numTads);
+}
+
+template <typename X, typename Y>
+void DoubleMethods<X, Y>::sortTadByValue(NDArray *xArr, NDArray *yArr,
+                                         NDArray *dimension, bool descending) {
+  auto x = reinterpret_cast<X *>(xArr->bufferAsT<X>());
+  auto y = reinterpret_cast<Y *>(yArr->bufferAsT<Y>());
+  auto dimensionData = dimension->bufferAsT<sd::LongType>();
+  auto len = dimension->lengthOf();
+  auto packX = ConstantTadHelper::getInstance().tadForDimensions(xArr->shapeInfo(), dimensionData, len);
+  auto packY = ConstantTadHelper::getInstance().tadForDimensions(yArr->shapeInfo(), dimensionData, len);
+
+  auto xLength = xArr->lengthOf();
+  auto xTadLength = shape::length(packX->primaryShapeInfo());
+  auto numTads = packX->numberOfTads();
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto r = start; r < stop; r++) {
+      NDArray *xView = packX->extractTadView(xArr,r);
+      NDArray *yView = packY->extractTadView(yArr,r);
+      quickSort_parallel_value<X, Y>(xView,
+                                   yView, xTadLength, 1,
+                                   descending);
+      delete xView;
+      delete yView;
+    }
+  };
+
+  samediff::Threads::parallel_tad(func, 0, numTads);
+}
+}  // namespace sd
+
+

libnd4j/include/ops/impl/specials_single.hpp

@@ -0,0 +1,441 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com, created on 07.10.2017.
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <array/NDArray.h>
+#include <helpers/Loops.h>
+
+#include <helpers/shape.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/specials.h>
+#include <types/types.h>
+
+namespace sd {
+
+/**
+ * @brief Checks if the shape of NDArray contains 1 before(order c) or after(order f) the specified axis
+ *
+ * @param input
+ * @param axis
+ * @return int
+ */
+SD_INLINE int isShapeExtendedWithOnes(NDArray&input, LongType axis) {
+  bool isAllOne = true;
+  auto shapes = shape::shapeOf(input.shapeInfo());
+  auto rank = input.rankOf();
+  if (rank > axis) {
+    if (input.ordering() == 'c') {
+      // check before the axis
+      for (sd::LongType i = 0; i < axis; i++) {
+        isAllOne = isAllOne && (shapes[i] == 1);
+      }
+    } else {
+      // check after the axis
+      for (int i = axis + 1; i < rank; i++) {
+        isAllOne = isAllOne && (shapes[i] == 1);
+      }
+    }
+    return isAllOne;
+  }
+
+  return true;
+}
+
+template <typename T>
+struct InputArgsCase2 {
+  const T *ptr;
+  int size;
+};
+
+template <typename T>
+void SpecialMethods<T>::concatCpuGeneric(const std::vector<NDArray *> &inArrs, NDArray &output,
+                                         const LongType axis) {
+  const sd::LongType numOfInArrs = inArrs.size();
+  const auto sizeofT = output.sizeOfT();
+
+  T *zBuff = output.bufferAsT<T>();
+
+  bool shapeExtendedWithOnes = isShapeExtendedWithOnes(output, axis);
+  bool followEws1 = false;
+  bool matchesOutputOrdering = true;
+  for (int i = 0; i < numOfInArrs; ++i) {
+    shapeExtendedWithOnes = shapeExtendedWithOnes && isShapeExtendedWithOnes(*inArrs[i], axis);
+    matchesOutputOrdering = matchesOutputOrdering && inArrs[i]->ordering() == output.ordering();
+  }
+
+  bool copyCaseEws1 = followEws1 & matchesOutputOrdering;
+  bool copyCase1 = numOfInArrs > 1 ? copyCaseEws1 & shapeExtendedWithOnes : copyCaseEws1;
+
+  if (copyCase1) {
+    // copyCase1:
+    // When NdArrays follow the same order and unit elementwise stride and
+    // the concantneation axis is 0th or has only 1 before it {1, 1, ..., axis} for "c"
+    // or axis is (rank-1)th or has only 1 after it {axis, 1, 1, ..., 1} for "f"
+    // we will concatenate them by sequential copying of the whole buffers
+
+    std::vector<T *> zPtrList;
+    T *z = output.bufferAsT<T>();
+    for (sd::LongType i = 0; i < numOfInArrs; i++) {
+      zPtrList.push_back(z);
+      z += inArrs[i]->lengthOf();
+    }
+    auto func = [&inArrs, &zPtrList](sd::LongType thread_id, sd::LongType start, sd::LongType stop,
+                                     sd::LongType increment) -> void {
+      for (sd::LongType i = start; i < stop; ++i) {
+        const auto memAmountToCopy = inArrs[i]->lengthOf();
+        const auto inputPtr = inArrs[i]->bufferAsT<T>();
+
+        auto zPtr = zPtrList[i];
+        for (int j = 0; j < memAmountToCopy; j++) {
+          zPtr[j] = inputPtr[j];
+        }
+      }
+    };
+
+    samediff::Threads::parallel_tad(func, 0, numOfInArrs, 1);
+    return;
+  }
+
+  // for one Array
+  if (numOfInArrs < 2) {
+    output.assign(inArrs[0]);
+    return;
+  }
+  bool copyCase2 = copyCaseEws1 && output.ordering() == 'c';
+  if (copyCase2) {
+    sd::LongType times = 1;
+    auto shapes = shape::shapeOf(output.shapeInfo());
+
+    T *z = output.bufferAsT<T>();
+    for (int i = 0; i < axis; i++) {
+      times = times * shapes[i];
+    }
+
+    sd::LongType totalCopySize = output.lengthOf() / times;
+
+    std::vector<InputArgsCase2<T>> inputArgs;
+    for (sd::LongType i = 0; i < numOfInArrs; i++) {
+      InputArgsCase2<T> input = {inArrs[i]->bufferAsT<T>(),
+                                 static_cast<int>(inArrs[i]->lengthOf()) / static_cast<int>(times)};
+      inputArgs.push_back(input);
+    }
+
+    auto func = [&inputArgs, z, totalCopySize](uint64_t thread_id, int64_t start, int64_t stop,
+                                               int64_t increment) -> void {
+      auto outPtr = &(z[start * totalCopySize]);
+      auto numOfInArrs = inputArgs.size();
+      for (int i = start; i < stop; i++) {
+        for (size_t j = 0; j < numOfInArrs; j++) {
+          auto inputCopySize = inputArgs[j].size;
+          const T *inputBasePtr = inputArgs[j].ptr;
+          auto inputPtr = &(inputBasePtr[i * inputCopySize]);
+          // copy
+          PRAGMA_OMP_SIMD
+          for (int k = 0; k < inputCopySize; k++) {
+            outPtr[k] = inputPtr[k];
+          }
+          outPtr += inputCopySize;
+        }
+      }
+    };
+    samediff::Threads::parallel_tad(func, 0, times, 1);
+    return;
+  }
+
+  // Cache shape and stride information for output
+  const sd::LongType zRank = shape::rank(output.shapeInfo());
+  const sd::LongType* zShape = shape::shapeOf(output.shapeInfo());
+  const sd::LongType* zStride = shape::stride(output.shapeInfo());
+
+  // Pre-cache input arrays' shape information
+  std::vector<const sd::LongType*> inShapes(numOfInArrs);
+  std::vector<const sd::LongType*> inStrides(numOfInArrs);
+  std::vector<sd::LongType> inRanks(numOfInArrs);
+
+  for (sd::LongType i = 0; i < numOfInArrs; i++) {
+    inRanks[i] = shape::rank(inArrs[i]->shapeInfo());
+    inShapes[i] = shape::shapeOf(inArrs[i]->shapeInfo());
+    inStrides[i] = shape::stride(inArrs[i]->shapeInfo());
+  }
+
+  // general case
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType coords[SD_MAX_RANK], temp;
+
+    for (sd::LongType i = start; i < stop; i += increment) {
+      INDEX2COORDS(i, zRank, zShape, coords);
+
+      sd::LongType zOffset;
+      COORDS2INDEX(zRank, zStride, coords, zOffset);
+
+      sd::LongType inArrIdx = 0;
+      sd::LongType xDim = inArrs[inArrIdx]->sizeAt(axis);
+
+      temp = coords[axis];
+      while (coords[axis] >= xDim) {
+        coords[axis] -= xDim;
+        xDim = inArrs[++inArrIdx]->sizeAt(axis);
+      }
+
+      const T *x = inArrs[inArrIdx]->bufferAsT<T>();
+      sd::LongType xOffset;
+      COORDS2INDEX(inRanks[inArrIdx], inStrides[inArrIdx], coords, xOffset);
+
+      zBuff[zOffset] = x[xOffset];
+
+      coords[axis] = temp;
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, output.lengthOf());
+}
+/**
+ * Concatneate multi array of the same shape together
+ * along a particular dimension
+ */
+template <typename T>
+void SpecialMethods<T>::concatCpuGeneric(LongType dimension, int numArrays,NDArray **data,
+                                         NDArray *vresult) {
+  auto result = reinterpret_cast<T *>(vresult);
+  std::vector<NDArray *> inputs(numArrays);
+
+
+  for (sd::LongType i = 0; i < numArrays; ++i)
+    inputs[i] =
+        new NDArray(static_cast<void *>(data[i]), data[i]->shapeInfo(), nullptr, false, 0);
+
+  sd::SpecialMethods<T>::concatCpuGeneric(inputs, *vresult, dimension);
+
+  for (sd::LongType i = 0; i < numArrays; ++i) {
+    delete inputs[i];
+  }
+}
+
+template <typename T>
+void SpecialMethods<T>::splitCpuGeneric(NDArray& input, const std::vector<NDArray*>& outArrs, const LongType axis) {
+  int numSplits = outArrs.size();
+  const auto sizeofT = input.sizeOfT();
+  auto xBuff = input.bufferAsT<T>();
+
+  bool luckCase1 = ((axis == 0 && input.ordering() == 'c') || (axis == input.rankOf() - 1 && input.ordering() == 'f'));
+
+  if (luckCase1) {
+    T* x = const_cast<T*>(xBuff);
+    for (sd::LongType i = 0; i < numSplits; ++i) {
+      const auto memAmountToCopy = outArrs[i]->lengthOf();
+      ops::safe_copy(x, static_cast<const T*>(outArrs[i]->buffer()), static_cast<size_t>(memAmountToCopy));
+      x += memAmountToCopy;
+    }
+    return;
+  }
+
+  // Cache shape and stride information
+  const sd::LongType xRank = shape::rank(input.shapeInfo());
+  const sd::LongType* xShape = shape::shapeOf(input.shapeInfo());
+  const sd::LongType* xStride = shape::stride(input.shapeInfo());
+
+  // Pre-cache output array ranks, shapes, and strides
+  std::vector<const sd::LongType*> outShapes(numSplits);
+  std::vector<const sd::LongType*> outStrides(numSplits);
+  std::vector<sd::LongType> outRanks(numSplits);
+
+  for (int i = 0; i < numSplits; i++) {
+    outRanks[i] = shape::rank(outArrs[i]->shapeInfo());
+    outShapes[i] = shape::shapeOf(outArrs[i]->shapeInfo());
+    outStrides[i] = shape::stride(outArrs[i]->shapeInfo());
+  }
+
+  sd::LongType zDim = outArrs[0]->sizeAt(axis);
+
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType coords[SD_MAX_RANK], temp;
+
+    for (sd::LongType i = start; i < stop; i += increment) {
+      INDEX2COORDS(i, xRank, xShape, coords);
+      sd::LongType xOffset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+
+      sd::LongType outArrIdx = 0;
+      temp = coords[axis];
+
+      while (coords[axis] >= zDim) {
+        coords[axis] -= zDim;
+        ++outArrIdx;
+      }
+
+      T* z = outArrs[outArrIdx]->bufferAsT<T>();
+      sd::LongType zOffset;
+      COORDS2INDEX(outRanks[outArrIdx], outStrides[outArrIdx], coords, zOffset);
+      z[zOffset] = xBuff[xOffset];
+
+      coords[axis] = temp;
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, input.lengthOf());
+}
+
+template <typename T>
+void SpecialMethods<T>::sortGeneric(NDArray *input, bool descending) {
+  auto x = input->bufferAsT<T>();
+  auto xShapeInfo = input->shapeInfo();
+  quickSort_parallel(input, Environment::getInstance().maxMasterThreads(), descending);
+}
+
+template <typename T>
+void SpecialMethods<T>::quickSort_parallel_internal(NDArray *x, int left, int right, int cutoff, bool descending) {
+  if (right - left <= cutoff) {
+    // Use insertion sort for small arrays
+    auto xBuff = x->bufferAsT<T>();
+    for (int i = left + 1; i <= right; i++) {
+      T key = xBuff[i];
+      int j = i - 1;
+      if (descending) {
+        while (j >= left && xBuff[j] < key) {
+          xBuff[j + 1] = xBuff[j];
+          j--;
+        }
+      } else {
+        while (j >= left && xBuff[j] > key) {
+          xBuff[j + 1] = xBuff[j];
+          j--;
+        }
+      }
+      xBuff[j + 1] = key;
+    }
+    return;
+  }
+
+  // Choose pivot as median of three
+  auto xBuff = x->bufferAsT<T>();
+  int mid = (left + right) / 2;
+  if (descending) {
+    if (xBuff[right] > xBuff[left]) std::swap(xBuff[right], xBuff[left]);
+    if (xBuff[mid] > xBuff[left]) std::swap(xBuff[mid], xBuff[left]);
+    if (xBuff[right] > xBuff[mid]) std::swap(xBuff[right], xBuff[mid]);
+  } else {
+    if (xBuff[right] < xBuff[left]) std::swap(xBuff[right], xBuff[left]);
+    if (xBuff[mid] < xBuff[left]) std::swap(xBuff[mid], xBuff[left]);
+    if (xBuff[right] < xBuff[mid]) std::swap(xBuff[right], xBuff[mid]);
+  }
+
+  // Partition
+  T pivot = xBuff[mid];
+  int i = left;
+  int j = right;
+
+  while (i <= j) {
+    if (descending) {
+      while (xBuff[i] > pivot) i++;
+      while (xBuff[j] < pivot) j--;
+    } else {
+      while (xBuff[i] < pivot) i++;
+      while (xBuff[j] > pivot) j--;
+    }
+
+    if (i <= j) {
+      std::swap(xBuff[i], xBuff[j]);
+      i++;
+      j--;
+    }
+  }
+
+  // Recursively sort sub-arrays
+  if (left < j) quickSort_parallel_internal(x, left, j, cutoff, descending);
+  if (i < right) quickSort_parallel_internal(x, i, right, cutoff, descending);
+}
+
+template <typename T>
+void SpecialMethods<T>::quickSort_parallel(NDArray *x, int numThreads, bool descending) {
+  const int CUTOFF = 32;  // Threshold for switching to insertion sort
+  auto length = x->lengthOf();
+
+  if (length <= 1) return;
+
+  // For very small arrays, just use the internal sort
+  if (length <= CUTOFF || numThreads <= 1) {
+    quickSort_parallel_internal(x, 0, length - 1, CUTOFF, descending);
+    return;
+  }
+
+  // For larger arrays, partition into segments and sort in parallel
+  int segmentSize = length / numThreads;
+  auto func = PRAGMA_THREADS_FOR {
+    int threadLeft = start * segmentSize;
+    int threadRight = (start == numThreads - 1) ? length - 1 : (start + 1) * segmentSize - 1;
+    quickSort_parallel_internal(x, threadLeft, threadRight, CUTOFF, descending);
+  };
+
+  samediff::Threads::parallel_for(func, 0, numThreads);
+
+  // Merge sorted segments if we used multiple threads
+  if (numThreads > 1) {
+    auto xBuff = x->bufferAsT<T>();
+    std::vector<T> temp(length);
+    for (int size = segmentSize; size < length; size *= 2) {
+      for (int left = 0; left < length; left += 2 * size) {
+        int mid = std::min(left + size, (int)length);
+        int right = std::min(left + 2 * size, (int)length);
+        int i = left, j = mid, k = left;
+
+        // Merge two segments
+        while (i < mid && j < right) {
+          if (descending) {
+            temp[k++] = (xBuff[i] >= xBuff[j]) ? xBuff[i++] : xBuff[j++];
+          } else {
+            temp[k++] = (xBuff[i] <= xBuff[j]) ? xBuff[i++] : xBuff[j++];
+          }
+        }
+        while (i < mid) temp[k++] = xBuff[i++];
+        while (j < right) temp[k++] = xBuff[j++];
+
+        // Copy back
+        for (i = left; i < right; i++) {
+          xBuff[i] = temp[i];
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+void SpecialMethods<T>::sortTadGeneric(NDArray *input, sd::LongType *dimension, int dimensionLength, bool descending) {
+  auto x = input->bufferAsT<T>();
+  sd::LongType xLength = input->lengthOf();
+  sd::LongType xTadLength = shape::tadLength(input->shapeInfo(), dimension, dimensionLength);
+  int numTads = xLength / xTadLength;
+
+  const std::vector<sd::LongType> dimVector(dimension, dimension + dimensionLength);
+  auto pack = sd::ConstantTadHelper::getInstance().tadForDimensions(
+      const_cast<sd::LongType *>(input->shapeInfo()), const_cast<sd::LongType *>(dimVector.data()), false);
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto r = start; r < stop; r++) {
+      NDArray *dx = pack->extractTadView(input, r);
+      quickSort_parallel(dx,  xTadLength, descending);
+      delete dx;
+    }
+  };
+  samediff::Threads::parallel_tad(func, 0, numTads);
+}
+
+}  // namespace sd

libnd4j/include/ops/impl/specials_sparse.cpp

@@ -0,0 +1,299 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver119@gmail.com
+//
+#include <helpers/shape.h>
+#include <ops/specials_sparse.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <types/float16.h>
+#include <types/types.h>
+
+namespace sd {
+namespace sparse {
+
+template <typename T>
+void SparseUtils<T>::printIndex(sd::LongType *indices, int rank, int x) {
+  printf(" [");
+  for (int e = 0; e < rank; e++) {
+    if (e > 0) printf(", ");
+
+    printf("%lld", (long long)indices[x * rank + e]);
+  }
+  printf("] ");
+}
+
+template <typename T>
+bool SparseUtils<T>::ltIndices(sd::LongType *indices, int rank, sd::LongType x, sd::LongType y) {
+  for (int e = 0; e < rank; e++) {
+    sd::LongType idxX = indices[x * rank + e];
+    sd::LongType idxY = indices[y * rank + e];
+    // we're comparing indices one by one, starting from outer dimension
+    if (idxX < idxY) {
+      return true;
+    } else if (idxX == idxY) {
+      // do nothing, continue to next dimension
+    } else
+      return false;
+  }
+
+  return false;
+}
+
+template <typename T>
+bool SparseUtils<T>::gtIndices(sd::LongType *indices, int rank, sd::LongType x, sd::LongType y) {
+  for (int e = 0; e < rank; e++) {
+    // we're comparing indices one by one, starting from outer dimension
+    sd::LongType idxX = indices[x * rank + e];
+    sd::LongType idxY = indices[y * rank + e];
+    if (idxX > idxY) {
+      return true;
+    } else if (idxX == idxY) {
+      // do nothing, continue to next dimension
+    } else
+      return false;
+  }
+  return false;
+}
+
+template <typename T>
+void SparseUtils<T>::swapEverything(sd::LongType *indices, T *array, int rank, sd::LongType x, sd::LongType y) {
+  // swap indices
+  for (int e = 0; e < rank; e++) {
+    sd::LongType tmp = indices[x * rank + e];
+    indices[x * rank + e] = indices[y * rank + e];
+    indices[y * rank + e] = tmp;
+  }
+
+  // swap values
+  T tmp = array[x];
+  array[x] = array[y];
+  array[y] = tmp;
+}
+
+template <typename T>
+sd::LongType SparseUtils<T>::coo_quickSort_findPivot(sd::LongType *indices, T *array, sd::LongType left,
+                                                     sd::LongType right, int rank) {
+  sd::LongType mid = (left + right) / 2;
+
+  // ensure left < mid
+  if (ltIndices(indices, rank, mid, left)) {  // ensure lo < mid
+    swapEverything(indices, array, rank, mid, left);
+  }
+
+  // ensure left < right
+  if (ltIndices(indices, rank, right, left)) {
+    swapEverything(indices, array, rank, right, left);
+  }
+
+  // ensure mid < right
+  if (ltIndices(indices, rank, right, mid)) {
+    swapEverything(indices, array, rank, right, mid);
+  }
+
+  // mid is the median of the 3, and is the optimal pivot point
+  return mid;
+}
+
+template <typename T>
+void SparseUtils<T>::coo_quickSort_parallel_internal(sd::LongType *indices, T *array, sd::LongType left,
+                                                     sd::LongType right, int cutoff, int rank) {
+  sd::LongType span = right - left;  // elements to be partitioned - 1
+
+  if (span == 1) {
+    // only 2 elements to partition. swap if needed and return directly without further sorting.
+    if (ltIndices(indices, rank, right, left)) {
+      swapEverything(indices, array, rank, left, right);
+    }
+    return;
+  }
+
+  // find optimal pivot and sort left < right < right
+  sd::LongType pvt = coo_quickSort_findPivot(indices, array, left, right, rank);
+
+  if (span == 2) {
+    // only 3 elements to partition. findPivot has already sorted them. no further sorting is needed.
+    return;
+  }
+
+  // index that is greater than pivot - leftmost element is already partitioned because of findPivot.
+  sd::LongType i = left + 1;
+
+  // index that is smaller than pivot - rightmost element is already partitioned because of findPivot.
+  sd::LongType j = right - 1;
+
+  {
+    // flag that indicates that pivot index lies between i and j and *could* be swapped.
+    bool checkPivot = true;
+    /* PARTITION PART */
+    while (i <= j) {
+      while (ltIndices(indices, rank, i, pvt)) i++;
+
+      while (gtIndices(indices, rank, j, pvt)) j--;
+
+      if (i <= j) {
+        if (i != j) {  // swap can be fairly expensive. don't swap i -> i
+          swapEverything(indices, array, rank, i, j);
+        }
+
+        // only check pivot if it hasn't already been swapped.
+        if (checkPivot) {
+          // check if we moved the pivot, if so, change pivot index accordingly
+          if (pvt == j) {
+            pvt = i;
+            checkPivot = false;
+          } else if (pvt == i) {
+            pvt = j;
+            checkPivot = false;
+          }
+        }
+
+        i++;
+        j--;
+      }
+    }
+  }
+
+  if ((span < cutoff)) {
+    if (left < j) {
+      coo_quickSort_parallel_internal(indices, array, left, j, cutoff, rank);
+    }
+    if (i < right) {
+      coo_quickSort_parallel_internal(indices, array, i, right, cutoff, rank);
+    }
+
+  } else {
+    PRAGMA_OMP_TASK { coo_quickSort_parallel_internal(indices, array, left, j, cutoff, rank); }
+    PRAGMA_OMP_TASK { coo_quickSort_parallel_internal(indices, array, i, right, cutoff, rank); }
+  }
+}
+
+template <typename T>
+void SparseUtils<T>::coo_quickSort_parallel(sd::LongType *indices, T *array, sd::LongType lenArray, int numThreads,
+                                            int rank) {
+  int cutoff = 1000;
+
+  PRAGMA_OMP_PARALLEL_THREADS(numThreads) {
+    PRAGMA_OMP_SINGLE_ARGS(nowait) { coo_quickSort_parallel_internal(indices, array, 0, lenArray - 1, cutoff, rank); }
+  }
+}
+
+template <typename T>
+void SparseUtils<T>::sortCooIndicesGeneric(sd::LongType *indices, void *vx, sd::LongType length, int rank) {
+  auto values = reinterpret_cast<T *>(vx);
+#ifdef _OPENMP
+  coo_quickSort_parallel(indices, values, length, omp_get_max_threads(), rank);
+#else
+  coo_quickSort_parallel(indices, values, length, 1, rank);
+#endif
+}
+
+BUILD_SINGLE_TEMPLATE( class SparseUtils, , SD_COMMON_TYPES);
+
+void IndexUtils::ravelMultiIndex(sd::LongType *indices, sd::LongType *flatIndices, sd::LongType length,
+                                 sd::LongType *shapeInfo, int mode) {
+  sd::LongType *shape = shape::shapeOf(shapeInfo);
+  sd::LongType *stride = shape::stride(shapeInfo);
+  sd::LongType rank = shape::rank(shapeInfo);
+  int errorCount = 0;
+
+  PRAGMA_OMP_PARALLEL_FOR
+  for (sd::LongType i = 0; i < length; ++i) {
+    sd::LongType raveledIndex = 0;
+    for (sd::LongType j = 0; j < rank; ++j) {
+      sd::LongType idx = indices[i * rank + j];
+      if (idx >= shape[j]) {
+        // index does not fit into shape at j dimension.
+        if (mode == ND4J_CLIPMODE_CLIP) {
+          // set idx to largest possible value (clip to shape)
+          idx = shape[j] - 1;
+        } else if (mode == ND4J_CLIPMODE_WRAP) {
+          idx %= shape[j];
+        } else {
+          // mode is ND4J_CLIPMODE_THROW or is unknown. either way. throw an error later.
+          // cannot throw here because of parallel region
+          sd_printf(
+              "sparse::IndexUtils::ravelMultiIndex Cannot ravel index at element %d, does not fit into specified "
+              "shape.\n",
+              i);
+          ++errorCount;
+        }
+      }
+      raveledIndex += idx * stride[j];
+    }
+    flatIndices[i] = raveledIndex;
+  }
+
+  if (errorCount > 0) {
+    // throw error if one ocurred in loop
+    THROW_EXCEPTION("sparse::IndexUtils::ravelMultiIndex Cannot ravel index");
+  }
+}
+
+void IndexUtils::unravelIndex(sd::LongType *indices, sd::LongType *flatIndices, sd::LongType length,
+                              sd::LongType *shapeInfo) {
+  sd::LongType *shape = shape::shapeOf(shapeInfo);
+  sd::LongType *stride = shape::stride(shapeInfo);
+  sd::LongType rank = shape::rank(shapeInfo);
+  int errorCount = 0;
+
+  // unravelOrder ensures that the dimensions with largest stride are unraveled first.
+  // create vector with elements 0..rank
+  int *unravelOrder = shape::range<int>(0, rank);
+
+  // sort order according to stride length.
+  std::sort(unravelOrder, unravelOrder + rank, [&](int i1, int i2) { return stride[i1] > stride[i2]; });
+
+  // calculate the largest raveled index that will fit into passed shape
+  sd::LongType maxRaveledIndex = shape[unravelOrder[0]] * stride[unravelOrder[0]] - 1;
+
+  PRAGMA_OMP_PARALLEL_FOR
+  for (sd::LongType i = 0; i < length; ++i) {
+    sd::LongType raveledIndex = flatIndices[i];
+    if (raveledIndex > maxRaveledIndex) {
+      // cannot throw here because of parallel region
+      sd_printf(
+          "sparse::IndexUtils::unravelIndex Cannot unravel index at element %d. raveled index of %d does not fit into "
+          "specified shape.\n",
+          i, raveledIndex);
+      ++errorCount;
+    }
+
+    for (int *it = unravelOrder; it != unravelOrder + rank; it++) {
+      int j = *it;
+      // how many strides of this size?
+      indices[i * rank + j] = raveledIndex / stride[j];
+
+      // remainder for subsequent smaller strides.
+      raveledIndex %= stride[j];
+    }
+  }
+
+  if (errorCount > 0) {
+    // throw error if one occurred in loop
+    sd_printf("Largest raveled index is: %d, ", maxRaveledIndex) std::vector<sd::LongType> v(shape, shape + rank);
+    sd_printv("Shape: ", v);
+    THROW_EXCEPTION("sparse::IndexUtils::unravelIndex Cannot unravel index");
+  }
+
+  delete[] unravelOrder;
+}
+}  // namespace sparse
+}  // namespace sd