chore: import upstream snapshot with attribution

libnd4j/include/helpers/impl/ArrayUtils.cpp

@@ -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 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/ArrayUtils.h>
+
+namespace sd {
+namespace ArrayUtils {
+void toIntPtr(std::initializer_list<int> list, int* target) {
+  std::vector<int> vec(list);
+  toIntPtr(vec, target);
+}
+
+void toIntPtr(std::vector<int>& list, int* target) { memcpy(target, list.data(), list.size() * sizeof(int)); }
+
+void toLongPtr(std::initializer_list<LongType> list, LongType* target) {
+  std::vector<LongType> vec(list);
+  toLongPtr(vec, target);
+}
+
+void toLongPtr(std::vector<LongType>& list, LongType* target) {
+  memcpy(target, list.data(), list.size() * sizeof(LongType));
+}
+
+std::vector<LongType> toLongVector(std::vector<int> vec) {
+  std::vector<LongType> result(vec.size());
+  LongType vecSize = vec.size();
+
+  for (LongType e = 0; e < vecSize; e++) result[e] = vec[e];
+
+  return result;
+}
+
+std::vector<LongType> toLongVector(std::vector<LongType> vec) { return vec; }
+}  // namespace ArrayUtils
+}  // namespace sd

libnd4j/include/helpers/impl/AttentionHelper.cpp

@@ -0,0 +1,513 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 Paul Dubs
+// @author Adam Gibson
+//
+
+#ifndef LIBND4J_ATTENTIONHELPER_CPP
+#define LIBND4J_ATTENTIONHELPER_CPP
+#include "../AttentionHelper.h"
+#include <indexing/NDIndexUtils.h>
+#include <helpers/AttentionHelper.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/batched_gemm.h>
+#if NOT_EXCLUDED(OP_multi_head_dot_product_attention)
+
+namespace sd {
+
+NDArray AttentionHelper::multiHeadProject(NDArray *input, NDArray *projectionMatrix,
+                                          LaunchContext *context) {
+  auto miniBatchSize = input->sizeAt(0);
+  auto seqLength = input->sizeAt(2);
+  auto numHeads = projectionMatrix->sizeAt(0);
+  auto projectedSize = projectionMatrix->sizeAt(1);
+
+  std::vector<sd::LongType> epsPermVec = {1, 0,2};
+  auto inputPerm = input->permute(epsPermVec, false, false);  //[batch, nIn, timeSteps] -> [nIn, batch, timeSteps]
+  std::vector<sd::LongType> inputPermShape = {input->sizeAt(1), (miniBatchSize * seqLength)};
+  auto inputPrep = inputPerm->reshape('c', inputPermShape);  //[nIn, batch*timeSteps]
+  std::vector<sd::LongType> projectionMatrixShape = {numHeads * projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)};
+  auto projectionPrep = projectionMatrix->reshape(
+      'c',
+      projectionMatrixShape);  //[nHeads, hS, nIn] -> [nHeads*hS, nIn]
+
+  std::vector<LongType> projectedShape = {numHeads * projectionMatrix->sizeAt(1), (miniBatchSize * seqLength)};
+  NDArray projected('c',projectedShape, input->dataType(),
+                    context);  //[nHeads*hS, batch*timeSteps]
+  ops::matmul mmul;
+  mmul.execute({&projectionPrep, &inputPrep}, {&projected});
+
+  projected.reshapei({numHeads, projectedSize, miniBatchSize, seqLength});
+  projected.permutei({2, 0, 1, 3}, false, false);  //[minibatch, numHeads, projectedSize, seqLength]
+
+  return projected;
+}
+
+
+/**
+ * @param shape
+ * @return
+ */
+NDArray * AttentionHelper::lowerTriangularMask(std::vector<LongType> *shape) {
+  auto rowIndexOnes = NDArrayFactory::valueOf(*shape,1,'c');
+  auto colIndexOnes = NDArrayFactory::valueOf(*shape, 1, 'c');
+  ops::cumsum cumsum;
+  auto rowCumSum = cumsum.evaluate({rowIndexOnes},{},{-2,0},{});
+  auto colsCumSum = cumsum.evaluate({colIndexOnes}, {}, {-1, 0}, {});
+  ops::greater_equal greaterEqual;
+  auto ret = greaterEqual.evaluate({rowCumSum.at(0),colsCumSum.at(0)});
+  return ret[0];
+}
+
+/**
+ * @param query
+ * @param value
+ * @return
+ */
+NDArray *AttentionHelper::computeCasualMask(NDArray *query, NDArray *value, bool multiHead) {
+  if(multiHead) {
+    auto qSeqLength = query->sizeAt(1);
+    auto vSeqLength = value != nullptr ? value->sizeAt(1) : qSeqLength;
+    ops::matrix_band_part matrixBandPart;
+    auto ones = NDArrayFactory::create('c',{1,qSeqLength,vSeqLength}, INT32);
+    int assignVal = 1;
+    ones->assign(assignVal);
+    auto lower = matrixBandPart.evaluate({ones},{},{-1,0});
+    auto ret = lower.at(0)->cast(BOOL);
+    delete ones;
+    return ret;
+
+  } else {
+    std::vector<LongType> causalMaskShape2;
+    causalMaskShape2.push_back(query->sizeAt(0));
+    //4d
+    if(query->rankOf() > 3)
+      causalMaskShape2.push_back(query->sizeAt(1));
+
+    causalMaskShape2.push_back(query->sizeAt(-2));
+    causalMaskShape2.push_back(value->sizeAt(-2));
+
+    auto ret  = lowerTriangularMask(&causalMaskShape2);
+    return ret;
+
+  }
+
+}
+
+
+/**
+ * @param query
+ * @param value
+ * @param attentionMask
+ * @param useCausalMask
+ * @return
+ */
+NDArray *AttentionHelper::computeAttentionMask(NDArray *query, NDArray *value, NDArray *queryMask, NDArray *valueMask,
+                                               NDArray *attentionMask, bool useCausalMask) {
+  auto internalQueryMask = queryMask;
+  auto internalValueMask = valueMask;
+  NDArray *autoMask = nullptr;
+  ops::create_view createView;
+  ops::boolean_and booleanAnd;
+  auto all = NDIndexUtils::createAll();
+  auto newAxis = NDIndexUtils::createNewAxis();
+
+  if (internalQueryMask != nullptr && !internalQueryMask->isEmpty()) {
+    internalQueryMask = queryMask->cast(BOOL);
+    if (autoMask != nullptr && !autoMask->isEmpty()) {
+      autoMask = createView.evaluate({internalQueryMask, all, all, newAxis}).at(0);
+    }
+  }
+
+  if (valueMask != nullptr && !valueMask->isEmpty()) {
+    internalValueMask = valueMask->cast(BOOL);
+    auto mask = createView.evaluate({internalValueMask, all, newAxis, all}).at(0);
+    if (autoMask == nullptr || autoMask->isEmpty()) {
+      autoMask = mask;
+    } else {
+      autoMask = booleanAnd.evaluate({autoMask, mask}).at(0);
+    }
+  }
+
+  if (useCausalMask) {
+    auto mask = computeCasualMask(query, value, false);
+    if (autoMask == nullptr) {
+      autoMask = mask;
+    } else {
+      autoMask = booleanAnd.evaluate({autoMask, mask}).at(0);
+    }
+  }
+
+  if (autoMask != nullptr && !autoMask->isEmpty()) {
+    if (attentionMask == nullptr || attentionMask->isEmpty()) {
+      return autoMask;
+    } else {
+      auto ret = booleanAnd.evaluate({attentionMask, autoMask}).at(0);
+      return ret;
+    }
+  }
+
+  delete all;
+  delete newAxis;
+
+  return autoMask;
+}
+
+NDArray * AttentionHelper::mergeMasks(NDArray *x, NDArray *y) {
+  if(x == nullptr || x->isEmpty()) {
+    return y;
+  }
+
+  if (y == nullptr || y->isEmpty()) {
+    return x;
+  }
+
+  ops::boolean_and booleanAnd;
+  auto ret = booleanAnd.evaluate({x,y});
+  return ret.at(0);
+}
+
+void AttentionHelper::applyAttentionScores(NDArray *scores, NDArray *value, NDArray *scoresMask,
+                                           double dropout, int randomSeed, NDArray *applyScoresOut, NDArray *attentionLogits,
+                                           NDArray *dropoutMask) {
+  ops::boolean_not booleanNot;
+  ops::softmax softmax;
+  ops::dropout dropoutOp;
+  ops::matmul matmul;
+
+  int softmaxDim = -1;
+  if (scoresMask != nullptr && !scoresMask->isEmpty()) {
+    REQUIRE_TRUE(scoresMask->sizeAt(-2) == 1 || scoresMask->sizeAt(-2) == scores->sizeAt(-2),0,
+                 "Scores mask must be either broadcastable or equal to scores shape. scores size at -2: was: %i scores size at -2 was: %i",scoresMask->sizeAt(-2),scores->sizeAt(-2));
+
+    REQUIRE_TRUE(scoresMask->sizeAt(-1) == scores->sizeAt(-1),0,
+                 "Scores mask must be either broadcastable or equal to scores shape. scores size at -1: was: %i scores size at -1 was: %i",scoresMask->sizeAt(-1),scores->sizeAt(-1));
+
+    auto castedScoresMask = scoresMask->cast(BOOL);
+    auto paddingMask = booleanNot.evaluate({castedScoresMask}).at(0);
+    auto paddingMaskCast = paddingMask->cast(scores->dataType());
+    if (attentionLogits->dataType() == BFLOAT16) {
+      auto minus =  65504 * *paddingMaskCast;
+      *attentionLogits -= *minus;
+      delete minus;
+    } else {
+      auto minus = 1.0e9 * *paddingMask;
+      *attentionLogits -= *minus;
+      delete minus;
+    }
+
+    if(paddingMaskCast != paddingMask) {
+      delete paddingMaskCast;
+    }
+
+    if(scoresMask != castedScoresMask) {
+      delete castedScoresMask;
+    }
+  }
+
+
+
+  softmax.execute({attentionLogits},{scores},{},{softmaxDim});
+  auto weights = scores;
+
+  if (dropout > 0) {
+    dropoutOp.execute({weights},{weights,dropoutMask},{dropout},{randomSeed});
+  }
+
+  //batch size, tq tv
+  //batch size tv dim
+  //output: batch size, tq dim
+  matmul.execute({weights,value},{applyScoresOut});
+
+}
+
+void AttentionHelper::dotProductAttentionBpHelper(NDArray *query, NDArray *key, NDArray *values,
+                                                  double scale,
+                                                  NDArray *dLdq, NDArray *dLdk, NDArray *dLdv, NDArray *eps, LongType dropoutSeed, NDArray *qMask, NDArray *vMask, bool useCausalMask, double dropout, bool training,
+                                                  NDArray *attentionScoresWeights, NDArray *attentionLogits,
+                                                  NDArray *dropoutMask) {
+  ops::matmul_bp matMulBp;
+  ops::softmax_bp softmaxBp;
+  NDArray dldW(attentionScoresWeights->shapeInfo());
+  NDArray dldS(attentionScoresWeights->shapeInfo());
+  NDArray * mask = nullptr;
+  NDArray *causalPointer = nullptr;
+
+  if(useCausalMask) {
+    std::vector<LongType> causalMaskShape2;
+    causalMaskShape2.push_back(attentionLogits->sizeAt(0));
+    //4d
+    if(attentionLogits->rankOf() > 3)
+      causalMaskShape2.push_back(attentionLogits->sizeAt(1));
+
+    for(int i = attentionLogits->rankOf() - 2; i < attentionLogits->rankOf(); i++) {
+      causalMaskShape2.push_back(attentionLogits->sizeAt(i));
+    }
+    causalPointer = lowerTriangularMask(&causalMaskShape2);
+  }
+
+  mask = mergeMasks(vMask,causalPointer);
+
+
+
+  matMulBp.execute({attentionScoresWeights,values,eps},{&dldW,dLdv},{},{});
+  if(dropout > 0.0 && training) {
+    ops::dropout_bp dropoutOp;
+    auto inputs = {attentionScoresWeights,dropoutMask,&dldW};
+    dropoutOp.execute(inputs,{&dldW},{dropout},{dropoutSeed},{false});
+  }
+
+
+  softmaxBp.execute({attentionLogits,&dldW,attentionScoresWeights},{&dldS},{},{-1},{});
+
+
+
+  if(scale != 0.0 && scale != 1.0) {
+    dldS *= scale;
+  }
+
+  // Initialize times as a scalar placeholder (will be reassigned if mask is present)
+  NDArray times(query->dataType(), query->getContext(), true);
+  if(mask != nullptr && !mask->isEmpty()) {
+    ops::expand_dims expandDims;
+    auto maskCast = mask->cast(query->dataType());
+    auto mask2 = *maskCast * 1e9;
+    times = *mask2;
+    dldS *= times;
+    delete mask2;
+    delete maskCast;
+
+  }
+
+
+
+  matMulBp.execute({query,key,&dldS},{dLdq,dLdk},{},{0,1,0});
+}
+
+
+
+
+/**
+   *
+   * @param query
+   * @param key
+   * @param scoreMode
+   * @param scale
+   * @return
+ */
+void AttentionHelper::attentionBpHelper(NDArray *query, NDArray *key, NDArray *values, double scale, NDArray *dLdq,
+                                        NDArray *dLdk, NDArray *dLdv, NDArray *eps,
+                                        LongType dropoutSeed,
+                                        NDArray *qMask, NDArray *vMask,
+                                        bool useCausalMask, double dropout, bool training, NDArray *attentionScoresOut,
+                                        NDArray *attentionScoresWeights,
+                                        NDArray *attentionScoresLogits,
+                                        NDArray *dropoutMask) {
+  dotProductAttentionBpHelper(query, key, values, scale, dLdq, dLdk, dLdv, eps, dropoutSeed, qMask, vMask,
+                              useCausalMask, dropout, training, attentionScoresWeights, attentionScoresLogits,
+                              dropoutMask);
+
+
+}
+
+/**
+   *
+   * @param query
+   * @param key
+   * @param scoreMode
+   * @param scale
+   * @return
+ */
+void AttentionHelper::attentionHelper(NDArray *query, NDArray *key, double scale, NDArray *attentionLogits) {
+  ops::matmul matmul3;
+  matmul3.execute({query,key},{attentionLogits},{},{0,1});
+  if(scale != 0.0 && scale != 1.0) {
+    *attentionLogits *= scale;
+  }
+
+  // Clamp attention logits to prevent numerical overflow in subsequent softmax
+  // Values beyond this range would produce Inf in exp() which leads to NaN
+  // Use clipbyvalue op for proper clamping
+  ops::clipbyvalue clipOp;
+  clipOp.execute({attentionLogits}, {attentionLogits}, {-1e4, 1e4}, {});
+}
+
+
+
+
+/**
+ * @param inputs
+ * @param mask
+ * @param training
+ * @param returnAttentionScores
+ * @param useCausalMask
+ */
+void AttentionHelper::doAttentionBp(std::vector<NDArray *> &inputs, std::vector<NDArray *> &masks, bool training,
+                                    bool useCausalMask, double dropout, double scale, std::vector<NDArray *> outputs,
+                                    LongType dropoutSeed) {
+  auto q = inputs[0];
+  auto v = inputs[1];
+  auto k = inputs[2];
+  auto attentionScoresOut = inputs[3];
+  auto attentionScoresWeights = inputs[4];
+  auto attentionScoresLogits = inputs[5];
+  auto eps = inputs[6];
+
+  auto dropoutMask = inputs.size() > 7 ? inputs[7] : inputs[7];
+
+  ops::expand_dims expandDims;
+  ops::ones_as onesAs;
+  ops::shape_of shapeOf;
+  ops::concat concatOp;
+  ops::create_view createView;
+  auto qMask = masks.size() > 0 ? masks[0] : nullptr;
+  auto vMask = masks.size() > 1 ? masks[1] : nullptr;
+  auto vmaskInternal = vMask;
+  auto qMaskInternal = qMask;
+  if(vMask != nullptr && !vMask->isEmpty() && vMask->rankOf() < v->rankOf()) {
+    vmaskInternal = expandDims.evaluate({vMask},{},{-2}).at(0);
+  }
+
+  if(qMask != nullptr && !qMask->isEmpty()) {
+    qMaskInternal = expandDims.evaluate({qMaskInternal},{},{-1}).at(0);
+  }
+
+
+  auto dLdq = outputs[0];
+  auto dLdv = outputs[1];
+  auto dLdk = outputs[2];
+  attentionBpHelper(q, k, v, scale, dLdq, dLdk, dLdv, eps, dropoutSeed, qMaskInternal, vmaskInternal, useCausalMask,
+                    dropout, training, attentionScoresOut, attentionScoresWeights, attentionScoresLogits, dropoutMask);
+
+}
+
+
+/**
+ * @param inputs
+ * @param mask
+ * @param training
+ * @param returnAttentionScores
+ * @param useCausalMask
+ */
+void AttentionHelper::doAttention(std::vector<NDArray *> &inputs, std::vector<NDArray *> &masks, bool training,
+                                  bool useCausalMask, double dropout, double scale, NDArray *attentionScores,
+                                  int dropoutSeed, NDArray *applyScoresOut, NDArray *attentionLogits,
+                                  NDArray *dropoutMask) {
+  auto q = inputs[0];
+  auto v = inputs[1];
+  auto k = inputs.size() > 2 ? inputs[2]  : v;
+  auto concatWeights = inputs.size() > 3 ? inputs[3] : nullptr;
+
+  ops::expand_dims expandDims;
+  ops::ones_as onesAs;
+  ops::shape_of shapeOf;
+  ops::concat concatOp;
+  ops::create_view createView;
+  auto qMask = masks.size() > 0 ? masks[0] : nullptr;
+  auto vMask = masks.size() > 1 ? masks[1] : nullptr;
+  auto vmaskInternal = vMask;
+  auto qMaskInternal = qMask;
+
+  NDArray *casualPointer = nullptr;
+  //inputs: query and value
+  //shape: batch_size Tq dim (batch_size Tv dim)
+  //note this does not apply softmax yet, we are just computing logits here
+  attentionHelper(q, k, scale, attentionLogits);
+
+  if(vMask != nullptr && !vMask->isEmpty() && vMask->rankOf() < v->rankOf()) {
+    vmaskInternal = expandDims.evaluate({vMask},{},{-2}).at(0);
+  }
+
+  if(useCausalMask) {
+    std::vector<LongType> causalMaskShape2;
+    causalMaskShape2.push_back(attentionScores->sizeAt(0));
+    //4d
+    if(attentionScores->rankOf() > 3)
+      causalMaskShape2.push_back(attentionScores->sizeAt(1));
+
+    for(int i = attentionScores->rankOf() - 2; i < attentionScores->rankOf(); i++) {
+      causalMaskShape2.push_back(attentionScores->sizeAt(i));
+    }
+    casualPointer = lowerTriangularMask(&causalMaskShape2);
+  }
+
+  auto scoresMask = mergeMasks(vmaskInternal,casualPointer);
+
+  //compute actual softmax now
+  if(training) {
+    applyAttentionScores(attentionScores, v, scoresMask, dropout, dropoutSeed, applyScoresOut, attentionLogits,
+                         dropoutMask);
+  } else {
+    applyAttentionScores(attentionScores, v, scoresMask, 0, dropoutSeed, applyScoresOut, attentionLogits, dropoutMask);
+  }
+  //inputs: scores:  batch size tq tv value:batch size, tv,dim scoresmask: batch size 1 tv or batch size tq tv
+  if(qMask != nullptr && !qMask->isEmpty()) {
+    qMaskInternal = expandDims.evaluate({qMaskInternal},{},{-1}).at(0);
+    auto casted = qMaskInternal->cast(attentionScores->dataType());
+    *attentionScores *= *casted;
+  }
+
+}
+
+
+void AttentionHelper::multiHeadProjectBp(NDArray *input, NDArray *projectionMatrix,
+                                         NDArray *eps,
+                                         NDArray *dLdInput, NDArray *dLdProjectionMatrix, LaunchContext *context) {
+  auto miniBatchSize = input->sizeAt(0);
+  auto seqLength = input->sizeAt(2);
+  auto numHeads = projectionMatrix->sizeAt(0);
+  auto projectedSize = projectionMatrix->sizeAt(1);
+
+  std::vector<sd::LongType> epsPermVec = {1, 2, 0, 3};
+  auto epsPerm = eps->permute(epsPermVec, false, false);
+  std::vector<sd::LongType> epsReshapeVec = {numHeads * projectedSize, miniBatchSize * seqLength};
+  auto epsReshaped = epsPerm->reshape('c', epsReshapeVec);
+
+  std::vector<sd::LongType> inputPermVec = {1, 0, 2};
+  auto inputPerm = input->permute(inputPermVec, false, false);
+  std::vector<sd::LongType> inputPermShape = {input->sizeAt(1), miniBatchSize * seqLength};
+  auto inputPrep = inputPerm->reshape('c',inputPermShape,false);
+  std::vector<sd::LongType> projectionMatrixShape = {numHeads * projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)};
+  auto projectionPrep =
+      projectionMatrix->reshape('c', projectionMatrixShape);
+
+  ops::matmul_bp mmulBp;
+  NDArray dLdProjectionPrep(projectionPrep->shapeInfo(), false, context);
+  NDArray dLdInputPrep(inputPrep->shapeInfo(), false, context);
+  mmulBp.execute({projectionPrep, inputPrep, epsReshaped}, std::vector<NDArray *>{&dLdProjectionPrep, &dLdInputPrep},
+                 {}, {}, {});
+
+  dLdProjectionPrep.reshapei({numHeads, projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)});
+  dLdProjectionMatrix->assign(&dLdProjectionPrep);
+
+  dLdInputPrep.reshapei({input->sizeAt(1), miniBatchSize, seqLength});
+  dLdInputPrep.permutei({1, 0, 2}, false, false);
+  dLdInput->assign(&dLdInputPrep);
+
+  delete epsReshaped;
+  delete projectionPrep;
+
+}
+}  // namespace sd
+#endif
+
+#endif

libnd4j/include/helpers/impl/BitwiseUtils.cpp

@@ -0,0 +1,75 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 10.11.2017.
+//
+#include <helpers/BitwiseUtils.h>
+#include <helpers/logger.h>
+#include <types/float16.h>
+
+namespace sd {
+
+bool BitwiseUtils::isBE() {
+  short int word = 0x0001;
+  char *byte = (char *)&word;
+  return (byte[0] ? false : true);
+}
+
+int BitwiseUtils::valueBit(int holder) {
+  if (holder == 0) return -1;
+
+#ifdef REVERSE_BITS
+  for (int e = 32; e >= 0; e--) {
+#else
+  for (int e = 0; e < 32; e++) {
+#endif
+    bool isOne = (holder & 1 << e) != 0;
+
+    if (isOne) return e;
+  }
+
+  return -1;
+}
+
+std::vector<LongType> BitwiseUtils::valueBits(int holder) {
+  std::vector<LongType> bits;
+  if (holder == 0) {
+    for (int e = 0; e < 32; e++) bits.emplace_back(0);
+
+    return bits;
+  }
+
+#ifdef REVERSE_BITS
+  for (int e = 32; e >= 0; e--) {
+#else
+  for (int e = 0; e < 32; e++) {
+#endif
+    bool isOne = (holder & 1 << e) != 0;
+
+    if (isOne)
+      bits.emplace_back(1);
+    else
+      bits.emplace_back(0);
+  }
+
+  return bits;
+}
+
+ByteOrder BitwiseUtils::asByteOrder() { return isBE() ? BE : LE; }
+}  // namespace sd

libnd4j/include/helpers/impl/BlasHelper.cpp

@@ -0,0 +1,523 @@
+/* ******************************************************************************
+ *
+ *
+ * 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/BlasHelper.h>
+#include <cstdlib>
+#include <string>
+
+// OpenBLAS thread control function declaration
+#if HAVE_OPENBLAS
+extern "C" void openblas_set_num_threads(int num_threads);
+extern "C" int openblas_get_num_threads(void);
+#endif
+
+namespace sd {
+
+BlasHelper::BlasHelper() {
+  // Initialize BLAS threading configuration from environment
+  initializeBlasThreading();
+}
+
+BlasHelper &BlasHelper::getInstance() {
+  static BlasHelper instance;
+  return instance;
+}
+
+void BlasHelper::initializeFunctions(Pointer *functions) {
+  sd_debug("Initializing BLAS\n", "");
+
+  _hasSgemv = functions[0] != nullptr;
+  _hasSgemm = functions[2] != nullptr;
+
+  _hasDgemv = functions[1] != nullptr;
+  _hasDgemm = functions[3] != nullptr;
+
+  _hasSgemmBatch = functions[4] != nullptr;
+  _hasDgemmBatch = functions[5] != nullptr;
+#if !defined(SD_CUDA)
+  this->cblasSgemv = (CblasSgemv)functions[0];
+  this->cblasDgemv = (CblasDgemv)functions[1];
+  this->cblasSgemm = (CblasSgemm)functions[2];
+  this->cblasDgemm = (CblasDgemm)functions[3];
+  this->cblasSgemmBatch = (CblasSgemmBatch)functions[4];
+  this->cblasDgemmBatch = (CblasDgemmBatch)functions[5];
+  this->lapackeSgesvd = (LapackeSgesvd)functions[6];
+  this->lapackeDgesvd = (LapackeDgesvd)functions[7];
+  this->lapackeSgesdd = (LapackeSgesdd)functions[8];
+  this->lapackeDgesdd = (LapackeDgesdd)functions[9];
+#endif
+}
+
+void BlasHelper::initializeDeviceFunctions(Pointer *functions) {
+  sd_debug("Initializing device BLAS\n", "");
+
+}
+
+#if defined(HAS_FLOAT32)
+template <>
+bool BlasHelper::hasGEMV<float>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return true;
+#else
+  return _hasSgemv;
+#endif
+}
+#endif
+
+#if defined(HAS_DOUBLE)
+template <>
+bool BlasHelper::hasGEMV<double>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return true;
+#else
+  return _hasDgemv;
+#endif
+}
+#endif
+
+#if defined(HAS_FLOAT16)
+template <>
+bool BlasHelper::hasGEMV<float16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BFLOAT16)
+template <>
+bool BlasHelper::hasGEMV<bfloat16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BOOL)
+template <>
+bool BlasHelper::hasGEMV<bool>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT32)
+template <>
+bool BlasHelper::hasGEMV<int>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT8)
+template <>
+bool BlasHelper::hasGEMV<int8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_UINT8)
+template <>
+bool BlasHelper::hasGEMV<uint8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT16)
+template <>
+bool BlasHelper::hasGEMV<int16_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_LONG)
+template <>
+bool BlasHelper::hasGEMV<LongType>() {
+  return false;
+}
+#endif
+
+bool BlasHelper::hasGEMV(const DataType dtype) {
+#if defined(HAS_FLOAT32)
+  if (dtype == FLOAT32) {
+    if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+    return true;
+#else
+    return _hasSgemv;
+#endif
+  }
+#endif
+#if defined(HAS_DOUBLE)
+  if (dtype == DOUBLE) {
+    if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+    return true;
+#else
+    return _hasDgemv;
+#endif
+  }
+#endif
+  return false;
+}
+
+#if defined(HAS_FLOAT32)
+template <>
+bool BlasHelper::hasGEMM<float>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return true;
+#else
+  return _hasSgemm;
+#endif
+}
+#endif
+
+#if defined(HAS_DOUBLE)
+template <>
+bool BlasHelper::hasGEMM<double>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return true;
+#else
+  return _hasDgemm;
+#endif
+}
+#endif
+
+#if defined(HAS_FLOAT16)
+template <>
+bool BlasHelper::hasGEMM<float16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BFLOAT16)
+template <>
+bool BlasHelper::hasGEMM<bfloat16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT32)
+template <>
+bool BlasHelper::hasGEMM<int>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_UINT8)
+template <>
+bool BlasHelper::hasGEMM<uint8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT8)
+template <>
+bool BlasHelper::hasGEMM<int8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT16)
+template <>
+bool BlasHelper::hasGEMM<int16_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BOOL)
+template <>
+bool BlasHelper::hasGEMM<bool>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_LONG)
+template <>
+bool BlasHelper::hasGEMM<LongType>() {
+  return false;
+}
+#endif
+
+bool BlasHelper::hasGEMM(const DataType dtype) {
+#if defined(HAS_FLOAT32)
+  if (dtype == FLOAT32) {
+    if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+    return true;
+#else
+    return _hasSgemm;
+#endif
+  }
+#endif
+#if defined(HAS_DOUBLE)
+  if (dtype == DOUBLE) {
+    if (Environment::getInstance().blasFallback()) return false;
+
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+    return true;
+#else
+    return _hasDgemm;
+#endif
+  }
+#endif
+  return false;
+}
+
+#if defined(HAS_FLOAT32)
+template <>
+bool BlasHelper::hasBatchedGEMM<float>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+  return _hasSgemmBatch;
+}
+#endif
+
+#if defined(HAS_DOUBLE)
+template <>
+bool BlasHelper::hasBatchedGEMM<double>() {
+  if (Environment::getInstance().blasFallback()) return false;
+
+  return _hasDgemmBatch;
+}
+#endif
+
+#if defined(HAS_FLOAT16)
+template <>
+bool BlasHelper::hasBatchedGEMM<float16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BFLOAT16)
+template <>
+bool BlasHelper::hasBatchedGEMM<bfloat16>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_LONG)
+template <>
+bool BlasHelper::hasBatchedGEMM<LongType>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT32)
+template <>
+bool BlasHelper::hasBatchedGEMM<int>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT8)
+template <>
+bool BlasHelper::hasBatchedGEMM<int8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_UINT8)
+template <>
+bool BlasHelper::hasBatchedGEMM<uint8_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_INT16)
+template <>
+bool BlasHelper::hasBatchedGEMM<int16_t>() {
+  return false;
+}
+#endif
+
+#if defined(HAS_BOOL)
+template <>
+bool BlasHelper::hasBatchedGEMM<bool>() {
+  return false;
+}
+#endif
+
+#if !defined(SD_CUDA)
+#if defined(HAS_FLOAT32)
+CblasSgemv BlasHelper::sgemv() {
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return (CblasSgemv)&cblas_sgemv;
+#else
+  return this->cblasSgemv;
+#endif
+}
+
+CblasSgemm BlasHelper::sgemm() {
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return (CblasSgemm)&cblas_sgemm;
+#else
+  return this->cblasSgemm;
+#endif
+}
+
+CblasSgemmBatch BlasHelper::sgemmBatched() { return this->cblasSgemmBatch; }
+
+LapackeSgesvd BlasHelper::sgesvd() { return this->lapackeSgesvd; }
+
+LapackeSgesdd BlasHelper::sgesdd() { return this->lapackeSgesdd; }
+#endif
+
+#if defined(HAS_DOUBLE)
+CblasDgemv BlasHelper::dgemv() {
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return (CblasDgemv)&cblas_dgemv;
+#else
+  return this->cblasDgemv;
+#endif
+}
+
+CblasDgemm BlasHelper::dgemm() {
+#if __EXTERNAL_BLAS__ || HAVE_OPENBLAS
+  return (CblasDgemm)&cblas_dgemm;
+#else
+  return this->cblasDgemm;
+#endif
+}
+
+CblasDgemmBatch BlasHelper::dgemmBatched() { return this->cblasDgemmBatch; }
+
+LapackeDgesvd BlasHelper::dgesvd() { return this->lapackeDgesvd; }
+
+LapackeDgesdd BlasHelper::dgesdd() { return this->lapackeDgesdd; }
+#endif
+#endif
+
+// BLAS call serialization implementation
+
+std::unique_lock<std::mutex> BlasHelper::lockBlas() const {
+  if (_serializeBlasCalls.load()) {
+    return std::unique_lock<std::mutex>(_blasMutex);
+  }
+  // Return an unlocked lock if serialization is disabled
+  return std::unique_lock<std::mutex>(_blasMutex, std::defer_lock);
+}
+
+bool BlasHelper::isSerializeBlasCalls() const {
+  return _serializeBlasCalls.load();
+}
+
+void BlasHelper::setSerializeBlasCalls(bool serialize) {
+  _serializeBlasCalls.store(serialize);
+}
+
+int BlasHelper::getOpenblasThreads() const {
+  return _openblasThreads.load();
+}
+
+void BlasHelper::setOpenblasThreads(int threads) {
+  _openblasThreads.store(threads);
+#if HAVE_OPENBLAS
+  if (threads > 0) {
+    openblas_set_num_threads(threads);
+    sd_debug("OpenBLAS threads set to %d\n", threads);
+  }
+#endif
+}
+
+void BlasHelper::initializeBlasThreading() {
+  // Check SD_BLAS_SERIALIZE environment variable
+  // Default is true (serialization enabled) for OpenBLAS safety
+  const char* serializeEnv = std::getenv("SD_BLAS_SERIALIZE");
+  if (serializeEnv != nullptr) {
+    std::string val(serializeEnv);
+    if (val == "0" || val == "false" || val == "FALSE" || val == "no" || val == "NO") {
+      _serializeBlasCalls.store(false);
+      sd_debug("BLAS call serialization DISABLED via SD_BLAS_SERIALIZE=%s\n", serializeEnv);
+    } else {
+      _serializeBlasCalls.store(true);
+      sd_debug("BLAS call serialization ENABLED via SD_BLAS_SERIALIZE=%s\n", serializeEnv);
+    }
+  } else {
+    // Default: enable serialization for OpenBLAS safety
+    _serializeBlasCalls.store(true);
+    sd_debug("BLAS call serialization ENABLED by default (set SD_BLAS_SERIALIZE=0 to disable)\n", "");
+  }
+
+  // Check SD_OPENBLAS_THREADS environment variable for OpenBLAS thread count
+  // This is separate from the serialization - you can have both:
+  // - Serialization ON + multi-threaded OpenBLAS = safe concurrent BLAS with internal parallelism
+  // - Serialization OFF + single-threaded OpenBLAS = original behavior
+  const char* threadsEnv = std::getenv("SD_OPENBLAS_THREADS");
+  if (threadsEnv != nullptr) {
+#ifdef __cpp_exceptions
+    try {
+      int threads = std::stoi(std::string(threadsEnv));
+      if (threads > 0) {
+        _openblasThreads.store(threads);
+#if HAVE_OPENBLAS
+        openblas_set_num_threads(threads);
+        sd_debug("OpenBLAS threads set to %d via SD_OPENBLAS_THREADS\n", threads);
+#endif
+      }
+    } catch (...) {
+      // Invalid value, ignore
+    }
+#else
+    int threads = std::atoi(threadsEnv);
+    if (threads > 0) {
+      _openblasThreads.store(threads);
+#if HAVE_OPENBLAS
+      openblas_set_num_threads(threads);
+      sd_debug("OpenBLAS threads set to %d via SD_OPENBLAS_THREADS\n", threads);
+#endif
+    }
+#endif
+  }
+
+  // Also check OPENBLAS_NUM_THREADS (standard OpenBLAS env var) if SD_OPENBLAS_THREADS not set
+  if (_openblasThreads.load() == 0) {
+    const char* openblasEnv = std::getenv("OPENBLAS_NUM_THREADS");
+    if (openblasEnv != nullptr) {
+#ifdef __cpp_exceptions
+      try {
+        int threads = std::stoi(std::string(openblasEnv));
+        if (threads > 0) {
+          _openblasThreads.store(threads);
+          sd_debug("OpenBLAS threads detected from OPENBLAS_NUM_THREADS=%d\n", threads);
+        }
+      } catch (...) {
+        // Invalid value, ignore
+      }
+#else
+      int threads = std::atoi(openblasEnv);
+      if (threads > 0) {
+        _openblasThreads.store(threads);
+        sd_debug("OpenBLAS threads detected from OPENBLAS_NUM_THREADS=%d\n", threads);
+      }
+#endif
+    }
+  }
+}
+
+// destructor
+BlasHelper::~BlasHelper() noexcept {}
+}  // namespace sd

libnd4j/include/helpers/impl/ConstantShapeHelper.cpp

@@ -0,0 +1,447 @@
+/* ******************************************************************************
+*
+* Copyright (c) 2024 Konduit K.K.
+* 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.
+*
+* 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
+******************************************************************************/
+
+#include "helpers/ConstantShapeHelper.h"
+
+#include "array/ConstantShapeBuffer.h"
+#include "system/common.h"
+#include <array/PrimaryPointerDeallocator.h>
+#include <helpers/ConstantShapeHelper.h>
+#include <helpers/ShapeBuilders.h>
+#include <helpers/ShapeUtils.h>
+#include <helpers/shape.h>
+#include <system/Environment.h>
+
+namespace sd {
+
+ConstantShapeHelper::~ConstantShapeHelper() {
+
+}
+
+ConstantShapeHelper::ConstantShapeHelper() {
+}
+
+ConstantShapeHelper& ConstantShapeHelper::getInstance() {
+ static ConstantShapeHelper instance;
+ return instance;
+}
+
+void ConstantShapeHelper::initializeEarly() {
+  ConstantShapeHelper& instance = getInstance();
+  instance._shapeTrie.waitForInitialization();
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::createConstBuffFromExisting(sd::LongType* shapeInfo) {
+ auto result = bufferForShapeInfo(shapeInfo);
+ return result;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(LongType* shapeInfo) {
+ if(shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+ if(shape::rank(shapeInfo) < 0 || shape::rank(shapeInfo) > SD_MAX_RANK) {
+   THROW_EXCEPTION("shapeInfo is not a valid rank.");
+ }
+
+ auto buffer = _shapeTrie.getOrCreate(shapeInfo);
+ if (buffer == nullptr || buffer->primary() == nullptr) {
+   THROW_EXCEPTION("Failed to get/create shape buffer");
+ }
+ return buffer;
+}
+ConstantShapeBuffer* ConstantShapeHelper::createSubArrShapeInfo( sd::LongType* inShapeInfo,  LongType* dims,
+                                                                 sd::LongType dimsSize) {
+ sd::LongType* newShapeInfo = ShapeBuilders::createSubArrShapeInfo(inShapeInfo, dims, dimsSize, nullptr);
+ auto ret = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return ret;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(DataType dataType, char order,
+                                                            const std::vector<LongType>& shape) {
+ auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, shape);
+ auto result = bufferForShapeInfo(descriptor);
+ delete[] descriptor;
+ return result;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(DataType dataType, char order,
+                                                            int rank,  LongType* shape) {
+ auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, rank, shape, nullptr, false);
+ auto result = bufferForShapeInfo(descriptor);
+ delete[] descriptor;
+ return result;
+}
+
+LongType* ConstantShapeHelper::emptyShapeInfoWithShape(DataType dataType, std::vector<LongType>& shape) {
+ auto descriptor = ShapeBuilders::createShapeInfo(dataType, 'c', shape, nullptr);
+ ArrayOptions::setPropertyBit(descriptor, ARRAY_EMPTY);
+ auto existing = createFromExisting(descriptor);
+ delete[] descriptor;
+ return existing;
+}
+
+LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, char order,
+                                              const std::vector<LongType>& shape) {
+ auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, shape);
+ auto result = bufferForShapeInfo(descriptor)->primary();
+ delete[] descriptor;
+ return result;
+}
+
+LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, char order, int rank,
+                                              LongType* shape, LongType extraProperties) {
+ if (extraProperties < 0) {
+   extraProperties = ArrayOptions::flagForDataType(dataType);
+ }
+
+ std::unique_ptr<LongType[]> strides(order == 'c' ? shape::calcStrides(shape, rank)
+                                                  : shape::calcStridesFortran(shape, rank));
+
+ auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, rank, shape, strides.get(),
+                                                  nullptr, extraProperties);
+ auto ret = bufferForShapeInfo(descriptor)->primary();
+ ArrayOptions::validateSingleDataType(ArrayOptions::dataType(ret));
+
+ delete[] descriptor;
+ return ret;
+}
+
+LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, LongType* shapeInfo) {
+ auto result = createShapeInfo(dataType, shape::order(shapeInfo), shape::rank(shapeInfo),
+                        shape::shapeOf(const_cast<LongType*>(shapeInfo)), -1);
+ return result;
+}
+
+LongType* ConstantShapeHelper::emptyShapeInfo(DataType dataType) {
+ auto descriptor = ShapeBuilders::emptyShapeInfo(dataType);
+ auto result = bufferForShapeInfo(descriptor)->primary();
+ delete[] descriptor;
+ return result;
+}
+
+
+LongType* ConstantShapeHelper::scalarShapeInfo(DataType dataType) {
+ auto descriptor = ShapeBuilders::createScalarShapeInfo(dataType);
+ return bufferForShapeInfo(descriptor)->primary();
+}
+
+LongType* ConstantShapeHelper::vectorShapeInfo(LongType length, DataType dataType) {
+ auto descriptor = ShapeBuilders::createVectorShapeInfo(dataType, length);
+ auto result = bufferForShapeInfo(descriptor)->primary();
+ delete[] descriptor;
+ return result;
+}
+
+
+LongType* ConstantShapeHelper::createShapeInfo(ShapeDescriptor* descriptor) {
+ auto shapeInfo = descriptor->toShapeInfo();
+ auto result = bufferForShapeInfo(shapeInfo)->primary();
+ delete[] shapeInfo;
+ return result;
+}
+
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithView(LongType* shapeInfo) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+
+
+
+ ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+
+ delete[] newShapeInfo;
+
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutView(LongType* shapeInfo) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ ArrayOptions::unsetPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithNeedsCopy(LongType* shapeInfo) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_NEEDS_COPY);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutNeedsCopy(LongType* shapeInfo) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ ArrayOptions::unsetPropertyBit(newShapeInfo, ARRAY_NEEDS_COPY);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithCopyOffset(LongType* shapeInfo, int inputIndex) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ if (inputIndex < 0 || inputIndex > 10) {
+   THROW_EXCEPTION("Input index out of range [0-10]");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
+ ArrayOptions::setPropertyBit(newShapeInfo, flag);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutCopyOffset(LongType* shapeInfo, int inputIndex) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ if (inputIndex < 0 || inputIndex > 10) {
+   THROW_EXCEPTION("Input index out of range [0-10]");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
+ ArrayOptions::unsetPropertyBit(newShapeInfo, flag);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutAllCopyOffsets(LongType* shapeInfo) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+ ArrayOptions::clearAllCopyOffsets(newShapeInfo);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithFlags(LongType* shapeInfo,
+                                                                      LongType flagsToSet,
+                                                                      LongType flagsToUnset) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+
+ // Unset flags first
+ if (flagsToUnset != 0) {
+   LongType extraIdx = ArrayOptions::extraIndex(newShapeInfo);
+   newShapeInfo[extraIdx] = newShapeInfo[extraIdx] & ~flagsToUnset;
+ }
+
+ // Then set flags
+ if (flagsToSet != 0) {
+   LongType extraIdx = ArrayOptions::extraIndex(newShapeInfo);
+   newShapeInfo[extraIdx] = newShapeInfo[extraIdx] | flagsToSet;
+ }
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoAsViewWithOffset(LongType* shapeInfo,
+                                                                             int inputIndex) {
+ if (shapeInfo == nullptr) {
+   THROW_EXCEPTION("shapeInfo is nullptr");
+ }
+
+ if (inputIndex < 0 || inputIndex > 10) {
+   THROW_EXCEPTION("Input index out of range [0-10]");
+ }
+
+ LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
+
+ // Set view flag
+ ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
+
+ // Set copy offset flag for specified input
+ LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
+ ArrayOptions::setPropertyBit(newShapeInfo, flag);
+
+ auto buffer = bufferForShapeInfo(newShapeInfo);
+ delete[] newShapeInfo;
+ return buffer;
+}
+
+LongType* ConstantShapeHelper::createFromExisting(LongType* shapeInfo) {
+ if (!shapeInfo) {
+   THROW_EXCEPTION("Null shape info");
+ }
+ auto buffer = bufferForShapeInfo(shapeInfo);
+ return buffer->primary();
+}
+
+
+LongType* ConstantShapeHelper::castToDataType(LongType* shapeInfo, DataType newType) {
+ if (!shapeInfo) {
+   THROW_EXCEPTION("Null shape info");
+ }
+ if (ArrayOptions::dataType(shapeInfo) == newType) {
+   return shapeInfo;
+ }
+
+ auto tempShapeInfo = ShapeBuilders::copyShapeInfoWithNewType(shapeInfo, newType);
+ if (!tempShapeInfo) {
+   THROW_EXCEPTION("Failed to create temp shape info");
+ }
+
+ auto buffer = bufferForShapeInfo(tempShapeInfo);
+ auto result = buffer->primary();
+ delete[] tempShapeInfo;
+ if(ArrayOptions::dataType(result) != newType) {
+   std::string errorMessage;
+   errorMessage += "castToDataType: new data type is ";
+   errorMessage += DataTypeUtils::asString(newType);
+   errorMessage += " data type from new constant created data type ";
+   errorMessage += DataTypeUtils::asString(ArrayOptions::dataType(result));
+   errorMessage += "\n";
+   THROW_EXCEPTION(errorMessage.c_str());
+ }
+ return result;
+}
+
+
+ConstantShapeBuffer* ConstantShapeHelper::createShapeInfoWithUnitiesForBroadcast(sd::LongType* maxShapeInfo,
+                                                                                sd::LongType* minShapeInfo,
+                                                                                sd::memory::Workspace* workspace,
+                                                                                const std::vector<LongType>& dimensions) {
+ sd::LongType* newShapeInfo = nullptr;
+ ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(shape::rank(maxShapeInfo)), sd::LongType);
+
+ newShapeInfo[0] = shape::rank(maxShapeInfo);
+ newShapeInfo[2 * shape::rank(maxShapeInfo) + 1] = 0;
+ sd::ArrayOptions::copyDataType(newShapeInfo, minShapeInfo);                      // type
+ newShapeInfo[2 * newShapeInfo[0] + 2] = shape::elementWiseStride(minShapeInfo);  // ews
+ newShapeInfo[2 * newShapeInfo[0] + 3] = shape::order(minShapeInfo);              // order
+
+ if (!dimensions.empty()) {
+   for (sd::LongType k = 0, j = 0, i = 0; i < shape::rank(maxShapeInfo); ++i) {
+     if (j < static_cast<sd::LongType>(dimensions.size()) && dimensions[j] == i) {
+       shape::shapeOf(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[k];
+       shape::stride(newShapeInfo)[i] = shape::stride(minShapeInfo)[k++];
+       ++j;
+     } else {
+       shape::shapeOf(newShapeInfo)[i] = 1;
+       shape::stride(newShapeInfo)[i] = 0;
+       if (shape::sizeAt(minShapeInfo, k) == 1 && static_cast<sd::LongType>(dimensions.size()) != shape::rank(minShapeInfo)) ++k;
+     }
+   }
+ } else {
+   for (int j = shape::rank(minShapeInfo) - 1, i = shape::rank(maxShapeInfo) - 1; i >= 0; --i) {
+     if (j >= 0) {
+       shape::shapeOf(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[j];
+       shape::stride(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[j] == 1 ? 0 : shape::stride(minShapeInfo)[j];
+       --j;
+     } else {
+       shape::shapeOf(newShapeInfo)[i] = 1;
+       shape::stride(newShapeInfo)[i] = 0;
+     }
+   }
+ }
+
+ auto ret = bufferForShapeInfo(newShapeInfo);
+ RELEASE(newShapeInfo, workspace);
+ return ret;
+}
+
+ConstantShapeBuffer* ConstantShapeHelper::createShapeInfoWithNoUnitiesForReduce(const sd::LongType* maxShapeInfo,
+                                                                               const std::vector<LongType>* dimsWithUnities,
+                                                                               sd::memory::Workspace* workspace) {
+ sd::LongType* newShapeInfo = nullptr;
+ ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(shape::rank(maxShapeInfo) - dimsWithUnities->size()),
+          sd::LongType);
+
+ sd::LongType temp;
+ if (dimsWithUnities->size() == 1 && shape::isCommonVector(maxShapeInfo, temp) && temp == dimsWithUnities->at(0)) {
+   auto dims = ShapeUtils::evalDimsToExclude(shape::rank(maxShapeInfo), 1,&temp);
+   shape::excludeUnitiesFromShapeInfo(maxShapeInfo, dims->data(), dims->size(), newShapeInfo);
+   delete dims;
+ } else {
+   shape::excludeUnitiesFromShapeInfo(maxShapeInfo, dimsWithUnities->data(), dimsWithUnities->size(), newShapeInfo);
+ }
+
+ auto ret = bufferForShapeInfo(newShapeInfo);
+ RELEASE(newShapeInfo, workspace);
+ return ret;
+}
+
+void ConstantShapeHelper::clearCache() {
+ std::lock_guard<std::mutex> lock(_mutex);
+ _shapeTrie.clearCache();
+}
+
+LongType ConstantShapeHelper::getCachedEntries() const {
+ return _shapeTrie.getCachedEntries();
+}
+
+LongType ConstantShapeHelper::getCachedBytes() const {
+ return _shapeTrie.getCachedBytes();
+}
+
+LongType ConstantShapeHelper::getPeakCachedEntries() const {
+ return _shapeTrie.getPeakCachedEntries();
+}
+
+LongType ConstantShapeHelper::getPeakCachedBytes() const {
+ return _shapeTrie.getPeakCachedBytes();
+}
+
+std::string ConstantShapeHelper::toString(int maxDepth, int maxEntries) const {
+ return _shapeTrie.toString(maxDepth, maxEntries);
+}
+
+void ConstantShapeHelper::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
+ _shapeTrie.getCachedPointers(out_pointers);
+}
+
+} // namespace sd
+

libnd4j/include/helpers/impl/ConstantTadHelper.cpp

@@ -0,0 +1,140 @@
+/* ******************************************************************************
+ *
+ * Copyright (c) 2024 Konduit K.K.
+ * 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.
+ *
+ * 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
+ ******************************************************************************/
+
+#include <array/TadDescriptor.h>
+#include <array/TadPack.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/ShapeUtils.h>
+
+namespace sd {
+
+ConstantTadHelper& ConstantTadHelper::getInstance() {
+  static ConstantTadHelper instance;
+  return instance;
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType dimension) {
+  return tadForDimensions(originalShape, &dimension, 1);
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, std::vector<LongType>* dimensions) {
+  if (dimensions == nullptr) {
+    THROW_EXCEPTION("Dimensions vector is null");
+  }
+  return tadForDimensions(originalShape, const_cast<LongType*>(dimensions->data()), dimensions->size());
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(TadDescriptor* descriptor) {
+  if (descriptor == nullptr) {
+    THROW_EXCEPTION("TadDescriptor is null");
+  }
+  return tadForDimensions(descriptor->originalShape(), descriptor->axis().data(),
+                          descriptor->axis().size());
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType* dimensions, LongType dimLength) {
+  if (originalShape == nullptr) {
+    THROW_EXCEPTION("Original shape is null");
+  }
+
+  if (dimensions == nullptr && dimLength > 0) {
+    THROW_EXCEPTION("Dimensions array is null but dimLength > 0");
+  }
+
+  // Check for empty array
+  if (shape::isEmptyConst(originalShape)) {
+    THROW_EXCEPTION("Cannot create TADs for empty array");
+  }
+
+  sd::LongType rank = shape::rank(originalShape);
+  if (rank < 0) {
+    THROW_EXCEPTION("Invalid shape rank");
+  }
+
+  // Check for zero-sized dimensions
+  for (LongType i = 0; i < rank; i++) {
+    if (shape::sizeAt(originalShape, i) == 0) {
+      THROW_EXCEPTION("Cannot create TADs for array with zero-sized dimensions");
+    }
+  }
+
+  // Handle zero dimension length case - treat entire array as single TAD
+  if (dimLength <= 0) {
+    // When no dimensions specified, create TAD along all dimensions
+    // This means the entire array is treated as a single TAD
+    std::vector<LongType> allDims;
+    for (LongType i = 0; i < rank; i++) {
+      allDims.push_back(i);
+    }
+
+    // Recursively call with all dimensions
+    return tadForDimensions(originalShape, allDims.data(), rank);
+  }
+
+  // Additional validation: check if dimensions are within valid range
+  for (LongType i = 0; i < dimLength; i++) {
+    LongType dim = dimensions[i];
+    if (dim < 0) dim += rank;  // Handle negative dimensions
+    if (dim < 0 || dim >= rank) {
+      THROW_EXCEPTION("Dimension index is out of bounds");
+    }
+  }
+
+  // Create non-temporary vector to satisfy the reference requirement
+  std::vector<LongType> dims(dimensions, dimensions + dimLength);
+
+  // The shared_ptr keeps the TadPack alive even if the cache tries to clear it
+  std::shared_ptr<TadPack> result = nullptr;
+  try {
+    result = _trie.getOrCreate(dims, originalShape);
+  } catch (const std::exception& e) {
+    THROW_EXCEPTION("Failed to create or retrieve TAD pack");
+  }
+
+  // DO NOT call checkAndCleanupCaches() here - would delete the pack we just created!
+  // Cleanup happens at NativeOps layer AFTER operations complete.
+  return result;
+}
+
+void ConstantTadHelper::clearCache() {
+  _trie.clear();
+}
+
+LongType ConstantTadHelper::getCachedEntries() const {
+  return _trie.getCachedEntries();
+}
+
+LongType ConstantTadHelper::getCachedBytes() const {
+  return _trie.getCachedBytes();
+}
+
+LongType ConstantTadHelper::getPeakCachedEntries() const {
+  return _trie.getPeakCachedEntries();
+}
+
+LongType ConstantTadHelper::getPeakCachedBytes() const {
+  return _trie.getPeakCachedBytes();
+}
+
+std::string ConstantTadHelper::toString(int maxDepth, int maxEntries) const {
+  return _trie.toString(maxDepth, maxEntries);
+}
+
+void ConstantTadHelper::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
+  _trie.getCachedPointers(out_pointers);
+}
+
+} // namespace sd

libnd4j/include/helpers/impl/CudaLaunchHelper.cpp

@@ -0,0 +1,35 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver on 4/5/2018.
+//
+#include <helpers/CudaLaunchHelper.h>
+#include <math/templatemath.h>
+
+namespace sd {
+
+int CudaLaunchHelper::getReductionBlocks(LongType xLength, int blockSize) {
+  int div = xLength / blockSize;
+  int can = sd::math::sd_max<int>(div, 1);
+  if (xLength % blockSize != 0 && xLength > blockSize) can++;
+
+  // not more then 512 blocks
+  return sd::math::sd_min<int>(can, 512);
+}
+}  // namespace sd

libnd4j/include/helpers/impl/DebugHelper.cpp

@@ -0,0 +1,109 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 20/04/18.
+//
+#include <array/NDArray.h>
+#include <array/NDArrayFactory.h>
+#include <execution/Threads.h>
+#include <helpers/DebugHelper.h>
+#include <helpers/DebugInfo.h>
+#include <ops/declarable/headers/parity_ops.h>
+
+namespace sd {
+DebugInfo DebugHelper::debugStatistics(NDArray * input) {
+  DebugInfo info;
+  retrieveDebugStatistics(&info, const_cast<NDArray*>(input));
+  return info;
+}
+void DebugHelper::retrieveDebugStatistics(DebugInfo* info, NDArray* input) {
+  if (nullptr == info) return;
+
+  info->_minValue = 0.;
+  info->_maxValue = -1;
+  info->_meanValue = 0.;
+  info->_stdDevValue = 1.;
+  info->_zeroCount = 0;
+  info->_positiveCount = 0;
+  info->_negativeCount = 0;
+  info->_infCount = 0;
+  info->_nanCount = 0;
+  if (input->lengthOf() == 1) {  // scalar case
+    info->_minValue = input->e<double>(0);
+    info->_maxValue = info->_minValue;
+    info->_meanValue = info->_minValue;
+    info->_stdDevValue = info->_minValue;
+    info->_zeroCount = math::sd_abs<double,double>(input->e<double>(0)) > 0.00001 ? 0 : 1;
+    info->_positiveCount = input->e<double>(0) > 0 ? 1 : 0;
+    info->_negativeCount = input->e<double>(0) < 0 ? 1 : 0;
+    info->_infCount = math::sd_isinf(input->e<double>(0));
+    info->_nanCount = math::sd_isnan(input->e<double>(0));
+  } else if (input->lengthOf() > 0) {
+    // TO DO: here processing for all elements with array
+    auto _minValue = input->e<double>(0);
+    auto _maxValue = input->e<double>(0);
+    auto _meanValue = input->e<double>(0);
+    auto _stdDevValue = 0.;  // info->_minValue;
+    auto _zeroCount = math::sd_abs<double,double>(input->e<double>(0)) > 0.00001 ? 0L : 1L;
+    auto _positiveCount = input->e<double>(0) > 0 ? 1L : 0L;
+    auto _negativeCount = input->e<double>(0) < 0 ? 1L : 0L;
+    auto _infCount = math::sd_isinf(input->e<double>(0)) ? 1L : 0L;
+    auto _nanCount = math::sd_isnan(input->e<double>(0)) ? 1L : 0L;
+
+    PRAGMA_OMP_PARALLEL_FOR_ARGS(schedule(guided) reduction(+:_nanCount,_infCount,_meanValue,_zeroCount,_positiveCount,_negativeCount) reduction(min:_minValue) reduction(max:_maxValue))
+    for (LongType e = 1; e < input->lengthOf(); e++) {
+      auto current = input->e<double>(e);
+      auto n = e + 1.;
+      //                auto delta = current - _meanValue;
+      //                auto delta2 = delta * delta;
+      _minValue = math::sd_min(current, _minValue);
+      _maxValue = math::sd_max(current, _maxValue);
+
+      _meanValue += current;
+      //_meanValue += delta / n; // this is a perfect formula but not working with omp in this notation
+      //_stdDevValue += delta2 * e / n;
+
+      _zeroCount += math::sd_abs<double,double>(current) > 0.00001 ? 0 : 1;
+      _positiveCount += current > 0 ? 1 : 0;
+      _negativeCount += current < 0 ? 1 : 0;
+      _infCount += math::sd_isinf(current);
+      _nanCount += math::sd_isnan(current);
+    }
+    *info = {_minValue,      _maxValue,  _meanValue / input->lengthOf(),
+             _stdDevValue,   _zeroCount, _positiveCount,
+             _negativeCount, _infCount,  _nanCount};
+    _stdDevValue = 0;  // math::sd_sqrt<double, double>(info->_stdDevValue / (input->lengthOf() - 1));
+
+    auto func = PRAGMA_REDUCE_DOUBLE {
+      auto _stdDevValue = 0.0;
+      for (auto e = start; e < stop; e++) {
+        double current = input->e<double>(e);
+        _stdDevValue += (info->_meanValue - current) * (info->_meanValue - current);  // info->_minValue;
+      }
+
+      return _stdDevValue;
+    };
+    _stdDevValue = samediff::Threads::parallel_double(
+        func, LAMBDA_AD { return _old + _new; }, 0, input->lengthOf());
+
+    info->_stdDevValue = math::sd_sqrt<double, double>(_stdDevValue / input->lengthOf());
+  }
+  // else - no statistics for empty
+}
+}  // namespace sd

libnd4j/include/helpers/impl/DirectShapeTrie.cpp

@@ -0,0 +1,867 @@
+/* ******************************************************************************
+*
+* 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
+******************************************************************************/
+
+#include <array/ArrayOptions.hXX>
+#include <array/ConstantShapeBuffer.h>
+#include <array/DataType.h>
+#include <array/PrimaryPointerDeallocator.h>
+#include <helpers/DirectShapeTrie.h>
+#include <helpers/shape.h>
+#include <system/common.h>
+
+#include <atomic>
+#include <chrono>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <thread>
+
+#include "helpers/ShapeBufferCreatorHelper.h"
+
+#if defined(SD_GCC_FUNCTRACE)
+#include <array/ShapeCacheLifecycleTracker.h>
+#endif
+
+namespace sd {
+
+void DirectShapeTrie::waitForInitialization() const {
+  if (_initialization_complete.load(std::memory_order_acquire)) {
+    return;
+  }
+
+  int attempts = 0;
+  while (_initialization_in_progress.load(std::memory_order_acquire)) {
+    if (attempts < 10) {
+      std::this_thread::yield();
+    } else {
+      auto delay = std::chrono::microseconds(std::min<int>(500, attempts * 10));
+      std::this_thread::sleep_for(delay);
+    }
+    attempts++;
+  }
+
+  if (!_initialization_complete.load(std::memory_order_acquire)) {
+    THROW_EXCEPTION("DirectShapeTrie initialization did not complete before use");
+  }
+}
+
+void ShapeTrieNode::setBuffer(ConstantShapeBuffer* buf) {
+  if (!buf) return;  // Nothing to do if buffer is null
+
+  // If we already have a buffer, don't replace it
+  if (_buffer != nullptr) {
+    // The existing buffer takes precedence
+    // Don't delete the new buffer - let the caller handle it
+    return;
+  }
+
+  // At this point, we know _buffer is null and buf is valid
+  // Set the buffer atomically
+  _buffer = buf;
+}
+
+
+#if defined(SD_GCC_FUNCTRACE)
+void ShapeTrieNode::collectStoreStackTrace() {
+  this->storeStackTrace = backward::StackTrace();
+  this->storeStackTrace.load_here(32);
+}
+#endif
+
+size_t DirectShapeTrie::computeHash(const LongType* shapeInfo) const {
+  size_t hash = 17; // Prime number starting point
+  const int rank = shape::rank(shapeInfo);
+
+  // Add rank first with high weight
+  hash = hash * 31 + rank * 19;
+
+  // Add shape elements to hash with position-dependent multipliers
+  const LongType* shape = shape::shapeOf(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    hash = hash * 13 + static_cast<size_t>(shape[i]) * (7 + i);
+  }
+
+  // Add stride elements to hash with position-dependent multipliers
+  const LongType* strides = shape::stride(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    hash = hash * 19 + static_cast<size_t>(strides[i]) * (11 + i);
+  }
+
+  // Add data type and order with higher weights
+  hash = hash * 23 + static_cast<size_t>(ArrayOptions::dataType(shapeInfo)) * 29;
+  hash = hash * 37 + static_cast<size_t>(shape::order(shapeInfo)) * 41;
+
+  // Add total element count
+  hash = hash * 43 + shape::length(shapeInfo);
+
+  // **NEW: Add property flags to distinguish views from non-views**
+  hash = hash * 47 + static_cast<size_t>(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+
+  return hash;
+}
+
+int DirectShapeTrie::calculateShapeSignature(const LongType* shapeInfo) const {
+  int signature = 17;
+  const int rank = shape::rank(shapeInfo);
+
+  // Incorporate rank with weight
+  signature = signature * 31 + rank * 13;
+
+  // Incorporate shape dimensions with position weights
+  const LongType* shapeValues = shape::shapeOf(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    signature = signature * 13 + static_cast<int>(shapeValues[i]) * (7 + i);
+  }
+
+  // Incorporate data type and order
+  signature = signature * 7 + static_cast<int>(ArrayOptions::dataType(shapeInfo)) * 11;
+  signature = signature * 17 + static_cast<int>(shape::order(shapeInfo)) * 19;
+
+  // Include element count
+  signature = signature * 23 + static_cast<int>(shape::length(shapeInfo) % 10000);
+
+  // **NEW: Include property flags**
+  signature = signature * 29 + static_cast<int>(shapeInfo[ArrayOptions::extraIndex(shapeInfo)] % 10000);
+
+  return signature;
+}
+
+size_t DirectShapeTrie::getStripeIndex(const LongType* shapeInfo) const {
+  return computeHash(shapeInfo) % NUM_STRIPES;
+}
+
+bool DirectShapeTrie::shapeInfoEqual(const LongType* a, const LongType* b) const {
+  if (a == b) return true;
+  if (a == nullptr || b == nullptr) return false;
+
+  const int rankA = shape::rank(a);
+  if (rankA != shape::rank(b)) return false;
+
+  const int len = shape::shapeInfoLength(rankA);
+  return std::memcmp(a, b, len * sizeof(LongType)) == 0;
+}
+
+void DirectShapeTrie::validateShapeInfo(const LongType* shapeInfo) const {
+  if (shapeInfo == nullptr) {
+    std::string msg = "Shape info cannot be null";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  const int rank = shape::rank(shapeInfo);
+  if (rank < 0 || rank > SD_MAX_RANK) {
+    std::string errorMessage = "Invalid rank: " + std::to_string(rank) +
+                               ". Valid range is 0 to " + std::to_string(SD_MAX_RANK);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (rank == 0) {
+    const int len = shape::shapeInfoLength(rank);
+    bool allZero = true;
+    for (int i = 0; i < len; i++) {
+      if (shapeInfo[i] != 0) {
+        allZero = false;
+        break;
+      }
+    }
+    if (allZero) {
+      std::string msg = "Found shape buffer with all zero values. Values likely unset.";
+      THROW_EXCEPTION(msg.c_str());
+    }
+  }
+
+  if (ArrayOptions::dataType(shapeInfo) == UNKNOWN) {
+    std::string msg = "Shape info created with invalid data type";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  char order = shape::order(shapeInfo);
+  if (order != 'c' && order != 'f') {
+    std::string errorMessage = "Invalid ordering in shape buffer: ";
+    errorMessage += order;
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+}
+
+const ShapeTrieNode* DirectShapeTrie::findChild(const ShapeTrieNode* node, LongType value,
+                                                int level, bool isShape, int shapeHash) const {
+  if (!node) return nullptr;
+
+  for (const auto& child : node->children()) {
+    if (child->value() == value &&
+        child->level() == level &&
+        child->isShape() == isShape &&
+        (shapeHash == 0 || child->shapeHash() == shapeHash)) {
+      return child;
+    }
+  }
+  return nullptr;
+}
+
+// Modified search method - still returns null when shape not found but with improved debugging
+ConstantShapeBuffer* DirectShapeTrie::search(const LongType* shapeInfo, size_t stripeIdx) const {
+  // Validate input
+  if (shapeInfo == nullptr) {
+    std::string msg = "Null shapeInfo passed to search method";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  if (stripeIdx >= NUM_STRIPES) {
+    std::string msg = "Invalid stripe index: " + std::to_string(stripeIdx) +
+                      " (max: " + std::to_string(NUM_STRIPES - 1) + ")";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  if (_roots == nullptr) {
+    std::string msg = "Root nodes array is null";
+    THROW_EXCEPTION(msg.c_str());
+  }
+  auto rootsRef = *_roots;
+
+
+  // No locks here - caller handles locking
+  const ShapeTrieNode* current = rootsRef[stripeIdx];
+  if (current == nullptr) {
+    // Cannot use createFallbackBuffer here as it's const method
+    // Caller should handle this case
+    return nullptr;
+  }
+
+  const int rank = shape::rank(shapeInfo);
+  const int shapeSignature = calculateShapeSignature(shapeInfo);
+
+  // Check rank
+  current = findChild(current, rank, 0, true, shapeSignature);
+  if (!current) {
+    return nullptr;  // Not found, but this is expected behavior
+  }
+
+  // Check datatype
+  current = findChild(current, ArrayOptions::dataType(shapeInfo), 1, true, shapeSignature);
+  if (!current) {
+    return nullptr;  // Not found, but this is expected behavior
+  }
+
+  // Check order
+  current = findChild(current, shape::order(shapeInfo), 2, true, shapeSignature);
+  if (!current) {
+    return nullptr;  // Not found, but this is expected behavior
+  }
+
+  // Check shape values
+  const LongType* shapeValues = shape::shapeOf(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    current = findChild(current, shapeValues[i], 3 + i, true, shapeSignature);
+    if (!current) {
+      return nullptr;  // Not found, but this is expected behavior
+    }
+  }
+
+  // Check stride values
+  const LongType* strides = shape::stride(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    current = findChild(current, strides[i], 3 + rank + i, false, shapeSignature);
+    if (!current) {
+      return nullptr;  // Not found, but this is expected behavior
+    }
+  }
+
+  return current ? current->buffer() : nullptr;
+}
+
+
+// Helper method to create a fallback buffer when the trie insertion fails
+ConstantShapeBuffer* DirectShapeTrie::createFallbackBuffer(const LongType* shapeInfo, int rank) {
+  if (shapeInfo == nullptr) {
+    std::string msg = "Null shapeInfo passed to createFallbackBuffer";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  if (rank < 0 || rank > SD_MAX_RANK) {
+    std::string msg = "Invalid rank in createFallbackBuffer: " + std::to_string(rank);
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  // Create a direct copy of the shape info
+  const int shapeInfoLength = shape::shapeInfoLength(rank);
+  LongType* shapeCopy = new LongType[shapeInfoLength];
+  if (shapeCopy == nullptr) {
+    std::string msg = "Failed to allocate memory for shape copy";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  std::memcpy(shapeCopy, shapeInfo, shapeInfoLength * sizeof(LongType));
+
+  // Create a deallocator for memory management
+  auto deallocator = std::shared_ptr<PrimaryPointerDeallocator>(
+      new PrimaryPointerDeallocator(),
+      [] (PrimaryPointerDeallocator* ptr) { delete ptr; });
+
+  // Create a pointer wrapper and buffer
+  auto hPtr = new PointerWrapper(shapeCopy, deallocator);
+  if (hPtr == nullptr) {
+    delete[] shapeCopy;
+    std::string msg = "Failed to create PointerWrapper";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  auto buffer = new ConstantShapeBuffer(hPtr);
+  if (buffer == nullptr) {
+    delete hPtr;
+    std::string msg = "Failed to create ConstantShapeBuffer";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // Track shape cache allocation
+  sd::array::ShapeCacheLifecycleTracker::getInstance().recordAllocation(shapeCopy);
+#endif
+
+  // Fallback buffer is NOT cached, so refCount stays at 1 (caller owns it)
+  // Caller will call deleteConstantShapeBuffer() which calls release()
+  return buffer;
+}
+
+// Updated getOrCreate method to ensure it always creates a shape buffer
+ConstantShapeBuffer* DirectShapeTrie::getOrCreate(const LongType* shapeInfo) {
+  waitForInitialization();
+
+  if (!shapeInfo) {
+    std::string msg = "Null shapeInfo passed to getOrCreate";
+    THROW_EXCEPTION(msg.c_str());
+  }
+
+  validateShapeInfo(shapeInfo);
+
+  size_t stripeIdx = getStripeIndex(shapeInfo);
+  int rank = shape::rank(shapeInfo);
+
+  // Validate stripe index
+  if (stripeIdx >= NUM_STRIPES) {
+    stripeIdx = NUM_STRIPES - 1;
+  }
+
+  int shapeSignature = calculateShapeSignature(shapeInfo);
+
+  // Check if mutex pointer is valid
+  if (_mutexes == nullptr || (*_mutexes)[stripeIdx] == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+
+  // First try a read-only lookup without obtaining a write lock
+  {
+    SHARED_LOCK_TYPE<MUTEX_TYPE> readLock(*(*_mutexes)[stripeIdx]);
+    ConstantShapeBuffer* existing = search(shapeInfo, stripeIdx);
+    if (existing != nullptr) {
+      if (shapeInfoEqual(existing->primary(), shapeInfo)) {
+        existing->addRef();  // Increment refcount before returning cached buffer
+        return existing;
+      }
+    }
+  }
+
+  // If not found or not matching, grab exclusive lock and try again
+  SHARED_LOCK_TYPE<MUTEX_TYPE> writeLock(*(*_mutexes)[stripeIdx]);
+
+  // Check again under the write lock
+  ConstantShapeBuffer* existing = search(shapeInfo, stripeIdx);
+  if (existing != nullptr) {
+    if (shapeInfoEqual(existing->primary(), shapeInfo)) {
+      existing->addRef();  // Increment refcount before returning cached buffer
+      return existing;
+    }
+  }
+
+  if (_roots == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+  // Not found, create a new entry
+  auto rootsRef = *_roots;
+
+
+  ShapeTrieNode* current = rootsRef[stripeIdx];
+  if (current == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+
+  if (rank < 0 || rank > SD_MAX_RANK) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+
+  // Safe pointer to track the current node through the insertion process
+  ShapeTrieNode* safeNodePtr = nullptr;
+
+  // Insert rank with signature
+  safeNodePtr = current->findOrCreateChild(rank, 0, true, shapeSignature);
+  if (safeNodePtr == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+  current = safeNodePtr;
+
+  // Insert datatype with signature
+  safeNodePtr = current->findOrCreateChild(ArrayOptions::dataType(shapeInfo), 1, true, shapeSignature);
+  if (safeNodePtr == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+  current = safeNodePtr;
+
+  // Insert order with signature
+  safeNodePtr = current->findOrCreateChild(shape::order(shapeInfo), 2, true, shapeSignature);
+  if (safeNodePtr == nullptr) {
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+  current = safeNodePtr;
+
+  // Insert shape values with signature
+  const LongType* shapeValues = shape::shapeOf(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    safeNodePtr = current->findOrCreateChild(shapeValues[i], 3 + i, true, shapeSignature);
+    if (safeNodePtr == nullptr) {
+      return createFallbackBuffer(shapeInfo, rank);
+    }
+    current = safeNodePtr;
+  }
+
+  // Insert stride values with signature
+  const LongType* strides = shape::stride(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    safeNodePtr = current->findOrCreateChild(strides[i], 3 + rank + i, false, shapeSignature);
+    if (safeNodePtr == nullptr) {
+      return createFallbackBuffer(shapeInfo, rank);
+    }
+    current = safeNodePtr;
+  }
+
+  // Check if another thread has already created the buffer
+  if (ConstantShapeBuffer* nodeBuffer = current->buffer()) {
+    if (shapeInfoEqual(nodeBuffer->primary(), shapeInfo)) {
+      nodeBuffer->addRef();  // Increment refcount before returning cached buffer
+      return nodeBuffer;
+    }
+  }
+
+  // Create the shape buffer
+  ConstantShapeBuffer* buffer = ShapeBufferCreatorHelper::getCurrentCreator().create(shapeInfo, rank);
+  if (buffer == nullptr || buffer->primary() == nullptr) {
+    // Use fallback if creator fails
+    if (buffer != nullptr) {
+      delete buffer;  // Clean up invalid buffer
+    }
+    return createFallbackBuffer(shapeInfo, rank);
+  }
+
+  // Set the buffer - setBuffer handles ownership properly
+  current->setBuffer(buffer);
+
+  // Return the buffer from the node (could be the one we just set or a pre-existing one)
+  ConstantShapeBuffer* resultBuffer = current->buffer();
+  if (resultBuffer == nullptr) {
+    // Rare case: setBuffer failed to store, return the buffer we created
+    // Caller owns it with refCount=1 (no addRef needed)
+    return buffer;
+  }
+
+  // Buffer is now cached, increment refcount for the caller
+  resultBuffer->addRef();
+  return resultBuffer;
+}
+
+bool DirectShapeTrie::exists(const LongType* shapeInfo) const {
+  waitForInitialization();
+
+  validateShapeInfo(shapeInfo);
+  size_t stripeIdx = getStripeIndex(shapeInfo);
+
+  // Validate stripe index
+  if (stripeIdx >= NUM_STRIPES) {
+    return false;
+  }
+
+  // Check if mutex pointer is valid
+  if (_mutexes == nullptr || (*_mutexes)[stripeIdx] == nullptr) {
+    return false;
+  }
+
+  int shapeSignature = calculateShapeSignature(shapeInfo);
+
+  SHARED_LOCK_TYPE<MUTEX_TYPE> lock(*(*_mutexes)[stripeIdx]);
+  ConstantShapeBuffer* found = search(shapeInfo, stripeIdx);
+  return found != nullptr && shapeInfoEqual(found->primary(), shapeInfo);
+}
+
+// Original insert method kept for compatibility, but getOrCreate should be used instead
+ConstantShapeBuffer* DirectShapeTrie::insert(const LongType* shapeInfo, size_t stripeIdx) {
+  auto rootsRef = *_roots;
+
+  ShapeTrieNode* current = rootsRef[stripeIdx];
+  const int rank = shape::rank(shapeInfo);
+  const int shapeSignature = calculateShapeSignature(shapeInfo);
+
+  // Insert rank
+  current = current->findOrCreateChild(rank, 0, true, shapeSignature);
+  if (!current) {
+    std::string msg = "Failed to create rank node";
+    THROW_EXCEPTION(msg.c_str());
+    return nullptr;
+  }
+
+  // Insert datatype
+  current = current->findOrCreateChild(ArrayOptions::dataType(shapeInfo), 1, true, shapeSignature);
+  if (!current) {
+    std::string msg = "Failed to create datatype node";
+    THROW_EXCEPTION(msg.c_str());
+    return nullptr;
+  }
+
+  // Insert order
+  current = current->findOrCreateChild(shape::order(shapeInfo), 2, true, shapeSignature);
+  if (!current) {
+    std::string msg = "Failed to create order node";
+    THROW_EXCEPTION(msg.c_str());
+    return nullptr;
+  }
+
+  // Insert shape values
+  const LongType* shape = shape::shapeOf(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    current = current->findOrCreateChild(shape[i], 3 + i, true, shapeSignature);
+    if (!current) {
+      std::string msg = "Failed to create shape value node at index " + std::to_string(i);
+      THROW_EXCEPTION(msg.c_str());
+      return nullptr;
+    }
+  }
+
+  // Insert stride values
+  const LongType* strides = shape::stride(shapeInfo);
+  for (int i = 0; i < rank; i++) {
+    current = current->findOrCreateChild(strides[i], 3 + rank + i, false, shapeSignature);
+    if (!current) {
+      std::string msg = "Failed to create stride value node at index " + std::to_string(i);
+      THROW_EXCEPTION(msg.c_str());
+      return nullptr;
+    }
+  }
+
+  if (!current->buffer()) {
+    try {
+      const int shapeInfoLength = shape::shapeInfoLength(rank);
+      LongType* shapeCopy = new LongType[shapeInfoLength];
+      std::memcpy(shapeCopy, shapeInfo, shapeInfoLength * sizeof(LongType));
+
+      auto deallocator = std::shared_ptr<PrimaryPointerDeallocator>(new PrimaryPointerDeallocator(),
+                                                                    [] (PrimaryPointerDeallocator* ptr) { delete ptr; });
+      auto hPtr = new PointerWrapper(shapeCopy, deallocator);
+      auto buffer = new ConstantShapeBuffer(hPtr);
+
+#if defined(SD_GCC_FUNCTRACE)
+      // Track shape cache allocation
+      sd::array::ShapeCacheLifecycleTracker::getInstance().recordAllocation(shapeCopy);
+#endif
+
+      current->setBuffer(buffer);
+      // Buffer is now cached (trie owns it with refCount=1)
+      // Increment refcount so caller also has a reference
+      buffer->addRef();
+      return buffer;
+    } catch (const std::exception& e) {
+      std::string msg = "Shape buffer creation failed: ";
+      msg += e.what();
+      THROW_EXCEPTION(msg.c_str());
+    } catch (...) {
+      std::string msg = "Shape buffer creation failed with unknown exception";
+      THROW_EXCEPTION(msg.c_str());
+    }
+  }
+
+  ConstantShapeBuffer* result = current->buffer();
+  if (result != nullptr) {
+    result->addRef();  // Increment refcount before returning cached buffer
+  }
+  return result;
+}
+
+void DirectShapeTrie::clearCache() {
+  waitForInitialization();
+
+  if (_roots == nullptr || _mutexes == nullptr) {
+    return;
+  }
+
+  // Clear each stripe
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    MUTEX_TYPE* mutex = (*_mutexes)[i];
+    if (mutex == nullptr) continue;
+
+    // Lock this stripe
+    std::lock_guard<MUTEX_TYPE> lock(*mutex);
+
+    // Delete the old root node (destructor recursively cleans up all children and buffers)
+    ShapeTrieNode* oldRoot = (*_roots)[i];
+    if (oldRoot != nullptr) {
+      delete oldRoot;
+    }
+
+    // Create a new empty root node
+    (*_roots)[i] = new ShapeTrieNode(0, 0, false);
+  }
+
+  // Reset current counters (but preserve peak values for diagnostics)
+  _current_entries.store(0);
+  _current_bytes.store(0);
+}
+
+void DirectShapeTrie::countEntriesAndBytes(const ShapeTrieNode* node, LongType& entries, LongType& bytes) const {
+  if (node == nullptr) return;
+
+  // If this node has a buffer, count it
+  ConstantShapeBuffer* buffer = node->buffer();
+  if (buffer != nullptr) {
+    entries++;
+    // Calculate buffer size: shapeInfo length is stored at index 0
+    const LongType* shapeInfo = buffer->primary();
+    if (shapeInfo != nullptr) {
+      LongType bufferLength = shape::shapeInfoLength(shapeInfo);
+      bytes += bufferLength * sizeof(LongType);
+    }
+  }
+
+  // Recursively count children
+  const std::vector<ShapeTrieNode*>& children = node->children();
+  for (const auto* child : children) {
+    countEntriesAndBytes(child, entries, bytes);
+  }
+}
+
+LongType DirectShapeTrie::getCachedEntries() const {
+  waitForInitialization();
+
+  LongType total_entries = 0;
+  LongType total_bytes = 0;
+
+  if (_roots == nullptr || _mutexes == nullptr) {
+    return 0;
+  }
+
+  // Count entries across all stripes
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    MUTEX_TYPE* mutex = (*_mutexes)[i];
+    if (mutex == nullptr) continue;
+
+    // Lock this stripe for reading
+    std::lock_guard<MUTEX_TYPE> lock(*mutex);
+
+    ShapeTrieNode* root = (*_roots)[i];
+    if (root != nullptr) {
+      countEntriesAndBytes(root, total_entries, total_bytes);
+    }
+  }
+
+  // Update current counters
+  _current_entries.store(total_entries);
+  _current_bytes.store(total_bytes);
+
+  // Update peak if current exceeds it
+  LongType current_peak = _peak_entries.load();
+  while (total_entries > current_peak) {
+    if (_peak_entries.compare_exchange_weak(current_peak, total_entries)) {
+      break;
+    }
+  }
+
+  current_peak = _peak_bytes.load();
+  while (total_bytes > current_peak) {
+    if (_peak_bytes.compare_exchange_weak(current_peak, total_bytes)) {
+      break;
+    }
+  }
+
+  return total_entries;
+}
+
+LongType DirectShapeTrie::getCachedBytes() const {
+  // getCachedEntries() updates both entries and bytes
+  getCachedEntries();
+  return _current_bytes.load();
+}
+
+LongType DirectShapeTrie::getPeakCachedEntries() const {
+  return _peak_entries.load();
+}
+
+LongType DirectShapeTrie::getPeakCachedBytes() const {
+  return _peak_bytes.load();
+}
+
+void DirectShapeTrie::buildStringRepresentation(const ShapeTrieNode* node, std::stringstream& ss,
+                                                const std::string& indent, int currentDepth,
+                                                int maxDepth, int& entriesShown, int maxEntries) const {
+  if (node == nullptr) return;
+  if (maxDepth != -1 && currentDepth > maxDepth) return;
+  if (maxEntries != -1 && entriesShown >= maxEntries) return;
+
+  // Check if this node has a buffer
+  ConstantShapeBuffer* buffer = node->buffer();
+  if (buffer != nullptr) {
+    const LongType* shapeInfo = buffer->primary();
+    if (shapeInfo != nullptr) {
+      entriesShown++;
+
+      // Display node info
+      ss << indent << "Node[level=" << node->level()
+         << ", value=" << node->value()
+         << ", isShape=" << (node->isShape() ? "true" : "false")
+         << "]\n";
+
+      // Display shape info details
+      int rank = shape::rank(shapeInfo);
+      ss << indent << "  Shape: rank=" << rank << ", order=" << shape::order(shapeInfo)
+         << ", dtype=" << DataTypeUtils::asString(ArrayOptions::dataType(shapeInfo)) << "\n";
+
+      // Display shape dimensions
+      ss << indent << "  Dims: [";
+      const LongType* dims = shape::shapeOf(shapeInfo);
+      for (int i = 0; i < rank; i++) {
+        if (i > 0) ss << ", ";
+        ss << dims[i];
+      }
+      ss << "]\n";
+
+      // Display strides
+      ss << indent << "  Strides: [";
+      const LongType* strides = shape::stride(shapeInfo);
+      for (int i = 0; i < rank; i++) {
+        if (i > 0) ss << ", ";
+        ss << strides[i];
+      }
+      ss << "]\n";
+
+      // Display total elements and buffer size
+      LongType length = shape::length(shapeInfo);
+      LongType bufferLength = shape::shapeInfoLength(shapeInfo);
+      ss << indent << "  Elements: " << length
+         << ", Buffer size: " << (bufferLength * sizeof(LongType)) << " bytes\n";
+
+      if (maxEntries != -1 && entriesShown >= maxEntries) {
+        ss << indent << "  ... (max entries reached)\n";
+        return;
+      }
+    }
+  }
+
+  // Recursively process children
+  const std::vector<ShapeTrieNode*>& children = node->children();
+  if (!children.empty() && (maxDepth == -1 || currentDepth < maxDepth)) {
+    for (const auto* child : children) {
+      if (maxEntries != -1 && entriesShown >= maxEntries) break;
+      buildStringRepresentation(child, ss, indent + "  ", currentDepth + 1,
+                               maxDepth, entriesShown, maxEntries);
+    }
+  }
+}
+
+std::string DirectShapeTrie::toString(int maxDepth, int maxEntries) const {
+  waitForInitialization();
+
+  std::stringstream ss;
+
+  if (_roots == nullptr || _mutexes == nullptr) {
+    ss << "DirectShapeTrie: [UNINITIALIZED]\n";
+    return ss.str();
+  }
+
+  // Get current statistics
+  LongType totalEntries = getCachedEntries();
+  LongType totalBytes = getCachedBytes();
+  LongType peakEntries = getPeakCachedEntries();
+  LongType peakBytes = getPeakCachedBytes();
+
+  // Header
+  ss << "DirectShapeTrie [" << NUM_STRIPES << " stripes]\n";
+  ss << "Current: " << totalEntries << " entries, " << totalBytes << " bytes\n";
+  ss << "Peak: " << peakEntries << " entries, " << peakBytes << " bytes\n";
+  ss << "Showing: max depth=" << (maxDepth == -1 ? "unlimited" : std::to_string(maxDepth))
+     << ", max entries=" << (maxEntries == -1 ? "unlimited" : std::to_string(maxEntries)) << "\n";
+  ss << "---\n";
+
+  int entriesShown = 0;
+
+  // Traverse each stripe
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    MUTEX_TYPE* mutex = (*_mutexes)[i];
+    if (mutex == nullptr) continue;
+
+    // Lock this stripe for reading
+    std::lock_guard<MUTEX_TYPE> lock(*mutex);
+
+    ShapeTrieNode* root = (*_roots)[i];
+    if (root != nullptr && !root->children().empty()) {
+      ss << "Stripe " << i << ":\n";
+      buildStringRepresentation(root, ss, "  ", 0, maxDepth, entriesShown, maxEntries);
+
+      if (maxEntries != -1 && entriesShown >= maxEntries) {
+        ss << "... (max entries limit reached, " << (totalEntries - entriesShown)
+           << " more entries not shown)\n";
+        break;
+      }
+    }
+  }
+
+  if (entriesShown == 0) {
+    ss << "(Cache is empty)\n";
+  }
+
+  return ss.str();
+}
+
+void DirectShapeTrie::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
+  waitForInitialization();
+
+  if (_roots == nullptr || _mutexes == nullptr) {
+    return;
+  }
+
+  // Traverse all stripes and collect ConstantShapeBuffer pointers
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    MUTEX_TYPE* mutex = (*_mutexes)[i];
+    if (mutex == nullptr) continue;
+
+    std::lock_guard<MUTEX_TYPE> lock(*mutex);
+
+    ShapeTrieNode* root = (*_roots)[i];
+    if (root != nullptr) {
+      collectCachedPointers(root, out_pointers);
+    }
+  }
+}
+
+void DirectShapeTrie::collectCachedPointers(const ShapeTrieNode* node, std::unordered_set<void*>& out_pointers) const {
+  if (node == nullptr) return;
+
+  // If this node has a ConstantShapeBuffer, add its primary pointer (LongType* shape_info) to the set
+  // This is what ShapeCacheLifecycleTracker uses to track allocations
+  ConstantShapeBuffer* buffer = node->buffer();
+  if (buffer != nullptr && buffer->primary() != nullptr) {
+    out_pointers.insert(buffer->primary());
+  }
+
+  // Recursively collect from all children
+  for (const auto* child : node->children()) {
+    collectCachedPointers(child, out_pointers);
+  }
+}
+
+}  // namespace sd

libnd4j/include/helpers/impl/DirectTadTrie.cpp

@@ -0,0 +1,563 @@
+/* ******************************************************************************
+ *
+ * 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
+ ******************************************************************************/
+
+#include "../DirectTadTrie.h"
+
+#include <array/TadPack.h>
+
+#include <algorithm>
+#include <memory>
+#include <atomic>
+#include <sstream>
+#include <string>
+#include <unordered_set>
+
+#include "array/TadCalculator.h"
+
+namespace sd {
+
+std::shared_ptr<TadPack> DirectTadTrie::enhancedSearch(const std::vector<LongType>& dimensions, LongType* originalShape, size_t stripeIdx) {
+  const TadTrieNode* current = _roots[stripeIdx].get();
+  int rank = shape::rank(originalShape);
+
+  // Navigate to dimension length node
+  current = findChild(current, dimensions.size(), 0, false, rank);
+  if (!current) return nullptr;
+
+  // Navigate through dimension nodes
+  for (size_t i = 0; i < dimensions.size(); i++) {
+    current = findChild(current, dimensions[i], i + 1, true, rank);
+    if (!current) return nullptr;
+  }
+
+  // Found a matching node, now verify TadPack compatibility
+  std::shared_ptr<TadPack> pack = current->pack();
+  if (!pack) return nullptr;
+
+  // Use cached signature for fast comparison - no TadCalculator needed!
+  const TadPackSignature* signature = current->packSignature();
+  if (!signature) {
+    // Signature not cached (shouldn't happen, but handle gracefully)
+    return nullptr;
+  }
+
+  // Fast comparison using cached signature instead of creating TadCalculator
+  if (!signature->matches(originalShape)) {
+    return nullptr;
+  }
+
+  return pack;
+}
+
+// Enhanced stride-aware hash computation
+size_t DirectTadTrie::computeStrideAwareHash(const std::vector<LongType>& dimensions, LongType* originalShape)  {
+  if (!originalShape) return 0;
+
+  size_t hash = 17; // Prime number starting point
+
+  // Handle empty dimensions specially
+  if (dimensions.empty()) {
+    // Empty dimensions case - hash based on shape only
+    hash = hash * 31 + 0; // Marker for empty dimensions
+  } else {
+    // Add dimension-specific hash contribution with position-dependence
+    for (size_t i = 0; i < dimensions.size(); i++) {
+      hash = hash * 31 + static_cast<size_t>(dimensions[i]) * (i + 1);
+    }
+  }
+
+  // Add rank - critical for distinguishing different dimension arrays
+  int rank = shape::rank(originalShape);
+  hash = hash * 13 + rank * 19;
+
+  // Add shape signature based on shape dimensions with position-dependence
+  LongType* shapeInfo = shape::shapeOf(originalShape);
+  for (int i = 0; i < rank; i++) {
+    hash = hash * 17 + static_cast<size_t>(shapeInfo[i]) * (11 + i);
+  }
+
+  // Add stride information to make the hash more specific
+  LongType* strides = shape::stride(originalShape);
+  for (int i = 0; i < rank; i++) {
+    hash = hash * 23 + static_cast<size_t>(strides[i]) * (7 + i);
+  }
+
+  // Add total element count to distinguish differently sized arrays
+  hash = hash * 41 + shape::length(originalShape);
+
+  // Add data type and order information
+  hash = hash * 29 + static_cast<size_t>(ArrayOptions::dataType(originalShape));
+  hash = hash * 37 + static_cast<size_t>(shape::order(originalShape));
+
+  // Compute the final stripe index
+  return hash % NUM_STRIPES;
+}
+
+bool DirectTadTrie::exists(const std::vector<LongType>& dimensions, LongType* originalShape)  {
+  if (!originalShape) return false;
+
+  const size_t stripeIdx = computeStripeIndex(dimensions, originalShape);
+  SHARED_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[stripeIdx]);
+
+  // Using the enhanced search method which verifies TadPack compatibility
+  return enhancedSearch(dimensions, originalShape, stripeIdx) != nullptr;
+}
+
+std::vector<LongType> DirectTadTrie::sortDimensions(const std::vector<LongType>& dimensions) const {
+  std::vector<LongType> sorted = dimensions;
+  std::sort(sorted.begin(), sorted.end());
+  return sorted;
+}
+
+std::shared_ptr<TadPack> DirectTadTrie::search(const std::vector<LongType>& dimensions, int originalShapeRank, size_t stripeIdx) const {
+  // No need for locking - caller handles locking (e.g., in getOrCreate)
+  const TadTrieNode* current = _roots[stripeIdx].get();
+
+  // First level: dimension length
+  current = findChild(current, dimensions.size(), 0, false, originalShapeRank);
+  if (!current) return nullptr;
+
+  // Second level: dimensions
+  for (size_t i = 0; i < dimensions.size(); i++) {
+    current = findChild(current, dimensions[i], i + 1, true, originalShapeRank);
+    if (!current) return nullptr;
+  }
+
+  return current->pack();
+}
+
+std::shared_ptr<TadPack> DirectTadTrie::getOrCreate(std::vector<LongType>& dimensions, LongType* originalShape) {
+  if (!originalShape) {
+    THROW_EXCEPTION("Original shape cannot be null in TAD calculation");
+  }
+
+  // Use the enhanced hash computation for better distribution
+  const size_t stripeIdx = computeStrideAwareHash(dimensions, originalShape);
+
+  // First try a read-only lookup
+  {
+    SHARED_LOCK_TYPE<MUTEX_TYPE> readLock(_mutexes[stripeIdx]);
+    std::shared_ptr<TadPack> existing = enhancedSearch(dimensions, originalShape, stripeIdx);
+    if (existing) {
+      return existing;
+    }
+  }
+
+  // If not found, use insert which will handle the write lock
+  return insert(dimensions, originalShape);
+}
+
+std::shared_ptr<TadPack> DirectTadTrie::insert(std::vector<LongType>& dimensions, LongType* originalShape) {
+  if (!originalShape) {
+    THROW_EXCEPTION("Original shape cannot be null in TAD calculation");
+  }
+
+  int rank = shape::rank(originalShape);
+  // Use the enhanced hash computation for better distribution
+  const size_t stripeIdx = computeStrideAwareHash(dimensions, originalShape);
+  // Use exclusive lock for write operation (inserting new TAD packs)
+  EXCLUSIVE_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[stripeIdx]);
+
+  // Check if a compatible TadPack already exists
+  std::shared_ptr<TadPack> existing = enhancedSearch(dimensions, originalShape, stripeIdx);
+  if (existing) {
+    return existing;
+  }
+
+  // No compatible TadPack found, create a new one
+  TadTrieNode* current = _roots[stripeIdx].get();
+
+  // First level: dimension length node with shape rank
+  current = current->findOrCreateChild(dimensions.size(), 0, false, rank);
+  if (!current) {
+    THROW_EXCEPTION("Failed to create dimension length node");
+  }
+
+  // Second level: dimension nodes with shape rank
+  for (size_t i = 0; i < dimensions.size(); i++) {
+    current = current->findOrCreateChild(dimensions[i], i + 1, true, rank);
+    if (!current) {
+      THROW_EXCEPTION("Failed to create dimension node");
+    }
+  }
+
+  // Create the TadPack only if it doesn't exist yet
+  if (!current->pack()) {
+    TadCalculator *calculator = nullptr;
+    std::shared_ptr<TadPack> newPack;
+
+    try {
+      calculator = new TadCalculator(originalShape);
+      calculator->createTadPack(dimensions);
+
+      // Create a new TadPack with full dimension information
+      // Use releaseOffsets() to transfer ownership of the offsets buffer to TadPack
+      // Wrap in shared_ptr for proper memory management
+      newPack = std::make_shared<TadPack>(
+          calculator->tadShape(),
+          calculator->releaseOffsets(),  // Transfer ownership
+          calculator->numberOfTads(),
+          dimensions.data(),
+          dimensions.size());
+
+      // Store the TadPack in the node
+      // setPack now also caches the signature for future fast comparisons
+      current->setPack(newPack);
+
+      // Clean up the calculator (safe now that offsets ownership was transferred)
+      delete calculator;
+      calculator = nullptr;
+
+    } catch (const std::exception& e) {
+      // Clean up on exception to prevent memory leaks
+      // shared_ptr will automatically clean up newPack
+      if (calculator != nullptr) {
+        delete calculator;
+        calculator = nullptr;
+      }
+      std::string msg = "TAD creation failed: ";
+      msg += e.what();
+      THROW_EXCEPTION(msg.c_str());
+    }
+  }
+
+  return current->pack();
+}
+
+
+const TadTrieNode* DirectTadTrie::findChild(const TadTrieNode* node, LongType value, int level, bool isDimension, int shapeRank) const {
+  if (!node) return nullptr;
+
+  for (const auto& child : node->children()) {
+    if (child->value() == value &&
+        child->level() == level &&
+        child->isDimension() == isDimension &&
+        child->shapeRank() == shapeRank) {
+      return child.get();
+    }
+  }
+
+  return nullptr;
+}
+
+// Helper function to recursively delete TadPacks from a node and its children
+// This ensures TadPack destructors are called, which triggers recordDeallocation()
+static void deleteTadPacksRecursive(TadTrieNode* node, int& deletedCount) {
+  if (!node) return;
+
+  // First, recursively delete from all children
+  const auto& children = node->children();
+  for (const auto& child : children) {
+    deleteTadPacksRecursive(child.get(), deletedCount);
+  }
+
+  // Then delete this node's TadPack if it exists
+  // shared_ptr will handle deletion automatically when we reset it
+  auto pack = node->pack();
+  if (pack) {
+    deletedCount++;
+    // Clear the shared_ptr to trigger TadPack destructor
+    // The destructor will call TADCacheLifecycleTracker::recordDeallocation()
+    // if SD_GCC_FUNCTRACE is defined during compilation
+    node->setPack(nullptr);
+  }
+}
+
+void DirectTadTrie::clear() {
+  // CRITICAL: Skip cleanup during shutdown to avoid SIGSEGV from corrupted memory
+  // During JVM/static destruction, memory allocators may have been destroyed,
+  // leaving corrupted pointers in the trie. Traversing the tree in this state
+  // causes crashes in deleteTadPacksRecursive.
+  if (_shutdownInProgress.load(std::memory_order_acquire)) {
+    return;  // Let the OS reclaim memory at exit - this is safe
+  }
+
+  // Clear all stripes
+  // NOTE: Removed #ifndef __JAVACPP_HACK__ guard to fix TAD cache memory leak
+  // The guard was preventing cache cleanup when JavaCPP is used (production mode)
+  // This caused indefinite accumulation of TADPack objects despite clearTADCache() calls
+
+  int totalDeleted = 0;
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    // Use exclusive lock for write operation (clearing the cache)
+    EXCLUSIVE_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[i]);
+
+    // This ensures TadPack destructors are called, which invokes recordDeallocation()
+    // for proper lifecycle tracking.
+    //
+    // IMPORTANT: We CANNOT rely on unique_ptr cascade deletion because:
+    // 1. TadTrieNode destructor deletes _tadPack only if SD_GCC_FUNCTRACE is defined
+    // 2. Functrace may be auto-disabled during build, causing guards to evaluate false
+    // 3. Even if guards pass, destructor might not run if roots are replaced before going out of scope
+    //
+    // By explicitly calling deleteTadPacksRecursive() BEFORE replacing roots,
+    // we guarantee that:
+    // - All TadPack objects are explicitly deleted via delete operator
+    // - Their destructors run and call recordDeallocation() (if tracking enabled)
+    // - Pointers are cleared to nullptr to prevent double-delete in node destructors
+    int deletedCount = 0;
+    deleteTadPacksRecursive(_roots[i].get(), deletedCount);
+    totalDeleted += deletedCount;
+
+    // Recreate the root node - this will delete the old tree structure
+    // (nodes are already cleaned of TadPacks above via deleteTadPacksRecursive)
+    // The old root's unique_ptr goes out of scope here, triggering node destructor cascade
+    // But TadPacks are already deleted and nulled out, so no double-delete occurs
+    _roots[i] = std::make_unique<TadTrieNode>(0, 0, false);
+    _stripeCounts[i].store(0);
+  }
+
+  // Reset current counters (but preserve peak values for diagnostics)
+  _current_entries.store(0);
+  _current_bytes.store(0);
+}
+
+void DirectTadTrie::countEntriesAndBytes(const TadTrieNode* node, LongType& entries, LongType& bytes) const {
+  if (node == nullptr) return;
+
+  // If this node has a TadPack, count it
+  auto pack = node->pack();
+  if (pack != nullptr) {
+    entries++;
+
+    // Calculate total bytes for this TadPack
+    // Shape info buffer
+    const LongType* shapeInfo = pack->primaryShapeInfo();
+    if (shapeInfo != nullptr) {
+      LongType shapeInfoLength = shape::shapeInfoLength(shapeInfo);
+      bytes += shapeInfoLength * sizeof(LongType);
+    }
+
+    // Offsets buffer
+    const LongType* offsets = pack->primaryOffsets();
+    if (offsets != nullptr) {
+      LongType numTads = pack->numberOfTads();
+      bytes += numTads * sizeof(LongType);
+    }
+  }
+
+  // Recursively count children
+  const std::vector<std::unique_ptr<TadTrieNode>>& children = node->children();
+  for (const auto& child : children) {
+    countEntriesAndBytes(child.get(), entries, bytes);
+  }
+}
+
+LongType DirectTadTrie::getCachedEntries() const {
+  LongType total_entries = 0;
+  LongType total_bytes = 0;
+
+  // Count entries across all stripes
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    // Lock this stripe for reading
+    SHARED_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[i]);
+
+    const TadTrieNode* root = _roots[i].get();
+    if (root != nullptr) {
+      countEntriesAndBytes(root, total_entries, total_bytes);
+    }
+  }
+
+  // Update current counters
+  _current_entries.store(total_entries);
+  _current_bytes.store(total_bytes);
+
+  // Update peak if current exceeds it
+  LongType current_peak = _peak_entries.load();
+  while (total_entries > current_peak) {
+    if (_peak_entries.compare_exchange_weak(current_peak, total_entries)) {
+      break;
+    }
+  }
+
+  current_peak = _peak_bytes.load();
+  while (total_bytes > current_peak) {
+    if (_peak_bytes.compare_exchange_weak(current_peak, total_bytes)) {
+      break;
+    }
+  }
+
+  return total_entries;
+}
+
+LongType DirectTadTrie::getCachedBytes() const {
+  // getCachedEntries() updates both entries and bytes
+  getCachedEntries();
+  return _current_bytes.load();
+}
+
+LongType DirectTadTrie::getPeakCachedEntries() const {
+  return _peak_entries.load();
+}
+
+LongType DirectTadTrie::getPeakCachedBytes() const {
+  return _peak_bytes.load();
+}
+
+void DirectTadTrie::buildStringRepresentation(const TadTrieNode* node, std::stringstream& ss,
+                                              const std::string& indent, int currentDepth,
+                                              int maxDepth, int& entriesShown, int maxEntries) const {
+  if (node == nullptr) return;
+  if (maxDepth != -1 && currentDepth > maxDepth) return;
+  if (maxEntries != -1 && entriesShown >= maxEntries) return;
+
+  // Check if this node has a TadPack
+  auto pack = node->pack();
+  if (pack != nullptr) {
+    entriesShown++;
+
+    // Display node info
+    ss << indent << "Node[level=" << node->level()
+       << ", value=" << node->value()
+       << ", isDim=" << (node->isDimension() ? "true" : "false")
+       << ", rank=" << node->shapeRank()
+       << "]\n";
+
+    // Display TAD pack details
+    const LongType* shapeInfo = pack->primaryShapeInfo();
+    if (shapeInfo != nullptr) {
+      int rank = shape::rank(shapeInfo);
+      ss << indent << "  TAD Shape: rank=" << rank
+         << ", order=" << shape::order(shapeInfo)
+         << ", dtype=" << DataTypeUtils::asString(ArrayOptions::dataType(shapeInfo)) << "\n";
+
+      // Display TAD dimensions
+      ss << indent << "  TAD Dims: [";
+      const LongType* dims = shape::shapeOf(shapeInfo);
+      for (int i = 0; i < rank; i++) {
+        if (i > 0) ss << ", ";
+        ss << dims[i];
+      }
+      ss << "]\n";
+
+      // Display TAD strides
+      ss << indent << "  TAD Strides: [";
+      const LongType* strides = shape::stride(shapeInfo);
+      for (int i = 0; i < rank; i++) {
+        if (i > 0) ss << ", ";
+        ss << strides[i];
+      }
+      ss << "]\n";
+    }
+
+    // Display number of TADs and offset info
+    LongType numTads = pack->numberOfTads();
+    ss << indent << "  Number of TADs: " << numTads << "\n";
+
+    // Display memory usage
+    LongType shapeInfoBytes = 0;
+    LongType offsetsBytes = 0;
+    if (shapeInfo != nullptr) {
+      LongType shapeInfoLength = shape::shapeInfoLength(shapeInfo);
+      shapeInfoBytes = shapeInfoLength * sizeof(LongType);
+    }
+    if (pack->primaryOffsets() != nullptr) {
+      offsetsBytes = numTads * sizeof(LongType);
+    }
+    ss << indent << "  Memory: shape_info=" << shapeInfoBytes
+       << " bytes, offsets=" << offsetsBytes
+       << " bytes, total=" << (shapeInfoBytes + offsetsBytes) << " bytes\n";
+
+    if (maxEntries != -1 && entriesShown >= maxEntries) {
+      ss << indent << "  ... (max entries reached)\n";
+      return;
+    }
+  }
+
+  // Recursively process children
+  const std::vector<std::unique_ptr<TadTrieNode>>& children = node->children();
+  if (!children.empty() && (maxDepth == -1 || currentDepth < maxDepth)) {
+    for (const auto& child : children) {
+      if (maxEntries != -1 && entriesShown >= maxEntries) break;
+      buildStringRepresentation(child.get(), ss, indent + "  ", currentDepth + 1,
+                               maxDepth, entriesShown, maxEntries);
+    }
+  }
+}
+
+std::string DirectTadTrie::toString(int maxDepth, int maxEntries) const {
+  std::stringstream ss;
+
+  // Get current statistics
+  LongType totalEntries = getCachedEntries();
+  LongType totalBytes = getCachedBytes();
+  LongType peakEntries = getPeakCachedEntries();
+  LongType peakBytes = getPeakCachedBytes();
+
+  // Header
+  ss << "DirectTadTrie [" << NUM_STRIPES << " stripes]\n";
+  ss << "Current: " << totalEntries << " entries, " << totalBytes << " bytes\n";
+  ss << "Peak: " << peakEntries << " entries, " << peakBytes << " bytes\n";
+  ss << "Showing: max depth=" << (maxDepth == -1 ? "unlimited" : std::to_string(maxDepth))
+     << ", max entries=" << (maxEntries == -1 ? "unlimited" : std::to_string(maxEntries)) << "\n";
+  ss << "---\n";
+
+  int entriesShown = 0;
+
+  // Traverse each stripe
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    // Lock this stripe for reading
+    SHARED_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[i]);
+
+    const TadTrieNode* root = _roots[i].get();
+    if (root != nullptr && !root->children().empty()) {
+      ss << "Stripe " << i << ":\n";
+      buildStringRepresentation(root, ss, "  ", 0, maxDepth, entriesShown, maxEntries);
+
+      if (maxEntries != -1 && entriesShown >= maxEntries) {
+        ss << "... (max entries limit reached, " << (totalEntries - entriesShown)
+           << " more entries not shown)\n";
+        break;
+      }
+    }
+  }
+
+  if (entriesShown == 0) {
+    ss << "(Cache is empty)\n";
+  }
+
+  return ss.str();
+}
+
+void DirectTadTrie::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
+  // Traverse all stripes and collect TadPack pointers
+  for (size_t i = 0; i < NUM_STRIPES; i++) {
+    SHARED_LOCK_TYPE<MUTEX_TYPE> lock(_mutexes[i]);
+
+    const TadTrieNode* root = _roots[i].get();
+    if (root != nullptr) {
+      collectCachedPointers(root, out_pointers);
+    }
+  }
+}
+
+void DirectTadTrie::collectCachedPointers(const TadTrieNode* node, std::unordered_set<void*>& out_pointers) const {
+  if (node == nullptr) return;
+
+  // If this node has a TadPack, add it to the set
+  auto pack = node->pack();
+  if (pack != nullptr) {
+    out_pointers.insert(pack.get());
+  }
+
+  // Recursively collect from all children
+  for (const auto& child : node->children()) {
+    collectCachedPointers(child.get(), out_pointers);
+  }
+}
+
+}  // namespace sd

libnd4j/include/helpers/impl/EigenValsAndVecs.cpp

@@ -0,0 +1,368 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/EigenValsAndVecs.h>
+#include <helpers/HessenbergAndSchur.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+EigenValsAndVecs<T>::EigenValsAndVecs(NDArray& matrix)
+    : _Vals(matrix.dataType(), matrix.getContext(), true),
+      _Vecs(matrix.dataType(), matrix.getContext(), true) {
+  if (matrix.rankOf() != 2)
+    THROW_EXCEPTION("ops::helpers::EigenValsAndVecs constructor: input matrix must be 2D !");
+
+  if (matrix.sizeAt(0) != matrix.sizeAt(1))
+    THROW_EXCEPTION("ops::helpers::EigenValsAndVecs constructor: input array must be 2D square matrix !");
+
+  Schur<T> schur(matrix);
+
+  NDArray* schurMatrixU = schur.u;
+  NDArray* schurMatrixT = schur.t;
+
+  std::vector<LongType> shape = {schurMatrixU->sizeAt(1), schurMatrixU->sizeAt(1), 2};
+  _Vecs = NDArray(matrix.ordering(), shape, matrix.dataType(),
+                  matrix.getContext());
+  std::vector<LongType> shape2 = {matrix.sizeAt(1), 2};
+  _Vals = NDArray(matrix.ordering(), shape2, matrix.dataType(), matrix.getContext());
+
+  // sequence of methods calls matters
+  calcEigenVals(*schurMatrixT);
+  calcPseudoEigenVecs(*schurMatrixT, *schurMatrixU);  // pseudo-eigenvectors are real and will be stored in schurMatrixU
+  calcEigenVecs(*schurMatrixU);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void calcEigenVals_(NDArray& schurMatrixT, NDArray& _Vals) {
+  const int numOfCols = schurMatrixT.sizeAt(1);
+
+  // calculate eigenvalues _Vals
+  int i = 0;
+  while (i < numOfCols) {
+    if (i == numOfCols - 1 || schurMatrixT.t<T>(i + 1, i) == T(0.f)) {
+      _Vals.r<T>(i, 0) = schurMatrixT.t<T>(i, i);  // real part
+      _Vals.r<T>(i, 1) = T(0);                     // imaginary part
+
+      if (!math::sd_isfin<T>(_Vals.t<T>(static_cast<T>(i), static_cast<T>(0)))) {
+        THROW_EXCEPTION("ops::helpers::igenValsAndVec::calcEigenVals: got infinite eigen value !");
+        return;
+      }
+
+      ++i;
+    } else {
+      T p = T(0.5) * (schurMatrixT.t<T>(i, i) - schurMatrixT.t<T>(i + 1, i + 1));
+      T z;
+      {
+        T t0 = schurMatrixT.t<T>(i + 1, i);
+        T t1 = schurMatrixT.t<T>(i, i + 1);
+        T maxval = math::sd_max<T>(math::sd_abs<T,T>(p), math::sd_max<T>(math::sd_abs<T,T>(t0), math::sd_abs<T,T>(t1)));
+        t0 /= maxval;
+        t1 /= maxval;
+        T p0 = p / maxval;
+        z = maxval * math::sd_sqrt<T, T>(math::sd_abs<T,T>(p0 * p0 + t0 * t1));
+      }
+
+      _Vals.r<T>(i, 0) = _Vals.r<T>(i + 1, 0) = schurMatrixT.t<T>(i + 1, i + 1) + p;
+      _Vals.r<T>(i, 1) = z;
+      _Vals.r<T>(i + 1, 1) = -z;
+
+      if (!(math::sd_isfin<T>(_Vals.t<T>(i, 0)) && math::sd_isfin<T>(_Vals.t<T>(i + 1, 0)) &&
+            math::sd_isfin<T>(_Vals.t<T>(i, 1))) &&
+          math::sd_isfin<T>(_Vals.t<T>(i + 1, 1))) {
+        THROW_EXCEPTION("ops::helpers::igenValsAndVec::calcEigenVals: got infinite eigen value !");
+        return;
+      }
+
+      i += 2;
+    }
+  }
+}
+
+template <typename T>
+void EigenValsAndVecs<T>::calcEigenVals(NDArray& schurMatrixT) {
+  calcEigenVals_<T>(schurMatrixT, _Vals);
+}
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void calcPseudoEigenVecs_(NDArray& schurMatrixT, NDArray& schurMatrixU, NDArray& _Vals) {
+  const int numOfCols = schurMatrixU.sizeAt(1);
+
+  T norm = static_cast<T>(0);
+  for (int j = 0; j < numOfCols; ++j) {
+    NDArray *viewPtr = schurMatrixT({j, j + 1, math::sd_max<LongType>(j - 1, 0), numOfCols});
+    auto* reduceResult = viewPtr->reduceNumber(reduce::ASum);
+    norm += reduceResult->template t<T>(0);
+    delete reduceResult;
+    delete viewPtr;
+  }
+
+  if (norm == T(0)) return;
+
+  for (int n = numOfCols - 1; n >= 0; n--) {
+    T p = _Vals.t<T>(n, 0);  // real part
+    T q = _Vals.t<T>(n, 1);  // imaginary part
+
+    if (q == (T)0) {  // not complex
+
+      T lastr((T)0), lastw((T)0);
+      int l = n;
+
+      schurMatrixT.r<T>(n, n) = T(1);
+
+      for (int i = n - 1; i >= 0; i--) {
+        T w = schurMatrixT.t<T>(i, i) - p;
+        
+        NDArray *view1Ptr = schurMatrixT({i, i + 1, l, n + 1}, true);
+        NDArray *view2Ptr = schurMatrixT({l, n + 1, n, n + 1}, true);
+        NDArray *dotResult = mmul(*view1Ptr, *view2Ptr);
+        T r = dotResult->template t<T>(0);
+        delete view1Ptr;
+        delete view2Ptr;
+        delete dotResult;
+
+        if (_Vals.t<T>(i, 1) < T(0)) {
+          lastw = w;
+          lastr = r;
+        } else {
+          l = i;
+          if (_Vals.t<T>(i, 1) == T(0)) {
+            if (w != T(0))
+              schurMatrixT.r<T>(i, n) = -r / w;
+            else
+              schurMatrixT.r<T>(i, n) = -r / (DataTypeUtils::eps<T>() * norm);
+          } else {
+            T x = schurMatrixT.t<T>(i, i + 1);
+            T y = schurMatrixT.t<T>(i + 1, i);
+            T denom = (_Vals.t<T>(i, 0) - p) * (_Vals.t<T>(i, 0) - p) + _Vals.t<T>(i, 1) * _Vals.t<T>(i, 1);
+            T t = (x * lastr - lastw * r) / denom;
+            schurMatrixT.r<T>(i, n) = t;
+
+            if (math::sd_abs<T,T>(x) > math::sd_abs<T,T>(lastw))
+              schurMatrixT.r<T>(i + 1, n) = (-r - w * t) / x;
+            else
+              schurMatrixT.r<T>(i + 1, n) = (-lastr - y * t) / lastw;
+          }
+
+          T t = math::sd_abs<T,T>(schurMatrixT.t<T>(i, n));
+          if ((DataTypeUtils::eps<T>() * t) * t > T(1)) {
+            NDArray *divViewPtr = schurMatrixT({schurMatrixT.sizeAt(0) - numOfCols + i, -1, n, n + 1});
+            *divViewPtr /= t;
+            delete divViewPtr;
+          }
+        }
+      }
+    } else if (q < T(0) && n > 0) {  // complex
+
+      T lastra(0), lastsa(0), lastw(0);
+      int l = n - 1;
+
+      if (math::sd_abs<T,T>(schurMatrixT.t<T>(n, n - 1)) > math::sd_abs<T,T>(schurMatrixT.t<T>(n - 1, n))) {
+        schurMatrixT.r<T>(n - 1, n - 1) = q / schurMatrixT.t<T>(n, n - 1);
+        schurMatrixT.r<T>(n - 1, n) = -(schurMatrixT.t<T>(n, n) - p) / schurMatrixT.t<T>(n, n - 1);
+      } else {
+        EigenValsAndVecs<T>::divideComplexNums(T(0), -schurMatrixT.t<T>(n - 1, n), schurMatrixT.t<T>(n - 1, n - 1) - p,
+                                               q, schurMatrixT.r<T>(n - 1, n - 1), schurMatrixT.r<T>(n - 1, n));
+      }
+
+      schurMatrixT.r<T>(n, n - 1) = T(0);
+      schurMatrixT.r<T>(n, n) = T(1);
+
+      for (int i = n - 2; i >= 0; i--) {
+        NDArray *raView1Ptr = schurMatrixT({i, i + 1, l, n + 1}, true);
+        NDArray *raView2Ptr = schurMatrixT({l, n + 1, n - 1, n}, true);
+        NDArray *raDotResult = mmul(*raView1Ptr, *raView2Ptr);
+        T ra = raDotResult->template t<T>(0);
+        delete raView1Ptr;
+        delete raView2Ptr;
+        delete raDotResult;
+        
+        NDArray *saView1Ptr = schurMatrixT({i, i + 1, l, n + 1}, true);
+        NDArray *saView2Ptr = schurMatrixT({l, n + 1, n, n + 1}, true);
+        NDArray *saDotResult = mmul(*saView1Ptr, *saView2Ptr);
+        T sa = saDotResult->template t<T>(0);
+        delete saView1Ptr;
+        delete saView2Ptr;
+        delete saDotResult;
+
+        T w = schurMatrixT.t<T>(i, i) - p;
+
+        if (_Vals.t<T>(i, 1) < T(0)) {
+          lastw = w;
+          lastra = ra;
+          lastsa = sa;
+        } else {
+          l = i;
+
+          if (_Vals.t<T>(i, 1) == T(0)) {
+            EigenValsAndVecs<T>::divideComplexNums(-ra, -sa, w, q, schurMatrixT.r<T>(i, n - 1),
+                                                   schurMatrixT.r<T>(i, n));
+          } else {
+            T x = schurMatrixT.t<T>(i, i + 1);
+            T y = schurMatrixT.t<T>(i + 1, i);
+            T vr = (_Vals.t<T>(i, 0) - p) * (_Vals.t<T>(i, 0) - p) + _Vals.t<T>(i, 1) * _Vals.t<T>(i, 1) - q * q;
+            T vi = (_Vals.t<T>(i, 0) - p) * T(2) * q;
+
+            if ((vr == T(0)) && (vi == T(0)))
+              vr = DataTypeUtils::eps<T>() * norm *
+                   (math::sd_abs<T,T>(w) + math::sd_abs<T,T>(q) + math::sd_abs<T,T>(x) + math::sd_abs<T,T>(y) +
+                    math::sd_abs<T,T>(lastw));
+
+            EigenValsAndVecs<T>::divideComplexNums(x * lastra - lastw * ra + q * sa, x * lastsa - lastw * sa - q * ra,
+                                                   vr, vi, schurMatrixT.r<T>(i, n - 1), schurMatrixT.r<T>(i, n));
+
+            if (math::sd_abs<T,T>(x) > (math::sd_abs<T,T>(lastw) + math::sd_abs<T,T>(q))) {
+              schurMatrixT.r<T>(i + 1, n - 1) =
+                  (-ra - w * schurMatrixT.t<T>(i, n - 1) + q * schurMatrixT.t<T>(i, n)) / x;
+              schurMatrixT.r<T>(i + 1, n) = (-sa - w * schurMatrixT.t<T>(i, n) - q * schurMatrixT.t<T>(i, n - 1)) / x;
+            } else
+              EigenValsAndVecs<T>::divideComplexNums(-lastra - y * schurMatrixT.t<T>(i, n - 1),
+                                                     -lastsa - y * schurMatrixT.t<T>(i, n), lastw, q,
+                                                     schurMatrixT.r<T>(i + 1, n - 1), schurMatrixT.r<T>(i + 1, n));
+          }
+
+          T t = math::sd_max<T>(math::sd_abs<T,T>(schurMatrixT.t<T>(i, n - 1)), math::sd_abs<T,T>(schurMatrixT.t<T>(i, n)));
+          if ((DataTypeUtils::eps<T>() * t) * t > T(1)) {
+            NDArray *divViewPtr = schurMatrixT({i, numOfCols, n - 1, n + 1});
+            *divViewPtr /= t;
+            delete divViewPtr;
+          }
+        }
+      }
+      n--;
+    } else
+      THROW_EXCEPTION("ops::helpers::EigenValsAndVecs::calcEigenVecs: internal bug !");
+  }
+
+  for (int j = numOfCols - 1; j >= 0; j--) {
+    NDArray *uViewPtr = schurMatrixU({0, 0, 0, j + 1}, true);
+    NDArray *tViewPtr = schurMatrixT({0, j + 1, j, j + 1}, true);
+    NDArray *assignResult = mmul(*uViewPtr, *tViewPtr);
+    delete uViewPtr;
+    delete tViewPtr;
+    
+    NDArray *uAssignPtr = schurMatrixU({0, 0, j, j + 1}, true);
+    uAssignPtr->assign(assignResult);
+    delete uAssignPtr;
+    delete assignResult;
+  }
+}
+
+template <typename T>
+void EigenValsAndVecs<T>::calcPseudoEigenVecs(NDArray& schurMatrixT, NDArray& schurMatrixU) {
+  calcPseudoEigenVecs_<T>(schurMatrixT, schurMatrixU, _Vals);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void calcEigenVecs_(NDArray& schurMatrixU, NDArray& _Vals, NDArray& _Vecs) {
+  const T precision = T(2) * DataTypeUtils::eps<T>();
+
+  const int numOfCols = schurMatrixU.sizeAt(1);
+
+  for (int j = 0; j < numOfCols; ++j) {
+    if (math::sd_abs<T,T>(_Vals.t<T>(j, 1)) <= math::sd_abs<T,T>(_Vals.t<T>(j, 0)) * precision ||
+        j + 1 == numOfCols) {  // real
+
+      _Vecs.syncToDevice();
+      NDArray *assignPtr = schurMatrixU({0, 0, j, j + 1});
+      NDArray *vecsViewPtr = _Vecs({0, 0, j, j + 1, 0, 1});
+      vecsViewPtr->assign(assignPtr);
+      delete assignPtr;
+      delete vecsViewPtr;
+      
+      NDArray *vecsView2Ptr = _Vecs({0, 0, j, j + 1, 1, 2});
+      *vecsView2Ptr = (T)0;
+      delete vecsView2Ptr;
+
+      // normalize
+      NDArray *norm2ViewPtr = _Vecs({0, 0, j, j + 1, 0, 1});
+      auto* norm2Result = norm2ViewPtr->reduceNumber(reduce::SquaredNorm);
+      const T norm2 = norm2Result->template t<T>(0);
+      delete norm2Result;
+      if (norm2 > (T)0) *norm2ViewPtr /= math::sd_sqrt<T, T>(norm2);
+      delete norm2ViewPtr;
+    } else {  // complex
+
+      for (int i = 0; i < numOfCols; ++i) {
+        _Vecs.r<T>(i, j, 0) = _Vecs.r<T>(i, j + 1, 0) = schurMatrixU.t<T>(i, j);
+        _Vecs.r<T>(i, j, 1) = schurMatrixU.t<T>(i, j + 1);
+        _Vecs.r<T>(i, j + 1, 1) = -schurMatrixU.t<T>(i, j + 1);
+      }
+
+      // normalize
+      NDArray *norm2View1Ptr = _Vecs({0, 0, j, j + 1, 0, 0});
+      auto* norm2Result1 = norm2View1Ptr->reduceNumber(reduce::SquaredNorm);
+      T norm2 = norm2Result1->template t<T>(0);
+      delete norm2Result1;
+      if (norm2 > (T)0) *norm2View1Ptr /= math::sd_sqrt<T, T>(norm2);
+      delete norm2View1Ptr;
+
+      // normalize
+      NDArray *norm2View2Ptr = _Vecs({0, 0, j + 1, j + 2, 0, 0});
+      auto* norm2Result2 = norm2View2Ptr->reduceNumber(reduce::SquaredNorm);
+      norm2 = norm2Result2->template t<T>(0);
+      delete norm2Result2;
+      if (norm2 > (T)0) *norm2View2Ptr /= math::sd_sqrt<T, T>(norm2);
+      delete norm2View2Ptr;
+
+      ++j;
+    }
+  }
+}
+
+template <typename T>
+void EigenValsAndVecs<T>::calcEigenVecs(NDArray& schurMatrixU) {
+  calcEigenVecs_<T>(schurMatrixU, _Vals, _Vecs);
+}
+
+template <typename T>
+void eig_(NDArray& input, NDArray& vals, NDArray& vecs) {
+  assert(input.rankOf() == 2 && "input is not a matrix");
+  assert(input.sizeAt(0) == input.sizeAt(1) && "input is not a square matrix");
+  assert(vals.rankOf() == 2 && vals.sizeAt(0) == input.sizeAt(0) && vals.sizeAt(1) == 2 &&
+         "incorrect shape for the eigenvalue results vals");
+  assert(vecs.rankOf() == 3 && vecs.sizeAt(0) == input.sizeAt(0) && vecs.sizeAt(1) == input.sizeAt(0) &&
+         vecs.sizeAt(2) == 2 && "incorrect shape for the eigenvector results vecs");
+
+  Schur<T> schur(input);
+  NDArray* schurMatrixU = schur.u;
+  NDArray* schurMatrixT = schur.t;
+  calcEigenVals_<T>(*schurMatrixT, vals);
+  calcPseudoEigenVecs_<T>(*schurMatrixT, *schurMatrixU, vals);
+  calcEigenVecs_<T>(*schurMatrixU, vals, vecs);
+}
+
+void eig(NDArray& input, NDArray& vals, NDArray& vecs) {
+  BUILD_SINGLE_SELECTOR(input.dataType(), eig_, (input, vals, vecs), SD_FLOAT_TYPES);
+}
+
+BUILD_SINGLE_TEMPLATE( void eig_, (NDArray& input, NDArray& vals, NDArray& vecs), SD_FLOAT_TYPES);
+
+BUILD_SINGLE_TEMPLATE( class EigenValsAndVecs, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/EnumUtils.cpp

@@ -0,0 +1,114 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <graph/VariableType.h>
+#include <helpers/EnumUtils.h>
+
+using namespace sd::graph;
+
+namespace sd {
+const char* EnumUtils::_VariableTypeToString(VariableType variableType) {
+  switch (variableType) {
+    case NDARRAY:
+      return "NDARRAY";
+    case ARRAY_LIST:
+      return "ARRAY_LIST";
+    case FLOW:
+      return "FLOW";
+    default:
+      return "UNKNOWN VariableType";
+  }
+}
+
+const char* EnumUtils::_OpTypeToString(::graph::OpType opType) {
+  switch (opType) {
+    case ::graph::OpType_REDUCE_SAME:
+      return "REDUCE_SAME";
+    case ::graph::OpType_REDUCE_BOOL:
+      return "REDUCE_BOOL";
+    case ::graph::OpType_REDUCE_LONG:
+      return "REDUCE_LONG";
+    case ::graph::OpType_REDUCE_FLOAT:
+      return "REDUCE_FLOAT";
+    case ::graph::OpType_BOOLEAN:
+      return "BOOLEAN";
+    case ::graph::OpType_BROADCAST:
+      return "BROADCAST";
+    case ::graph::OpType_BROADCAST_BOOL:
+      return "BROADCAST_BOOL";
+    case ::graph::OpType_PAIRWISE:
+      return "PAIRWISE";
+    case ::graph::OpType_PAIRWISE_BOOL:
+      return "PAIRWISE_BOOL";
+    case ::graph::OpType_CUSTOM:
+      return "CUSTOM";
+    case ::graph::OpType_LOGIC:
+      return "LOGIC";
+    case ::graph::OpType_TRANSFORM_SAME:
+      return "TRANSFORM_SAME";
+    case ::graph::OpType_TRANSFORM_FLOAT:
+      return "TRANSFORM_FLOAT";
+    case ::graph::OpType_TRANSFORM_BOOL:
+      return "TRANSFORM_BOOL";
+    case ::graph::OpType_TRANSFORM_STRICT:
+      return "TRANSFORM_STRICT";
+    case ::graph::OpType_TRANSFORM_ANY:
+      return "TRANSFORM_ANY";
+    case ::graph::OpType_INDEX_REDUCE:
+      return "INDEX_ACCUMULATION";
+    case ::graph::OpType_SCALAR:
+      return "SCALAR";
+    case ::graph::OpType_SCALAR_BOOL:
+      return "SCALAR_BOOL";
+    case ::graph::OpType_SHAPE:
+      return "SHAPE";
+    default:
+      return "UNKNOWN OpType";
+  }
+}
+
+const char* EnumUtils::_LogicOpToString(int opNum) {
+  switch (opNum) {
+    case 0:
+      return "WHILE";
+    case 10:
+      return "SCOPE";
+    case 20:
+      return "CONDITIONAL";
+    case 30:
+      return "SWITCH";
+    case 40:
+      return "RETURN";
+    case 60:
+      return "MERGE";
+    case 70:
+      return "LOOP_COND";
+    case 80:
+      return "NEXT_ITERATION";
+    case 90:
+      return "EXIT";
+    case 100:
+      return "ENTER";
+    default:
+      return "UNKNOWN OPERATION";
+  }
+}
+}  // namespace sd

libnd4j/include/helpers/impl/FullPivLU.cpp

@@ -0,0 +1,226 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/FullPivLU.h>
+#include <ops/declarable/helpers/triangular_solve.h>
+
+#include <numeric>
+
+#if NOT_EXCLUDED(OP_triangular_solve)
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+// A{M,K} * x{K,N} = b{M,N}
+template <typename T>
+void FullPivLU<T>::solve(NDArray &A, NDArray &b, NDArray& x) {
+  if (A.rankOf() != 2) THROW_EXCEPTION("FullPivLU::solve: input matrix A must be 2D !");
+
+  if (A.sizeAt(0) != b.sizeAt(0))
+    THROW_EXCEPTION("FullPivLU::solve: A and b must have the same number of rows !");
+
+  if (A.sizeAt(1) != x.sizeAt(0))
+    THROW_EXCEPTION("FullPivLU::solve: number of A columns must be equal to number of x rows !");
+
+  NDArray *LU = A.dup(A.ordering());
+  NDArray luRef = *LU;
+  const int rows = LU->sizeAt(0);
+  const int cols = LU->sizeAt(1);
+  const int diagLen = math::sd_min<int>(rows, cols);
+
+  std::vector<int> rowsInds(rows), colsInds(cols);
+
+  int nonZeroPivots1 = diagLen;
+
+  T maxPivot = T(0);
+
+  for (int k = 0; k < diagLen; ++k) {
+    NDArray *bottomRightCornerPtr = luRef({k, rows, k, cols}, true);
+    NDArray bottomRightCorner = *bottomRightCornerPtr;
+    delete bottomRightCornerPtr;
+    
+    NDArray *indexNum = bottomRightCorner.indexReduceNumber(indexreduce::IndexAbsoluteMax);
+    const int indPivot = static_cast<int>(indexNum->t<LongType>(0));
+
+    int colPivot = indPivot % (cols - k);
+    int rowPivot = indPivot / (cols - k);
+
+    T currentMax = math::sd_abs<T,T>(bottomRightCorner.t<T>(rowPivot, colPivot));
+
+    // take into account that this was calculated in corner, not in whole LU
+    rowPivot += k;
+    colPivot += k;
+
+    if (currentMax == T(0)) {
+      nonZeroPivots1 = k;
+
+      for (int i = k; i < diagLen; ++i) rowsInds[i] = colsInds[i] = i;
+
+      delete indexNum;
+      break;
+    }
+
+    if (currentMax > maxPivot) maxPivot = currentMax;
+
+    rowsInds[k] = rowPivot;
+    colsInds[k] = colPivot;
+
+    if (k != rowPivot) {
+      NDArray *row1Ptr = luRef({k, k + 1, 0, 0}, true);
+      NDArray *row2Ptr = luRef({rowPivot, rowPivot + 1, 0, 0}, true);
+      row1Ptr->swapUnsafe(*row2Ptr);
+      delete row1Ptr;
+      delete row2Ptr;
+    }
+    if (k != colPivot) {
+      NDArray *col1Ptr = luRef({0, 0, k, k + 1}, true);
+      NDArray *col2Ptr = luRef({0, 0, colPivot, colPivot + 1}, true);
+      col1Ptr->swapUnsafe(*col2Ptr);
+      delete col1Ptr;
+      delete col2Ptr;
+    }
+
+    if (k < rows - 1) {
+      NDArray *divViewPtr = luRef({k + 1, rows, k, k + 1}, true);
+      *divViewPtr /= luRef.t<T>(k, k);
+      delete divViewPtr;
+    }
+
+    if (k < diagLen - 1) {
+      NDArray *leftPtr = luRef({k + 1, rows, k, k + 1}, true);
+      NDArray *rightPtr = luRef({k, k + 1, k + 1, cols}, true);
+      NDArray *targetPtr = luRef({k + 1, rows, k + 1, cols}, true);
+      NDArray left = *leftPtr;
+      NDArray right = *rightPtr;
+      NDArray *mulResult = mmul(left, right);
+      *targetPtr -= *mulResult;
+      delete mulResult;
+      delete leftPtr;
+      delete rightPtr;
+      delete targetPtr;
+    }
+
+    delete indexNum;
+  }
+  //***************************************************//
+
+  const T threshold = maxPivot * DataTypeUtils::eps<T>() * (T)diagLen;
+
+  int nonZeroPivots2 = 0;
+  for (int i = 0; i < nonZeroPivots1; ++i)
+    nonZeroPivots2 += static_cast<int>(math::sd_abs<T,T>(luRef.t<T>(i, i)) > threshold);
+
+  if (nonZeroPivots2 == 0) {
+    x.nullify();
+    delete LU;
+    return;
+  }
+
+  //***************************************************//
+
+  std::vector<int> rowsPermut1(rows), rowsPermut2(rows), colsPermut(cols);
+  std::iota(rowsPermut1.begin(), rowsPermut1.end(), 0);
+  std::iota(colsPermut.begin(), colsPermut.end(), 0);
+
+  for (int k = diagLen - 1; k >= 0; --k) math::sd_swap<int>(rowsPermut1[k], rowsPermut1[rowsInds[k]]);
+
+  for (int k = 0; k < diagLen; ++k) math::sd_swap<int>(colsPermut[k], colsPermut[colsInds[k]]);
+
+  for (int i = 0; i < rows; ++i)
+    for (int j = 0; j < rows; ++j)
+      if (i == rowsPermut1[j]) {
+        rowsPermut2[i] = j;
+        break;
+      }
+
+  //***************************************************//
+
+  NDArray *bUlike = b.ulike();
+  NDArray c = *bUlike;
+
+  for (int i = 0; i < rows; ++i) {
+    NDArray *cAssignPtr = b({rowsPermut2[i], rowsPermut2[i] + 1, 0, 0}, true);
+    NDArray cAssign = *cAssignPtr;
+    delete cAssignPtr;
+    
+    NDArray *cTargetPtr = c({i, i + 1, 0, 0}, true);
+    cTargetPtr->assign(&cAssign);
+    delete cTargetPtr;
+  }
+  
+  NDArray *cTopRows1Ptr = c({0, diagLen, 0, 0}, true);
+  NDArray cTopRows1 = *cTopRows1Ptr;
+  delete cTopRows1Ptr;
+  
+  NDArray *luDiagPtr = luRef({0, diagLen, 0, diagLen}, true);
+  // TriangularSolver<T>::solve(LU({0,diagLen, 0,diagLen}, true), cTopRows1, true, true, cTopRows1);
+  helpers::triangularSolve2D<T>(nullptr, *luDiagPtr, cTopRows1, true, true, cTopRows1);
+  delete luDiagPtr;
+
+  if (rows > cols) {
+    NDArray *leftPtr = luRef({cols, -1, 0, 0}, true);
+    NDArray *rightPtr = c({0, cols, 0, 0}, true);
+    NDArray *targetPtr = c({cols, -1, 0, 0}, true);
+    NDArray left = *leftPtr;
+    NDArray right = *rightPtr;
+    NDArray *mulResult = mmul(left, right);
+    *targetPtr -= *mulResult;
+    delete mulResult;
+    delete leftPtr;
+    delete rightPtr;
+    delete targetPtr;
+  }
+  
+  NDArray *cTopRows2Ptr = c({0, nonZeroPivots2, 0, 0}, true);
+  NDArray cTopRows2 = *cTopRows2Ptr;
+  delete cTopRows2Ptr;
+  
+  NDArray *luNonZeroPtr = luRef({0, nonZeroPivots2, 0, nonZeroPivots2}, true);
+  helpers::triangularSolve2D<T>(nullptr, *luNonZeroPtr, cTopRows2, false, false, cTopRows2);
+  delete luNonZeroPtr;
+
+  for (int i = 0; i < nonZeroPivots2; ++i) {
+    NDArray *cAssignPtr = c({i, i + 1, 0, 0}, true);
+    NDArray cAssign = *cAssignPtr;
+    delete cAssignPtr;
+    
+    NDArray *xTargetPtr = x({colsPermut[i], colsPermut[i] + 1, 0, 0}, true);
+    xTargetPtr->assign(&cAssign);
+    delete xTargetPtr;
+  }
+  
+  for (int i = nonZeroPivots2; i < cols; ++i) {
+    NDArray *xNullifyPtr = x({colsPermut[i], colsPermut[i] + 1, 0, 0}, true);
+    xNullifyPtr->nullify();
+    delete xNullifyPtr;
+  }
+
+  delete LU;
+  delete bUlike;
+}
+
+BUILD_SINGLE_TEMPLATE( class FullPivLU, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd
+#endif

libnd4j/include/helpers/impl/GradCheck.cpp

@@ -0,0 +1,151 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com), created on 16.07.2018
+//
+#include <array/NDArrayFactory.h>
+#include <helpers/GradCheck.h>
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////
+void GradCheck::fillGradArrays(const LossFunc loss, const std::vector<NDArray*>& gradArrs) {
+  const int numInGradArrs = gradArrs.size();
+
+  // fill input gradient arrays in accordance to type of loss function
+  switch (loss) {
+    case MEAN:
+      for (int i = 0; i < numInGradArrs; ++i) *gradArrs[i] = 1. / gradArrs[i]->lengthOf();
+      break;
+
+    case SUM:
+      for (int i = 0; i < numInGradArrs; ++i) *gradArrs[i] = 1.;
+      break;
+
+    default:
+      THROW_EXCEPTION("GradCheck::fillGradArrays: invalid type of loss function !");
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+bool GradCheck::checkGrad(ops::DeclarableOp& opFF, ops::DeclarableOp& opBP, const OpArgsHolder& argsHolderFF,
+                          const OpArgsHolder& argsHolderBP, const std::vector<bool>& whatArrsToCheck,
+                          const std::vector<double>& idxRange, const LossFunc loss) {
+  const int numInArrsFF =
+      argsHolderFF.getNumInArrs();  // at the same time numInArrsFF = number of output arrays in opBP
+  const int numInGradArrsBP =
+      argsHolderBP.getNumInArrs() - numInArrsFF;  // because argsHolderBP.getNumInArrs() = numInArrsFF + numInGradArrsBP
+  const std::vector<NDArray*>& inArrsFF = argsHolderFF.getInArrs();
+  const std::vector<NDArray*>& inArrsBP = argsHolderBP.getInArrs();
+
+  // fill input gradient arrays in accordance to kind of loss function
+  fillGradArrays(loss, std::vector<NDArray*>(&inArrsBP[numInArrsFF], &inArrsBP[numInArrsFF + numInGradArrsBP]));
+
+  // back prop pass
+  ResultSet outArrsBP = opBP.execute(argsHolderBP);  // number of output arrays in back prop = numInArrsFF;
+
+  NDArray tmpScalar(DOUBLE, inArrsFF[0]->getContext());  // scalar = 0
+
+  for (int i = 0; i < numInArrsFF; ++i) {  // loop through input array
+
+    if (!whatArrsToCheck.empty() && static_cast<bool>(whatArrsToCheck[i]) == false) continue;
+
+    const LongType idxStart = static_cast<LongType>(idxRange[0] * inArrsFF[i]->lengthOf());
+    const LongType idxEnd = static_cast<LongType>(idxRange[1] * inArrsFF[i]->lengthOf());
+
+    for (LongType j = idxStart; j < idxEnd; ++j) {  // loop through all elements for current array
+
+      const double orig = inArrsFF[i]->e<double>(j);
+
+      // add epsilon, feed forward
+      inArrsFF[i]->p<double>(j, orig + EPSILON);
+      ResultSet outArrsFF = opFF.execute(argsHolderFF);
+      int numOutArrs = outArrsFF.size();
+      double scorePlus = 0.;
+
+      for (int k = 0; k < numOutArrs; ++k) {  // loop through output arrays
+        if (loss == SUM)
+          outArrsFF.at(k)->reduceNumber(reduce::Sum, &tmpScalar);
+        else
+          outArrsFF.at(k)->reduceNumber(reduce::Mean, &tmpScalar);
+        scorePlus += tmpScalar.e<double>(0);
+      }
+
+      // subtract epsilon, feed forward
+      inArrsFF[i]->p<double>(j, orig - EPSILON);
+      outArrsFF = opFF.execute(argsHolderFF);
+      double scoreMinus = 0.;
+
+      for (int k = 0; k < numOutArrs; ++k) {  // loop through output arrays
+        if (loss == SUM)
+          outArrsFF.at(k)->reduceNumber(reduce::Sum, &tmpScalar);
+        else
+          outArrsFF.at(k)->reduceNumber(reduce::Mean, &tmpScalar);
+        scoreMinus += tmpScalar.e<double>(0);
+      }
+
+      // restore initial element value
+      inArrsFF[i]->p<double>(j, orig);
+
+      // calculate numerical gradient
+      const double numericalGrad = (scorePlus - scoreMinus) / (2 * EPSILON);
+      if (std::isnan(numericalGrad) || std::isinf(numericalGrad)) {
+        printf(
+            "GradCheck::checkGrad: got wrong value for numerical gradient for input array # %i and its element at "
+            "position %lld ! \n",
+            i, j);
+        THROW_EXCEPTION("");
+      }
+
+      // get analytical gradient
+      const double analyticGrad = outArrsBP.at(i)->e<double>(j);
+      if (std::isnan(analyticGrad) || std::isinf(analyticGrad)) {
+        printf(
+            "GradCheck::checkGrad: got wrong value for analytical gradient for input array # %i and its element at "
+            "position %lld ! \n",
+            i, j);
+        THROW_EXCEPTION("");
+      }
+
+
+      // calculate relative error
+      double relError;
+      if (numericalGrad == 0. && analyticGrad == 0.)
+        relError = 0.;
+      else
+        relError = math::sd_abs<double,double>(analyticGrad - numericalGrad) /
+                   (math::sd_abs<double,double>(analyticGrad) + math::sd_abs<double,double>(numericalGrad));
+
+      // verify result
+      if (relError > MAXRELERR || std::isnan(relError)) {
+        if (math::sd_abs<double,double>(analyticGrad - numericalGrad) < MINABSERR) continue;
+        printf("numericalGrad = %.15f,  analyticGrad = %.15f \n", numericalGrad, analyticGrad);
+        printf(
+            "GradCheck::checkGrad: got RELERROR = %f > MAXRELERROR(%f) for input array # %i and its element at "
+            "position %lld ! \n",
+            relError, MAXRELERR, i, j);
+        return false;
+      }
+    }
+  }
+
+  return true;
+}
+
+}  // namespace sd

libnd4j/include/helpers/impl/HessenbergAndSchur.cpp

@@ -0,0 +1,405 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/HessenbergAndSchur.h>
+#include <helpers/hhSequence.h>
+#include <helpers/householder.h>
+#include <helpers/jacobiSVD.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+Hessenberg<T>::Hessenberg(NDArray* matrix) {
+  if (matrix->rankOf() != 2) THROW_EXCEPTION("ops::helpers::Hessenberg constructor: input matrix must be 2D !");
+
+  if (matrix->sizeAt(0) == 1) {
+    std::vector<LongType> qShape = {1, 1};
+    _Q = new NDArray(matrix->ordering(),qShape, matrix->dataType(), matrix->getContext());
+    *_Q = 1;
+    _H = matrix->dup(matrix->ordering());
+    return;
+  }
+
+  if (matrix->sizeAt(0) != matrix->sizeAt(1))
+    THROW_EXCEPTION("ops::helpers::Hessenberg constructor: input array must be 2D square matrix !");
+
+  _H = matrix->dup(matrix->ordering());
+  _Q = matrix->ulike();
+
+  evalData();
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Hessenberg<T>::evalData() {
+  const int rows = _H->sizeAt(0);
+
+  std::vector<LongType> coeffsShape = {rows - 1};
+  NDArray hhCoeffs(_H->ordering(), coeffsShape, _H->dataType(), _H->getContext());
+
+  // calculate _H
+  for (LongType i = 0; i < rows - 1; ++i) {
+    T coeff, norm;
+
+    NDArray hRef = *_H;
+    NDArray *tail1Ptr = hRef({i + 1, -1, i, i + 1});
+    NDArray tail1 = *tail1Ptr;
+    delete tail1Ptr;
+    
+    NDArray *tail2Ptr = hRef({i + 2, -1, i, i + 1}, true);
+    NDArray tail2 = *tail2Ptr;
+    delete tail2Ptr;
+
+    Householder<T>::evalHHmatrixDataI(tail1, coeff, norm);
+
+    NDArray *hViewPtr = hRef({0, 0, i, i + 1});
+    hViewPtr->template r<T>(i + 1) = norm;
+    delete hViewPtr;
+    
+    hhCoeffs.template r<T>(i) = coeff;
+
+    NDArray *bottomRightCornerPtr = hRef({i + 1, -1, i + 1, -1}, true);
+    NDArray bottomRightCorner = *bottomRightCornerPtr;
+    delete bottomRightCornerPtr;
+    
+    Householder<T>::mulLeft(bottomRightCorner, tail2, coeff);
+    
+    NDArray *tail2Trans = tail2.transpose();
+    NDArray *rightColsPtr = hRef({0, 0, i + 1, -1}, true);
+    NDArray rightCols = *rightColsPtr;
+    delete rightColsPtr;
+    
+    Householder<T>::mulRight(rightCols, *tail2Trans, coeff);
+    delete tail2Trans;
+  }
+
+  // calculate _Q
+  HHsequence hhSeq(_H, &hhCoeffs, 'u');
+  hhSeq._diagSize = rows - 1;
+  hhSeq._shift = 1;
+  hhSeq.applyTo_<T>(_Q);
+
+  // fill down with zeros starting at first subdiagonal
+  _H->fillAsTriangular<T>(0, -1, -1, *_H, 'l',false);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+Schur<T>::Schur(NDArray& matrix) {
+  if (matrix.rankOf() != 2) THROW_EXCEPTION("ops::helpers::Schur constructor: input matrix must be 2D !");
+
+  if (matrix.sizeAt(0) != matrix.sizeAt(1))
+    THROW_EXCEPTION("ops::helpers::Schur constructor: input array must be 2D square matrix !");
+
+  evalData(matrix);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::evalData(NDArray& matrix) {
+  auto res = matrix.reduceNumber(reduce::AMax);
+  const T scale = res->template t<T>(0);
+  delete res;
+
+  if (scale < DataTypeUtils::min_positive<T>()) {
+    t = matrix.ulike();
+    u = matrix.ulike();
+
+    t->nullify();
+    u->setIdentity();
+
+    return;
+  }
+
+  // perform Hessenberg decomposition
+  NDArray *matrixScale = matrix / scale;
+  Hessenberg<T> hess(matrixScale);
+
+  t = hess._H;
+  u = hess._Q;
+
+  calcFromHessenberg();
+
+  *t *= scale;
+  delete matrixScale;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::splitTwoRows(const int ind, const T shift) {
+  const int numCols = t->sizeAt(1);
+
+  T p = (T)0.5 * (t->t<T>(ind - 1, ind - 1) - t->t<T>(ind, ind));
+
+  T q = p * p + t->t<T>(ind, ind - 1) * t->t<T>(ind - 1, ind);
+
+  t->r<T>(ind, ind) += shift;
+  t->r<T>(ind - 1, ind - 1) += shift;
+
+  if (q >= (T)0) {
+    T z = math::sd_sqrt<T, T>(math::sd_abs<T,T>(q));
+
+    std::vector<LongType> rotShape = {2, 2};
+    NDArray rotation(t->ordering(), rotShape, t->dataType(), t->getContext());
+
+    if (p >= (T)0)
+      JacobiSVD<T>::createJacobiRotationGivens(p + z, t->t<T>(ind, ind - 1), rotation);
+    else
+      JacobiSVD<T>::createJacobiRotationGivens(p - z, t->t<T>(ind, ind - 1), rotation);
+
+    NDArray tRef = *t;
+    NDArray *rightColsPtr = tRef({0, 0, ind - 1, -1});
+    NDArray rightCols = *rightColsPtr;
+    delete rightColsPtr;
+    
+    NDArray *rotT = rotation.transpose();
+    JacobiSVD<T>::mulRotationOnLeft(ind - 1, ind, rightCols, *rotT);
+
+    NDArray *topRowsPtr = tRef({0, ind + 1, 0, 0});
+    NDArray topRows = *topRowsPtr;
+    delete topRowsPtr;
+    
+    JacobiSVD<T>::mulRotationOnRight(ind - 1, ind, topRows, rotation);
+
+    JacobiSVD<T>::mulRotationOnRight(ind - 1, ind, *u, rotation);
+
+    t->r<T>(ind, ind - 1) = (T)0;
+    delete rotT;
+  }
+
+  if (ind > 1) t->r<T>(ind - 1, ind - 2) = (T)0;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::calcShift(const int ind, const int iter, T& shift, NDArray& shiftVec) {
+  // shiftVec has length = 3
+
+  shiftVec.r<T>(0) = t->t<T>(ind, ind);
+  shiftVec.r<T>(1) = t->t<T>(ind - 1, ind - 1);
+  shiftVec.r<T>(2) = t->t<T>(ind, ind - 1) * t->t<T>(ind - 1, ind);
+
+  if (iter == 10) {
+    shift += shiftVec.t<T>(0);
+
+    for (int i = 0; i <= ind; ++i) t->r<T>(i, i) -= shiftVec.t<T>(0);
+
+    T s = math::sd_abs<T,T>(t->t<T>(ind, ind - 1)) + math::sd_abs<T,T>(t->t<T>(ind - 1, ind - 2));
+
+    shiftVec.r<T>(0) = T(0.75) * s;
+    shiftVec.r<T>(1) = T(0.75) * s;
+    shiftVec.r<T>(2) = T(-0.4375) * s * s;
+  }
+
+  if (iter == 30) {
+    T s = (shiftVec.t<T>(1) - shiftVec.t<T>(0)) / T(2.0);
+    s = s * s + shiftVec.t<T>(2);
+
+    if (s > T(0)) {
+      s = math::sd_sqrt<T, T>(s);
+
+      if (shiftVec.t<T>(1) < shiftVec.t<T>(0)) s = -s;
+
+      s = s + (shiftVec.t<T>(1) - shiftVec.t<T>(0)) / T(2.0);
+      s = shiftVec.t<T>(0) - shiftVec.t<T>(2) / s;
+      shift += s;
+
+      for (int i = 0; i <= ind; ++i) t->r<T>(i, i) -= s;
+
+      shiftVec = T(0.964);
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::initFrancisQR(const int ind1, const int ind2, NDArray& shiftVec, int& ind3,
+                             NDArray& householderVec) {
+  // shiftVec has length = 3
+
+  for (ind3 = ind2 - 2; ind3 >= ind1; --ind3) {
+    const T mm = t->t<T>(ind3, ind3);
+    const T r = shiftVec.t<T>(0) - mm;
+    const T s = shiftVec.t<T>(1) - mm;
+
+    householderVec.r<T>(0) = (r * s - shiftVec.t<T>(2)) / t->t<T>(ind3 + 1, ind3) + t->t<T>(ind3, ind3 + 1);
+    householderVec.r<T>(1) = t->t<T>(ind3 + 1, ind3 + 1) - mm - r - s;
+    householderVec.r<T>(2) = t->t<T>(ind3 + 2, ind3 + 1);
+
+    if (ind3 == ind1) break;
+
+    const T lhs =
+        t->t<T>(ind3, ind3 - 1) * (math::sd_abs<T,T>(householderVec.t<T>(1)) + math::sd_abs<T,T>(householderVec.t<T>(2)));
+    const T rhs = householderVec.t<T>(0) * (math::sd_abs<T,T>(t->t<T>(ind3 - 1, ind3 - 1)) + math::sd_abs<T,T>(mm) +
+                                            math::sd_abs<T,T>(t->t<T>(ind3 + 1, ind3 + 1)));
+
+    if (math::sd_abs<T,T>(lhs) < DataTypeUtils::eps<T>() * rhs) break;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::doFrancisQR(const int ind1, const int ind2, const int ind3, NDArray& householderVec) {
+  if (!(ind2 >= ind1))
+    THROW_EXCEPTION(
+        "ops::helpers::Schur:doFrancisQR: wrong input indexes, condition ind2 >= ind1 must be true !");
+  if (!(ind2 <= ind3 - 2))
+    THROW_EXCEPTION(
+        "ops::helpers::Schur:doFrancisQR: wrong input indexes, condition iind2 <= ind3-2 must be true !");
+
+  const int numCols = t->sizeAt(1);
+  NDArray tRef = *t;
+  NDArray uRef = *u;
+  for (int k = ind2; k <= ind3 - 2; ++k) {
+    const bool firstIter = (k == ind2);
+
+    T coeff, normX;
+    std::vector<LongType> tailShape = {2,1};
+    NDArray tail(t->ordering(),tailShape, t->dataType(), t->getContext());
+    
+    NDArray *firstPtr = firstIter ? &householderVec : tRef({k, k + 3, k - 1, k});
+    NDArray first = *firstPtr;
+    if (!firstIter) delete firstPtr;
+    
+    Householder<T>::evalHHmatrixData(first, tail, coeff, normX);
+
+    if (normX != T(0)) {
+      if (firstIter && k > ind1)
+        t->r<T>(k, k - 1) = -t->t<T>(k, k - 1);
+      else if (!firstIter)
+        t->r<T>(k, k - 1) = normX;
+
+      NDArray *block1Ptr = tRef({k, k + 3, k, numCols}, true);
+      NDArray block1 = *block1Ptr;
+      delete block1Ptr;
+      Householder<T>::mulLeft(block1, tail, coeff);
+
+      NDArray *block2Ptr = tRef({0, math::sd_min<int>(ind3, k + 3) + 1, k, k + 3}, true);
+      NDArray block2 = *block2Ptr;
+      delete block2Ptr;
+      Householder<T>::mulRight(block2, tail, coeff);
+
+      NDArray *block3Ptr = uRef({0, numCols, k, k + 3}, true);
+      NDArray block3 = *block3Ptr;
+      delete block3Ptr;
+      Householder<T>::mulRight(block3, tail, coeff);
+    }
+  }
+
+  T coeff, normX;
+  std::vector<LongType> tailShape = {1,1};
+  NDArray tail(t->ordering(), tailShape, t->dataType(), t->getContext());
+  NDArray *firstPtr = tRef({ind3 - 1, ind3 + 1, ind3 - 2, ind3 - 1});
+  NDArray first = *firstPtr;
+  delete firstPtr;
+  
+  Householder<T>::evalHHmatrixData(first, tail, coeff, normX);
+
+  if (normX != T(0)) {
+    t->r<T>(ind3 - 1, ind3 - 2) = normX;
+
+    NDArray *block1Ptr = tRef({ind3 - 1, ind3 + 1, ind3 - 1, numCols}, true);
+    NDArray block1 = *block1Ptr;
+    delete block1Ptr;
+    Householder<T>::mulLeft(block1, tail, coeff);
+
+    NDArray *block2Ptr = tRef({0, ind3 + 1, ind3 - 1, ind3 + 1}, true);
+    NDArray block2 = *block2Ptr;
+    delete block2Ptr;
+    Householder<T>::mulRight(block2, tail, coeff);
+
+    NDArray *block3Ptr = uRef({0, numCols, ind3 - 1, ind3 + 1}, true);
+    NDArray block3 = *block3Ptr;
+    delete block3Ptr;
+    Householder<T>::mulRight(block3, tail, coeff);
+  }
+
+  for (int i = ind2 + 2; i <= ind3; ++i) {
+    t->r<T>(i, i - 2) = T(0);
+    if (i > ind2 + 2) t->r<T>(i, i - 3) = T(0);
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Schur<T>::calcFromHessenberg() {
+  const int maxIters = _maxItersPerRow * t->sizeAt(0);
+
+  const int numCols = t->sizeAt(1);
+  int iu = numCols - 1;
+  int iter = 0;
+  int totalIter = 0;
+
+  T shift = T(0);
+  NDArray tRef = *t;
+  NDArray uRef = *u;
+  T norm = static_cast<T>(0);
+  for (int j = 0; j < numCols; ++j) {
+    NDArray *viewPtr = tRef({0, math::sd_min<int>(numCols, j + 2), j, j + 1});
+    auto sum = viewPtr->reduceNumber(reduce::ASum);
+    norm += sum->template t<T>(0);
+    delete viewPtr;
+    delete sum;
+  }
+
+  if (norm != T(0)) {
+    while (iu >= 0) {
+      const int il = getSmallSubdiagEntry(iu);
+
+      if (il == iu) {
+        t->r<T>(iu, iu) = t->t<T>(iu, iu) + shift;
+        if (iu > 0) t->r<T>(iu, iu - 1) = T(0);
+        iu--;
+        iter = 0;
+
+      } else if (il == iu - 1) {
+        splitTwoRows(iu, shift);
+        iu -= 2;
+        iter = 0;
+      } else {
+        std::vector<LongType> shiftVecShape = {3};
+        NDArray householderVec(t->ordering(), shiftVecShape, t->dataType(), t->getContext());
+        NDArray shiftVec(t->ordering(), shiftVecShape, t->dataType(), t->getContext());
+
+        calcShift(iu, iter, shift, shiftVec);
+
+        ++iter;
+        ++totalIter;
+
+        if (totalIter > maxIters) break;
+
+        int im;
+        initFrancisQR(il, iu, shiftVec, im, householderVec);
+        doFrancisQR(il, im, iu, householderVec);
+      }
+    }
+  }
+}
+
+BUILD_SINGLE_TEMPLATE( class Hessenberg, , SD_FLOAT_TYPES);
+BUILD_SINGLE_TEMPLATE( class Schur, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/MmulHelper.cpp

@@ -0,0 +1,616 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com), created on 05.06.2018
+//
+
+#ifndef LIBND4J_MMULHELPER_CPP
+#define LIBND4J_MMULHELPER_CPP
+#include "../MmulHelper.h"
+
+#include <array/NDArrayFactory.h>
+#include <helpers/BlasHelper.h>
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/headers/shape.h>
+#include <ops/declarable/helpers/batched_gemm.h>
+
+#include <algorithm>
+#include <iterator>
+#include <numeric>
+#include <vector>
+
+#include "ops/declarable/headers/blas.h"
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////
+NDArray* MmulHelper::tensorDot(NDArray* A, NDArray* B,
+                               const std::initializer_list<LongType>& axesA,
+                               const std::initializer_list<LongType>& axesB) {
+  std::vector<LongType> aA(axesA);
+  std::vector<LongType> aB(axesB);
+  return tensorDot(A, B, aA, aB);
+}
+
+//////////////////////////////////////////////////////////////////////////
+NDArray* MmulHelper::tensorDot(NDArray* A, NDArray* B, const std::vector<LongType>& axesA,
+                               const std::vector<LongType>& axesB) {
+  std::vector<LongType> permutAt, permutBt;
+  std::vector<LongType> shapeAt, shapeBt;
+
+  auto outShape = ShapeUtils::evalShapeForTensorDot(A, B, axesA, axesB, permutAt, permutBt, shapeAt, shapeBt);
+
+  // check whether permutation is necessary
+  NDArray* aP = permutAt.empty() ? A : A->permute(permutAt, false, false);
+  NDArray* bP = permutBt.empty() ? B : B->permute(permutBt, false, false);
+
+  // check whether reshape is necessary
+  NDArray* aPR = aP->isSameShape(shapeAt) ? aP : aP->reshape(aP->ordering(), shapeAt);
+  NDArray* bPR = bP->isSameShape(shapeAt) ? bP : bP->reshape(bP->ordering(), shapeBt);
+
+  NDArray* c = mmul(aPR, bPR, nullptr, 1.0, 0.0);
+
+  c->reshapei(outShape);
+
+  // Delete reshaped arrays first
+  if(aPR != A && aPR != aP) {
+    delete aPR;
+  }
+  if(bPR != B && bPR != bP) {
+    delete bPR;
+  }
+
+  // Then delete permuted arrays
+  if(aP != A) {
+    delete aP;
+  }
+  if(bP != B) {
+    delete bP;
+  }
+
+  return c;
+}
+
+
+void MmulHelper::computeNewShapesAndAxes(
+    NDArray& as_, const std::vector<LongType>& axes_a,
+    NDArray& bs, const std::vector<LongType>& axes_b,
+    std::vector<LongType>& newshape_a, std::vector<LongType>& newaxes_a,
+    std::vector<LongType>& newshape_b, std::vector<LongType>& newaxes_b
+) {
+
+
+  std::vector<LongType> *as_shape = as_.getShapeAsVector();
+  std::vector<LongType> *bs_shape = bs.getShapeAsVector();
+
+  std::vector<LongType> notin_a;
+  for(size_t k = 0; k < as_shape->size(); ++k) {
+    if(std::find(axes_a.begin(), axes_a.end(), k) == axes_a.end())
+      notin_a.push_back(k);
+  }
+
+
+
+  newaxes_a.clear();
+  std::copy(notin_a.begin(), notin_a.end(), std::back_inserter(newaxes_a));
+  std::copy(axes_a.begin(), axes_a.end(), std::back_inserter(newaxes_a));
+
+  LongType N2_a = std::accumulate(axes_a.begin(), axes_a.end(), 1L, [&](LongType product, LongType i){
+    return product * (*as_shape)[i];
+  });
+
+  newshape_a.clear();
+  newshape_a.push_back(std::accumulate(notin_a.begin(), notin_a.end(), 1L, [&](LongType product, LongType i){
+    return product * (*as_shape)[i];
+  }));
+  newshape_a.push_back(N2_a);
+
+
+
+  std::vector<LongType> notin_b;
+  for(size_t k = 0; k < bs_shape->size(); ++k) {
+    if(std::find(axes_b.begin(), axes_b.end(), k) == axes_b.end())
+      notin_b.push_back(k);
+  }
+
+
+  newaxes_b.clear();
+  std::copy(axes_b.begin(), axes_b.end(), std::back_inserter(newaxes_b));
+  std::copy(notin_b.begin(), notin_b.end(), std::back_inserter(newaxes_b));
+
+
+
+  LongType N2_b = std::accumulate(axes_b.begin(), axes_b.end(), 1L, [&](LongType product, LongType i){
+    return product * (*bs_shape)[i];
+  });
+
+
+
+  newshape_b.clear();
+  newshape_b.push_back(N2_b);
+  newshape_b.push_back(std::accumulate(notin_b.begin(), notin_b.end(), 1L, [&](LongType product, LongType i){
+    return product * (*bs_shape)[i];
+  }));
+
+
+}
+
+//////////////////////////////////////////////////////////////////////////
+void MmulHelper::tensorDot2(NDArray* a, NDArray* b, NDArray* c, const std::vector<LongType>& axes_a,
+                            const std::vector<LongType>& axes_b, std::vector<LongType>& permutAt,
+                            std::vector<LongType>& permuteBt, std::vector<LongType>& permuteCt,
+                            NDArray* realFinalResult) {
+
+  // check whether permutation is required
+  NDArray* cP  =permuteCt.empty() ? c : c->permute(permuteCt, false, false);
+
+  std::vector<LongType> shapeAt, shapeBt;
+  std::vector<LongType> permutAtDummy, permuteBtDummy;
+
+  std::vector<LongType> newshape_a, newaxes_a, newshape_b, newaxes_b;
+  computeNewShapesAndAxes(*a, axes_a, *b, axes_b, newshape_a, newaxes_a, newshape_b, newaxes_b);
+
+  NDArray* aP = permutAt.empty() ? a : a->permute(permutAt, false, false);
+  NDArray* bP = permuteBt.empty() ? b :b->permute(permuteBt, false, false);
+
+  NDArray* aPermuted = aP->permute(newaxes_a, false, false);
+  NDArray* aPR = aPermuted->reshape('c', newshape_a, true);
+
+  NDArray* bPermuted = bP->permute(newaxes_b, false, false);
+  NDArray* bPR = bPermuted->reshape('c', newshape_b, true);
+
+  std::vector<LongType> requiredCshape  = {aPR->sizeAt(0), bPR->sizeAt(1)};
+  NDArray *cP2 = cP->reshape('f', requiredCshape, false);
+  NDArray* cPR = cP2;
+
+  NDArray * ret = mmul(aPR, bPR, cPR, 1.0, 0.0);
+
+  if (cPR->buffer() != cP->buffer() ||
+      cPR->specialBuffer() != cP->specialBuffer()) {  // this means both permute and reshape have been performed on c, cP
+    if(c->buffer() == cP->buffer()) {
+      auto copyFromBuff = cP->dataBuffer();
+      cP->dataBuffer()->copyBufferFrom(*copyFromBuff);
+    } else {
+      auto copyFromBuff = cP->dataBuffer();
+      c->dataBuffer()->copyBufferFrom(*copyFromBuff);
+    }
+  }
+
+  if(realFinalResult != c) {
+    realFinalResult->dataBuffer()->copyBufferFrom(*c->dataBuffer());
+  }
+
+  if(cP != c) {
+    delete cP;
+  }
+  if(cPR != c) {
+    delete cPR;
+  }
+
+  if(aP != a && !aP->isView()) {
+    delete aP;
+  }
+  if(bP != b && !bP->isView()) {
+    delete bP;
+  }
+
+  // Delete in reverse order of creation to avoid use-after-free
+  if(bPR != b && bPR != bP && bPR != bPermuted && !bPR->isView()) {
+    delete bPR;
+  }
+  if(bPermuted != b && bPermuted != bP && !bPermuted->isView()) {
+    delete bPermuted;
+  }
+
+  if(aPR != a && aPR != aP && aPR != aPermuted && !aPR->isView()) {
+    delete aPR;
+  }
+  if(aPermuted != a && aPermuted != aP && !aPermuted->isView()) {
+    delete aPermuted;
+  }
+
+}
+
+
+void MmulHelper::tensorDot(NDArray* a, NDArray* b, NDArray* c,
+                           std::vector<LongType>& axes_a, std::vector<LongType>& axes_b,
+                           std::vector<LongType>& permutForC) {
+
+  std::vector<LongType> permutAt, permutBt;
+  std::vector<LongType> shapeAt, shapeBt;
+  ShapeUtils::evalShapeForTensorDot(a, b, axes_a, axes_b, permutAt, permutBt, shapeAt, shapeBt);
+
+
+  // check whether permutation is required
+  NDArray* cP = permutForC.empty() ? c :c->permute(permutForC, false, false);
+  // check whether permutation is necessary
+  NDArray* aP = permutAt.empty() ? a :a->permute(permutAt, false, false);
+  NDArray* bP = permutBt.empty() ? b : b->permute(permutBt, false, false);
+
+  // check whether reshape is necessary
+  NDArray* aPR = aP->isSameShape(shapeAt) ? aP : aP->reshape(aP->ordering(), shapeAt);
+  NDArray* bPR = bP->isSameShape(shapeAt) ? bP : bP->reshape(bP->ordering(), shapeBt);
+
+  std::vector<LongType> requiredCshape = {aPR->sizeAt(0), bPR->sizeAt(1)};
+
+
+  NDArray* cPR = cP->isSameShape(requiredCshape) ? cP : cP->reshape(cP->ordering(), requiredCshape, false);
+  NDArray *ret = mmul(aPR, bPR, cPR, 1.0, 0.0);
+
+  if (c != ret) {  // this means both permute and reshape have been performed on c, cP
+    // always points on c->buffer()
+    NDArray *assign2 = ret->reshape(c->ordering(),requiredCshape);
+    c->assign(assign2);
+    delete assign2;
+  }
+
+
+  if(c != cP && !cP->isView()) {
+    delete cP;
+  }
+
+  if(aP != a && !aP->isView()) {
+    delete aP;
+  }
+
+  if(bP != b && !bP->isView()) {
+    delete bP;
+  }
+
+  if(aPR != a && aPR != aP && !aPR->isView()) {
+    delete aPR;
+  }
+  if(bPR != b && bPR != bP && !bPR->isView()) {
+    delete bPR;
+  }
+
+  if(cPR != c && cPR != cP && !cPR->isView()) {
+    delete cPR;
+  }
+}
+
+#ifndef __JAVACPP_HACK__
+//////////////////////////////////////////////////////////////////////////
+void MmulHelper::tensorDot(NDArray* a, NDArray* b, NDArray* c,
+                           std::vector<std::vector<LongType>>& modifA,
+                           std::vector<std::vector<LongType>>& modifB,
+                           std::vector<std::vector<LongType>>& modifC) {
+  NDArray *aPR(const_cast<NDArray*>(a)), *bPR(const_cast<NDArray*>(b));
+  std::string whatToDoWithA, whatToDoWithB,
+      whatToDoWithC;  // "" - nothing; "p" - permutation; "r" - reshaping; "pr" - permutation+reshaping; "rp" -
+  // reshaping/permutation, and so on; if another string is produced - throw exception
+
+  for (const auto& arr : modifA)
+    whatToDoWithA =
+        (std::find(arr.begin(), arr.end(), 0) != arr.end())
+        ? whatToDoWithA + "p"
+        : whatToDoWithA +
+          "r";  // when 0 is present in arr then it is permutation array, otherwise - it is reshaping array
+  for (const auto& arr : modifB)
+    whatToDoWithB = (std::find(arr.begin(), arr.end(), 0) != arr.end()) ? whatToDoWithB + "p" : whatToDoWithB + "r";
+  for (const auto& arr : modifC)
+    whatToDoWithC = (std::find(arr.begin(), arr.end(), 0) != arr.end()) ? whatToDoWithC + "p" : whatToDoWithC + "r";
+
+  // first step for a array
+
+  if (!whatToDoWithA.empty())
+    aPR = (whatToDoWithA[0] == 'p') ? a->permute(modifA[0], false, false)
+                                    :a->reshape(a->ordering(), modifA[0]);
+  // first step for b array
+  if (!whatToDoWithB.empty())
+    bPR = (whatToDoWithB[0] == 'p') ? b->permute(modifB[0], false, false)
+                                    : b->reshape(b->ordering(), modifB[0]);
+  // rest steps for a array
+  for (size_t i = 1; i < whatToDoWithA.size(); ++i)
+    if (whatToDoWithA[i] == 'p')
+      aPR->permutei(modifA[i], false, false);
+    else
+      aPR->reshapei(modifA[i]);
+  // rest steps for b array
+  for (size_t i = 1; i < whatToDoWithB.size(); ++i)
+    if (whatToDoWithB[i] == 'p')
+      bPR->permutei(modifB[i], false, false);
+    else
+      bPR->reshapei(modifB[i]);
+
+  // now work with c array
+  std::vector<NDArray*> cArrs = {c};
+  if (!whatToDoWithC.empty()) {
+    cArrs = std::vector<NDArray*>(whatToDoWithC.size() + 1, c);
+    for (size_t i = 0; i < cArrs.size() - 1; ++i)
+      cArrs[i + 1] =
+          (whatToDoWithC[i] == 'p')
+          ? cArrs[i]->permute(modifC[i], false, false)
+          : cArrs[i]->reshape(
+              c->ordering(), modifC[i],
+              false);  // since we ignore first element in cArrs (that is cArrs[0]) then it is always equal to c
+  }
+
+  mmul(aPR, bPR, cArrs[cArrs.size() - 1], 1.0, 0.0);
+
+  // check whether new buffer allocation was happened for c array
+  if (!whatToDoWithC.empty()) {
+    for (int i = cArrs.size() - 1; i > 0; --i) {
+      if (cArrs[i]->buffer() != cArrs[i - 1]->buffer() || cArrs[i]->specialBuffer() != cArrs[i - 1]->specialBuffer())
+        cArrs[i - 1]->assign(cArrs[i]);
+      delete cArrs[i];
+    }
+  }
+
+  if (aPR != a) delete aPR;
+  if (bPR != b) delete bPR;
+}
+
+//////////////////////////////////////////////////////////////////////////
+NDArray* MmulHelper::tensorDot(NDArray* a, NDArray* b,
+                               std::vector<std::vector<LongType>>& modifA,
+                               std::vector<std::vector<LongType>>& modifB) {
+  NDArray *aPR(const_cast<NDArray*>(a)), *bPR(const_cast<NDArray*>(b));
+  std::string whatToDoWithA,
+      whatToDoWithB;  // "" - nothing; "p" - permutation only; "r" - reshaping only; "pr" - permutation+reshaping; "rp"
+  // - reshaping/permutation; another string - throw exception
+
+  for (const auto& arr : modifA)
+    whatToDoWithA =
+        (std::find(arr.begin(), arr.end(), 0) != arr.end())
+        ? whatToDoWithA + "p"
+        : whatToDoWithA +
+          "r";  // when 0 is present in arr then it is permutation array, otherwise - it is reshaping array
+  for (const auto& arr : modifB)
+    whatToDoWithB = (std::find(arr.begin(), arr.end(), 0) != arr.end()) ? whatToDoWithB + "p" : whatToDoWithB + "r";
+
+  // first step for a array
+  if (!whatToDoWithA.empty())
+    aPR = (whatToDoWithA[0] == 'p') ?a->permute(modifA[0], false, false)
+                                    : a->reshape(a->ordering(), modifA[0]);
+  // first step for b array
+  if (!whatToDoWithB.empty())
+    bPR = (whatToDoWithB[0] == 'p') ? b->permute(modifB[0], false, false)
+                                    : b->reshape(b->ordering(), modifB[0]);
+  // rest steps for a array
+  for (size_t i = 1; i < whatToDoWithA.size(); ++i)
+    if (whatToDoWithA[i] == 'p')
+      aPR->permutei(modifA[i], false, false);
+    else
+      aPR->reshapei(modifA[i]);
+  // rest steps for b array
+  for (size_t i = 1; i < whatToDoWithB.size(); ++i)
+    if (whatToDoWithB[i] == 'p')
+      bPR->permutei(modifB[i], false, false);
+    else
+      bPR->reshapei(modifB[i]);
+
+  NDArray* result = mmul(aPR, bPR, nullptr, 1.0, 0.0);
+
+  return result;
+}
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+NDArray* MmulHelper::mmul(NDArray* A, NDArray* B, NDArray* C, const double alpha,
+                          const double beta, const char outOrder) {
+  LongType lenDim;
+  const LongType aRank = A->rankOf();
+  const LongType bRank = B->rankOf();
+  const bool isAVector = shape::isCommonVector(A->shapeInfo(), lenDim);
+  const bool isBVector = shape::isCommonVector(B->shapeInfo(), lenDim);
+  // dot product of 2 vectors
+  if (A->lengthOf() == B->lengthOf() && isAVector && isBVector &&
+      (aRank != 2 ||
+       (aRank == 2 && (A->isSameShape(B) ||
+                      (bRank == 1 && A->sizeAt(1) == 1))))) {  // (1x1x1 * 1x1) or (1x4 * 1*4) or (4x1 * 4x1) or (4x1 * 4)
+
+
+    return dot(A, B, C, alpha, beta);
+  }
+  // matrix x matrix
+  if (aRank == 2 && bRank == 2) {
+    return mmulMxM(A, B, C, alpha, beta, outOrder);
+  }
+
+  // matrix x vector
+  if (aRank == 2 && isBVector) {
+    return mmulMxV(A, B, C, alpha, beta, outOrder);
+  }
+
+  // vector x matrix, A{M} x B{M,N} = C{N} -> reduce to matrix x matrix A2{1,M} x B{M,N} = C2{1,N}, since there is no
+  // corresponding blas operation sgevm
+  if (isAVector && bRank == 2) {
+    std::vector<sd::LongType> aShape = {1, A->lengthOf()};
+    std::vector<sd::LongType> cShape = {1, C->lengthOf()};
+
+
+    NDArray* A2 = A->reshape(A->ordering(), aShape);                       // A{M} -> A2{1,M}
+    NDArray* C2 = C ? C->reshape(C->ordering(), cShape, false) : nullptr;  // C{N} -> C2{1,N}
+    auto result = mmulMxM(A2, B, C2, alpha, beta, outOrder);                                        // result{1,N}
+
+    // Cleanup reshaped arrays
+    if (A2 != A) delete A2;
+    if (C2 != nullptr && C2 != C) delete C2;
+
+    if (!C) {
+      result->reshapei({result->lengthOf()});  // result{1,N} -> result{N}
+      return result;
+    }
+    return C;
+  }
+
+  // batched matrix multiplication
+  return mmulNxN(A, B, C, alpha, beta, outOrder);
+}
+
+bool MmulHelper::resolveTranspose(sd::NDArray& a, sd::NDArray& b, bool& transA, bool& transB) {
+  int rowsA = a.sizeAt(-2);
+  int colsA = a.sizeAt(-1);
+  int rowsB = b.sizeAt(-2);
+  int colsB = b.sizeAt(-1);
+
+  transA = false;
+  transB = false;
+
+
+  if (colsA == rowsB) {
+    // No transpose needed
+    return true;
+  } else if (rowsA == rowsB) {
+    // Transpose A
+    transA = true;
+    return true;
+  } else if (colsA == colsB) {
+    // Transpose B
+    transB = true;
+    return true;
+  } else {
+    // Dimensions do not match for matrix multiply
+    return false;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+void MmulHelper::matmul(NDArray* x, NDArray* y, NDArray* z, const bool transX, const bool transY, double alpha,
+                        double beta, NDArray* realFinalResult) {
+  int xRank = x->rankOf();
+  int yRank = y->rankOf();
+
+  auto outShape = ShapeUtils::evalShapeForMatmul(x->shapeInfo(), y->shapeInfo(), transX, transY);
+  if (!z->isSameShape(outShape)) {
+    std::string errorMessage;
+    errorMessage = "NDArrayFactory::matmul static method: input shape of output array is wrong, actual is";
+    errorMessage += ShapeUtils::shapeAsString(z).c_str();
+    errorMessage += " and expected is ";
+    errorMessage += ShapeUtils::shapeAsString(outShape).c_str();
+    errorMessage += " ! \n";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (z->isEmpty()) return;
+
+  NDArray *xT = const_cast<NDArray *>(x);
+  NDArray *yT = const_cast<NDArray *>(y);
+  NDArray *zT = z;
+
+  // Handle transpose via permute + dup for contiguous data
+  // permute creates a view with swapped strides, dup() makes a contiguous copy
+  if ((transX && xRank > 1) || (transY && yRank > 1)) {
+    const int rank = xRank >= yRank ? xRank : yRank;
+    std::vector<LongType> permut(rank);
+    for (int i = 0; i < rank - 2; ++i) permut[i] = i;
+    permut[rank - 2] = rank - 1;
+    permut[rank - 1] = rank - 2;
+
+    if (transX) {
+      NDArray *permutedView = x->permute(permut, false, false);  // Create view (non-contiguous)
+      xT = permutedView->dup();  // Make contiguous copy with proper data layout
+      delete permutedView;
+    }
+    if (transY) {
+      NDArray *permutedView = y->permute(permut, false, false);  // Create view (non-contiguous)
+      yT = permutedView->dup();  // Make contiguous copy with proper data layout
+      delete permutedView;
+    }
+  }
+
+  if (xRank <= 2 && yRank <= 2) {
+    // dot (1Dx1D), vector-matrix (1Dx2D), matrix-vector (2Dx1D), matrix-matrix (2Dx2D) product cases
+    NDArray* xReshaped = nullptr;
+    NDArray* zReshaped = nullptr;
+    
+    if (xRank == 1 && yRank == 2) {
+      // reduce vector-matrix to matrix-matrix case
+      std::vector<sd::LongType> xShape = {1, xT->lengthOf()};
+      std::vector<sd::LongType> zShape = {1, z->lengthOf()};
+
+      // Remember if we need to delete the permuted versions
+      NDArray* xPermuted = (xT != x) ? xT : nullptr;
+      NDArray* zPermuted = (zT != z) ? zT : nullptr;
+
+      xReshaped = xT->reshape(xT->ordering(), xShape, false);
+      xT = xReshaped;
+      zReshaped = z->reshape(z->ordering(), zShape, false);
+      zT = zReshaped;
+
+      // Clean up permuted versions if they exist
+      if(xPermuted != nullptr && !xPermuted->isView()) {
+        delete xPermuted;
+      }
+      if(zPermuted != nullptr && !zPermuted->isView()) {
+        delete zPermuted;
+      }
+    }
+
+    mmul(xT, yT, zT, alpha, beta);
+
+    // Copy back result and clean up reshaped output
+    if(zT != z) {
+      z->dataBuffer()->copyBufferFrom(*zT->dataBuffer(), zT->lengthOf() * zT->sizeOfT());
+      delete zT;
+      zT = z;  // Reset to original to prevent double-free at end of function
+    }
+
+    // Clean up reshaped input
+    if(xReshaped != nullptr && xReshaped != x) {
+      delete xReshaped;
+      xT = x;  // Reset to original to prevent double-free at end of function
+    }
+
+  } else {
+    // Batched matmul: loop over batch dimensions and call 2D gemm for each slice
+    // This is more reliable than mmulNxN which has bugs in batch index calculation
+
+    // For 3D arrays [batch, M, K] x [batch, K, N] = [batch, M, N]
+    // We iterate over batch dimension and call 2D mmul for each slice
+    const int xRankT = xT->rankOf();
+    const int yRankT = yT->rankOf();
+    const int zRankT = zT->rankOf();
+
+    if (xRankT == 3 && yRankT == 3 && zRankT == 3) {
+      // Simple case: all 3D with matching batch dimension
+      const LongType batchSize = xT->sizeAt(0);
+      const LongType M = xT->sizeAt(1);
+      const LongType K = xT->sizeAt(2);
+      const LongType N = yT->sizeAt(2);
+
+      for (LongType b = 0; b < batchSize; ++b) {
+        // Get 2D slices for this batch using subarray
+        auto xSlice = (*xT)(b, {0});  // [M, K]
+        auto ySlice = (*yT)(b, {0});  // [K, N]
+        auto zSlice = (*zT)(b, {0});  // [M, N]
+
+        // Call 2D matmul - no transpose flags since we already handled them via permute+dup
+        mmul(xSlice, ySlice, zSlice, alpha, beta);
+      }
+    } else {
+      // Fall back to mmulNxN for other cases (4D+, mixed ranks, etc.)
+      mmulNxN(xT, yT, zT, alpha, beta, z->ordering());
+    }
+  }
+
+  // Clean up permuted arrays (works for both cases)
+  if (xT != x && xT != nullptr) delete xT;
+  if (yT != y && yT != nullptr) delete yT;
+
+  if(realFinalResult != nullptr && realFinalResult != z) {
+    realFinalResult->dataBuffer()->copyBufferFrom(*z->dataBuffer());
+  }
+
+
+}
+}  // namespace sd
+
+#endif

libnd4j/include/helpers/impl/ModularHasher.cpp

@@ -0,0 +1,121 @@
+#include <helpers/ModularHasher.h>
+#include <cstring>
+
+namespace sd {
+namespace helpers {
+namespace detail {
+
+const uint64_t GOLDEN_RATIO = 0x9e3779b97f4a7c15ULL;
+const uint64_t INITIAL_HASH = 14695981039346656037ULL;
+
+// Specialization for uint64_t
+template<> uint64_t SIMDHasher<uint64_t>::hash_chunk(const uint64_t* data, size_t size, uint64_t initial_hash) {
+  uint64_t hash = initial_hash;
+
+#if defined(__ARM_NEON)
+  uint64x2_t hash_vec = vdupq_n_u64(initial_hash);
+  const uint64x2_t golden = vdupq_n_u64(GOLDEN_RATIO);
+
+  for (size_t i = 0; i < size - 1; i += 2) {
+    uint64x2_t val = vld1q_u64(data + i);
+    hash_vec = veorq_u64(hash_vec, val);
+    // Extract lower 32 bits of each 64-bit lane
+    uint32x2_t low_hash = vmovn_u64(hash_vec);
+    uint32x2_t low_golden = vmovn_u64(golden);
+    // Perform 32x32 -> 64 bit widening multiply
+    hash_vec = vmull_u32(low_hash, low_golden);
+  }
+
+  uint64_t tmp[2];
+  vst1q_u64(tmp, hash_vec);
+  hash = tmp[0] ^ tmp[1];
+
+#elif defined(__AVX2__)
+  __m256i hash_vec = _mm256_set1_epi64x(initial_hash);
+  const __m256i golden_vec = _mm256_set1_epi64x(GOLDEN_RATIO);
+
+  for (size_t i = 0; i < size - 3; i += 4) {
+    __m256i val = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i));
+    hash_vec = _mm256_xor_si256(hash_vec, val);
+    hash_vec = _mm256_mul_epi32(hash_vec, golden_vec);
+  }
+
+  uint64_t tmp[4];
+  _mm256_storeu_si256(reinterpret_cast<__m256i*>(tmp), hash_vec);
+  hash = tmp[0] ^ tmp[1] ^ tmp[2] ^ tmp[3];
+
+#elif defined(__SSE4_2__)
+  __m128i hash_vec = _mm_set1_epi64x(initial_hash);
+  const __m128i golden_vec = _mm_set1_epi64x(GOLDEN_RATIO);
+
+  for (size_t i = 0; i < size - 1; i += 2) {
+    __m128i val = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i));
+    hash_vec = _mm_xor_si128(hash_vec, val);
+    hash_vec = _mm_mul_epi32(hash_vec, golden_vec);
+  }
+
+  uint64_t tmp[2];
+  _mm_storeu_si128(reinterpret_cast<__m128i*>(tmp), hash_vec);
+  hash = tmp[0] ^ tmp[1];
+
+#else
+  if(size >= 4) {
+    // Scalar fallback with unrolling
+    for (size_t i = 0; i < size - 3; i += 4) {
+      hash ^= data[i];
+      hash = (hash * GOLDEN_RATIO) ^ (hash >> 32);
+      hash ^= data[i+1];
+      hash = (hash * GOLDEN_RATIO) ^ (hash >> 32);
+      hash ^= data[i+2];
+      hash = (hash * GOLDEN_RATIO) ^ (hash >> 32);
+      hash ^= data[i+3];
+      hash = (hash * GOLDEN_RATIO) ^ (hash >> 32);
+    }
+  }
+
+#endif
+
+  // Handle remaining elements
+  size_t remainder = size % 4;
+  if(size >= 4) {
+    size_t start = size - remainder;
+    for (size_t i = start; i < size; i++) {
+      hash ^= data[i];
+      hash = (hash * GOLDEN_RATIO) ^ (hash >> 32);
+    }
+  }
+  return hash;
+}
+
+// Specialization for double
+uint64_t DataChunkHasher<double>::hash_data(const double* data, size_t size, uint64_t initial_hash) {
+  return SIMDHasher<uint64_t>::hash_chunk(
+      reinterpret_cast<const uint64_t*>(data),
+      size,
+      initial_hash
+  );
+}
+
+uint64_t ModularHasher::combine_hashes(std::initializer_list<uint64_t> hashes) {
+  uint64_t result = INITIAL_HASH;
+  for (uint64_t h : hashes) {
+    result ^= h;
+    result = (result * GOLDEN_RATIO) ^ (result >> 32);
+  }
+  return result;
+}
+
+uint64_t ModularHasher::hash_scalar(uint64_t value, uint64_t initial_hash) {
+  uint64_t hash = initial_hash;
+  hash ^= value;
+  return (hash * GOLDEN_RATIO) ^ (hash >> 32);
+}
+
+// Explicit template instantiations
+template uint64_t ModularHasher::hash_vector<uint64_t>(const std::vector<uint64_t>&, uint64_t);
+template uint64_t ModularHasher::hash_vector<double>(const std::vector<double>&, uint64_t);
+template uint64_t ModularHasher::hash_vector<int64_t>(const std::vector<int64_t>&, uint64_t);
+
+} // namespace detail
+} // namespace helpers
+} // namespace sd

libnd4j/include/helpers/impl/OmpLaunchHelper.cpp

@@ -0,0 +1,103 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 6/30/2018
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/OmpLaunchHelper.h>
+#include <math/templatemath.h>
+#include <system/Environment.h>
+
+namespace sd {
+
+////////////////////////////////////////////////////////////////////////////////
+OmpLaunchHelper::OmpLaunchHelper(const LongType N, float desiredNumThreads) {
+  auto maxItersPerThread = Environment::getInstance().elementwiseThreshold();
+
+  if (N < maxItersPerThread)
+    _numThreads = 1;
+  else {
+#ifdef _OPENMP
+    if (desiredNumThreads == -1)
+      desiredNumThreads = omp_get_max_threads();
+    else if (desiredNumThreads < 1)
+      desiredNumThreads = 1;
+    else
+      desiredNumThreads = sd::math::sd_min<int>(omp_get_max_threads(), desiredNumThreads);
+#else
+    desiredNumThreads = Environment::getInstance().maxThreads();
+#endif
+    _numThreads = sd::math::sd_min<int>(N / maxItersPerThread, desiredNumThreads);
+  }
+
+  _itersPerThread = N / _numThreads;
+  _remainder = N % _numThreads;  // last thread may contain bigger number of iterations
+}
+
+LongType OmpLaunchHelper::betterSpan(LongType N) { return betterSpan(N, betterThreads(N)); }
+
+LongType OmpLaunchHelper::betterSpan(LongType N, LongType numThreads) {
+  auto r = N % numThreads;
+  auto t = N / numThreads;
+
+  if (r == 0)
+    return t;
+  else {
+    // breaks alignment
+    return t + 1;
+  }
+}
+
+int OmpLaunchHelper::betterThreads(LongType N) {
+#ifdef _OPENMP
+  return betterThreads(N, omp_get_max_threads());
+#else
+  return betterThreads(N, Environment::getInstance().maxThreads());
+  ;
+#endif
+}
+
+int OmpLaunchHelper::betterThreads(LongType N, int maxThreads) {
+  auto t = Environment::getInstance().elementwiseThreshold();
+  if (N < t)
+    return 1;
+  else {
+    return static_cast<int>(sd::math::sd_min<LongType>(N / t, maxThreads));
+  }
+}
+
+int OmpLaunchHelper::tadThreads(LongType tadLength, LongType numTads) {
+#ifdef _OPENMP
+  auto maxThreads = omp_get_max_threads();
+#else
+  auto maxThreads = Environment::getInstance().maxThreads();
+#endif
+
+  // if there's only 1 thread allowed - nothing to do here
+  if (maxThreads <= 1) return 1;
+
+  auto totalLength = tadLength * numTads;
+
+  // if array is tiny - no need to spawn any threeds
+  if (totalLength < Environment::getInstance().elementwiseThreshold()) return 1;
+
+  // by default we're spawning as many threads we can, but not more than number of TADs
+  return sd::math::sd_min<int>(numTads, maxThreads);
+}
+}  // namespace sd

libnd4j/include/helpers/impl/OpArgsHolder.cpp

@@ -0,0 +1,141 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com), created on 15.07.2018
+//
+#include <helpers/OpArgsHolder.h>
+
+namespace sd {
+
+////////////////////////////////////////////////////////////////////////
+// default constructor
+OpArgsHolder::OpArgsHolder() {
+  _inArrs = std::vector<NDArray*>();
+  _tArgs = std::vector<double>();
+  _iArgs = std::vector<sd::LongType>();
+  _bArgs = std::vector<bool>();
+
+  _isArrAlloc = std::vector<bool>();
+
+  _numInArrs = 0;
+  _numTArgs = 0;
+  _numIArgs = 0;
+  _numBArgs = 0;
+}
+
+////////////////////////////////////////////////////////////////////////
+// copy constructor
+OpArgsHolder::OpArgsHolder(const OpArgsHolder& other) {
+  THROW_EXCEPTION("OpArgsHolder::OpArgsHolder copy constructor: don't use me !");
+}
+
+////////////////////////////////////////////////////////////////////////
+// constructor
+OpArgsHolder::OpArgsHolder(const std::vector<NDArray*>& inArrs, const std::vector<double>& tArgs,
+                           const std::vector<sd::LongType>& iArgs, const std::vector<bool>& bArgs) {
+  _inArrs = inArrs;
+  _tArgs = tArgs;
+  _iArgs = iArgs;
+  _bArgs = bArgs;
+
+  _isArrAlloc = std::vector<bool>();
+
+  _numInArrs = _inArrs.size();
+  _numTArgs = _tArgs.size();
+  _numIArgs = _iArgs.size();
+  _numBArgs = _bArgs.size();
+}
+
+////////////////////////////////////////////////////////////////////////
+// move constructor
+OpArgsHolder::OpArgsHolder(OpArgsHolder&& other) noexcept
+    : _inArrs(std::move(other._inArrs)),
+      _tArgs(std::move(other._tArgs)),
+      _iArgs(std::move(other._iArgs)),
+      _bArgs(std::move(other._bArgs)),
+      _isArrAlloc(std::move(other._isArrAlloc)) {
+  other._isArrAlloc = std::vector<bool>();
+
+  _numInArrs = _inArrs.size();
+  _numTArgs = _tArgs.size();
+  _numIArgs = _iArgs.size();
+  _numBArgs = _bArgs.size();
+}
+
+////////////////////////////////////////////////////////////////////////
+// assignment operator
+OpArgsHolder& OpArgsHolder::operator=(const OpArgsHolder& other) {
+  return *this;
+}
+
+////////////////////////////////////////////////////////////////////////
+// move assignment operator
+OpArgsHolder& OpArgsHolder::operator=(OpArgsHolder&& other) noexcept {
+  if (this == &other) return *this;
+
+  for (size_t i = 0; i < _isArrAlloc.size(); ++i)  // delete arrays if necessary
+    if (_isArrAlloc[i]) delete _inArrs[i];
+
+  _inArrs = std::move(other._inArrs);
+  _tArgs = std::move(other._tArgs);
+  _iArgs = std::move(other._iArgs);
+  _bArgs = std::move(other._bArgs);
+  _isArrAlloc = std::move(other._isArrAlloc);
+
+  other._isArrAlloc = std::vector<bool>();
+
+  _numInArrs = _inArrs.size();
+  _numTArgs = _tArgs.size();
+  _numIArgs = _iArgs.size();
+  _numBArgs = _bArgs.size();
+
+  return *this;
+}
+
+////////////////////////////////////////////////////////////////////////
+OpArgsHolder OpArgsHolder::createArgsHolderForBP(const std::vector<NDArray*>& inGradArrs, const bool isInPlace) const {
+  const int numInGradArrs = inGradArrs.size();
+
+  OpArgsHolder result(std::vector<NDArray*>(_numInArrs + numInGradArrs, nullptr), _tArgs, _iArgs);
+
+  if (isInPlace) result._isArrAlloc = std::vector<bool>(_numInArrs + numInGradArrs, false);
+
+  for (int i = 0; i < _numInArrs; ++i) {
+    if (isInPlace) {
+      NDArray &arr2 = *_inArrs[i];
+      result._inArrs[i] = new NDArray(arr2);  // make copy
+      result._isArrAlloc[i] = true;
+    } else
+      result._inArrs[i] = _inArrs[i];
+  }
+
+  // input gradients
+  for (int i = 0; i < numInGradArrs; ++i) result._inArrs[_numInArrs + i] = inGradArrs[i];
+
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+// default destructor
+OpArgsHolder::~OpArgsHolder() noexcept {
+  for (size_t i = 0; i < _isArrAlloc.size(); ++i)
+    if (_isArrAlloc[i]) delete _inArrs[i];
+}
+
+}  // namespace sd

libnd4j/include/helpers/impl/OpTracker.cpp

@@ -0,0 +1,115 @@
+/* ******************************************************************************
+ *
+ *
+ * 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/OpTracker.h>
+#include <helpers/logger.h>
+#include <legacy/NativeOps.h>
+
+#include <sstream>
+
+using namespace sd::ops;
+using namespace sd::graph;
+
+namespace sd {
+
+OpTracker& OpTracker::getInstance() {
+  static OpTracker instance;
+  return instance;
+}
+
+void OpTracker::storeOperation(::graph::OpType opType, const OpDescriptor& descriptor) {
+  // check out CPU features
+  if (!isMinimalRequirementsMet()) {
+    auto binaryLevel = ::binaryLevel();
+    auto optimalLevel = ::optimalLevel();
+
+    switch (binaryLevel) {
+      case 3: {
+        sd_printf(
+            "libnd4j binary was built with AVX512 support, but current CPU doesn't have this instruction set. Exiting "
+            "now...",
+            "");
+      } break;
+      case 2: {
+        sd_printf(
+            "libnd4j binary was built with AVX/AVX2 support, but current CPU doesn't have this instruction set. "
+            "Exiting now...",
+            "");
+      } break;
+      default: {
+        sd_printf("Unknown binary validation error. Exiting now...", "");
+      } break;
+    }
+
+    // we're exiting now
+    exit(119);
+  }
+  //
+  if (_map.count(opType) < 1) {
+    std::vector<OpDescriptor> vec;
+    _map[opType] = vec;
+  }
+
+  _operations++;
+
+  auto vec = _map[opType];
+
+  if (std::find(vec.begin(), vec.end(), descriptor) == vec.end()) _map[opType].emplace_back(descriptor);
+}
+
+void OpTracker::storeOperation(::graph::OpType opType, const char* opName, const LongType opNum) {
+  OpDescriptor descriptor(0, opName, false);
+  descriptor.setOpNum((int)opNum);
+  descriptor.setHash(-1);
+
+  storeOperation(opType, descriptor);
+}
+
+template <typename T>
+std::string OpTracker::local_to_string(T value) {
+  std::ostringstream os;
+  os << value;
+  return os.str();
+}
+
+int OpTracker::totalGroups() { return (int)_map.size(); }
+
+int OpTracker::totalOperations() { return _operations; }
+
+const char* OpTracker::exportOperations() {
+  if (_export.length() == 0) {
+    for (auto& v : _map) {
+      std::string block = local_to_string(v.first) + " ";
+
+      for (auto& i : v.second) {
+        block += local_to_string(i.getHash()) + ":";
+        block += local_to_string(i.getOpNum()) + ":";
+        block += *i.getOpName() + "<<";
+      }
+
+      block += ">>";
+      _export += block;
+    }
+  }
+
+  return _export.c_str();
+}
+}  // namespace sd

libnd4j/include/helpers/impl/RandomLauncher.cpp

@@ -0,0 +1,191 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <graph/RandomGenerator.h>
+#include <helpers/RandomLauncher.h>
+#include <types/float16.h>
+#include <helpers/PointersManager.h>
+#include <legacy/NativeOpExecutioner.h>
+namespace sd {
+void RandomLauncher::applyDropOut(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                  double retainProb, NDArray* z) {
+  if (z == nullptr) z = array;
+
+  ExtraArguments arguments({retainProb});
+  PointersManager pm(context, "applyDropOut");
+
+  NDArray::prepareSpecialUse({z}, {array});
+
+  NativeOpExecutioner::execRandom(context, random::DropOut, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), z->buffer(), z->shapeInfo(),
+                                  z->specialBuffer(), z->specialShapeInfo(), arguments.argumentsAsT(z->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({z}, {array});
+}
+
+void RandomLauncher::applyInvertedDropOut(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                          double retainProb, NDArray* z) {
+  if (z == nullptr) z = array;
+
+  ExtraArguments arguments({retainProb});
+  PointersManager pm(context, "applyInvertedDropOut");
+
+  NDArray::prepareSpecialUse({z}, {array});
+
+  NativeOpExecutioner::execRandom(context, random::DropOutInverted, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), z->buffer(), z->shapeInfo(),
+                                  z->specialBuffer(), z->specialShapeInfo(), arguments.argumentsAsT(z->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({z}, {array});
+}
+
+void RandomLauncher::applyAlphaDropOut(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                       double retainProb, double alpha, double beta, double alphaPrime, NDArray* z) {
+  if (z == nullptr) z = array;
+
+  ExtraArguments arguments({retainProb, alpha, beta, alphaPrime});
+  PointersManager pm(context, "applyAlphaDropOut");
+
+  NDArray::prepareSpecialUse({z}, {array});
+
+  NativeOpExecutioner::execRandom(context, random::AlphaDropOut, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), z->buffer(), z->shapeInfo(),
+                                  z->specialBuffer(), z->specialShapeInfo(), arguments.argumentsAsT(z->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({z}, {array});
+}
+
+void RandomLauncher::fillBernoulli(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                   double prob) {
+  ExtraArguments arguments({prob});
+  PointersManager pm(context, "fillBernoulli");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::BernoulliDistribution, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(),
+                                  arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillUniform(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                 double from, double to) {
+  ExtraArguments arguments({from, to});
+  PointersManager pm(context, "fillUniform");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::UniformDistribution, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(),
+                                  arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillGaussian(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                  double mean, double stdev) {
+  ExtraArguments arguments({mean, stdev});
+  PointersManager pm(context, "fillGaussian");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::GaussianDistribution, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), array->buffer(),
+                                  array->shapeInfo(), array->specialBuffer(), array->specialShapeInfo(),
+                                  array->buffer(), array->shapeInfo(), array->specialBuffer(),
+                                  array->specialShapeInfo(), arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillExponential(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                     double lambda) {
+  ExtraArguments arguments({lambda});
+  PointersManager pm(context, "fillExponential");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::ExponentialDistribution, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(),
+                                  arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillLogNormal(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                   double mean, double stdev) {
+  ExtraArguments arguments({mean, stdev});
+  PointersManager pm(context, "fillLogNormal");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::GaussianDistribution, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), array->buffer(),
+                                  array->shapeInfo(), array->specialBuffer(), array->specialShapeInfo(),
+                                  array->buffer(), array->shapeInfo(), array->specialBuffer(),
+                                  array->specialShapeInfo(), arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillTruncatedNormal(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                         double mean, double stdev) {
+  ExtraArguments arguments({mean, stdev});
+  PointersManager pm(context, "fillTruncatedNormal");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(
+      context, random::TruncatedNormalDistribution, &rng, array->buffer(), array->shapeInfo(), array->specialBuffer(),
+      array->specialShapeInfo(), array->buffer(), array->shapeInfo(), array->specialBuffer(), array->specialShapeInfo(),
+      array->buffer(), array->shapeInfo(), array->specialBuffer(), array->specialShapeInfo(),
+      arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+
+void RandomLauncher::fillBinomial(LaunchContext* context, graph::RandomGenerator& rng, NDArray* array,
+                                  int trials, double prob) {
+  ExtraArguments arguments({(double)trials, prob});
+  PointersManager pm(context, "fillBinomial");
+
+  NDArray::prepareSpecialUse({array}, {});
+
+  NativeOpExecutioner::execRandom(context, random::BinomialDistributionEx, &rng, array->buffer(), array->shapeInfo(),
+                                  array->specialBuffer(), array->specialShapeInfo(), array->buffer(),
+                                  array->shapeInfo(), array->specialBuffer(), array->specialShapeInfo(),
+                                  array->buffer(), array->shapeInfo(), array->specialBuffer(),
+                                  array->specialShapeInfo(), arguments.argumentsAsT(array->dataType()));
+  pm.synchronize();
+
+  NDArray::registerSpecialUse({array}, {});
+}
+}  // namespace sd

libnd4j/include/helpers/impl/ShapeBufferCreatorHelper.cpp

@@ -0,0 +1,54 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+#include <helpers/ShapeBufferCreatorHelper.h>
+#include <helpers/ShapeBufferPlatformHelper.h>
+#include <helpers/cpu/CpuShapeBufferCreator.h>
+#include <stdexcept>
+
+namespace sd {
+
+// Initialize static member
+ShapeBufferCreator* ShapeBufferCreatorHelper::currentCreator_ = nullptr;
+
+ShapeBufferCreator& ShapeBufferCreatorHelper::getCurrentCreator() {
+    // This prevents SIGSEGV crash when getCurrentCreator() is called before initialization
+    if (currentCreator_ == nullptr) {
+        // Trigger platform initialization which will call setCurrentCreator()
+        ShapeBufferPlatformHelper::initialize();
+
+        // Double-check after initialization attempt
+        if (currentCreator_ == nullptr) {
+            THROW_EXCEPTION("FATAL: ShapeBufferCreator not initialized! "
+                          "ShapeBufferPlatformHelper::initialize() failed to set currentCreator_. "
+                          "This indicates a critical initialization order bug.");
+        }
+    }
+    return *currentCreator_;
+}
+
+void ShapeBufferCreatorHelper::setCurrentCreator(ShapeBufferCreator* creator) {
+    if (creator == nullptr) {
+        THROW_EXCEPTION("ShapeBufferCreator cannot be null");
+    }
+    currentCreator_ = creator;
+}
+
+} // namespace sd

libnd4j/include/helpers/impl/ShapeBufferPlatformHelper.cpp

@@ -0,0 +1,81 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+#include <helpers/ShapeBufferPlatformHelper.h>
+#include <helpers/cpu/CpuShapeBufferCreator.h>
+#include <mutex>
+
+// Include platform-specific headers conditionally
+#if defined(SD_CUDA)
+#include <helpers/cuda/CudaShapeBufferCreator.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+// Forward declare Environment class if it's used for platform detection
+
+namespace sd {
+
+// This ensures ShapeBufferPlatformHelper is initialized BEFORE DirectShapeTrie or any other
+// code tries to use it. The dummy struct with static member forces initialization at program startup.
+struct ShapeBufferInitializer {
+  ShapeBufferInitializer() {
+    ShapeBufferPlatformHelper::initialize();
+  }
+};
+static ShapeBufferInitializer _force_early_init;
+
+void ShapeBufferPlatformHelper::initialize() {
+  // Thread-safe initialization using static local mutex
+  // This prevents race conditions when multiple threads call initialize() simultaneously
+  static std::mutex init_mutex;
+  static bool init_done = false;
+
+  // Fast path: if already initialized, return immediately without locking
+  if (init_done) {
+    return;
+  }
+
+  // Slow path: acquire lock and check again
+  std::lock_guard<std::mutex> lock(init_mutex);
+  if (init_done) {
+    return;  // Another thread completed initialization while we were waiting
+  }
+
+#if defined(SD_CUDA)
+  printf("Initializing CUDA platform\n");
+  fflush(stdout);
+  // Switch to CUDA implementation
+  ShapeBufferCreatorHelper::setCurrentCreator(&CudaShapeBufferCreator::getInstance());
+#else
+  printf("Initializing CPU platform\n");
+  fflush(stdout);
+
+  ShapeBufferCreatorHelper::setCurrentCreator(&CpuShapeBufferCreator::getInstance());
+
+#endif
+
+  // Mark as complete - must be last line after all initialization
+  init_done = true;
+
+  // Add other platforms as needed (ROCm, OpenCL, etc.)
+}
+
+} // namespace sd

libnd4j/include/helpers/impl/ShapeBuilders.cpp

@@ -0,0 +1,308 @@
+/* ******************************************************************************
+ *
+ *
+ * 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/ShapeBuilders.h>
+
+#include "array/ShapeDescriptor.h"
+
+namespace sd {
+
+LongType* ShapeBuilders::createShapeInfoFrom(ShapeDescriptor* descriptor) {
+  LongType bufferLen = shape::shapeInfoLength(descriptor->rank());
+  auto ret = new LongType[bufferLen];
+  ret[0] = descriptor->rank();
+  if(descriptor->rank() > 0) {
+    shape::setShape(ret, descriptor->shape_strides());
+    shape::setStrideConst(ret, descriptor->stridesPtr());
+    shape::setOrder(ret, descriptor->order());
+  } else {
+    std::vector<LongType> shape = {0};
+    std::vector<LongType> strides = {1};
+    shape::setShape(ret,shape.data());
+    shape::setStrideConst(ret, strides.data());
+    shape::setOrder(ret,'c');
+  }
+
+  shape::setExtra(ret, descriptor->extra());
+  if(ArrayOptions::dataType(ret) != descriptor->dataType()) {
+    ArrayOptions::setDataType(ret, descriptor->dataType());
+  }
+  return ret;
+}
+
+LongType* ShapeBuilders::createScalarShapeInfo(const DataType dataType, memory::Workspace* workspace) {
+  // there is no reason for shape info to use workspaces. we have constant shape helper for this
+  // workspaces with shapebuffers also appears to cause issues when reused elsewhere.
+  LongType lenOfShapeInfo = 6;
+  auto newShape = new LongType[lenOfShapeInfo];
+  newShape[0] = 0;
+  newShape[1] = 0;
+  newShape[2] = 1;
+  newShape[3] = ArrayOptions::setDataTypeValue(ArrayOptions::defaultFlag(), dataType);
+  newShape[4] = 1;
+  newShape[5] = 99;
+
+
+  DataType actualType = ArrayOptions::dataType(newShape);
+  if (actualType != dataType) {
+    printf("ERROR: Data type mismatch in scalarShapeInfo - requested %d but got %d\n",
+           DataTypeUtils::asInt(dataType), DataTypeUtils::asInt(actualType));
+  }
+  return newShape;
+}
+LongType* ShapeBuilders::createVectorShapeInfo(const DataType dataType, const LongType length,
+                                               memory::Workspace* workspace) {
+  //there is no reason for shape info to use workspaces. we have constant shape helper for this
+  // workspaces with shapebuffers also appears to cause issues when reused elsewhere.
+  LongType* newShape = new LongType[shape::shapeInfoLength(static_cast<LongType>(1))];
+
+  newShape[0] = 1;
+  newShape[1] = length;
+  newShape[2] = 1;
+  newShape[3] =  ArrayOptions::setDataTypeValue(ArrayOptions::defaultFlag(), dataType);
+  newShape[4] = 1;
+  newShape[5] = 99;
+  return newShape;
+}
+
+
+LongType* ShapeBuilders::createShapeInfo(const DataType dataType, const char order, int rank,
+                                         const LongType* shapeOnly,
+                                         const LongType *strideOnly,
+                                         memory::Workspace* workspace, sd::LongType extras) {
+  LongType* shapeInfo = nullptr;
+
+  if (rank == 0) {  // scalar case
+    shapeInfo = createScalarShapeInfo(dataType, workspace);
+  } else {
+    shapeInfo = new LongType[shape::shapeInfoLength(rank)];
+
+    // Initialize entire buffer to zero first
+    memset(shapeInfo, 0, shape::shapeInfoLength(rank) * sizeof(LongType));
+
+    shapeInfo[0] = rank;
+
+    // Set shape values
+    for (int i = 0; i < rank; i++) {
+      shapeInfo[i + 1] = shapeOnly[i];
+    }
+
+    // Set stride values
+    for (int i = 0; i < rank; i++) {
+      shapeInfo[i + 1 + rank] = strideOnly[i];
+    }
+
+    // Explicitly set EWS to -1 (unused) at position length-2
+    shapeInfo[shape::shapeInfoLength(rank) - 2] = -1;
+
+    // Set order (at position length-1)
+    shapeInfo[shape::shapeInfoLength(rank) - 1] = order;
+  }
+
+  // The 'extras' parameter may not have data type flags set, which would cause
+  // ArrayOptions::dataType() to return UNKNOWN, triggering validation errors.
+  // We must call setDataType() AFTER setExtra() to ensure the data type is correct.
+  ArrayOptions::setExtra(shapeInfo, extras);
+  ArrayOptions::setDataType(shapeInfo, dataType);  // Ensure data type is set from the dataType parameter
+  shape::setOrder(shapeInfo, order);
+  return shapeInfo;
+}
+
+LongType* ShapeBuilders::copyShapeInfoWithNewType(const LongType* inShapeInfo, const DataType newType) {
+  int rank = shape::rank(inShapeInfo);
+  LongType* newShapeInfo = new LongType[shape::shapeInfoLength(rank)];
+
+  // Copy the basic shape structure
+  memcpy(newShapeInfo, inShapeInfo, shape::shapeInfoByteLength(inShapeInfo));
+
+  // Update the data type while preserving other properties
+  LongType currentExtra = ArrayOptions::extra(inShapeInfo);
+  LongType newExtra = ArrayOptions::setDataTypeValue(
+      ArrayOptions::propertyWithoutDataTypeValue(currentExtra),
+      newType
+  );
+  ArrayOptions::setExtra(newShapeInfo, newExtra);
+
+  return newShapeInfo;
+}
+
+
+
+////////////////////////////////////////////////////////////////////////////////
+LongType  * ShapeBuilders::createShapeInfo(const DataType dataType, const char order, int rank, const LongType* shapeOnly,
+                                           memory::Workspace* workspace, bool empty)  {
+  LongType* shapeInfo = nullptr;
+
+  if (rank == 0) {  // scalar case
+    shapeInfo = createScalarShapeInfo(dataType, workspace);
+  } else {
+    shapeInfo = new LongType[shape::shapeInfoLength(rank)];
+    shapeInfo[0] = rank;
+    for (int i = 0; i < rank; i++) {
+      shapeInfo[i + 1] = shapeOnly[i];
+    }
+
+    ArrayOptions::resetFlags(shapeInfo);
+    shape::updateStrides(shapeInfo, order, false);
+  }
+
+  ArrayOptions::setDataType(shapeInfo, dataType);
+
+  if (empty) {
+    ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  }
+
+  return shapeInfo;
+}
+
+LongType* ShapeBuilders::emptyShapeInfoWithShape(const DataType dataType, std::vector<LongType>& shape,
+                                                 memory::Workspace* workspace) {
+  auto shapeInfo = createShapeInfo(dataType, 'c', shape, workspace);
+  ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  return shapeInfo;
+}
+
+LongType* ShapeBuilders::emptyShapeInfo(const DataType dataType, memory::Workspace* workspace) {
+  auto shapeInfo = createScalarShapeInfo(dataType, workspace);
+  ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  return shapeInfo;
+}
+
+LongType* ShapeBuilders::emptyShapeInfo(const DataType dataType, const char order,
+                                        const std::vector<LongType>& shape, memory::Workspace* workspace) {
+  auto shapeInfo = createShapeInfo(dataType, order, shape.size(), shape.data(), workspace, true);
+  return shapeInfo;
+}
+
+LongType* ShapeBuilders::emptyShapeInfo(const DataType dataType, const char order, int rank,
+                                        const LongType* shapeOnly, memory::Workspace* workspace) {
+  auto shapeInfo2 = new LongType[shape::shapeInfoLength(rank)];
+  shapeInfo2[0] = rank;
+
+  for(int i = 0; i < rank; i++) {
+    shapeInfo2[i + 1] = shapeOnly[i];
+    //all empty strides are zero
+    shapeInfo2[i + 1 + rank] = 0;
+  }
+
+  shape::setOrder(shapeInfo2, order);
+
+
+  ArrayOptions::setPropertyBits(shapeInfo2, {ARRAY_EMPTY,ArrayOptions::flagForDataType(dataType)});
+  return shapeInfo2;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::createShapeInfo(const DataType dataType, const char order,
+                                         const std::vector<LongType>& shapeOnly, memory::Workspace* workspace) {
+  bool isEmpty = false;
+  //shape size 1 but 0 can be scalar
+  if(shapeOnly.size() > 1)
+    for(size_t i = 0; i < shapeOnly.size(); i++) {
+      if(shapeOnly[i] == 0) {
+        isEmpty = true;
+        break;
+      }
+    }
+  auto ret = createShapeInfo(dataType, order, shapeOnly.size(), shapeOnly.data(), workspace, isEmpty);
+  if(isEmpty && !ArrayOptions::hasPropertyBitSet(ret, ARRAY_EMPTY)) {
+    THROW_EXCEPTION("Shape builders: empty was specified was true but shape info returned false");
+  } else if(!isEmpty && ArrayOptions::hasPropertyBitSet(ret, ARRAY_EMPTY)) {
+    THROW_EXCEPTION("Shape builders: empty was specified was false but shape info returned true");
+  }
+  return ret;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::createShapeInfo(const DataType dataType, const char order,
+                                         const std::initializer_list<LongType>& shapeOnly,
+                                         memory::Workspace* workspace) {
+  return createShapeInfo(dataType, order, std::vector<LongType>(shapeOnly), workspace);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::copyShapeInfo(const LongType* inShapeInfo, const bool copyStrides,
+                                       memory::Workspace* workspace) {
+  LongType* outShapeInfo = new LongType[shape::shapeInfoLength(shape::rank(inShapeInfo))];
+  memcpy(outShapeInfo, inShapeInfo, shape::shapeInfoByteLength(inShapeInfo));
+
+  if (!copyStrides) shape::updateStrides(outShapeInfo, shape::order(outShapeInfo), false);
+
+  return outShapeInfo;
+}
+
+
+LongType* ShapeBuilders::setAsView(const LongType* inShapeInfo) {
+  LongType* outShapeInfo = copyShapeInfo(inShapeInfo, true, nullptr);
+  ArrayOptions::toggleIsView(outShapeInfo);
+  return outShapeInfo;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::copyShapeInfoAndType(const LongType* inShapeInfo, const DataType dtype,
+                                              const bool copyStrides, memory::Workspace* workspace) {
+  LongType* outShapeInfo = copyShapeInfo(inShapeInfo, copyStrides, workspace);
+  ArrayOptions::setExtra(outShapeInfo, ArrayOptions::propertyWithoutDataTypeValue(ArrayOptions::extra(inShapeInfo)));  // set extra value to 0 (like in DataTypeEx::TypeEx
+  ArrayOptions::setDataType(outShapeInfo, dtype);
+  return outShapeInfo;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::copyShapeInfoAndType(const LongType* inShapeInfo,
+                                              const LongType* shapeInfoToGetTypeFrom, const bool copyStrides,
+                                              memory::Workspace* workspace) {
+  return copyShapeInfoAndType(inShapeInfo, ArrayOptions::dataType(shapeInfoToGetTypeFrom), copyStrides,
+                              workspace);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+LongType* ShapeBuilders::createSubArrShapeInfo(const LongType* inShapeInfo, const LongType* dims, const int dimsSize,
+                                               memory::Workspace* workspace) {
+  LongType* subArrShapeInfo = nullptr;
+  ALLOCATE(subArrShapeInfo, workspace, shape::shapeInfoLength(dimsSize), LongType);
+
+  subArrShapeInfo[0] = dimsSize;  // rank
+  subArrShapeInfo[2 * dimsSize + 1] = 0;
+  ArrayOptions::copyDataType(subArrShapeInfo, inShapeInfo);   // type
+  subArrShapeInfo[2 * dimsSize + 3] = shape::order(inShapeInfo);  // order
+
+  LongType* shape = shape::shapeOf(subArrShapeInfo);
+  LongType* strides = shape::stride(subArrShapeInfo);
+
+  bool isEmpty = false;
+  for (int i = 0; i < dimsSize; ++i) {
+
+    shape[i] = shape::sizeAt(inShapeInfo, dims[i]);
+    if(shape[i] == 0) {
+      isEmpty = true;
+    }
+    strides[i] = shape::strideAt(inShapeInfo, dims[i]);
+  }
+
+
+
+  shape::checkStridesEwsAndOrder(subArrShapeInfo);
+  if(isEmpty)
+    ArrayOptions::togglePropertyBit(subArrShapeInfo, ARRAY_EMPTY);
+  return subArrShapeInfo;
+}
+
+}  // namespace sd

libnd4j/include/helpers/impl/SimpleReadWriteLock.cpp

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 raver on 8/29/2018.
+//
+#include <helpers/SimpleReadWriteLock.h>
+
+namespace sd {
+SimpleReadWriteLock::SimpleReadWriteLock(const SimpleReadWriteLock& other) {
+  _read_locks.store(other._read_locks.load());
+  _write_locks.store(other._write_locks.load());
+}
+
+SimpleReadWriteLock::SimpleReadWriteLock() {
+  _read_locks.store(0);
+  _write_locks.store(0);
+}
+
+void SimpleReadWriteLock::lockRead() {
+  _mutex.lock();
+  _read_locks++;
+  while (_write_locks.load() > 0) {
+    // just loop
+  }
+  _mutex.unlock();
+}
+
+void SimpleReadWriteLock::unlockRead() { _read_locks--; }
+
+// write lock
+void SimpleReadWriteLock::lockWrite() {
+  _mutex.lock();
+  _write_locks++;
+  while (_read_locks.load() > 0) {
+    // just loop
+  }
+  _mutex.unlock();
+}
+
+void SimpleReadWriteLock::unlockWrite() { _write_locks--; }
+
+SimpleReadWriteLock& SimpleReadWriteLock::operator=(const SimpleReadWriteLock& other) {
+  if (this == &other) return *this;
+
+  this->_write_locks.store(other._write_locks.load());
+  this->_read_locks.store(other._read_locks.load());
+
+  return *this;
+}
+}  // namespace sd

libnd4j/include/helpers/impl/Sqrtm.cpp

@@ -0,0 +1,335 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/EigenValsAndVecs.h>
+#include <helpers/FullPivLU.h>
+#include <helpers/HessenbergAndSchur.h>
+#include <helpers/MmulHelper.h>
+#include <helpers/Sqrtm.h>
+#include <ops/declarable/helpers/lup.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+static void sqrtmQuasiTrianDiag(NDArray& matrixT, NDArray& sqrtT) {
+  const int rows = matrixT.sizeAt(0);
+
+  for (int i = 0; i < rows; i++) {
+    if (i == rows - 1 || matrixT.t<T>(i + 1, i) == (T)0) {
+      const auto elemT = matrixT.t<T>(i, i);
+      if (elemT < (T)0)
+        THROW_EXCEPTION(
+            "ops::helpers::Sqrtm::sqrtmQuasiTrianDiag: can't take sqrt of negative diagonal element of T matrix !");
+      sqrtT.r<T>(i, i) = math::sd_sqrt<T, T>(elemT);
+    } else {
+      NDArray *esViewPtr = matrixT({i, i + 2, i, i + 2}, true);
+      EigenValsAndVecs<T> es(*esViewPtr);  // es._Vecs {2,2,2}, es._Vals{2,2}
+      delete esViewPtr;
+
+       NDArray& vecs = es._Vecs;
+       NDArray& vals = es._Vals;
+
+      const T& vecsReal00 = vecs.t<T>(0, 0, 0);
+      const T& vecsImag00 = vecs.t<T>(0, 0, 1);
+      const T& vecsReal01 = vecs.t<T>(0, 1, 0);
+      const T& vecsImag01 = vecs.t<T>(0, 1, 1);
+      const T& vecsReal10 = vecs.t<T>(1, 0, 0);
+      const T& vecsImag10 = vecs.t<T>(1, 0, 1);
+      const T& vecsReal11 = vecs.t<T>(1, 1, 0);
+      const T& vecsImag11 = vecs.t<T>(1, 1, 1);
+
+      // es.eigenvalues().cwiseSqrt().asDiagonal()
+      T eigenValsSqrt[2][2];
+      eigenValsSqrt[0][0] = vals.t<T>(0, 0);
+      eigenValsSqrt[0][1] = vals.t<T>(0, 1);
+      eigenValsSqrt[1][0] = vals.t<T>(1, 0);
+      eigenValsSqrt[1][1] = vals.t<T>(1, 1);
+      EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[0][0], eigenValsSqrt[0][1]);
+      EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[1][0], eigenValsSqrt[1][1]);
+
+      // es.eigenvectors() * es.eigenvalues().cwiseSqrt().asDiagonal()
+      T vecsElem[2][2][2];
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
+                                               vecsElem[0][0][0], vecsElem[0][0][1]);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
+                                               vecsElem[0][1][0], vecsElem[0][1][1]);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal10, vecsImag10, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
+                                               vecsElem[1][0][0], vecsElem[1][0][1]);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal11, vecsImag11, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
+                                               vecsElem[1][1][0], vecsElem[1][1][1]);
+
+      // es.eigenvectors().inverse()
+      T vecsElemInv[2][2][2];
+
+      T tempReal, tempImag, divisorReal, divisorImag;
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, vecsReal11, vecsImag11, divisorReal,
+                                               divisorImag);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, vecsReal10, vecsImag10, tempReal, tempImag);
+      divisorReal -= tempReal;
+      divisorImag -= tempImag;
+
+      EigenValsAndVecs<T>::divideComplexNums(vecsReal11, vecsImag11, divisorReal, divisorImag, vecsElemInv[0][0][0],
+                                             vecsElemInv[0][0][1]);
+      EigenValsAndVecs<T>::divideComplexNums(-vecsReal01, -vecsImag01, divisorReal, divisorImag, vecsElemInv[0][1][0],
+                                             vecsElemInv[0][1][1]);
+      EigenValsAndVecs<T>::divideComplexNums(-vecsReal10, -vecsImag10, divisorReal, divisorImag, vecsElemInv[1][0][0],
+                                             vecsElemInv[1][0][1]);
+      EigenValsAndVecs<T>::divideComplexNums(vecsReal00, vecsImag00, divisorReal, divisorImag, vecsElemInv[1][1][0],
+                                             vecsElemInv[1][1][1]);
+
+      // result
+      T result[2][2][2];
+
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][0][0],
+                                               vecsElemInv[0][0][1], tempReal, tempImag);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][0][0],
+                                               vecsElemInv[1][0][1], result[0][0][0], result[0][0][1]);
+      result[0][0][0] += tempReal;
+
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][1][0],
+                                               vecsElemInv[0][1][1], tempReal, tempImag);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][1][0],
+                                               vecsElemInv[1][1][1], result[0][1][0], result[0][1][1]);
+      result[0][1][0] += tempReal;
+
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][0][0],
+                                               vecsElemInv[0][0][1], tempReal, tempImag);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][0][0],
+                                               vecsElemInv[1][0][1], result[1][0][0], result[1][0][1]);
+      result[1][0][0] += tempReal;
+
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][1][0],
+                                               vecsElemInv[0][1][1], tempReal, tempImag);
+      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][1][0],
+                                               vecsElemInv[1][1][1], result[1][1][0], result[1][1][1]);
+      result[1][1][0] += tempReal;
+
+      sqrtT.r<T>(i, i) = result[0][0][0];
+      sqrtT.r<T>(i, i + 1) = result[0][1][0];
+      sqrtT.r<T>(i + 1, i) = result[1][0][0];
+      sqrtT.r<T>(i + 1, i + 1) = result[1][1][0];
+
+      ++i;
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+// all matrices are {2,2} here
+template <typename T>
+static void sqrtmQuasiTrianAuxEq(NDArray& A, NDArray& B, NDArray& C, NDArray& X) {
+  std::vector<LongType> tempShape = {4,4};
+  NDArray tempMatrix(A.ordering(),tempShape, A.dataType(), A.getContext());
+
+  tempMatrix.r<T>(0, 0) = A.t<T>(0, 0) + B.t<T>(0, 0);
+  tempMatrix.r<T>(1, 1) = A.t<T>(0, 0) + B.t<T>(1, 1);
+  tempMatrix.r<T>(2, 2) = A.t<T>(1, 1) + B.t<T>(0, 0);
+  tempMatrix.r<T>(3, 3) = A.t<T>(1, 1) + B.t<T>(1, 1);
+  tempMatrix.r<T>(0, 1) = B.t<T>(1, 0);
+  tempMatrix.r<T>(0, 2) = A.t<T>(0, 1);
+  tempMatrix.r<T>(1, 0) = B.t<T>(0, 1);
+  tempMatrix.r<T>(1, 3) = A.t<T>(0, 1);
+  tempMatrix.r<T>(2, 0) = A.t<T>(1, 0);
+  tempMatrix.r<T>(2, 3) = B.t<T>(1, 0);
+  tempMatrix.r<T>(3, 1) = A.t<T>(1, 0);
+  tempMatrix.r<T>(3, 2) = B.t<T>(0, 1);
+  tempMatrix.r<T>(0, 3) = (T)0;
+  tempMatrix.r<T>(1, 2) = (T)0;
+  tempMatrix.r<T>(2, 1) = (T)0;
+  tempMatrix.r<T>(3, 0) = (T)0;
+
+  std::vector<LongType> resultShape = {4,1};
+  NDArray result(A.ordering(), resultShape, A.dataType(), A.getContext());
+  result.r<T>(0, 0) = C.t<T>(0, 0);
+  result.r<T>(1, 0) = C.t<T>(0, 1);
+  result.r<T>(2, 0) = C.t<T>(1, 0);
+  result.r<T>(3, 0) = C.t<T>(1, 1);
+
+  FullPivLU<T>::solve(tempMatrix, result, result);
+
+  X.r<T>(0, 0) = result.t<T>(0);
+  X.r<T>(0, 1) = result.t<T>(1);
+  X.r<T>(1, 0) = result.t<T>(2);
+  X.r<T>(1, 1) = result.t<T>(3);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+static void sqrtmQuasiTrianOffDiag(NDArray& matrixT, NDArray& sqrtT) {
+  const int rows = matrixT.sizeAt(0);
+
+  for (int j = 1; j < rows; j++) {
+    if (matrixT.t<T>(j, j - 1) != (T)0) continue;
+
+    for (int i = j - 1; i >= 0; i--) {
+      if (i > 0 && matrixT.t<T>(i, i - 1) != (T)0) continue;
+
+      const bool iBlockIs2x2 = (i < rows - 1) && (matrixT.t<T>(i + 1, i) != (T)0);
+      const bool jBlockIs2x2 = (j < rows - 1) && (matrixT.t<T>(j + 1, j) != (T)0);
+
+      if (iBlockIs2x2 && jBlockIs2x2) {
+        NDArray *APtr = sqrtT({i, i + 2, i, i + 2}, true);
+        NDArray A = *APtr;
+        delete APtr;
+        
+        NDArray *BPtr = sqrtT({j, j + 2, j, j + 2}, true);
+        NDArray B = *BPtr;
+        delete BPtr;
+        
+        NDArray *XPtr = matrixT({i, i + 2, j, j + 2}, true);
+        NDArray X = *XPtr;
+        delete XPtr;
+
+        if (j - i > 2) {
+          NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
+          NDArray *rightPtr = sqrtT({i + 2, j, j, j + 2}, true);
+          auto mul = mmul(*leftPtr, *rightPtr);
+          X -= *mul;
+          delete leftPtr;
+          delete rightPtr;
+          delete mul;
+        }
+
+        sqrtmQuasiTrianAuxEq<T>(A, B, X, X);
+
+        sqrtT.syncToDevice();
+        NDArray *assignPtr = sqrtT({i, i + 2, j, j + 2}, true);
+        assignPtr->assign(&X);
+        delete assignPtr;
+      } else if (iBlockIs2x2 && !jBlockIs2x2) {
+        NDArray *rhsPtr = matrixT({i, i + 2, j, j + 1}, true);
+        NDArray rhs = *rhsPtr;
+        delete rhsPtr;
+
+        if (j - i > 2) {
+          NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
+          NDArray *rightPtr = sqrtT({i + 2, j, j, j + 1}, true);
+          auto mul = mmul(*leftPtr, *rightPtr);
+          rhs -= *mul;
+          delete leftPtr;
+          delete rightPtr;
+          delete mul;
+        }
+
+        std::vector<LongType> aShape = {2,2};
+        NDArray A(matrixT.ordering(), aShape, matrixT.dataType(), matrixT.getContext());
+        A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(j, j);
+        A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
+        
+        NDArray *addPtr = sqrtT({i, i + 2, i, i + 2}, true);
+        A += *addPtr;
+        delete addPtr;
+
+        FullPivLU<T>::solve(A, rhs, rhs);
+
+        // sqrtT.syncToDevice();
+        NDArray *assignPtr = sqrtT({i, i + 2, j, j + 1}, true);
+        assignPtr->assign(&rhs);
+        delete assignPtr;
+      } else if (!iBlockIs2x2 && jBlockIs2x2) {
+        NDArray *rhsPtr = matrixT({i, i + 1, j, j + 2}, true);
+        NDArray rhs = *rhsPtr;
+        delete rhsPtr;
+
+        if (j - i > 1) {
+          NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j}, true);
+          NDArray *rightPtr = sqrtT({i + 1, j, j, j + 2}, true);
+          auto mul =  mmul(*leftPtr, *rightPtr);
+          rhs -= *mul;
+          delete leftPtr;
+          delete rightPtr;
+          delete mul;
+        }
+
+        std::vector<LongType> aShape = {2,2};
+        NDArray A(matrixT.ordering(),aShape, matrixT.dataType(), matrixT.getContext());
+        A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(i, i);
+        A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
+        
+        NDArray *addPtr = sqrtT({j, j + 2, j, j + 2}, true);
+        NDArray *add = addPtr->transpose();
+        delete addPtr;
+        A += *add;
+        delete add;
+
+        NDArray *rhsT = rhs.transpose();
+        FullPivLU<T>::solve(A, *rhsT, *rhsT);
+
+        // sqrtT.syncToDevice();
+        NDArray *assignPtr = sqrtT({i, i + 1, j, j + 2}, true);
+        assignPtr->assign(&rhs);
+        delete assignPtr;
+        delete rhsT;
+      } else if (!iBlockIs2x2 && !jBlockIs2x2) {
+        NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j});
+        NDArray *rightPtr = sqrtT({i + 1, j, j, j + 1});
+        auto mul = mmul(*leftPtr, *rightPtr);
+        T temp = mul->t<T>(0);  // dot
+        delete leftPtr;
+        delete rightPtr;
+        delete mul;
+        
+        sqrtT.r<T>(i, j) = (matrixT.t<T>(i, j) - temp) / (sqrtT.t<T>(i, i) + sqrtT.t<T>(j, j));
+      }
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Sqrtm<T>::calc(NDArray& in, NDArray& out) {
+  if (in.rankOf() != 2 || in.sizeAt(0) != in.sizeAt(1))
+    THROW_EXCEPTION("ops::helpers::Sqrtm::calc: input matrix must have rank 2 and be square !");
+  if (!out.isSameShape(in))
+    THROW_EXCEPTION("ops::helpers::Sqrtm::calc: output matrix must have the same shape as input one!");
+
+  if (in.lengthOf() == 1) {
+    out.r<T>(0) = math::sd_sqrt<T, T>(in.t<T>(0));
+    return;
+  }
+
+  Schur<T> schur(in);
+
+
+  NDArray *inULike = in.ulike();
+  NDArray sqrtT = *inULike;
+  sqrtT.nullify();
+
+  sqrtmQuasiTrianDiag<T>(*schur.t, sqrtT);
+  sqrtmQuasiTrianOffDiag<T>(*schur.t, sqrtT);
+
+  NDArray *second = schur.u->transpose();
+  // out = U * sqrtT * U^T;
+  NDArray *temp = mmul(sqrtT, *second);
+  MmulHelper::mmul(schur.u, temp, &out);
+  delete inULike;
+  delete second;
+  delete temp;
+}
+
+BUILD_SINGLE_TEMPLATE( class Sqrtm, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/StringUtils.cpp

@@ -0,0 +1,357 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 20/04/18.
+// @author Oleg Semeniv <oleg.semeniv@gmail.com>
+//
+#include <exceptions/datatype_exception.h>
+#include <helpers/BitwiseUtils.h>
+#include <helpers/StringUtils.h>
+
+#include <bitset>
+
+#include "execution/Threads.h"
+#include "helpers/ShapeUtils.h"
+
+namespace sd {
+
+void StringUtils::setValueForDifferentDataType(NDArray* arr, LongType idx, NDArray* input, DataType zType) {
+  switch(zType) {
+#if HAS_UTF8
+    case UTF8: {
+      switch(input->dataType()) {
+        case UTF8:
+          arr->p<std::string>(idx, input->e<std::string>(idx));
+          break;
+        case UTF16:
+          arr->p<std::string>(idx, std::string(input->e<std::u16string>(idx).begin(), input->e<std::u16string>(idx).end()));
+          break;
+        case UTF32:
+          arr->p<std::string>(idx, std::string(input->e<std::u32string>(idx).begin(), input->e<std::u32string>(idx).end()));
+          break;
+        default:
+         THROW_EXCEPTION("Unsupported DataType for source string.");
+      }
+      break;
+    }
+#endif
+#if HAS_UTF16
+    case UTF16: {
+      switch(input->dataType()) {
+        case UTF8:
+          arr->p<std::u16string>(idx, std::u16string(input->e<std::string>(idx).begin(), input->e<std::string>(idx).end()));
+          break;
+        case UTF16:
+          arr->p<std::u16string>(idx, input->e<std::u16string>(idx));
+          break;
+        case UTF32:
+          arr->p<std::u16string>(idx, std::u16string(input->e<std::u32string>(idx).begin(), input->e<std::u32string>(idx).end()));
+          break;
+        default:
+          THROW_EXCEPTION("Unsupported DataType for source string.");
+      }
+      break;
+    }
+#endif
+#if HAS_UTF32
+    case UTF32: {
+      switch(input->dataType()) {
+        case UTF8:
+          arr->p<std::u32string>(idx, std::u32string(input->e<std::string>(idx).begin(), input->e<std::string>(idx).end()));
+          break;
+        case UTF16:
+          arr->p<std::u32string>(idx, std::u32string(input->e<std::u16string>(idx).begin(), input->e<std::u16string>(idx).end()));
+          break;
+        case UTF32:
+          arr->p<std::u32string>(idx, input->e<std::u32string>(idx));
+          break;
+        default:
+          THROW_EXCEPTION("Unsupported DataType for source string.");
+      }
+      break;
+    }
+#endif
+    default:
+      THROW_EXCEPTION("Unsupported DataType for destination string.");
+  }
+}
+
+void StringUtils::broadcastStringAssign(NDArray* x, NDArray* z) {
+  if (!x->isBroadcastableTo(*z)) {
+    THROW_EXCEPTION("Shapes of x and z are not broadcastable.");
+  }
+
+  auto zType = z->dataType();
+  auto xCasted = x->cast(zType);
+
+  std::vector<LongType> zeroVec = {0};
+  std::vector<LongType> *restDims = ShapeUtils::evalDimsToExclude(x->rankOf(), 1, zeroVec.data());
+
+  auto xTensors = xCasted->allTensorsAlongDimension(*restDims);
+  auto zTensors = z->allTensorsAlongDimension(*restDims);
+
+  delete restDims;
+
+  if (xCasted->isScalar()) {
+    for (int e = 0; e < zTensors.size(); e++) {
+      for (int f = 0; f < zTensors.at(e)->lengthOf(); f++) {
+        setValueForDifferentDataType(zTensors.at(e), f, xCasted, zType);
+      }
+    }
+  } else {
+    for (int e = 0; e < xTensors.size(); e++) {
+      auto tensor = xTensors.at(e);
+      for (int f = 0; f < tensor->lengthOf(); f++) {
+        setValueForDifferentDataType(zTensors.at(e), f, tensor, zType);
+      }
+    }
+  }
+}
+
+template <typename T>
+void StringUtils::convertStringsForDifferentDataType(NDArray* sourceArray, NDArray* targetArray) {
+  if (!sourceArray->isS() || !targetArray->isS()) THROW_EXCEPTION("Source or target array is not a string array!");
+
+  int numStrings = sourceArray->isScalar() ? 1 : sourceArray->lengthOf();
+
+  auto inData = sourceArray->bufferAsT<int8_t>() + ShapeUtils::stringBufferHeaderRequirements(sourceArray->lengthOf());
+  auto outData = targetArray->bufferAsT<int8_t>() + ShapeUtils::stringBufferHeaderRequirements(targetArray->lengthOf());
+
+  const auto nInputoffsets = sourceArray->bufferAsT<LongType>();
+  const auto nOutputoffsets = targetArray->bufferAsT<LongType>();
+
+  for (int e = 0; e < numStrings; e++) {
+    auto idata = inData + nInputoffsets[e];
+    auto cdata = outData + nOutputoffsets[e];
+
+    auto start = nInputoffsets[e];
+    auto end = nInputoffsets[e + 1];
+
+    // Convert based on target type (using UTF conversions)
+    if (DataTypeUtils::fromT<T>() == UTF16) {
+      if (sourceArray->dataType() == UTF8) {
+        unicode::utf8to16(idata, cdata, end);
+      } else if(sourceArray->dataType() == UTF32) {
+        unicode::utf32to16(idata, cdata, (end / sizeof(char32_t)));
+      }
+    } else if (DataTypeUtils::fromT<T>() == UTF32) {
+      if (sourceArray->dataType() == UTF8) {
+        unicode::utf8to32(idata, cdata, end);
+      } else if(sourceArray->dataType() == UTF16) {
+        unicode::utf16to32(idata, cdata, (end / sizeof(char16_t)));
+      }
+    } else {
+      if (sourceArray->dataType() == UTF16) {
+        unicode::utf16to8(idata, cdata, (end / sizeof(char16_t)));
+      } else if(sourceArray->dataType() == UTF32) {
+        unicode::utf32to8(idata, cdata, (end / sizeof(char32_t)));
+      }
+    }
+  }
+}
+
+#define DEFINE_CONVERT(T) template void StringUtils::convertStringsForDifferentDataType<GET_SECOND(T)>(NDArray* sourceArray, NDArray* targetArray);
+ITERATE_LIST((SD_STRING_TYPES),DEFINE_CONVERT)
+
+
+template <typename T>
+std::vector<LongType> StringUtils::calculateOffsetsForTargetDataType(NDArray* sourceArray) {
+  if (!sourceArray->isS()) THROW_EXCEPTION("Source array is not a string array!");
+
+  LongType offsetsLength = ShapeUtils::stringBufferHeaderRequirements(sourceArray->lengthOf());
+
+  std::vector<LongType> offsets(sourceArray->lengthOf() + 1);
+
+  const auto nInputoffsets = sourceArray->bufferAsT<LongType>();
+
+  LongType start = 0, stop = 0;
+  LongType dataLength = 0;
+
+  int numStrings = sourceArray->isScalar() ? 1 : sourceArray->lengthOf();
+  auto data = sourceArray->bufferAsT<int8_t>() + offsetsLength;
+  for (LongType e = 0; e < numStrings; e++) {
+    offsets[e] = dataLength;
+    start = nInputoffsets[e];
+    stop = nInputoffsets[e + 1];
+
+    // Determine size difference based on the target type (using UTF conversions)
+    if (sourceArray->dataType() == UTF8) {
+      dataLength += (DataTypeUtils::fromT<T>() == UTF16)
+                        ? unicode::offsetUtf8StringInUtf16(data + start, stop)
+                        : unicode::offsetUtf8StringInUtf32(data + start, stop);
+    } else if (sourceArray->dataType() == UTF16) {
+      dataLength += (DataTypeUtils::fromT<T>() == UTF32)
+                        ? unicode::offsetUtf16StringInUtf32(data + start, (stop / sizeof(char16_t)))
+                        : unicode::offsetUtf16StringInUtf8(data + start, (stop / sizeof(char16_t)));
+    } else if (sourceArray->dataType() == UTF32) {
+      dataLength += (DataTypeUtils::fromT<T>() == UTF16)
+                        ? unicode::offsetUtf32StringInUtf16(data + start, (stop / sizeof(char32_t)))
+                        : unicode::offsetUtf32StringInUtf8(data + start, (stop / sizeof(char32_t)));
+    }
+  }
+
+  offsets[numStrings] = dataLength;
+
+  return offsets;
+}
+#define DEFINE_OFFSET(T) template std::vector<LongType> StringUtils::calculateOffsetsForTargetDataType<GET_SECOND(T)>(NDArray* sourceArray);
+ITERATE_LIST((SD_STRING_TYPES),DEFINE_OFFSET)
+
+static SD_INLINE bool match(const LongType* haystack, const LongType* needle, LongType length) {
+  for (int e = 0; e < length; e++)
+    if (haystack[e] != needle[e]) return false;
+
+  return true;
+}
+
+template <typename T>
+std::string StringUtils::bitsToString(T value) {
+  return std::bitset<sizeof(T) * 8>(value).to_string();
+}
+
+template std::string StringUtils::bitsToString(int value);
+template std::string StringUtils::bitsToString(uint32_t value);
+template std::string StringUtils::bitsToString(LongType value);
+template std::string StringUtils::bitsToString(uint64_t value);
+
+LongType StringUtils::countSubarrays(const void* haystack, LongType haystackLength, const void* needle,
+                                     LongType needleLength) {
+  auto haystack2 = reinterpret_cast<const LongType*>(haystack);
+  auto needle2 = reinterpret_cast<const LongType*>(needle);
+
+  LongType number = 0;
+
+  for (LongType e = 0; e < haystackLength - needleLength; e++) {
+    if (match(&haystack2[e], needle2, needleLength)) number++;
+  }
+
+  return number;
+}
+
+LongType StringUtils::byteLength(NDArray& array) {
+  if (!array.isS())
+    THROW_EXCEPTION(datatype_exception::build("StringUtils::byteLength expects one of String types;", array.dataType()).what());
+
+  auto buffer = array.bufferAsT<LongType>();
+  return buffer[array.lengthOf()];
+}
+
+std::vector<std::string> StringUtils::split(const std::string& haystack, const std::string& delimiter) {
+  std::vector<std::string> output;
+
+  std::string::size_type prev_pos = 0, pos = 0;
+
+  // iterating through the haystack till the end
+  while ((pos = haystack.find(delimiter, pos)) != std::string::npos) {
+    output.emplace_back(haystack.substr(prev_pos, pos - prev_pos));
+    prev_pos = ++pos;
+  }
+
+  output.emplace_back(haystack.substr(prev_pos, pos - prev_pos));  // Last word
+
+  return output;
+}
+
+bool StringUtils::u8StringToU16String(const std::string& u8, std::u16string& u16) {
+  if (u8.empty()) return false;
+
+  u16.resize(unicode::offsetUtf8StringInUtf16(u8.data(), u8.size()) / sizeof(char16_t));
+  if (u8.size() == u16.size())
+    u16.assign(u8.begin(), u8.end());
+  else
+    return unicode::utf8to16(u8.data(), &u16[0], u8.size());
+
+  return true;
+}
+
+bool StringUtils::u8StringToU32String(const std::string& u8, std::u32string& u32) {
+  if (u8.empty()) return false;
+
+  u32.resize(unicode::offsetUtf8StringInUtf32(u8.data(), u8.size()) / sizeof(char32_t));
+  if (u8.size() == u32.size())
+    u32.assign(u8.begin(), u8.end());
+  else
+    return unicode::utf8to32(u8.data(), &u32[0], u8.size());
+
+  return true;
+}
+
+bool StringUtils::u16StringToU32String(const std::u16string& u16, std::u32string& u32) {
+  if (u16.empty()) return false;
+
+  u32.resize(unicode::offsetUtf16StringInUtf32(u16.data(), u16.size()) / sizeof(char32_t));
+  if (u16.size() == u32.size())
+    u32.assign(u16.begin(), u16.end());
+  else
+    return unicode::utf16to32(u16.data(), &u32[0], u16.size());
+
+  return true;
+}
+
+bool StringUtils::u16StringToU8String(const std::u16string& u16, std::string& u8) {
+  if (u16.empty()) return false;
+
+  u8.resize(unicode::offsetUtf16StringInUtf8(u16.data(), u16.size()));
+  if (u16.size() == u8.size())
+    u8.assign(u16.begin(), u16.end());
+  else
+    return unicode::utf16to8(u16.data(), &u8[0], u16.size());
+
+  return true;
+}
+
+bool StringUtils::u32StringToU16String(const std::u32string& u32, std::u16string& u16) {
+  if (u32.empty()) return false;
+
+  u16.resize(unicode::offsetUtf32StringInUtf16(u32.data(), u32.size()) / sizeof(char16_t));
+  if (u32.size() == u16.size())
+    u16.assign(u32.begin(), u32.end());
+  else
+    return unicode::utf32to16(u32.data(), &u16[0], u32.size());
+
+  return true;
+}
+
+bool StringUtils::u32StringToU8String(const std::u32string& u32, std::string& u8) {
+  if (u32.empty()) return false;
+
+  u8.resize(unicode::offsetUtf32StringInUtf8(u32.data(), u32.size()));
+  if (u32.size() == u8.size())
+    u8.assign(u32.begin(), u32.end());
+  else
+    return unicode::utf32to8(u32.data(), &u8[0], u32.size());
+
+  return true;
+}
+
+template <typename T>
+std::string StringUtils::vectorToString(const std::vector<T>& vec) {
+  std::string result;
+  for (auto v : vec) result += valueToString<T>(v);
+
+  return result;
+}
+
+template std::string StringUtils::vectorToString(const std::vector<int>& vec);
+template std::string StringUtils::vectorToString(const std::vector<LongType>& vec);
+template std::string StringUtils::vectorToString(const std::vector<int16_t>& vec);
+template std::string StringUtils::vectorToString(const std::vector<uint32_t>& vec);
+}  // namespace sd

libnd4j/include/helpers/impl/TAD.cpp

@@ -0,0 +1,25 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Adam Gibson
+//
+
+
+
+namespace shape {}

libnd4j/include/helpers/impl/biDiagonalUp.cpp

@@ -0,0 +1,180 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma on 18.12.2017
+//
+
+#include <helpers/biDiagonalUp.h>
+#include <helpers/householder.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+
+BiDiagonalUp::BiDiagonalUp(NDArray& matrix)
+    : _HHmatrix(matrix.dataType(), matrix.getContext(), true),
+      _HHbidiag(matrix.dataType(), matrix.getContext(), true),
+      _hhCoeffs(matrix.dataType(), matrix.getContext(), true) {
+      // input validation
+  if (matrix.rankOf() != 2 || matrix.isScalar())
+    THROW_EXCEPTION("ops::helpers::biDiagonalizeUp constructor: input array must be 2D matrix !");
+
+  std::vector<LongType> shape = {matrix.sizeAt(0), matrix.sizeAt(1)};
+  _HHmatrix = NDArray(matrix.ordering(), shape, matrix.dataType(), matrix.getContext());
+  std::vector<sd::LongType> shape2 = {matrix.sizeAt(1), matrix.sizeAt(1)};
+  _HHbidiag = NDArray(matrix.ordering(),shape2, matrix.dataType(), matrix.getContext());
+  _HHmatrix.assign(&matrix);
+  double zeroAssign = 0.;
+  _HHbidiag.assign(zeroAssign);
+
+  evalData();
+}
+
+template <typename T>
+void BiDiagonalUp::_evalData() {
+  const auto rows = _HHmatrix.sizeAt(0);
+  const auto cols = _HHmatrix.sizeAt(1);
+
+  if (rows < cols)
+    THROW_EXCEPTION(
+        "ops::helpers::BiDiagonalizeUp::evalData method: this procedure is applicable only for input matrix with rows "
+        ">= cols !");
+
+  T coeff, normX;
+
+  T x, y;
+
+  for (LongType i = 0; i < cols - 1; ++i) {
+    // evaluate Householder matrix nullifying columns
+    NDArray *column1Ptr = _HHmatrix({i, rows, i, i + 1});
+    NDArray column1 = *column1Ptr;
+    delete column1Ptr;
+
+    x = _HHmatrix.t<T>(i, i);
+    y = _HHbidiag.t<T>(i, i);
+
+    Householder<T>::evalHHmatrixDataI(column1, x, y);
+
+    _HHmatrix.r<T>(i, i) = x;
+    _HHbidiag.r<T>(i, i) = y;
+
+    // multiply corresponding matrix block on householder matrix from the left: P * bottomRightCorner
+    NDArray *bottomRightCorner1Ptr = _HHmatrix({i, rows, i + 1, cols}, true);  // {i, cols}
+    NDArray bottomRightCorner1 = *bottomRightCorner1Ptr;
+    delete bottomRightCorner1Ptr;
+    
+    NDArray *hhViewPtr = _HHmatrix({i + 1, rows, i, i + 1}, true);
+    Householder<T>::mulLeft(bottomRightCorner1, *hhViewPtr, _HHmatrix.t<T>(i, i));
+    delete hhViewPtr;
+
+    if (i == cols - 2) continue;  // do not apply right multiplying at last iteration
+
+    // evaluate Householder matrix nullifying rows
+    NDArray *row1Ptr = _HHmatrix({i, i + 1, i + 1, cols});
+    NDArray row1 = *row1Ptr;
+    delete row1Ptr;
+
+    x = _HHmatrix.t<T>(i, i + 1);
+    y = _HHbidiag.t<T>(i, i + 1);
+
+    Householder<T>::evalHHmatrixDataI(row1, x, y);
+
+    _HHmatrix.r<T>(i, i + 1) = x;
+    _HHbidiag.r<T>(i, i + 1) = y;
+
+    // multiply corresponding matrix block on householder matrix from the right: bottomRightCorner * P
+    NDArray *bottomRightCorner2Ptr = _HHmatrix({i + 1, rows, i + 1, cols}, true);  // {i, rows}
+    NDArray bottomRightCorner2 = *bottomRightCorner2Ptr;
+    delete bottomRightCorner2Ptr;
+
+    NDArray *hhView2Ptr = _HHmatrix({i, i + 1, i + 2, cols}, true);
+    Householder<T>::mulRight(bottomRightCorner2, *hhView2Ptr, _HHmatrix.t<T>(i, i + 1));
+    delete hhView2Ptr;
+  }
+
+  NDArray *row2Ptr = _HHmatrix({cols - 2, cols - 1, cols - 1, cols});
+  NDArray row2 = *row2Ptr;
+  delete row2Ptr;
+
+  x = _HHmatrix.t<T>(cols - 2, cols - 1);
+  y = _HHbidiag.t<T>(cols - 2, cols - 1);
+
+  Householder<T>::evalHHmatrixDataI(row2, x, y);
+
+  _HHmatrix.r<T>(cols - 2, cols - 1) = x;
+  _HHbidiag.r<T>(cols - 2, cols - 1) = y;
+
+  NDArray *column2Ptr = _HHmatrix({cols - 1, rows, cols - 1, cols});
+  NDArray column2 = *column2Ptr;
+  delete column2Ptr;
+
+  x = _HHmatrix.t<T>(cols - 1, cols - 1);
+  y = _HHbidiag.t<T>(cols - 1, cols - 1);
+
+  Householder<T>::evalHHmatrixDataI(column2, x, y);
+
+  _HHmatrix.r<T>(cols - 1, cols - 1) = x;
+  _HHbidiag.r<T>(cols - 1, cols - 1) = y;
+}
+
+//////////////////////////////////////////////////////////////////////////
+void BiDiagonalUp::evalData() {
+  auto xType = _HHmatrix.dataType();
+  BUILD_SINGLE_SELECTOR(xType, _evalData, ();, SD_FLOAT_TYPES);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+HHsequence BiDiagonalUp::makeHHsequence_(const char type) {
+  const int diagSize = type == 'u' ? _HHbidiag.sizeAt(0) : _HHbidiag.sizeAt(0) - 1;
+
+  std::vector<LongType> shape = {diagSize};
+  _hhCoeffs = NDArray(_HHmatrix.ordering(),shape, _HHmatrix.dataType(), _HHmatrix.getContext());
+
+  if (type == 'u')
+    for (int i = 0; i < diagSize; ++i) _hhCoeffs.r<T>(i) = _HHmatrix.t<T>(i, i);
+  else
+    for (int i = 0; i < diagSize; ++i) _hhCoeffs.r<T>(i) = _HHmatrix.t<T>(i, i + 1);
+
+  HHsequence result(&_HHmatrix, &_hhCoeffs, type);
+
+  if (type != 'u') {
+    result._diagSize = diagSize;
+    result._shift = 1;
+  }
+
+  return result;
+}
+
+//////////////////////////////////////////////////////////////////////////
+HHsequence BiDiagonalUp::makeHHsequence(const char type) {
+  auto xType = _HHmatrix.dataType();
+  BUILD_SINGLE_SELECTOR(xType, return makeHHsequence_, (type);, SD_FLOAT_TYPES);
+  // This should never be reached - BUILD_SINGLE_SELECTOR covers all SD_FLOAT_TYPES
+  THROW_EXCEPTION("BiDiagonalUp::makeHHsequence: unsupported data type");
+}
+
+BUILD_SINGLE_TEMPLATE( void BiDiagonalUp::_evalData, (), SD_FLOAT_TYPES);
+BUILD_SINGLE_TEMPLATE( HHsequence BiDiagonalUp::makeHHsequence_, (const char type), SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/generic/TypedTrie.cpp

@@ -0,0 +1,253 @@
+/* ******************************************************************************
+ *
+ * Copyright (c) 2024 Konduit K.K.
+ * 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.
+ *
+ * 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
+ ******************************************************************************/
+
+#include <helpers/generic/TypedTrie.h>
+
+#include "array/ConstantShapeBuffer.h"
+#include "array/TadPack.h"
+#include "helpers/DirectTadTrie.h"
+
+namespace sd {
+namespace generic {
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+struct TypedTrie<KeyType, ValueType, NUM_STRIPES>::TrieNode {
+  std::shared_ptr<ValueType> value;
+  std::unordered_map<typename KeyType::value_type, std::unique_ptr<TrieNode>> children;
+  std::atomic<bool> isComplete{false};
+  std::chrono::steady_clock::time_point lastAccess;
+
+  TrieNode() : lastAccess(std::chrono::steady_clock::now()) {}
+};
+
+
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+TypedTrie<KeyType, ValueType, NUM_STRIPES>::TypedTrie()
+ {
+#ifdef __cpp_exceptions
+  try {
+    for (auto& root : _roots) {
+      root = std::make_unique<TrieNode>();
+      _resourceManager.registerNode();
+    }
+  } catch (...) {
+    cleanup();
+    throw;
+  }
+#else
+  // Exceptions disabled - direct initialization without try/catch
+  for (auto& root : _roots) {
+    root = std::make_unique<TrieNode>();
+    _resourceManager.registerNode();
+  }
+#endif
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+TypedTrie<KeyType, ValueType, NUM_STRIPES>::~TypedTrie() {
+
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+size_t TypedTrie<KeyType, ValueType, NUM_STRIPES>::getStripeIndex(const KeyType& key) const {
+  size_t h = 0;
+  for (const auto& elem : key) {
+    h = h * 31 + std::hash<typename KeyType::value_type>{}(elem);
+  }
+  return h & (NUM_STRIPES - 1);
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+typename TypedTrie<KeyType, ValueType, NUM_STRIPES>::TrieNode*
+TypedTrie<KeyType, ValueType, NUM_STRIPES>::findOrCreateNode(TrieNode* root,
+                                                             const KeyType& key,
+                                                             bool createIfMissing) const {
+  if (!root) return nullptr;
+
+  TrieNode* current = root;
+  for (const auto& k : key) {
+    auto it = current->children.find(k);
+    if (it == current->children.end()) {
+      if (!createIfMissing) return nullptr;
+      auto newNode = std::make_unique<TrieNode>();
+      if (!newNode) return nullptr;
+      current->children[k] = std::move(newNode);
+    }
+    current = current->children[k].get();
+    if (!current) return nullptr;
+  }
+  return current;
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+void TypedTrie<KeyType, ValueType, NUM_STRIPES>::cleanupNode(TrieNode* node) {
+  if (!node) return;
+
+  auto now = std::chrono::steady_clock::now();
+  auto age = std::chrono::duration_cast<std::chrono::minutes>(
+      now - node->lastAccess).count();
+
+  // Clean up nodes older than 30 minutes
+  if (age > 30) {
+    node->value.reset();
+
+    for (auto it = node->children.begin(); it != node->children.end();) {
+      auto* child = it->second.get();
+      if (child) {
+        cleanupNode(child);
+        if (!child->value && child->children.empty()) {
+          _resourceManager.unregisterNode();
+          it = node->children.erase(it);
+          continue;
+        }
+      }
+      ++it;
+    }
+  }
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+std::shared_ptr<ValueType> TypedTrie<KeyType, ValueType, NUM_STRIPES>::get(const KeyType& key) const {
+  if (!key.empty()) {
+    auto scope = _resourceManager.createScope();
+    size_t stripe = getStripeIndex(key);
+
+    _locks.lockStripe(stripe, false);
+    auto node = findOrCreateNode(_roots[stripe].get(), key, false);
+    if (node && node->isComplete.load(std::memory_order_acquire)) {
+      auto result = node->value;
+      if (result) {
+        node->lastAccess = std::chrono::steady_clock::now();
+        _locks.unlockStripe(stripe, false);
+        return result;
+      }
+    }
+    _locks.unlockStripe(stripe, false);
+  }
+  return nullptr;
+}
+
+
+// Previously instantiated TypedTrie with ConstantShapeBuffer* (raw pointer) as ValueType
+// This caused undefined behavior by mixing raw pointer storage with shared_ptr returns
+// The correct instantiation uses std::shared_ptr<ConstantShapeBuffer> (line 246)
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+bool TypedTrie<KeyType, ValueType, NUM_STRIPES>::insert(const KeyType& key,
+                                                        std::shared_ptr<ValueType> value) {
+  if (!value || key.empty()) return false;
+
+  auto scope = _resourceManager.createScope();
+  size_t stripe = getStripeIndex(key);
+
+  _locks.lockStripe(stripe, true);
+  auto node = findOrCreateNode(_roots[stripe].get(), key, true);
+  if (!node) {
+    _locks.unlockStripe(stripe, true);
+    return false;
+  }
+
+  if (node->value || node->isComplete.load(std::memory_order_acquire)) {
+    _locks.unlockStripe(stripe, true);
+    return false;
+  }
+
+  node->value = value;
+  node->lastAccess = std::chrono::steady_clock::now();
+  node->isComplete.store(true, std::memory_order_release);
+  _locks.unlockStripe(stripe, true);
+  return true;
+}
+
+// template bool sd::generic::TypedTrie<..., ConstantShapeBuffer*, ...>::insert(...)
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+bool TypedTrie<KeyType, ValueType, NUM_STRIPES>::remove(const KeyType& key) {
+  auto scope = _resourceManager.createScope();
+  size_t stripe = getStripeIndex(key);
+
+  _locks.lockStripe(stripe, true);
+
+  auto node = findOrCreateNode(_roots[stripe].get(), key, false);
+  if (!node || !node->value) {
+    _locks.unlockStripe(stripe, true);
+    return false;
+  }
+
+  node->value.reset();
+  node->isComplete.store(false, std::memory_order_release);
+
+  _locks.unlockStripe(stripe, true);
+  return true;
+}
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+void TypedTrie<KeyType, ValueType, NUM_STRIPES>::cleanup() {
+  auto scope = _resourceManager.createScope();
+  std::vector<size_t> stripes;
+  for (size_t i = 0; i < NUM_STRIPES; ++i) {
+    stripes.push_back(i);
+  }
+  auto guard = _locks.acquireMultiLock(stripes, true);
+
+  for (auto& root : _roots) {
+    if (root) cleanupNode(root.get());
+  }
+}
+
+// template void sd::generic::TypedTrie<..., ConstantShapeBuffer*, ...>::cleanup();
+
+template void
+sd::generic::TypedTrie<std::vector<long long, std::allocator<long long>>,
+    sd::TadPack*,
+    32>::cleanup();
+
+// template sd::generic::TypedTrie<..., ConstantShapeBuffer*, ...>::~TypedTrie();
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES>
+typename TypedTrie<KeyType, ValueType, NUM_STRIPES>::Stats
+TypedTrie<KeyType, ValueType, NUM_STRIPES>::getStats() const {
+  Stats stats;
+  for (size_t i = 0; i < NUM_STRIPES; ++i) {
+    stats.stripeCounts[i] = _locks.getStripeCount(i);
+  }
+  return stats;
+}
+
+}  // namespace generic
+}  // namespace sd
+
+template class sd::generic::TypedTrie<std::vector<unsigned char>, std::shared_ptr<sd::ConstantShapeBuffer>, 32>;
+template class sd::generic::TypedTrie<std::vector<sd::LongType>, std::shared_ptr<sd::TadPack>, 32>;
+
+// NOTE: TadPack raw pointer instantiations kept for now (may need similar cleanup in future)
+template std::shared_ptr<sd::TadPack*>
+sd::generic::TypedTrie<std::vector<long long, std::allocator<long long>>,
+                       sd::TadPack*,
+                       32>::get(const std::vector<long long, std::allocator<long long>>& key) const;
+template bool
+sd::generic::TypedTrie<std::vector<long long, std::allocator<long long>>,
+                       sd::TadPack*,
+                       32>::insert(const std::vector<long long, std::allocator<long long>>& key,
+                                   std::shared_ptr<sd::TadPack*> value);
+template
+    sd::generic::TypedTrie<std::vector<long long, std::allocator<long long>>,
+                           sd::TadPack*,
+                           32>::~TypedTrie();
+
+// template sd::generic::TypedTrie<..., ConstantShapeBuffer*, ...>::TypedTrie();
+// This was mixing raw pointer storage with shared_ptr-based lifecycle management

libnd4j/include/helpers/impl/helper_hash.cpp

@@ -0,0 +1,71 @@
+/* ******************************************************************************
+ *
+ *
+ * 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/helper_hash.h>
+#include <helpers/logger.h>
+
+namespace sd {
+namespace ops {
+
+HashHelper& HashHelper::getInstance() {
+  static HashHelper instance;
+  return instance;
+}
+
+LongType HashHelper::getLongHash(std::string& str) {
+  _locker.lock();
+  if (!_isInit) {
+    sd_verbose("Building HashUtil table\n", "");
+
+    unsigned long long h = 0x544B2FBACAAF1684L;
+    for (int i = 0; i < 256; i++) {
+      for (int j = 0; j < 31; j++) {
+        h = (((unsigned long long)h) >> 7) ^ h;
+        h = (h << 11) ^ h;
+        h = (((unsigned long long)h) >> 10) ^ h;
+      }
+      _byteTable[i] = h;
+    }
+
+    _isInit = true;
+  }
+
+  _locker.unlock();
+
+  //note: DO NOT change this type.
+  //when something like thread sanitizer + cuda is used
+  //the offsets can get absurdly big.
+  //you get errors like: left shift of 11 places cannot be represented in type
+  unsigned long long h = HSTART;
+  unsigned long long hmult = HMULT;
+
+  LongType len = str.size();
+  for (int i = 0; i < len; i++) {
+    char ch = str.at(i);
+    auto uch = (unsigned char)ch;
+    h = (h * hmult) ^ _byteTable[ch & 0xff];
+    h = (h * hmult) ^ _byteTable[(uch >> 8) & 0xff];
+  }
+
+  return h;
+}
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/hhColPivQR.cpp

@@ -0,0 +1,175 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma on 11.01.2018
+//
+#include <helpers/hhColPivQR.h>
+#include <helpers/householder.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+HHcolPivQR::HHcolPivQR(NDArray &matrix) {
+  _qr = matrix.dup(matrix.ordering());
+  std::vector<LongType> coeffsShape = {1,_diagSize};
+  _diagSize = math::sd_min<int>(matrix.sizeAt(0), matrix.sizeAt(1));
+  std::vector<LongType> permShape = {matrix.sizeAt(1), matrix.sizeAt(1)};
+  _coeffs = new NDArray(matrix.ordering(),coeffsShape, matrix.dataType(), matrix.getContext());
+
+  _permut = new NDArray(matrix.ordering(), permShape, matrix.dataType(), matrix.getContext());
+
+  evalData();
+}
+
+void HHcolPivQR::evalData() { BUILD_SINGLE_SELECTOR(_qr->dataType(), _evalData, (), SD_FLOAT_TYPES); }
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void HHcolPivQR::_evalData() {
+  const int rows = _qr->sizeAt(0);
+  const int cols = _qr->sizeAt(1);
+
+  std::vector<LongType> colsShape = {cols};
+  NDArray transp(_qr->ordering(), colsShape, _qr->dataType(), _qr->getContext());
+  NDArray normsUpd(_qr->ordering(), colsShape , _qr->dataType(), _qr->getContext());
+  NDArray normsDir(_qr->ordering(),colsShape , _qr->dataType(), _qr->getContext());
+
+  auto qRDeRef = *_qr;
+  for (int k = 0; k < cols; ++k) {
+    NDArray *colViewPtr = qRDeRef({0, 0, k, k + 1});
+    auto norm =  colViewPtr->reduceNumber(reduce::Norm2);
+    normsDir.r<T>(k) = normsUpd.r<T>(k) = norm->t<T>(0);
+    delete norm;
+    delete colViewPtr;
+  }
+
+  auto max = (normsUpd.reduceNumber(reduce::Max));
+  T normScaled = max->t<T>(0) * DataTypeUtils::eps<T>();
+  T threshold1 = normScaled * normScaled / (T)rows;
+  T threshold2 = math::sd_sqrt<T, T>(DataTypeUtils::eps<T>());
+
+  T nonZeroPivots = static_cast<T>(_diagSize);
+  T maxPivot = static_cast<T>(0.);
+ delete max;
+  for (int k = 0; k < _diagSize; ++k) {
+    NDArray *normsUpdViewPtr = normsUpd({k, -1});
+    NDArray *indexNum = normsUpdViewPtr->indexReduceNumber(indexreduce::IndexMax);
+    int biggestColIndex = indexNum->e<int>(0);
+    auto max2 = normsUpdViewPtr->reduceNumber(reduce::Max);
+    T biggestColNorm = max2->t<T>(0);
+    delete normsUpdViewPtr;
+    delete max2;
+    
+    T biggestColSqNorm = biggestColNorm * biggestColNorm;
+    biggestColIndex += k;
+
+    if (nonZeroPivots == (T)_diagSize && biggestColSqNorm < threshold1 * (T)(rows - k)) nonZeroPivots = k;
+
+    transp.r<T>(k) = (T)biggestColIndex;
+
+    if (k != biggestColIndex) {
+      NDArray *temp1Ptr = qRDeRef({0, 0, k, k + 1});
+      NDArray temp1 = *temp1Ptr;
+      delete temp1Ptr;
+      
+      NDArray *temp2Ptr = qRDeRef({0, 0, biggestColIndex, biggestColIndex + 1});
+      NDArray temp2 = *temp2Ptr;
+      delete temp2Ptr;
+      
+      temp1.swapUnsafe(temp2);
+
+      math::sd_swap<T>(normsUpd.r<T>(k), normsUpd.r<T>(biggestColIndex));
+      math::sd_swap<T>(normsDir.r<T>(k), normsDir.r<T>(biggestColIndex));
+    }
+
+    T normX, c;
+    NDArray *qrBlockPtr = qRDeRef({k, rows, k, k + 1});
+    NDArray qrBlock = *qrBlockPtr;
+    delete qrBlockPtr;
+    
+    Householder<T>::evalHHmatrixDataI(qrBlock, c, normX);
+
+    _coeffs->r<T>(k) = c;
+
+    _qr->r<T>(k, k) = normX;
+
+    T max = math::sd_abs<T,T>(normX);
+    if (max > maxPivot) maxPivot = max;
+
+    if (k < rows && (k + 1) < cols) {
+      NDArray *qrBlock2Ptr = qRDeRef({k, rows, k + 1, cols}, true);
+      NDArray qrBlock2 = *qrBlock2Ptr;
+      delete qrBlock2Ptr;
+      
+      NDArray *tailPtr = qRDeRef({k + 1, rows, k, k + 1}, true);
+      NDArray tail = *tailPtr;
+      delete tailPtr;
+      
+      Householder<T>::mulLeft(qrBlock2, tail, _coeffs->t<T>(k));
+    }
+
+    for (int j = k + 1; j < cols; ++j) {
+      if (normsUpd.t<T>(j) != (T)0.f) {
+        T temp = math::sd_abs<T,T>(_qr->t<T>(k, j)) / normsUpd.t<T>(j);
+        temp = ((T)1. + temp) * ((T)1. - temp);
+        temp = temp < (T)0. ? (T)0. : temp;
+        T temp2 = temp * normsUpd.t<T>(j) * normsUpd.t<T>(j) / (normsDir.t<T>(j) * normsDir.t<T>(j));
+
+        if (temp2 <= threshold2) {
+          if (k + 1 < rows && j < cols) {
+            NDArray *normViewPtr = qRDeRef({k + 1, rows, j, j + 1});
+            auto reduce =  normViewPtr->reduceNumber(reduce::Norm2);
+            normsDir.r<T>(j) = reduce->t<T>(0);
+            delete normViewPtr;
+            delete reduce;
+          }
+
+          normsUpd.r<T>(j) = normsDir.t<T>(j);
+        } else
+          normsUpd.r<T>(j) = normsUpd.t<T>(j) * math::sd_sqrt<T, T>(temp);
+      }
+    }
+
+    delete indexNum;
+  }
+
+  _permut->setIdentity();
+
+  auto permuteRef = *_permut;
+  for (int k = 0; k < _diagSize; ++k) {
+    int idx = transp.e<int>(k);
+    NDArray *temp1Ptr = permuteRef({0, 0, k, k + 1});
+    NDArray temp1 = *temp1Ptr;
+    delete temp1Ptr;
+    
+    NDArray *temp2Ptr = permuteRef({0, 0, idx, idx + 1});
+    NDArray temp2 = *temp2Ptr;
+    delete temp2Ptr;
+    
+    temp1.swapUnsafe(temp2);
+  }
+}
+
+BUILD_SINGLE_TEMPLATE( void HHcolPivQR::_evalData, (), SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/hhSequence.cpp

@@ -0,0 +1,137 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma on 02.01.2018
+//
+#include <helpers/hhSequence.h>
+#include <helpers/householder.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+HHsequence::HHsequence(NDArray* vectors, NDArray* coeffs, const char type)
+    : _vectors(vectors), _coeffs(coeffs) {
+  _diagSize = math::sd_min(_vectors->sizeAt(0), _vectors->sizeAt(1));
+  _shift = 0;
+  _type = type;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void HHsequence::mulLeft_(NDArray* matrix) {
+  const int rows = _vectors->sizeAt(0);
+  const int cols = _vectors->sizeAt(1);
+  const int inRows = matrix->sizeAt(0);
+  NDArray matrixRef = *matrix;
+  NDArray vectorsRef =  *_vectors;
+  for (int i = _diagSize - 1; i >= 0; --i) {
+    if (_type == 'u') {
+      NDArray *blockPtr = matrixRef({inRows - rows + _shift + i, inRows, 0, 0}, true);
+      NDArray block = *blockPtr;
+      
+      NDArray *vectorPtr = vectorsRef({i + 1 + _shift, rows, i, i + 1}, true);
+      NDArray vector = *vectorPtr;
+      
+      Householder<T>::mulLeft(block, vector, _coeffs->t<T>(i));
+      
+      delete blockPtr;
+      delete vectorPtr;
+    } else {
+      NDArray *blockPtr = matrixRef({inRows - cols + _shift + i, inRows, 0, 0}, true);
+      NDArray block = *blockPtr;
+      
+      NDArray *vectorPtr = vectorsRef({i, i + 1, i + 1 + _shift, cols}, true);
+      NDArray vector = *vectorPtr;
+      
+      Householder<T>::mulLeft(block, vector, _coeffs->t<T>(i));
+      
+      delete blockPtr;
+      delete vectorPtr;
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+NDArray HHsequence::getTail(const int idx) const {
+  int first = idx + 1 + _shift;
+  NDArray vectorsRef =  *_vectors;
+
+  if (_type == 'u') {
+    NDArray *tailPtr = vectorsRef({first, -1, idx, idx + 1}, true);
+    NDArray tail = *tailPtr;
+    delete tailPtr;
+    return tail;
+  } else {
+    NDArray *tailPtr = vectorsRef({idx, idx + 1, first, -1}, true);
+    NDArray tail = *tailPtr;
+    delete tailPtr;
+    return tail;
+  }
+}
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void HHsequence::applyTo_(NDArray* dest) {
+  int size = _type == 'u' ? _vectors->sizeAt(0) : _vectors->sizeAt(1);
+ NDArray *originalDest = dest;
+ NDArray destRef = *dest;
+  std::vector<LongType> sizeShape = {size,size};
+  if (dest->rankOf() != 2 || (dest->sizeAt(0) != size && dest->sizeAt(1) != size)) {
+    dest = new NDArray(dest->ordering(), sizeShape, dest->dataType(), dest->getContext());
+    destRef = *dest;
+  }
+  dest->setIdentity();
+
+  for (int k = _diagSize - 1; k >= 0; --k) {
+    int curNum = size - k - _shift;
+    if (curNum < 1 || (k + 1 + _shift) >= size) continue;
+    
+    NDArray *blockPtr = destRef({dest->sizeAt(0) - curNum, dest->sizeAt(0), dest->sizeAt(1) - curNum, dest->sizeAt(1)}, true);
+    NDArray block = *blockPtr;
+
+    NDArray tailK = getTail(k);
+    Householder<T>::mulLeft(block, tailK, _coeffs->t<T>(k));
+    
+    delete blockPtr;
+  }
+
+  if(originalDest != dest) {
+    delete dest;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+void HHsequence::applyTo(NDArray* dest) {
+  auto xType = _coeffs->dataType();
+  BUILD_SINGLE_SELECTOR(xType, applyTo_, (dest), SD_FLOAT_TYPES);
+}
+
+//////////////////////////////////////////////////////////////////////////
+void HHsequence::mulLeft(NDArray* matrix) {
+  auto xType = _coeffs->dataType();
+  BUILD_SINGLE_SELECTOR(xType, mulLeft_, (matrix), SD_FLOAT_TYPES);
+}
+
+BUILD_SINGLE_TEMPLATE( void HHsequence::applyTo_, (sd::NDArray * dest), SD_FLOAT_TYPES);
+BUILD_SINGLE_TEMPLATE( void HHsequence::mulLeft_, (NDArray * matrix), SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/householder.cpp

@@ -0,0 +1,230 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma on 18.12.2017
+//
+#include <helpers/householder.h>
+
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Householder<T>::evalHHmatrixData(NDArray& x, NDArray& tail, T& coeff, T& normX) {
+  // input validation
+  if (x.rankOf() != 1 && !x.isScalar())
+    THROW_EXCEPTION(
+        "ops::helpers::Householder::evalHHmatrixData method: input array must have rank = 1 or to be scalar!");
+
+  if (!x.isScalar() && x.lengthOf() != tail.lengthOf() + 1)
+    THROW_EXCEPTION(
+        "ops::helpers::Householder::evalHHmatrixData method: input tail vector must have length less than unity "
+        "compared to input x vector!");
+
+  const auto xLen = x.lengthOf();
+
+  NDArray *xTail = xLen > 1 ? x({1, -1}) : nullptr;
+
+  T tailXnorm;
+  if (xLen > 1) {
+    auto* tailNormPtr = xTail->reduceNumber(reduce::SquaredNorm);
+    tailXnorm = tailNormPtr->t<T>(0);
+    delete tailNormPtr;
+  } else {
+    tailXnorm = (T)0;
+  }
+
+  const auto xFirstElem = x.t<T>(0);
+
+  if (tailXnorm <= DataTypeUtils::min_positive<T>()) {
+    normX = xFirstElem;
+    coeff = (T)0.f;
+    tail = (T)0.f;
+  } else {
+    normX = math::sd_sqrt<T, T>(xFirstElem * xFirstElem + tailXnorm);
+
+    if (xFirstElem >= (T)0.f) normX = -normX;  // choose opposite sign to lessen roundoff error
+
+    coeff = (normX - xFirstElem) / normX;
+    T divisor = xFirstElem - normX;
+    NDArray *tailAssign = (*xTail) / divisor;
+    tail.assign(tailAssign);
+    delete tailAssign;
+  }
+
+  if (xTail != nullptr) delete xTail;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Householder<T>::evalHHmatrixDataI(NDArray& x, T& coeff, T& normX) {
+  // input validation
+  if (x.rankOf() != 1 && !x.isScalar())
+    THROW_EXCEPTION(
+        "ops::helpers::Householder::evalHHmatrixDataI method: input array must have rank = 1 or to be scalar!");
+
+  int rows = (int)x.lengthOf() - 1;
+  int num = 1;
+
+  if (rows == 0) {
+    rows = 1;
+    num = 0;
+  }
+
+  NDArray *tailPtr = x({num, -1});
+  NDArray tail = *tailPtr;
+  delete tailPtr;
+
+  evalHHmatrixData(x, tail, coeff, normX);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Householder<T>::mulLeft(NDArray& matrix, NDArray& tail, const T coeff) {
+  if (matrix.sizeAt(0) == 1 && coeff != (T)0) {
+    NDArray *scaledResult = matrix * ((T)1.f - coeff);
+    matrix.assign(scaledResult);
+    delete scaledResult;
+  } else if (coeff != (T)0.f) {
+    NDArray *bottomPartPtr = matrix({1, matrix.sizeAt(0), 0, 0}, true);
+    NDArray bottomPart = *bottomPartPtr;
+    delete bottomPartPtr;
+    
+    NDArray *fistRowPtr = matrix({0, 1, 0, 0}, true);
+    NDArray fistRow = *fistRowPtr;
+    delete fistRowPtr;
+    
+    NDArray *tailTranspose = tail.transpose();
+    if (tail.isColumnVector()) {
+      NDArray *resultingRow = mmul(*tailTranspose, bottomPart);
+      NDArray *rowPlusFirst = (*resultingRow) + fistRow;
+      delete resultingRow;
+      resultingRow = rowPlusFirst;
+      
+      NDArray *scaledRow = (*resultingRow) * coeff;
+      delete resultingRow;
+      resultingRow = scaledRow;
+      
+      NDArray *firstMinusRow = fistRow - (*resultingRow);
+      fistRow.assign(firstMinusRow);
+      delete firstMinusRow;
+      
+      NDArray *tailMulRow = mmul(tail, *resultingRow);
+      NDArray *bottomMinusTailMul = bottomPart - (*tailMulRow);
+      bottomPart.assign(bottomMinusTailMul);
+      delete tailMulRow;
+      delete bottomMinusTailMul;
+      delete resultingRow;
+    } else {
+      NDArray *resultingRow = mmul(tail, bottomPart);
+      NDArray *rowPlusFirst = (*resultingRow) + fistRow;
+      delete resultingRow;
+      resultingRow = rowPlusFirst;
+      
+      NDArray *scaledRow = (*resultingRow) * coeff;
+      delete resultingRow;
+      resultingRow = scaledRow;
+      
+      NDArray *firstMinusRow = fistRow - (*resultingRow);
+      fistRow.assign(firstMinusRow);
+      delete firstMinusRow;
+      
+      NDArray *transTailMulRow = mmul(*tailTranspose, *resultingRow);
+      NDArray *bottomMinusTrans = bottomPart - (*transTailMulRow);
+      bottomPart.assign(bottomMinusTrans);
+      delete transTailMulRow;
+      delete bottomMinusTrans;
+      delete resultingRow;
+    }
+    delete tailTranspose;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void Householder<T>::mulRight(NDArray& matrix, NDArray& tail, const T coeff) {
+  if (matrix.sizeAt(1) == 1 && coeff != (T)0) {
+    NDArray *scaledResult = matrix * ((T)1.f - coeff);
+    matrix.assign(scaledResult);
+    delete scaledResult;
+  } else if (coeff != (T)0.f) {
+    NDArray *rightPartPtr = matrix({0, 0, 1, matrix.sizeAt(1)}, true);
+    NDArray rightPart = *rightPartPtr;
+    delete rightPartPtr;
+    
+    NDArray *fistColPtr = matrix({0, 0, 0, 1}, true);
+    NDArray fistCol = *fistColPtr;
+    delete fistColPtr;
+
+    NDArray *transposedTail = tail.transpose();
+    if (tail.isColumnVector()) {
+      NDArray *resultingCol = mmul(rightPart, tail);
+      NDArray *colPlusFirst = (*resultingCol) + fistCol;
+      delete resultingCol;
+      resultingCol = colPlusFirst;
+      
+      NDArray *scaledCol = (*resultingCol) * coeff;
+      delete resultingCol;
+      resultingCol = scaledCol;
+      
+      NDArray *firstMinusCol = fistCol - (*resultingCol);
+      fistCol.assign(firstMinusCol);
+      delete firstMinusCol;
+      
+      NDArray *colMulTransTail = mmul(*resultingCol, *transposedTail);
+      NDArray *rightMinusColMul = rightPart - (*colMulTransTail);
+      rightPart.assign(rightMinusColMul);
+      delete colMulTransTail;
+      delete rightMinusColMul;
+      delete resultingCol;
+    } else {
+      NDArray *resultingCol = mmul(rightPart, *transposedTail);
+      NDArray *colPlusFirst = (*resultingCol) + fistCol;
+      delete resultingCol;
+      resultingCol = colPlusFirst;
+      
+      NDArray *scaledCol = (*resultingCol) * coeff;
+      delete resultingCol;
+      resultingCol = scaledCol;
+      
+      NDArray *firstMinusCol = fistCol - (*resultingCol);
+      fistCol.assign(firstMinusCol);
+      delete firstMinusCol;
+      
+      NDArray *colMulTail = mmul(*resultingCol, tail);
+      NDArray *rightMinusColMul = rightPart - (*colMulTail);
+      rightPart.assign(rightMinusColMul);
+      delete colMulTail;
+      delete rightMinusColMul;
+      delete resultingCol;
+    }
+
+    delete transposedTail;
+  }
+}
+
+BUILD_SINGLE_TEMPLATE( class Householder, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd

libnd4j/include/helpers/impl/jacobiSVD.cpp

@@ -0,0 +1,503 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma on 11.01.2018
+//
+#include <helpers/MmulHelper.h>
+#include <helpers/hhColPivQR.h>
+#include <helpers/jacobiSVD.h>
+#if NOT_EXCLUDED(svd)
+namespace sd {
+namespace ops {
+namespace helpers {
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+JacobiSVD<T>::JacobiSVD(NDArray& matrix, const bool calcU, const bool calcV, const bool fullUV)
+    : _m(matrix.dataType(), matrix.getContext(), true),
+      _s(matrix.dataType(), matrix.getContext(), true),
+      _u(matrix.dataType(), matrix.getContext(), true),
+      _v(matrix.dataType(), matrix.getContext(), true) {
+  if (matrix.rankOf() != 2 || matrix.isScalar())
+    THROW_EXCEPTION("ops::helpers::JacobiSVD constructor: input array must be 2D matrix !");
+
+  _rows = static_cast<int>(matrix.sizeAt(0));
+  _cols = static_cast<int>(matrix.sizeAt(1));
+  _diagSize = math::sd_min<int>(_rows, _cols);
+
+  _calcU = calcU;
+  _calcV = calcV;
+  _fullUV = fullUV;
+
+  std::vector<LongType> sShape = {_diagSize,1};
+  _s = NDArray(matrix.ordering(),sShape, matrix.dataType(), matrix.getContext());
+
+  if (_calcU) {
+    std::vector<LongType> rowsShape = {_rows,_rows};
+    std::vector<LongType> rowsShape2 = {_rows,_diagSize};
+    if (_fullUV)
+      _u = NDArray(matrix.ordering(), rowsShape, matrix.dataType(), matrix.getContext());
+    else
+      _u = NDArray(matrix.ordering(), rowsShape2, matrix.dataType(), matrix.getContext());
+  } else {
+    std::vector<LongType> rowsShape = {_rows, 1};
+    _u = NDArray(matrix.ordering(), rowsShape, matrix.dataType(), matrix.getContext());
+  }
+  if (_calcV) {
+    if (_fullUV) {
+      std::vector<LongType> colsShape = {_cols, _cols};
+      _v = NDArray(matrix.ordering(), colsShape, matrix.dataType(), matrix.getContext());
+    }
+    else {
+      std::vector<LongType> shape = {_cols, _diagSize};
+      _v = NDArray(matrix.ordering(),shape, matrix.dataType(), matrix.getContext());
+    }
+  } else {
+    std::vector<LongType> vShape = {_cols, 1};
+    _v = NDArray(matrix.ordering(), vShape, matrix.dataType(), matrix.getContext());
+  }
+  std::vector<LongType> mShape = {_diagSize, _diagSize};
+  _m = NDArray(matrix.ordering(), mShape, matrix.dataType(), matrix.getContext());
+
+
+  evalData(matrix);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void JacobiSVD<T>::mulRotationOnLeft(const int i, const int j, NDArray& block, NDArray& rotation) {
+  if (i < j) {
+    if (j + 1 > block.sizeAt(0))
+      THROW_EXCEPTION(
+          "ops::helpers::JacobiSVD mulRotationOnLeft: second arguments is out of array row range !");
+
+    NDArray *tempPtr = block({i, j + 1, j - i, 0, 0, 0}, true, true);
+    NDArray temp = *tempPtr;
+    NDArray *tempAssignResult = mmul(rotation, temp);
+    temp.assign(tempAssignResult);
+    delete tempAssignResult;
+    delete tempPtr;
+
+  } else {
+    if (j + 1 > block.sizeAt(0) || i + 1 > block.sizeAt(0))
+      THROW_EXCEPTION(
+          "ops::helpers::JacobiSVD mulRotationOnLeft: some or both integer arguments are out of array row range !");
+
+    std::vector<LongType> tempShape = {2, block.sizeAt(1)};
+    NDArray temp(block.ordering(),tempShape, block.dataType(), block.getContext());
+    
+    NDArray *row1Ptr = block({i, i + 1, 0, 0}, true);
+    NDArray row1 = *row1Ptr;
+    
+    NDArray *row2Ptr = block({j, j + 1, 0, 0}, true);
+    NDArray row2 = *row2Ptr;
+    
+    NDArray *rowTemp1Ptr = temp({0, 1, 0, 0}, true);
+    NDArray rowTemp1 = *rowTemp1Ptr;
+    
+    NDArray *rowTemp2Ptr = temp({1, 2, 0, 0}, true);
+    NDArray rowTemp2 = *rowTemp2Ptr;
+    
+    rowTemp1.assign(&row1);
+    rowTemp2.assign(&row2);
+    NDArray *tempAssignResult = mmul(rotation, temp);
+    temp.assign(tempAssignResult);
+    delete tempAssignResult;
+    row1.assign(&rowTemp1);
+    row2.assign(&rowTemp2);
+    
+    delete row1Ptr;
+    delete row2Ptr;
+    delete rowTemp1Ptr;
+    delete rowTemp2Ptr;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void JacobiSVD<T>::mulRotationOnRight(const int i, const int j, NDArray& block, NDArray& rotation) {
+  if (i < j) {
+    if (j + 1 > block.sizeAt(1))
+      THROW_EXCEPTION(
+          "ops::helpers::JacobiSVD mulRotationOnRight: second argument is out of array column range !");
+
+    NDArray *tempPtr = block({0, 0, 0, i, j + 1, j - i}, true, true);
+    NDArray temp = *tempPtr;
+    NDArray *tempAssignResult = mmul(temp, rotation);
+    temp.assign(tempAssignResult);
+    delete tempAssignResult;
+    delete tempPtr;
+  } else {
+    if (j + 1 > block.sizeAt(1) || i + 1 > block.sizeAt(1))
+      THROW_EXCEPTION(
+          "ops::helpers::JacobiSVD mulRotationOnRight: some or both integer arguments are out of array column range !");
+
+    std::vector<LongType> tempShape = {block.sizeAt(0), 2};
+    NDArray temp(block.ordering(), tempShape, block.dataType(), block.getContext());
+    
+    NDArray *col1Ptr = block({0, 0, i, i + 1}, true);
+    NDArray col1 = *col1Ptr;
+    
+    NDArray *col2Ptr = block({0, 0, j, j + 1}, true);
+    NDArray col2 = *col2Ptr;
+    
+    NDArray *colTemp1Ptr = temp({0, 0, 0, 1}, true);
+    NDArray colTemp1 = *colTemp1Ptr;
+    
+    NDArray *colTemp2Ptr = temp({0, 0, 1, 2}, true);
+    NDArray colTemp2 = *colTemp2Ptr;
+    
+    colTemp1.assign(&col1);
+    colTemp2.assign(&col2);
+    NDArray *tempAssignResult = mmul(temp, rotation);
+    temp.assign(tempAssignResult);
+    delete tempAssignResult;
+    col1.assign(&colTemp1);
+    col2.assign(&colTemp2);
+    
+    delete col1Ptr;
+    delete col2Ptr;
+    delete colTemp1Ptr;
+    delete colTemp2Ptr;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+bool JacobiSVD<T>::isBlock2x2NotDiag(NDArray& block, int p, int q, T& maxElem) {
+  std::vector<LongType> shape = {2, 2};
+  NDArray rotation(_m.ordering(), shape, _m.dataType(), _m.getContext());
+
+  T n = math::sd_sqrt<T, T>(block.t<T>(p, p) * block.t<T>(p, p) + block.t<T>(q, p) * block.t<T>(q, p));
+
+  const T almostZero = DataTypeUtils::min_positive<T>();
+  const T precision = DataTypeUtils::eps<T>();
+
+  if (n == (T)0.f) {
+    block.r<T>(p, p) = (T)0;
+    block.r<T>(q, p) = (T)0;
+  } else {
+    T v = block.t<T>(p, p) / n;
+
+    rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = v;
+
+    v = block.t<T>(q, p) / n;
+    rotation.r<T>(0, 1) = v;
+
+    rotation.r<T>(1, 0) = -rotation.template t<T>(0, 1);
+    mulRotationOnLeft(p, q, block, rotation);
+    NDArray *rotT = rotation.transpose();
+    if (_calcU) mulRotationOnRight(p, q, _u, *rotT);
+    delete rotT;
+  }
+
+  maxElem =
+      math::sd_max<T>(maxElem, math::sd_max<T>(math::sd_abs<T,T>(block.t<T>(p, p)), math::sd_abs<T,T>(block.t<T>(q, q))));
+  T threshold = math::sd_max<T>(almostZero, precision * maxElem);
+
+  return math::sd_abs<T,T>(block.t<T>(p, q)) > threshold || math::sd_abs<T,T>(block.t<T>(q, p)) > threshold;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+bool JacobiSVD<T>::createJacobiRotation(const T& x, const T& y, const T& z, NDArray& rotation) {
+  T denom = (T)(2.f) * math::sd_abs<T,T>(y);
+
+  if (denom < DataTypeUtils::min_positive<T>()) {
+    rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = (T)1.f;
+    rotation.r<T>(0, 1) = rotation.r<T>(1, 0) = (T)0.f;
+
+    return false;
+  } else {
+    T tau = (x - z) / denom;
+    T w = math::sd_sqrt<T, T>(tau * tau + (T)1.f);
+    T t;
+
+    if (tau > (T)0.)
+      t = (T)1.f / (tau + w);
+    else
+      t = (T)1.f / (tau - w);
+
+    T sign = t > (T)0. ? (T)1.f : (T)-1.f;
+
+    T cos = (T)1.f / math::sd_sqrt<T, T>(t * t + (T)1.f);
+    T sin = -sign * (y / math::sd_abs<T,T>(y)) * math::sd_abs<T,T>(t) * cos;
+
+    rotation.r<T>(0, 1) = sin;
+    rotation.r<T>(1, 0) = -sin;
+    rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = cos;
+
+    return true;
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void JacobiSVD<T>::createJacobiRotationGivens(const T& p, const T& q, NDArray& rotation) {
+  T cos, sin;
+
+  if (q == (T)0) {
+    cos = p < (T)0 ? (T)-1 : (T)1;
+    sin = (T)0;
+  } else if (p == (T)0) {
+    cos = (T)0;
+    sin = q < (T)0 ? (T)1 : (T)-1;
+  } else if (math::sd_abs<T,T>(p) > math::sd_abs<T,T>(q)) {
+    T t = q / p;
+    T u = math::sd_sqrt<T, T>((T)1 + t * t);
+    if (p < (T)0) u = -u;
+    cos = (T)1 / u;
+    sin = -t * cos;
+  } else {
+    T t = p / q;
+    T u = math::sd_sqrt<T, T>((T)1 + t * t);
+    if (q < (T)0) u = -u;
+    sin = -(T)1 / u;
+    cos = -t * sin;
+  }
+
+  rotation.r<T>(0, 1) = sin;
+  rotation.r<T>(1, 0) = -sin;
+  rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = cos;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void JacobiSVD<T>::svd2x2(NDArray& block, int p, int q, NDArray& left, NDArray& right) {
+  std::vector<LongType> shape = {2, 2};
+  NDArray m(block.ordering(), shape, block.dataType(), block.getContext());
+  m.r<T>(0, 0) = block.t<T>(p, p);
+  m.r<T>(0, 1) = block.t<T>(p, q);
+  m.r<T>(1, 0) = block.t<T>(q, p);
+  m.r<T>(1, 1) = block.t<T>(q, q);
+
+  NDArray rotation(block.ordering(),shape, block.dataType(), block.getContext());
+  T t = m.t<T>(0, 0) + m.t<T>(1, 1);
+  T d = m.t<T>(1, 0) - m.t<T>(0, 1);
+
+  if (math::sd_abs<T,T>(d) < DataTypeUtils::min<T>()) {
+    rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = (T)1;
+    rotation.r<T>(0, 1) = rotation.r<T>(1, 0) = (T)0;
+  } else {
+    T u = t / d;
+    T tmp = math::sd_sqrt<T, T>((T)1.f + u * u);
+    rotation.r<T>(0, 0) = rotation.r<T>(1, 1) = u / tmp;
+    rotation.r<T>(0, 1) = (T)1.f / tmp;
+    rotation.r<T>(1, 0) = -rotation.t<T>(0, 1);
+  }
+  NDArray *mAssignResult = mmul(rotation, m);
+  m.assign(mAssignResult);
+  delete mAssignResult;
+
+  createJacobiRotation(m.t<T>(0, 0), m.t<T>(0, 1), m.t<T>(1, 1), right);
+  NDArray *rightT = right.transpose();
+  NDArray *leftAssignResult = mmul(rotation, *rightT);
+  left.assign(leftAssignResult);
+  delete leftAssignResult;
+  delete rightT;
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void JacobiSVD<T>::evalData(NDArray& matrix) {
+  const T precision = (T)2.f * DataTypeUtils::eps<T>();
+  const T almostZero = DataTypeUtils::min_positive<T>();
+
+  auto* scaleResult = matrix.reduceNumber(reduce::AMax);
+  T scale = scaleResult->template t<T>(0);
+  delete scaleResult;
+  if (scale <   (T)1.f) scale = (T)1.f;
+
+  if (_rows > _cols) {
+    NDArray *scaled = matrix / scale;
+    HHcolPivQR qr(*scaled);
+    delete scaled;
+    NDArray qrRef = *qr._qr;
+    
+    NDArray *mAssignPtr = qrRef({0, _cols, 0, _cols});
+    NDArray mAssign = *mAssignPtr;
+    _m.assign(&mAssign);
+    delete mAssignPtr;
+    
+    _m.fillAsTriangular<T>(0., 0, 0, _m, 'l',false);
+
+    HHsequence hhSeg(qr._qr, qr._coeffs, 'u');
+
+    if (_fullUV)
+      hhSeg.applyTo(&_u);
+    else if (_calcU) {
+      _u.setIdentity();
+      hhSeg.mulLeft(&_u);
+    }
+
+
+    if (_calcV) _v.assign(qr._permut);
+  } else if (_rows < _cols) {
+    NDArray *matrixT = matrix.transpose();
+    NDArray *scaled = (*matrixT) / scale;
+    HHcolPivQR qr(*scaled);
+    delete scaled;
+    NDArray qrRef = *qr._qr;
+    
+    NDArray *mAssignPtr = qrRef({0, _rows, 0, _rows});
+    NDArray mAssign = *mAssignPtr;
+    _m.assign(&mAssign);
+    delete mAssignPtr;
+    
+    _m.fillAsTriangular<T>(0., 0, 0, _m, 'l',false);
+    _m.transposei();
+
+    HHsequence hhSeg(qr._qr, qr._coeffs, 'u');  // type = 'u' is not mistake here !
+
+    if (_fullUV)
+      hhSeg.applyTo(&_v);
+    else if (_calcV) {
+      _v.setIdentity();
+      hhSeg.mulLeft(&_v);
+    }
+
+
+    if (_calcU) _u.assign(qr._permut);
+
+    delete matrixT;
+  } else {
+    NDArray *mAssignPtr = matrix({0, _diagSize, 0, _diagSize});
+    NDArray *mAssignDiv = (*mAssignPtr) / scale;
+    _m.assign(mAssignDiv);
+    delete mAssignDiv;
+    delete mAssignPtr;
+
+    if (_calcU) _u.setIdentity();
+
+    if (_calcV) _v.setIdentity();
+  }
+
+
+  T maxDiagElem = static_cast<T>(0.);
+  for (int i = 0; i < _diagSize; ++i) {
+    T current = math::sd_abs<T,T>(_m.t<T>(i, i));
+    if (maxDiagElem < current) maxDiagElem = current;
+  }
+
+  bool stop = false;
+
+  while (!stop) {
+    stop = true;
+
+    for (int p = 1; p < _diagSize; ++p) {
+      for (int q = 0; q < p; ++q) {
+        T threshold = math::sd_max<T>(almostZero, precision * maxDiagElem);
+
+        if (math::sd_abs<T,T>(_m.t<T>(p, q)) > threshold || math::sd_abs<T,T>(_m.t<T>(q, p)) > threshold) {
+          stop = false;
+
+          std::vector<LongType> shape = {2, 2};
+          NDArray rotLeft(_m.ordering(), shape, _m.dataType(), _m.getContext());
+          NDArray rotRight(_m.ordering(), shape, _m.dataType(), _m.getContext());
+          svd2x2(_m, p, q, rotLeft, rotRight);
+
+          mulRotationOnLeft(p, q, _m, rotLeft);
+          NDArray *rotLeftTranspose = rotLeft.transpose();
+          if (_calcU) mulRotationOnRight(p, q, _u, *rotLeftTranspose);
+
+          mulRotationOnRight(p, q, _m, rotRight);
+
+          if (_calcV) mulRotationOnRight(p, q, _v, rotRight);
+
+          maxDiagElem = math::sd_max<T>(
+              maxDiagElem, math::sd_max<T>(math::sd_abs<T,T>(_m.t<T>(p, p)), math::sd_abs<T,T>(_m.t<T>(q, q))));
+
+          delete rotLeftTranspose;
+        }
+      }
+    }
+  }
+
+
+  for (int i = 0; i < _diagSize; ++i) {
+    _s.r<T>(i) = math::sd_abs<T,T>(_m.t<T>(i, i));
+
+    if (_calcU && _m.t<T>(i, i) < (T)0.) {
+      NDArray *tempPtr = _u({0, 0, i, i + 1}, true);
+      NDArray temp = *tempPtr;
+      temp.applyTransform(transform::Neg, &temp, nullptr);
+      delete tempPtr;
+    }
+  }
+
+
+
+  _s *= scale;
+  for (int i = 0; i < _diagSize; i++) {
+    NDArray *sSlicePtr = _s({i, -1, 0, 0});
+    NDArray sSlice = *sSlicePtr;
+    
+    NDArray *indexNum = sSlice.indexReduceNumber(indexreduce::IndexMax, nullptr);
+    int pos = indexNum->template e<int>(0);
+    auto* maxResult = sSlice.reduceNumber(reduce::Max);
+    T maxSingVal = maxResult->template t<T>(0);
+    delete maxResult;
+    
+    delete sSlicePtr;
+    delete indexNum;
+
+    if (maxSingVal == (T)0.) break;
+
+    if (pos) {
+      pos += i;
+
+      math::sd_swap<T>(_s.r<T>(i), _s.r<T>(pos));
+
+      if (_calcU) {
+        NDArray *temp1Ptr = _u({0, 0, pos, pos + 1}, true);
+        NDArray temp1 = *temp1Ptr;
+        
+        NDArray *temp2Ptr = _u({0, 0, i, i + 1}, true);
+        NDArray temp2 = *temp2Ptr;
+        
+        temp1.swapUnsafe(temp2);
+        
+        delete temp1Ptr;
+        delete temp2Ptr;
+      }
+
+      if (_calcV) {
+        NDArray *temp1Ptr = _v({0, 0, pos, pos + 1}, true);
+        NDArray temp1 = *temp1Ptr;
+        
+        NDArray *temp2Ptr = _v({0, 0, i, i + 1}, true);
+        NDArray temp2 = *temp2Ptr;
+        
+        temp1.swapUnsafe(temp2);
+        
+        delete temp1Ptr;
+        delete temp2Ptr;
+      }
+    }
+  }
+
+}
+
+BUILD_SINGLE_TEMPLATE( class JacobiSVD, , SD_FLOAT_TYPES);
+
+}  // namespace helpers
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/helpers/impl/logger.cpp

@@ -0,0 +1,68 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 31.10.2017.
+//
+#include <helpers/logger.h>
+
+namespace sd {
+
+SD_HOST void Logger::info(const char *format, ...) {
+  va_list args;
+  va_start(args, format);
+  vprintf(format, args);
+  va_end(args);
+  fflush(stdout);
+}
+
+SD_HOST void Logger::infoEmpty(const char *format) {
+ if(format != nullptr)
+  printf("%s",format);
+}
+
+
+SD_HOST void Logger::printv(const char *format, const std::vector<int> &vec) {
+  printf("%s: {", format);
+  for (size_t e = 0; e < vec.size(); e++) {
+    auto v = vec[e];
+    printf("%i", v);
+    if (e < vec.size() - 1) printf(", ");
+  }
+  printf("}\n");
+  fflush(stdout);
+}
+
+SD_HOST void Logger::printv(const char *format, const std::vector<LongType> &vec) {
+  printf("%s: {", format);
+  for (size_t e = 0; e < vec.size(); e++) {
+    auto v = vec[e];
+    printf("%lld", (long long)v);
+    if (e < vec.size() - 1) printf(", ");
+  }
+  printf("}\n");
+  fflush(stdout);
+}
+
+SD_HOST_DEVICE Status Logger::logStatusMsg(Status code, const char *msg) {
+  if (msg != nullptr) sd_printf("%s\n", msg);
+  return code;
+}
+
+SD_HOST_DEVICE Status Logger::logKernelFailureMsg(const char *msg) { return logStatusMsg(Status::KERNEL_FAILURE, msg); }
+}  // namespace sd

libnd4j/include/helpers/impl/reshapeNoCopy.cpp

@@ -0,0 +1,147 @@
+//
+// Created by agibsonccc on 8/30/24.
+//
+
+#ifndef LIBND4J_RESHAPENOCOPY_H
+#define LIBND4J_RESHAPENOCOPY_H
+#include <system/op_boilerplate.h>
+#include <helpers/reshapeNoCopy.h>
+#include <helpers/shape.h>
+#include <array/ArrayOptions.hXX>
+namespace sd {
+namespace ops {
+namespace helpers {
+bool reshapeNoAlloc(const sd::LongType* inShape,
+                    const std::vector<sd::LongType>& newShape,
+                    char order,
+                    sd::LongType* outShape) {
+  LongType oldnd = shape::rank(inShape);
+  std::vector<sd::LongType> olddims(oldnd);
+  std::vector<sd::LongType> oldstrides(oldnd);
+  sd::LongType np, op, last_stride;
+  int oi, oj, ok, ni, nj, nk;
+  std::vector<sd::LongType> newStrides(newShape.size());
+
+  int newnd = newShape.size();
+  bool isFOrder = order == 'f';
+
+  // FIX: Set data type early, before any return statements
+  // This ensures data type is preserved even for empty arrays
+  if(ArrayOptions::numDataTypesSet(ArrayOptions::extra(outShape)) < 1) {
+    ArrayOptions::setDataType(outShape, ArrayOptions::dataType(inShape));
+  }
+
+  // Remove axes with dimension 1 from the old array
+  int actual_oldnd = 0;
+  for (oi = 0; oi < oldnd; oi++) {
+    if (shape::shapeOf(inShape)[oi] != 1) {
+      olddims[actual_oldnd] = shape::shapeOf(inShape)[oi];
+      oldstrides[actual_oldnd] = shape::stride(inShape)[oi];
+      actual_oldnd++;
+    }
+  }
+
+  oldnd = actual_oldnd;
+
+  np = 1;
+  for (ni = 0; ni < newnd; ni++) {
+    np *= newShape[ni];
+  }
+  op = 1;
+  for (oi = 0; oi < oldnd; oi++) {
+    op *= olddims[oi];
+  }
+  if (np != op) {
+    return false;  // total sizes must match
+  }
+
+  if (np == 0) {
+    // FIX: Data type has already been set above, so empty arrays will have correct type
+    return false;  // don't support empty arrays
+  }
+
+
+
+  // oi to oj and ni to nj give the axis ranges currently worked with
+  oi = 0;
+  oj = 1;
+  ni = 0;
+  nj = 1;
+  while (ni < newnd && oi < oldnd) {
+    np = newShape[ni];
+    op = olddims[oi];
+
+    while (np != op) {
+      if (np < op) {
+        np *= newShape[nj++];
+      } else {
+        op *= olddims[oj++];
+      }
+    }
+
+    // Check whether the original axes can be combined
+    for (ok = oi; ok < oj - 1; ok++) {
+      if (isFOrder) {
+        if (oldstrides[ok + 1] != olddims[ok] * oldstrides[ok]) {
+          return false;  // not contiguous enough
+        }
+      } else {
+        // C order
+        if (oldstrides[ok] != olddims[ok + 1] * oldstrides[ok + 1]) {
+          return false;  // not contiguous enough
+        }
+      }
+    }
+
+    // Calculate new strides for all axes currently worked with
+    if (isFOrder) {
+      newStrides[ni] = oldstrides[oi];
+      for (nk = ni + 1; nk < nj; nk++) {
+        newStrides[nk] = newStrides[nk - 1] * newShape[nk - 1];
+      }
+    } else {
+      // C order
+      newStrides[nj - 1] = oldstrides[oj - 1];
+      for (nk = nj - 1; nk > ni; nk--) {
+        newStrides[nk - 1] = newStrides[nk] * newShape[nk];
+      }
+    }
+    ni = nj++;
+    oi = oj++;
+  }
+
+
+
+  // Set strides corresponding to trailing 1s of the new shape
+  if (ni >= 1) {
+    last_stride = newStrides[ni - 1];
+  } else {
+    last_stride = 1;
+  }
+  if (isFOrder && ni >= 1) {
+    last_stride *= newShape[ni - 1];
+  }
+  for (nk = ni; nk < newnd; nk++) {
+    newStrides[nk] = last_stride;
+  }
+
+  // Update the output shape info
+  outShape[0] = newnd;  // Set rank
+
+  shape::setShape(outShape, const_cast<sd::LongType*>(newShape.data()));
+  shape::setStride(outShape, newStrides.data());
+
+  // Set order first
+  shape::setOrder(outShape, order);
+
+  // Data type was set early (lines 28-32) but we set it again here as a defensive measure
+  // to ensure it's preserved even if other shape operations modified the extra field
+  ArrayOptions::setDataType(outShape, ArrayOptions::dataType(inShape));
+
+  return true;
+}
+}
+}
+}
+
+#endif