chore: import upstream snapshot with attribution

libnd4j/include/array/ArrayOptions.h

@@ -0,0 +1,196 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef ND4J_ARRAY_OPTIONS_H
+#define ND4J_ARRAY_OPTIONS_H
+
+#pragma once
+
+#include <system/common.h>
+#include <system/op_boilerplate.h>
+#include <array/ArrayType.h>
+#include <array/DataType.h>
+#include <array/SpaceType.h>
+#include <array/SparseType.h>
+
+#include <initializer_list>
+#include <vector>
+
+
+//notice how each flag value is multiplied by 2
+//if they are too close in value values will clash.
+#define ARRAY_SPARSE 2
+#define ARRAY_COMPRESSED 4
+#define ARRAY_EMPTY 8
+#define ARRAY_RAGGED 16
+
+#define ARRAY_CSR 32
+#define ARRAY_CSC 64
+#define ARRAY_COO 128
+
+// complex values
+#define ARRAY_COMPLEX 512
+
+// quantized values
+#define ARRAY_QUANTIZED 1024
+
+//  16 bit float FP16
+#define ARRAY_HALF 4096
+
+//  16 bit bfloat16
+#define ARRAY_BHALF 2048
+
+// regular 32 bit float
+#define ARRAY_FLOAT 8192
+
+// regular 64 bit float
+#define ARRAY_DOUBLE 16384
+
+// 8 bit integer
+#define ARRAY_CHAR 32768
+
+// 16 bit integer
+#define ARRAY_SHORT 65536
+
+// 32 bit integer
+#define ARRAY_INT 131072
+
+// 64 bit integer
+#define ARRAY_LONG 262144
+
+// boolean values
+#define ARRAY_BOOL 524288
+
+// UTF values
+#define ARRAY_UTF8 1048576
+#define ARRAY_UTF16 4194304
+#define ARRAY_UTF32 16777216
+
+// flag for extras
+#define ARRAY_EXTRAS 2097152
+
+// flag for signed/unsigned integers
+#define ARRAY_UNSIGNED 8388608
+
+// flag for arrays with padded buffer
+#define ARRAY_HAS_PADDED_BUFFER (1 << 25)
+//flags for when array has a view or not
+#define ARRAY_IS_VIEW 33554432
+
+//flag for when array needs a copy
+//this is mainly used in the reshape_no_copy op but could be used elsewhere
+#define ARRAY_NEEDS_COPY 67108864
+
+
+//we need this in order to preserve the offset of the original buffer when creating the output array
+//when views are created, sometimes we need to use the original offset of the array
+//we don't need this very often and we don't store the offset in the shape info
+//this preserves the offsets only being in the ndarray but allowing us to pass information
+//when creating views, each flag is for an input in to an op, most of the time we only need the first 3
+//but may need more. For now only the first one is used but may be needed elsewhere.
+#define ARRAY_COPY_OFFSET_INPUT_0 134217728
+#define ARRAY_COPY_OFFSET_INPUT_1 268435456
+#define ARRAY_COPY_OFFSET_INPUT_2 536870912
+#define ARRAY_COPY_OFFSET_INPUT_3 1073741824
+#define ARRAY_COPY_OFFSET_INPUT_4 2147483648
+#define ARRAY_COPY_OFFSET_INPUT_5 4294967296
+#define ARRAY_COPY_OFFSET_INPUT_6 8589934592
+#define ARRAY_COPY_OFFSET_INPUT_7 17179869184
+#define ARRAY_COPY_OFFSET_INPUT_8 34359738368
+#define ARRAY_COPY_OFFSET_INPUT_9 68719476736
+#define ARRAY_COPY_OFFSET_INPUT_10 137438953472
+
+
+#define DEFAULT_FLAG 0
+
+
+
+
+namespace sd {
+class SD_LIB_EXPORT ArrayOptions {
+ public:
+  static SD_HOST SD_INLINE  LongType extra(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void setExtra(LongType *shapeInfo, LongType value);
+  static SD_HOST SD_INLINE  bool isNewFormat(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  bool hasPropertyBitSet(const LongType *shapeInfo, LongType property);
+  static SD_HOST SD_INLINE  bool togglePropertyBit(LongType *shapeInfo, LongType property);
+  static SD_HOST SD_INLINE  void unsetPropertyBit(LongType *shapeInfo, LongType property);
+  static SD_HOST SD_INLINE  void validateSingleDataType(LongType property);
+  static SD_HOST SD_INLINE  void setPropertyBit(LongType *shapeInfo, LongType property);
+  static SD_HOST SD_INLINE  void setPropertyBits(LongType *shapeInfo, std::initializer_list<LongType> properties);
+  static SD_HOST SD_INLINE  sd::LongType numDataTypesSet(sd::LongType property);
+  static SD_HOST SD_INLINE  bool isUnsigned(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  bool isSparseArray(sd::LongType *shapeInfo);
+  static SD_HOST SD_INLINE  DataType dataType(const LongType *shapeInfo);
+
+  static SD_HOST SD_INLINE  SpaceType spaceType(LongType *shapeInfo);
+  static  SD_INLINE SD_HOST_DEVICE SpaceType spaceType(const LongType *shapeInfo);
+
+  static SD_HOST SD_INLINE  ArrayType arrayType(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  ArrayType arrayType(const LongType *shapeInfo);
+
+  static SD_HOST SD_INLINE  bool isView(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void toggleIsView(LongType *shapeInfo);
+
+  static  SD_INLINE SD_HOST_DEVICE SparseType sparseType(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  SparseType sparseType(const LongType *shapeInfo);
+
+  static  SD_INLINE SD_HOST_DEVICE bool hasExtraProperties(LongType *shapeInfo);
+
+  static SD_HOST SD_INLINE  bool hasPaddedBuffer(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void flagAsPaddedBuffer(LongType *shapeInfo);
+
+  static SD_HOST SD_INLINE  void resetDataType(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  LongType propertyWithoutDataType(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void setDataType(LongType *shapeInfo, const DataType dataType);
+  static SD_HOST SD_INLINE  LongType setDataTypeValue(LongType extraStorage, const DataType dataType);
+  static SD_HOST SD_INLINE  LongType flagForDataType(const DataType dataType);
+  static SD_HOST SD_INLINE  void copyDataType(LongType *to, const LongType *from);
+  static SD_HOST SD_INLINE  const char *enumerateSetFlags(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void unsetAllFlags(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  int enumerateSetFlags(const LongType *shapeInfo, const char **setFlagsOutput, int maxFlags);
+  static SD_HOST SD_INLINE  const char *findFlagString(int flag);
+  static SD_HOST SD_INLINE  LongType extraIndex(const LongType *shapeInfo);
+  static SD_HOST SD_INLINE  LongType extraIndex(LongType *shapeInfo);
+  static SD_HOST SD_INLINE  void unsetAllFlags(LongType &flagStorage);
+  static SD_HOST SD_INLINE  const char *enumerateSetFlagsForFlags(const LongType flagStorage);
+  static SD_HOST SD_INLINE  SpaceType spaceTypeForFlags(const LongType &flagStorage);
+  static SD_HOST SD_INLINE  ArrayType arrayTypeForFlags(const LongType &flagStorage);
+  static SD_HOST SD_INLINE  bool togglePropertyBitForFlags(LongType &flagStorage, LongType property);
+  static SD_HOST SD_INLINE  LongType unsetPropertyBitForFlags(LongType &flagStorage, LongType property);
+  static SD_HOST SD_INLINE  SparseType sparseTypeForFlags(const LongType &flagStorage);
+  static SD_INLINE LongType setPropertyBitForFlagsValue(LongType extraStorage, LongType property);
+  static SD_HOST SD_INLINE  bool hasPropertyBitSet(const LongType extra, LongType property);
+  static SD_HOST SD_INLINE  void resetFlags(LongType *to);
+  static SD_HOST SD_INLINE  LongType defaultFlag();
+
+  static SD_HOST SD_INLINE   LongType propertyWithoutDataTypeValue(LongType extra);
+  static SD_HOST SD_INLINE  DataType dataTypeValue(LongType property);
+  static SD_INLINE bool isEmpty(LongType *shapeInfo);
+  static SD_INLINE void toggleIsEmpty(LongType *shapeInfo);
+
+  static SD_INLINE bool arrayNeedsCopy(LongType *shapeInfo);
+  static SD_INLINE void toggleArrayNeedsCopy(LongType *shapeInfo);
+};
+
+}
+#endif  // ND4J_ARRAY_OPTIONS_H :)

libnd4j/include/array/ArrayOptions.hXX

@@ -0,0 +1,799 @@
+#ifndef ND4J_ARRAY_OPTIONS_HXX
+#define ND4J_ARRAY_OPTIONS_HXX
+#include <array/ArrayOptions.h>
+#include <system/op_boilerplate.h>
+#include <helpers/shape.h>
+#include <array/DataTypeUtils.h>
+#pragma  once
+
+namespace sd {
+
+
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::extraIndex(const sd::LongType *shapeInfo) {
+  return ArrayOptions::extraIndex(const_cast<sd::LongType *>(shapeInfo));
+}
+
+
+SD_HOST  SD_INLINE sd::LongType ArrayOptions::extraIndex(sd::LongType *shapeInfo) {
+  if(shapeInfo == nullptr)
+    THROW_EXCEPTION("Shape info was null!");
+  sd::LongType  rank = shapeInfo[0];
+
+  sd::LongType  idx = 0;
+  //rank takes up 1 element + usual elements
+  if(rank == 0)
+    idx = 3;
+  else
+    // FIXME magic numbers
+    idx = rank + rank + 1;
+
+  return idx;
+}
+
+
+SD_HOST SD_INLINE void ArrayOptions::setExtra(sd::LongType *shapeInfo, sd::LongType value) {
+  sd::LongType  idx = ArrayOptions::extraIndex(shapeInfo);
+  shapeInfo[idx] = value;
+}
+
+SD_HOST SD_INLINE LongType ArrayOptions::extra(const LongType *shapeInfo) {
+  sd::LongType  idx = ArrayOptions::extraIndex(shapeInfo);
+  if(idx >  shape::shapeInfoLength(shape::rank(shapeInfo)))
+    THROW_EXCEPTION("Extra index is out of bounds!");
+  return shapeInfo[idx];
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::isNewFormat(const sd::LongType *shapeInfo) {
+  return (extra(const_cast<sd::LongType *>(shapeInfo)) != 0);
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::isSparseArray(sd::LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, ARRAY_SPARSE);
+}
+
+SD_HOST_DEVICE SD_INLINE bool ArrayOptions::hasExtraProperties(sd::LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, ARRAY_EXTRAS);
+}
+
+SD_HOST  SD_INLINE bool ArrayOptions::hasPropertyBitSet(const sd::LongType extra, LongType property) {
+  return ((static_cast<int>(extra) & static_cast<int>(property)) == static_cast<int>(property));
+}
+
+SD_HOST SD_INLINE bool ArrayOptions:: hasPropertyBitSet(const sd::LongType *shapeInfo, LongType property) {
+  if (!isNewFormat(shapeInfo)) return false;
+
+  return ((static_cast<int>(extra(const_cast<sd::LongType *>(shapeInfo)) & property)) == static_cast<int>(property));
+}
+
+
+SD_HOST_DEVICE SD_INLINE bool hasPropertyBitSetForFlags(const sd::LongType& flagStorage, LongType property) {
+  return static_cast<sd::LongType>(flagStorage & (property)) == (property);
+}
+
+SD_HOST  SD_INLINE void unsetPropertyBitForFlags(sd::LongType& flagStorage, LongType property) {
+  flagStorage &= ~property;
+}
+
+SD_HOST SD_INLINE const char *ArrayOptions::enumerateSetFlagsForFlags(const LongType flagStorage) {
+  int maxFlags = 24;
+  int flagsArray[] = {
+      ARRAY_SPARSE,
+      ARRAY_COMPRESSED,
+      ARRAY_EMPTY,
+      ARRAY_RAGGED,
+      ARRAY_CSR,
+      ARRAY_CSC,
+      ARRAY_COO,
+      ARRAY_COMPLEX,
+      ARRAY_QUANTIZED,
+      ARRAY_HALF,
+      ARRAY_BHALF,
+      ARRAY_FLOAT,
+      ARRAY_DOUBLE,
+      ARRAY_CHAR,
+      ARRAY_SHORT,
+      ARRAY_INT,
+      ARRAY_LONG,
+      ARRAY_BOOL,
+      ARRAY_UTF8,
+      ARRAY_UTF16,
+      ARRAY_UTF32,
+      ARRAY_EXTRAS,
+      ARRAY_UNSIGNED,
+      ARRAY_HAS_PADDED_BUFFER,
+      ARRAY_IS_VIEW
+  };
+
+  const char* flagsStrings[] = {
+      "ARRAY_SPARSE",
+      "ARRAY_COMPRESSED",
+      "ARRAY_EMPTY",
+      "ARRAY_RAGGED",
+      "ARRAY_CSR",
+      "ARRAY_CSC",
+      "ARRAY_COO",
+      "ARRAY_COMPLEX",
+      "ARRAY_QUANTIZED",
+      "ARRAY_HALF",
+      "ARRAY_BHALF",
+      "ARRAY_FLOAT",
+      "ARRAY_DOUBLE",
+      "ARRAY_CHAR",
+      "ARRAY_SHORT",
+      "ARRAY_INT",
+      "ARRAY_LONG",
+      "ARRAY_BOOL",
+      "ARRAY_UTF8",
+      "ARRAY_UTF16",
+      "ARRAY_UTF32",
+      "ARRAY_EXTRAS",
+      "ARRAY_UNSIGNED",
+      "ARRAY_HAS_PADDED_BUFFER",
+      "ARRAY_IS_VIEW"
+  };
+
+  std::string flags;
+  for (int i = 0; i <  maxFlags; i++) {
+    if (hasPropertyBitSetForFlags(flagStorage, flagsArray[i])) {
+      flags += flagsStrings[i];
+      flags += " ";
+    }
+
+  }
+
+  auto ret =  new std::string(flags);  // Returns the number of set flags found
+  return ret->c_str();
+}
+
+SD_HOST SD_INLINE void ArrayOptions::unsetAllFlags(sd::LongType& flagStorage) {
+
+  int flagsArray[] = {
+      ARRAY_SPARSE,
+      ARRAY_COMPRESSED,
+      ARRAY_EMPTY,
+      ARRAY_RAGGED,
+      ARRAY_CSR,
+      ARRAY_CSC,
+      ARRAY_COO,
+      ARRAY_COMPLEX,
+      ARRAY_QUANTIZED,
+      ARRAY_HALF,
+      ARRAY_BHALF,
+      ARRAY_FLOAT,
+      ARRAY_DOUBLE,
+      ARRAY_CHAR,
+      ARRAY_SHORT,
+      ARRAY_INT,
+      ARRAY_LONG,
+      ARRAY_BOOL,
+      ARRAY_UTF8,
+      ARRAY_UTF16,
+      ARRAY_UTF32,
+      ARRAY_EXTRAS,
+      ARRAY_UNSIGNED,
+      ARRAY_HAS_PADDED_BUFFER
+  };
+
+  for (int i = 0; i < 24; i++) {
+    unsetPropertyBitForFlags(flagStorage, flagsArray[i]);
+  }
+}
+
+SD_HOST SD_INLINE const char *ArrayOptions::enumerateSetFlags(const sd::LongType *shapeInfo) {
+  return enumerateSetFlagsForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+
+SD_HOST  SD_INLINE void ArrayOptions::unsetAllFlags(sd::LongType *shapeInfo) {
+  ArrayOptions::unsetAllFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+
+
+SD_HOST  SD_INLINE bool ArrayOptions::isUnsigned(sd::LongType *shapeInfo) {
+  if (!isNewFormat(shapeInfo)) return false;
+
+  return hasPropertyBitSet(shapeInfo, ARRAY_UNSIGNED);
+}
+
+
+
+#define DATA_TYPE_FLAGS { \
+    ARRAY_FLOAT, \
+    ARRAY_DOUBLE, \
+    ARRAY_HALF, \
+    ARRAY_BHALF, \
+    ARRAY_BOOL, \
+    ARRAY_CHAR, \
+    ARRAY_SHORT, \
+    ARRAY_INT, \
+    ARRAY_LONG, \
+    ARRAY_UTF8, \
+    ARRAY_UTF16, \
+    ARRAY_UTF32 \
+}
+
+#define DATA_TYPES { \
+    sd::DataType::FLOAT32, \
+    sd::DataType::DOUBLE, \
+    sd::DataType::HALF, \
+    sd::DataType::BFLOAT16, \
+    sd::DataType::BOOL, \
+    sd::DataType::INT8, \
+    sd::DataType::INT16, \
+    sd::DataType::INT32, \
+    sd::DataType::INT64, \
+    sd::DataType::UTF8, \
+    sd::DataType::UTF16, \
+    sd::DataType::UTF32 \
+}
+
+#define ARRAY_UNSIGNED_TYPES { \
+    ARRAY_CHAR, \
+    ARRAY_SHORT, \
+    ARRAY_INT, \
+    ARRAY_LONG, \
+    ARRAY_UTF8, \
+    ARRAY_UTF16, \
+    ARRAY_UTF32 \
+}
+
+#define UNSIGNED_DATA_TYPES { \
+    sd::DataType::UINT8, \
+    sd::DataType::UINT16, \
+    sd::DataType::UINT32, \
+    sd::DataType::UINT64, \
+    sd::DataType::UTF8, \
+    sd::DataType::UTF16, \
+    sd::DataType::UTF32 \
+}
+
+
+SD_HOST SD_INLINE sd::DataType ArrayOptions::dataTypeValue(sd::LongType property) {
+  const sd::LongType dataTypeFlags[] = DATA_TYPE_FLAGS;
+  const sd::DataType dataTypes[] = DATA_TYPES;
+  const size_t numTypes = sizeof(dataTypeFlags) / sizeof(sd::LongType);
+
+
+  if (hasPropertyBitSetForFlags(property, ARRAY_UNSIGNED)) {
+    const sd::LongType unsignedTypeFlags[] = ARRAY_UNSIGNED_TYPES;
+    const sd::DataType unsignedDataTypes[] = UNSIGNED_DATA_TYPES;
+    const size_t numUnsignedTypes = sizeof(unsignedTypeFlags) / sizeof(sd::LongType);
+
+    for (size_t i = 0; i < numUnsignedTypes; ++i) {
+      if (hasPropertyBitSetForFlags(property, unsignedTypeFlags[i])) {
+        return unsignedDataTypes[i];
+      }
+    }
+  } else {
+    for (size_t i = 0; i < numTypes; ++i) {
+      if (hasPropertyBitSetForFlags(property, dataTypeFlags[i])) {
+        return dataTypes[i];
+      }
+    }
+  }
+
+  return sd::DataType::UNKNOWN;
+}
+
+
+
+SD_HOST  SD_INLINE void validateFlags(sd::LongType property, const sd::LongType flags[], size_t numFlags) {
+  LongType *flagIndices = new LongType[numFlags];
+  int numFlagsSet = 0;
+  for (size_t i = 0; i < numFlags; ++i) {
+    if (hasPropertyBitSetForFlags(property, flags[i])) {
+      flagIndices[i] = 1;
+      numFlagsSet++;
+    } else {
+      flagIndices[i] = 0;
+    }
+  }
+
+  if (numFlagsSet > 1) {
+    std::ostringstream errorMsg;
+    errorMsg << "Multiple data types are set for the given property: ";
+    for (size_t i = 0; i < numFlags; i++) {
+      if(flagIndices[i] == 1) {
+        errorMsg << "Flag index " << i << " (flag value: " << flags[i] << "), ";
+      }
+    }
+    errorMsg << "Total: " << numFlagsSet << " data types set.";
+    THROW_EXCEPTION(errorMsg.str().c_str());
+  }
+
+  delete[] flagIndices;
+}
+
+SD_HOST SD_INLINE void ArrayOptions::validateSingleDataType(sd::LongType property) {
+  const sd::LongType dataTypeFlags[] = DATA_TYPE_FLAGS;
+  const size_t numDataTypeFlags = sizeof(dataTypeFlags) / sizeof(sd::LongType);
+  validateFlags(property, dataTypeFlags, numDataTypeFlags);
+
+  if (hasPropertyBitSetForFlags(property, ARRAY_UNSIGNED)) {
+    const sd::LongType unsignedTypeFlags[] = ARRAY_UNSIGNED_TYPES;
+    const size_t numUnsignedTypeFlags = sizeof(unsignedTypeFlags) / sizeof(sd::LongType);
+    validateFlags(property, unsignedTypeFlags, numUnsignedTypeFlags);
+  }
+}
+
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::numDataTypesSet(sd::LongType property) {
+  const sd::LongType dataTypeFlags[] = DATA_TYPE_FLAGS;
+  const size_t numDataTypeFlags = sizeof(dataTypeFlags) / sizeof(sd::LongType);
+
+  sd::LongType numFlagsSet = 0;
+  for (size_t i = 0; i < numDataTypeFlags; ++i) {
+    if (hasPropertyBitSetForFlags(property, dataTypeFlags[i])) {
+      numFlagsSet++;
+    }
+  }
+
+  return numFlagsSet;
+}
+
+SD_HOST SD_INLINE sd::DataType ArrayOptions::dataType(const sd::LongType *shapeInfo) {
+  // Return safe default instead of throwing - prevents SIGABRT crashes
+  if(shapeInfo == nullptr)
+    return DataType::FLOAT32;  // Safe default for uninitialized arrays
+
+  auto extra = ArrayOptions::extra(shapeInfo);
+  return ArrayOptions::dataTypeValue(extra);
+}
+
+
+
+SD_HOST SD_INLINE SpaceType ArrayOptions::spaceTypeForFlags(const sd::LongType& flagStorage) {
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_QUANTIZED)) return SpaceType::QUANTIZED;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_COMPLEX)) return SpaceType::COMPLEX;
+  return SpaceType::CONTINUOUS;  // by default we return continuous type here
+}
+
+SD_HOST SD_INLINE ArrayType ArrayOptions::arrayTypeForFlags(const sd::LongType& flagStorage) {
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_SPARSE)) return ArrayType::SPARSE;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_COMPRESSED)) return ArrayType::COMPRESSED;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_EMPTY)) return ArrayType::EMPTY;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_RAGGED)) return ArrayType::RAGGED;
+  return ArrayType::DENSE;  // by default we return DENSE type here
+}
+
+SD_HOST SD_INLINE  bool ArrayOptions::togglePropertyBitForFlags(sd::LongType& flagStorage, LongType property) {
+  flagStorage ^= property;
+  return hasPropertyBitSetForFlags(flagStorage, property);
+}
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::unsetPropertyBitForFlags(sd::LongType& flagStorage, LongType property) {
+  return flagStorage & ~property;
+}
+
+SD_HOST  SD_INLINE SparseType ArrayOptions::sparseTypeForFlags(const sd::LongType& flagStorage) {
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_CSC)) return SparseType::CSC;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_CSR)) return SparseType::CSR;
+  if (hasPropertyBitSetForFlags(flagStorage, ARRAY_COO)) return SparseType::COO;
+  return SparseType::LIL;
+}
+
+// Existing function that works with shapeInfo:
+SD_HOST_DEVICE  SD_INLINE SpaceType ArrayOptions::spaceType(const sd::LongType *shapeInfo) {
+  return spaceTypeForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+
+SD_HOST  SD_INLINE ArrayType ArrayOptions::arrayType(const sd::LongType *shapeInfo) {
+  return arrayTypeForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+
+SD_HOST SD_INLINE ArrayType ArrayOptions::arrayType(sd::LongType *shapeInfo) {
+  return arrayTypeForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::isEmpty(sd::LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, EMPTY);
+}
+
+
+ SD_HOST SD_INLINE bool ArrayOptions::arrayNeedsCopy(LongType *shapeInfo) {
+  return hasPropertyBitSetForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)], ARRAY_NEEDS_COPY);
+}
+
+ SD_HOST SD_INLINE void ArrayOptions::toggleArrayNeedsCopy(LongType *shapeInfo) {
+    togglePropertyBit(shapeInfo, ARRAY_NEEDS_COPY);
+}
+
+SD_HOST SD_INLINE void ArrayOptions::toggleIsEmpty(sd::LongType *shapeInfo) {
+  togglePropertyBit(shapeInfo, EMPTY);
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::isView(sd::LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, ARRAY_IS_VIEW);
+}
+
+SD_HOST SD_INLINE void ArrayOptions::toggleIsView(sd::LongType *shapeInfo) {
+  togglePropertyBit(shapeInfo, ARRAY_IS_VIEW);
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::togglePropertyBit(sd::LongType *shapeInfo, LongType property) {
+  return togglePropertyBitForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)], property);
+}
+
+SD_HOST SD_INLINE void ArrayOptions::setPropertyBit(sd::LongType *shapeInfo, LongType property) {
+  shapeInfo[ArrayOptions::extraIndex(shapeInfo)] = setPropertyBitForFlagsValue(shapeInfo[ArrayOptions::extraIndex(shapeInfo)], property);
+}
+
+SD_HOST SD_INLINE void ArrayOptions::unsetPropertyBit(sd::LongType *shapeInfo, LongType property) {
+  shapeInfo[ArrayOptions::extraIndex(shapeInfo)] = unsetPropertyBitForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)], property);
+}
+
+SD_HOST SD_INLINE SparseType ArrayOptions::sparseType(const sd::LongType *shapeInfo) {
+  return sparseTypeForFlags(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+}
+SD_HOST SD_INLINE void ArrayOptions::setPropertyBits(sd::LongType *shapeInfo, std::initializer_list<LongType> properties) {
+  for (auto v : properties) {
+    if (!hasPropertyBitSet(shapeInfo, v)) setPropertyBit(shapeInfo, v);
+  }
+}
+
+SD_HOST SD_INLINE void ArrayOptions::flagAsPaddedBuffer(sd::LongType *shapeInfo) {
+  if (!isNewFormat(shapeInfo)) return;
+
+  return setPropertyBit(shapeInfo, ARRAY_HAS_PADDED_BUFFER);
+}
+
+SD_HOST SD_INLINE bool ArrayOptions::hasPaddedBuffer(const sd::LongType *shapeInfo) {
+  if (!isNewFormat(shapeInfo)) return false;
+
+  return hasPropertyBitSet(shapeInfo, ARRAY_HAS_PADDED_BUFFER);
+}
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::propertyWithoutDataTypeValue(sd::LongType extra) {
+  sd::LongType property = extra;
+  property = property & (~ARRAY_BOOL);
+  property = property & (~ARRAY_HALF);
+  property = property & (~ARRAY_BHALF);
+  property = property & (~ARRAY_FLOAT);
+  property = property & (~ARRAY_DOUBLE);
+  property = property & (~ARRAY_INT);
+  property = property & (~ARRAY_LONG);
+  property = property & (~ARRAY_CHAR);
+  property = property & (~ARRAY_SHORT);
+  property = property & (~ARRAY_UNSIGNED);
+  return property;
+}
+
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::propertyWithoutDataType(const sd::LongType *shapeInfo) {
+  auto newCast = const_cast<sd::LongType *>(shapeInfo);
+  sd::LongType property = extra(newCast);
+  return propertyWithoutDataTypeValue(property);
+}
+
+SD_HOST SD_INLINE void ArrayOptions::resetDataType(sd::LongType *shapeInfo) {
+  setExtra(shapeInfo, propertyWithoutDataType(shapeInfo));
+}
+
+SD_HOST SD_INLINE LongType ArrayOptions::flagForDataType(const sd::DataType dataType) {
+  switch (dataType) {
+    case sd::DataType::BOOL:
+      return ARRAY_BOOL;
+    case sd::DataType::HALF:
+      return ARRAY_HALF;
+    case sd::DataType::BFLOAT16:
+      return ARRAY_BHALF;
+    case sd::DataType::FLOAT32:
+      return ARRAY_FLOAT;
+    case sd::DataType::DOUBLE:
+      return ARRAY_DOUBLE;
+    case sd::DataType::INT8:
+      return ARRAY_CHAR;
+    case sd::DataType::INT16:
+      return ARRAY_SHORT;
+    case sd::DataType::INT32:
+      return ARRAY_INT;
+    case sd::DataType::INT64:
+      return ARRAY_LONG;
+    case sd::DataType::UINT8:
+      return ARRAY_CHAR | ARRAY_UNSIGNED;
+    case sd::DataType::UINT16:
+      return ARRAY_SHORT | ARRAY_UNSIGNED;
+    case sd::DataType::UINT32  :
+      return ARRAY_INT | ARRAY_UNSIGNED;
+    case sd::DataType::UINT64 :
+      return ARRAY_LONG | ARRAY_UNSIGNED;
+    case sd::DataType::UTF8:
+      return ARRAY_UTF8;
+    case sd::DataType::UTF16:
+      return ARRAY_UTF16;
+    case sd::DataType::UTF32:
+      return ARRAY_UTF32;
+    default:
+#ifndef __CUDA_ARCH__
+      std::string errorMessage;
+      errorMessage += "Can't set unknown data type ArrayOptions: ";
+      errorMessage += sd::DataTypeUtils::asString(dataType);
+      THROW_EXCEPTION(errorMessage.c_str());
+      return 0;  // Never reached, but satisfies compiler warning
+
+#else
+      printf("Can't set unknown data type");
+      return 0;  // Return sentinel value for unknown types in CUDA code
+#endif
+  }
+}
+
+SD_HOST SD_INLINE void ArrayOptions::setDataType(sd::LongType *shapeInfo, const sd::DataType dataType) {
+  switch (dataType) {
+    case sd::DataType::BOOL:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_BOOL);
+      break;
+    case sd::DataType::HALF:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_HALF);
+      break;
+    case sd::DataType::BFLOAT16:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_BHALF);
+      break;
+    case sd::DataType::FLOAT32:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_FLOAT);
+      break;
+    case sd::DataType::DOUBLE:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_DOUBLE);
+      break;
+    case sd::DataType::INT8:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_CHAR);
+      break;
+    case sd::DataType::INT16:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_SHORT);
+      break;
+    case sd::DataType::INT32:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_INT);
+      break;
+    case sd::DataType::INT64:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_LONG);
+      break;
+    case sd::DataType::UINT8:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_CHAR | ARRAY_UNSIGNED);
+      break;
+    case sd::DataType::UINT16:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_SHORT | ARRAY_UNSIGNED);
+      break;
+    case sd::DataType::UINT32:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_INT | ARRAY_UNSIGNED);
+      break;
+    case sd::DataType::UINT64:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_LONG | ARRAY_UNSIGNED);
+      break;
+    case sd::DataType::UTF8:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_UTF8);
+      break;
+    case sd::DataType::UTF16:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_UTF16);
+      break;
+    case sd::DataType::UTF32:
+      ArrayOptions::resetDataType(shapeInfo);
+      setPropertyBit(shapeInfo, ARRAY_UTF32);
+      break;
+    default:
+#ifndef __CUDA_ARCH__
+      std::string errorMessage;
+      errorMessage += "Can't set unknown data type: ";
+      errorMessage += DataTypeUtils::asString(dataType);
+      THROW_EXCEPTION(errorMessage.c_str());
+#else
+      printf("Can't set unknown data type");
+#endif
+
+      validateSingleDataType(shapeInfo[ArrayOptions::extraIndex(shapeInfo)]);
+
+  }
+
+#ifndef __CUDA_ARCH__
+  if(ArrayOptions::dataType(shapeInfo) != dataType) {
+    std::string errorMessage;
+    errorMessage += "setDataType:  Data type set was not correct one. Expected ";
+    errorMessage += DataTypeUtils::asString(dataType);
+    errorMessage += " but got ";
+    errorMessage += sd::DataTypeUtils::asString(dataType);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+#else
+  printf("setDataType: Data type set was incorrect.");
+#endif
+
+}
+
+SD_HOST SD_INLINE sd::LongType ArrayOptions::setDataTypeValue(sd::LongType extraStorage, const sd::DataType dataType) {
+  sd::LongType  ret = extraStorage;
+  if (dataType == sd::DataType::UINT8
+      || dataType == sd::DataType::UINT16
+      || dataType == sd::DataType::UINT32 ||
+      dataType == sd::DataType::UINT64) {
+    ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_UNSIGNED);
+    extraStorage = ret;
+    return extraStorage;
+  }
+
+
+  switch (dataType) {
+    case sd::DataType::BOOL:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_BOOL);
+      break;
+    case sd::DataType::HALF:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_HALF);
+      break;
+    case sd::DataType::BFLOAT16:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_BHALF);
+      break;
+    case sd::DataType::FLOAT32:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_FLOAT);
+      break;
+    case sd::DataType::DOUBLE:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_DOUBLE);
+      break;
+    case sd::DataType::INT8:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_CHAR);
+      break;
+    case sd::DataType::INT16:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_SHORT);
+      break;
+    case sd::DataType::INT32:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_INT);
+      break;
+    case sd::DataType::INT64:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_LONG);
+      break;
+    case sd::DataType::UINT8:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_CHAR);
+      break;
+    case sd::DataType::UINT16:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_SHORT);
+      break;
+    case sd::DataType::UINT32:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_INT);
+      break;
+    case sd::DataType::UINT64:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_LONG);
+      break;
+    case sd::DataType::UTF8:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_UTF8);
+      break;
+    case sd::DataType::UTF16:
+      ret =  setPropertyBitForFlagsValue(extraStorage, ARRAY_UTF16);
+      break;
+    case sd::DataType::UTF32:
+      ret = setPropertyBitForFlagsValue(extraStorage, ARRAY_UTF32);
+      break;
+    default:
+#ifndef __CUDA_ARCH__
+      THROW_EXCEPTION("Can't set unknown data type");
+#else
+      printf("Can't set unknown data type");
+#endif
+  }
+
+
+
+  return ret;
+
+}
+
+SD_HOST sd::LongType ArrayOptions::defaultFlag() {
+  return DEFAULT_FLAG;
+}
+
+SD_HOST void ArrayOptions::resetFlags(sd::LongType *to) {
+  to[ArrayOptions::extraIndex(to)] = 1;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+SD_HOST SD_INLINE  void ArrayOptions::copyDataType(sd::LongType *to, const sd::LongType *from) {
+  if(to == nullptr)
+    THROW_EXCEPTION("To Shape info was null!");
+  if(from == nullptr)
+    THROW_EXCEPTION("From Shape info was null!");
+  DataType dt = dataType(from);
+  sd::LongType flagForDt = flagForDataType(dt);
+  if(hasPropertyBitSet(to, flagForDt))
+    return;
+  setDataType(to, dt);
+}
+SD_HOST sd::LongType ArrayOptions::setPropertyBitForFlagsValue(LongType extraStorage, LongType property) {
+  if(hasPropertyBitSet(extraStorage, property))
+    return extraStorage;
+  return extraStorage | property;
+}
+
+// New method implementations to add to ArrayOptions.hXX
+
+// needsCopy - const overload to check if array needs copy
+SD_HOST SD_INLINE bool ArrayOptions::needsCopy(const LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, ARRAY_NEEDS_COPY);
+}
+
+// hasCopyOffset - check if specific input index has copy offset flag set
+SD_HOST SD_INLINE bool ArrayOptions::hasCopyOffset(const LongType *shapeInfo, int inputIndex) {
+  if (inputIndex < 0 || inputIndex > 10) {
+    return false;
+  }
+  LongType flag = copyOffsetFlagForInput(inputIndex);
+  return hasPropertyBitSet(shapeInfo, flag);
+}
+
+// toggleCopyOffset - toggle copy offset flag for specific input index
+SD_HOST SD_INLINE void ArrayOptions::toggleCopyOffset(LongType *shapeInfo, int inputIndex) {
+  if (inputIndex < 0 || inputIndex > 10) {
+#ifndef __CUDA_ARCH__
+    THROW_EXCEPTION("Input index out of range [0-10]");
+#else
+    printf("Input index out of range [0-10]\n");
+#endif
+    return;
+  }
+  LongType flag = copyOffsetFlagForInput(inputIndex);
+  togglePropertyBit(shapeInfo, flag);
+}
+
+// copyOffsetFlagForInput - get the flag constant for a specific input index
+SD_HOST SD_INLINE LongType ArrayOptions::copyOffsetFlagForInput(int inputIndex) {
+  switch (inputIndex) {
+    case 0: return ARRAY_COPY_OFFSET_INPUT_0;
+    case 1: return ARRAY_COPY_OFFSET_INPUT_1;
+    case 2: return ARRAY_COPY_OFFSET_INPUT_2;
+    case 3: return ARRAY_COPY_OFFSET_INPUT_3;
+    case 4: return ARRAY_COPY_OFFSET_INPUT_4;
+    case 5: return ARRAY_COPY_OFFSET_INPUT_5;
+    case 6: return ARRAY_COPY_OFFSET_INPUT_6;
+    case 7: return ARRAY_COPY_OFFSET_INPUT_7;
+    case 8: return ARRAY_COPY_OFFSET_INPUT_8;
+    case 9: return ARRAY_COPY_OFFSET_INPUT_9;
+    case 10: return ARRAY_COPY_OFFSET_INPUT_10;
+    default:
+#ifndef __CUDA_ARCH__
+      THROW_EXCEPTION("Invalid input index for copy offset flag");
+#else
+      printf("Invalid input index for copy offset flag\n");
+#endif
+      return 0;
+  }
+}
+
+// clearAllCopyOffsets - clear all 11 copy offset flags at once
+SD_HOST SD_INLINE void ArrayOptions::clearAllCopyOffsets(LongType *shapeInfo) {
+  for (int i = 0; i <= 10; i++) {
+    LongType flag = copyOffsetFlagForInput(i);
+    if (hasPropertyBitSet(shapeInfo, flag)) {
+      unsetPropertyBit(shapeInfo, flag);
+    }
+  }
+}
+
+// getActiveCopyOffsets - return count or bitmask of which inputs have offset flags set
+SD_HOST SD_INLINE int ArrayOptions::getActiveCopyOffsets(const LongType *shapeInfo) {
+  int count = 0;
+  for (int i = 0; i <= 10; i++) {
+    LongType flag = copyOffsetFlagForInput(i);
+    if (hasPropertyBitSet(shapeInfo, flag)) {
+      count++;
+    }
+  }
+  return count;
+}
+
+// isView - const overload for consistency
+SD_HOST SD_INLINE bool ArrayOptions::isView(const LongType *shapeInfo) {
+  return hasPropertyBitSet(shapeInfo, ARRAY_IS_VIEW);
+}
+}
+#endif

libnd4j/include/array/ArrayType.h

@@ -0,0 +1,36 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef ND4J_ARRAY_TYPE_H
+#define ND4J_ARRAY_TYPE_H
+
+namespace sd {
+enum ArrayType {
+  DENSE = 1,
+  SPARSE = 2,
+  COMPRESSED = 3,
+  EMPTY = 4,
+  RAGGED = 5,
+};
+}
+
+#endif

libnd4j/include/array/ByteOrder.h

@@ -0,0 +1,33 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 21.11.17.
+//
+
+#ifndef LIBND4J_BYTEORDER_H
+#define LIBND4J_BYTEORDER_H
+
+namespace sd {
+enum ByteOrder {
+  LE = 0,
+  BE = 1,
+};
+}
+
+#endif  // LIBND4J_BYTEORDER_H

libnd4j/include/array/ByteOrderUtils.h

@@ -0,0 +1,36 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 21.11.17.
+//
+
+#ifndef LIBND4J_BYTEORDERUTILS_H
+#define LIBND4J_BYTEORDERUTILS_H
+#include <array/ByteOrder.h>
+#include <graph/generated/array_generated.h>
+#include <system/common.h>
+
+namespace sd {
+class SD_LIB_EXPORT ByteOrderUtils {
+ public:
+  static ByteOrder fromFlatByteOrder(graph::ByteOrder order);
+};
+}  // namespace sd
+
+#endif  // LIBND4J_BYTEORDERUTILS_H

libnd4j/include/array/ConstantDataBuffer.h

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+
+#ifndef LIBND4J_CONSTANTDATABUFFER_H
+#define LIBND4J_CONSTANTDATABUFFER_H
+
+#include <array/DataType.h>
+#include <array/PointerWrapper.h>
+#include <system/common.h>
+
+#include <memory>
+
+namespace sd {
+class SD_LIB_EXPORT ConstantDataBuffer {
+ private:
+  std::shared_ptr<PointerWrapper> _primaryBuffer;
+  std::shared_ptr<PointerWrapper> _specialBuffer = nullptr;
+  uint64_t _length = 0;
+  uint8_t _sizeOf = 0;
+
+ public:
+  ConstantDataBuffer(const std::shared_ptr<PointerWrapper> &primary, uint64_t numEelements, DataType dype);
+  ConstantDataBuffer(const std::shared_ptr<PointerWrapper> &primary, const std::shared_ptr<PointerWrapper> &special,
+                     uint64_t numEelements, DataType dype);
+  ConstantDataBuffer(const ConstantDataBuffer &other);
+  ConstantDataBuffer() = default;
+  ~ConstantDataBuffer() = default;
+
+  uint8_t sizeOf() const;
+  uint64_t length() const;
+
+  void *primary() const;
+  void *special() const;
+
+  ConstantDataBuffer &operator=(const ConstantDataBuffer &other) = default;
+  ConstantDataBuffer &operator=(ConstantDataBuffer &&other) noexcept = default;
+
+  template <typename T>
+  T *primaryAsT() const;
+
+  template <typename T>
+  T *specialAsT() const;
+};
+}  // namespace sd
+
+#endif  // DEV_TESTS_CONSTANTDATABUFFER_H

libnd4j/include/array/ConstantDescriptor.h

@@ -0,0 +1,77 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef DEV_TESTS_CONSTANTDESCRIPTOR_H
+#define DEV_TESTS_CONSTANTDESCRIPTOR_H
+
+#include <array/ConstantDataBuffer.h>
+#include <array/DataType.h>
+#include <system/common.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace sd {
+class SD_LIB_EXPORT ConstantDescriptor {
+ private:
+  std::vector<LongType> _integerValues;
+  std::vector<double> _floatValues;
+
+ public:
+  ConstantDescriptor(double *values, int length);
+  ConstantDescriptor(LongType const *values, int length);
+  ConstantDescriptor(std::initializer_list<double> values);
+
+  explicit ConstantDescriptor(std::vector<LongType> &values);
+  explicit ConstantDescriptor(std::vector<double> &values);
+
+  ~ConstantDescriptor() = default;
+
+  // equal to operator
+  bool operator==(const ConstantDescriptor &other) const;
+
+  // less than operator
+  bool operator<(const ConstantDescriptor &other) const;
+
+  bool isInteger() const;
+  bool isFloat() const;
+
+  LongType length() const;
+
+  const std::vector<LongType> &integerValues() const;
+  const std::vector<double> &floatValues() const;
+};
+}  // namespace sd
+
+#ifndef __JAVACPP_HACK__
+
+namespace std {
+template <>
+class SD_LIB_EXPORT hash<sd::ConstantDescriptor> {
+ public:
+  size_t operator()(const sd::ConstantDescriptor &k) const;
+};
+}  // namespace std
+
+#endif
+
+#endif  // DEV_TESTS_CONSTANTDESCRIPTOR_H

libnd4j/include/array/ConstantHolder.h

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com
+//
+
+#ifndef LIBND4J_CONSTANTHOLDER_H
+#define LIBND4J_CONSTANTHOLDER_H
+#include <array/ConstantDataBuffer.h>
+#include <array/ConstantDescriptor.h>
+
+#include <map>
+#include <mutex>
+
+namespace sd {
+class ConstantHolder {
+ private:
+  int _deviceId = 0;
+  std::mutex _mutex;
+
+  std::map<DataType, ConstantDataBuffer> _buffers;
+
+ public:
+  ConstantHolder(const ConstantHolder &other);
+  ConstantHolder() = default;
+  ~ConstantHolder() = default;
+
+  ConstantHolder &operator=(const ConstantHolder &other) = delete;
+  ConstantHolder &operator=(ConstantHolder &&other) = delete;
+
+  bool hasBuffer(DataType dataType);
+
+  template <typename T>
+  bool hasBuffer();
+
+  void addBuffer(ConstantDataBuffer &pointer, DataType dataType);
+
+  template <typename T>
+  void addBuffer(ConstantDataBuffer &pointer);
+
+  ConstantDataBuffer *getConstantDataBuffer(DataType dataType);
+
+  template <typename T>
+  ConstantDataBuffer *getConstantDataBuffer();
+
+  std::mutex *mutex();
+};
+}  // namespace sd
+
+#endif  // DEV_TESTS_CONSTANTHOLDER_H

libnd4j/include/array/ConstantOffsetsBuffer.h

@@ -0,0 +1,57 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+//
+
+#ifndef SD_ARRAY_CONSTANTOFFSETSBUFFER_H_
+#define SD_ARRAY_CONSTANTOFFSETSBUFFER_H_
+
+#include <array/PointerWrapper.h>
+#include <system/common.h>
+
+#include <memory>
+
+namespace sd {
+
+class SD_LIB_EXPORT ConstantOffsetsBuffer {
+ private:
+  static constexpr uint32_t MAGIC_VALID = 0xC0FFE575;  // "COFFSETS" - magic number for validity check
+  uint32_t _magic = 0;  // Validity marker to detect use-after-free
+  std::shared_ptr<PointerWrapper> _primaryOffsets;
+  std::shared_ptr<PointerWrapper> _specialOffsets;
+
+ public:
+  ConstantOffsetsBuffer(const std::shared_ptr<PointerWrapper> &primary);
+  ConstantOffsetsBuffer(const std::shared_ptr<PointerWrapper> &primary, const std::shared_ptr<PointerWrapper> &special);
+  ConstantOffsetsBuffer();  // Default constructor - sets magic number
+  ~ConstantOffsetsBuffer();  // Need custom destructor to clear magic
+
+  inline bool isValid() const { return _magic == MAGIC_VALID; }
+
+  LongType *primary();
+  LongType *special();
+  LongType *platform();
+};
+
+}  // namespace sd
+
+#endif  // SD_ARRAY_CONSTANTOFFSETSBUFFER_H_

libnd4j/include/array/ConstantShapeBuffer.h

@@ -0,0 +1,98 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+//
+
+#ifndef SD_ARRAY_CONSTANTSHAPEBUFFER_H_
+#define SD_ARRAY_CONSTANTSHAPEBUFFER_H_
+
+#include <array/PointerWrapper.h>
+#include <system/common.h>
+
+#include <atomic>
+#include <memory>
+#include <string>
+#ifndef  __JAVACPP_HACK__
+#if defined(SD_GCC_FUNCTRACE)
+#include <exceptions/backward.hpp>
+#endif
+#endif
+namespace sd {
+
+class SD_LIB_EXPORT ConstantShapeBuffer {
+ private:
+  static constexpr uint32_t MAGIC_VALID = 0x58A9EB0F;  // Magic number for validity check
+  uint32_t _magic;  // Validity marker to detect use-after-free/garbage pointers
+  PointerWrapper* _primaryShapeInfo;
+  PointerWrapper*  _specialShapeInfo;
+  std::atomic<int> _refCount;
+
+
+ public:
+  ConstantShapeBuffer( PointerWrapper* primary);
+  ConstantShapeBuffer( PointerWrapper* primary, PointerWrapper* special);
+  ConstantShapeBuffer();
+  ~ConstantShapeBuffer();
+
+  /**
+   * Check if this buffer is valid (not garbage/use-after-free).
+   * Uses magic number validation.
+   */
+  inline bool isValid() const { return _magic == MAGIC_VALID; }
+#ifndef  __JAVACPP_HACK__
+#if defined(SD_GCC_FUNCTRACE)
+  backward::StackTrace st;
+#endif
+#endif
+  LongType *primary() ;
+  LongType *special() ;
+  LongType *platform() ;
+
+  /**
+   * Get the stack trace as a formatted string.
+   * Returns empty string if functrace is not enabled.
+   */
+  std::string getStackTraceAsString() const;
+
+  /**
+   * Manual reference counting for safe cross-JNI ownership.
+   * Increments reference count - call when handing out a pointer.
+   */
+  void addRef();
+
+  /**
+   * Manual reference counting for safe cross-JNI ownership.
+   * Decrements reference count and deletes when reaching zero.
+   * Call when done with a pointer (e.g., in deleteConstantShapeBuffer).
+   */
+  void release();
+
+  /**
+   * Get current reference count (for debugging).
+   */
+  int getRefCount() const;
+};
+
+
+}  // namespace sd
+
+#endif  // SD_ARRAY_CONSTANTSHAPEBUFFER_H_

libnd4j/include/array/CudaPointerDeallocator.h

@@ -0,0 +1,41 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+//
+
+#ifndef SD_CUDAYPOINTERDEALLOCATOR_H_
+#define SD_CUDAYPOINTERDEALLOCATOR_H_
+
+#include <array/PointerDeallocator.h>
+#include <system/common.h>
+
+namespace sd {
+class SD_LIB_EXPORT CudaPointerDeallocator : public PointerDeallocator {
+ public:
+  CudaPointerDeallocator() = default;
+  ~CudaPointerDeallocator() = default;
+
+  void release(void *ptr) override;
+};
+}  // namespace sd
+
+#endif  // SD_CUDAYPOINTERDEALLOCATOR_H_

libnd4j/include/array/DataBuffer.h

@@ -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 raver119@gmail.com
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#ifndef DEV_TESTS_DATABUFFER_H
+#define DEV_TESTS_DATABUFFER_H
+
+#include <array/DataType.h>
+#include <execution/LaunchContext.h>
+#include <memory/Workspace.h>
+#include <system/common.h>
+#include <system/op_boilerplate.h>
+
+#include <cstring>
+#include <mutex>
+namespace sd {
+
+class SD_LIB_EXPORT DataBuffer {
+ private:
+  // Magic number for validity checking (pattern from DirectShapeTrie validation)
+  // Set in constructor, cleared in destructor, checked before use
+  // Helps detect use-after-free and corrupted pointers
+  static constexpr uint32_t MAGIC_NUMBER = 0xDA7ABF01;  // "DA7ABF01" (DataBuffer v01)
+  uint32_t _magicNumber = MAGIC_NUMBER;
+
+  void *_primaryBuffer = nullptr;
+  void *_specialBuffer = nullptr;
+  LongType _lenInBytes = 0;
+  memory::Workspace *_workspace = nullptr;
+
+  std::atomic<int> _deviceId;
+  std::mutex _deleteMutex;
+#ifndef __JAVACPP_HACK__
+#if defined(SD_CUDA)
+  mutable std::atomic<LongType> _counter;
+  mutable std::atomic<LongType> _writePrimary;
+  mutable std::atomic<LongType> _writeSpecial;
+  mutable std::atomic<LongType> _readPrimary;
+  mutable std::atomic<LongType> _readSpecial;
+#endif
+
+#if defined(SD_GCC_FUNCTRACE)
+  StackTrace *allocationStackTracePrimary = nullptr;
+  StackTrace *allocationStackTraceSpecial = nullptr;
+  StackTrace *creationStackTrace = nullptr;
+
+#endif
+
+
+#endif
+
+  bool closed = false;
+
+  // Helper template function for printing host buffer content (implementation in .cpp)
+  template <typename T>
+  void printHostBufferContent(void* buffer, sd::LongType offset, sd::LongType length);
+
+
+
+  void setCountersToZero();
+  void copyCounters(const DataBuffer &other);
+  void deleteSpecial();
+  void deletePrimary();
+  void deleteBuffers();
+  void setAllocFlags(const bool isOwnerPrimary, const bool isOwnerSpecial = false);
+  void allocateBuffers(const bool allocBoth = false);
+
+  void setSpecial(void *special, const bool isOwnerSpecial);
+
+  void copyBufferFromHost(const void *hostBuffer, size_t sizeToCopyinBytes = 0, const LongType offsetThis = 0,
+                          const LongType offsetHostBuffer = 0);
+
+ public:
+
+  bool _isOwnerPrimary;
+  bool _isOwnerSpecial;
+  bool isConstant = false;
+  DataType _dataType;
+
+  DataBuffer(void *primary, void *special, const size_t lenInBytes, const DataType dataType,
+             const bool isOwnerPrimary = false, const bool isOwnerSpecial = false,
+             memory::Workspace *workspace = nullptr);
+
+  DataBuffer(void *primary, const size_t lenInBytes, const DataType dataType, const bool isOwnerPrimary = false,
+             memory::Workspace *workspace = nullptr);
+
+  DataBuffer(const void *hostBuffer,  // copies data from hostBuffer to own memory buffer
+             const DataType dataType, const size_t lenInBytes, memory::Workspace *workspace = nullptr);
+
+  DataBuffer(const sd::LongType lenInBytes, const DataType dataType, memory::Workspace *workspace = nullptr,
+             const bool allocBoth = false);
+
+  DataBuffer(const DataBuffer &other);
+  DataBuffer(DataBuffer &&other);
+  explicit DataBuffer();
+  ~DataBuffer();
+
+  DataBuffer &operator=(const DataBuffer &other);
+  DataBuffer &operator=(DataBuffer &&other) noexcept;
+
+  DataType getDataType();
+  void setDataType(DataType dataType);
+  size_t getLenInBytes() const;
+
+  size_t getNumElements();
+
+  template <typename T>
+  void *primaryAtOffset(const LongType offset);
+  template <typename T>
+  void *specialAtOffset(const LongType offset);
+
+  void *primary();
+  void *special();
+  void printAllocationTrace();
+
+  /**
+   * Validate that this DataBuffer object is in a sane state.
+   * Following DirectShapeTrie validation pattern: check magic number, closed flag, etc.
+   * Throws exception with detailed message if validation fails.
+   * Call this before accessing any member in methods that might be called
+   * on dangling/corrupted pointers (like special(), primary(), etc.)
+   */
+  void validateIntegrity() const;
+
+  void allocatePrimary();
+  void allocateSpecial();
+
+  void writePrimary() const;
+  void writeSpecial() const;
+  void readPrimary() const;
+  void readSpecial() const;
+  bool isPrimaryActual() const;
+  bool isSpecialActual() const;
+
+  void expand(const uint64_t size);
+
+  int deviceId() const;
+  void setDeviceId(int deviceId);
+  void migrate();
+
+  template <typename T>
+  SD_INLINE T *primaryAsT();
+  template <typename T>
+  SD_INLINE T *specialAsT();
+
+  void markConstant(bool reallyConstant);
+
+  void syncToPrimary(const LaunchContext *context, const bool forceSync = false);
+  void syncToSpecial(const bool forceSync = false);
+
+  void setToZeroBuffers(const bool both = false);
+
+  void copyBufferFrom(const DataBuffer &other, size_t sizeToCopyinBytes = 0, const LongType offsetThis = 0,
+                      const LongType offsetOther = 0);
+
+
+  void setPrimaryBuffer(void *buffer, size_t length);
+  void setSpecialBuffer(void *buffer, size_t length);
+
+
+  void  showBufferLimited();
+  //for Debug purposes
+  void showCounters(const char* msg1, const char* msg2);
+
+  /**
+   * This method deletes buffers, if we're owners
+   */
+  void close();
+  bool isClosed() { return closed; }
+  void printPrimaryAllocationStackTraces();
+  void printSpecialAllocationTraces();
+  DataBuffer  dup();
+
+  /**
+   * Helper method to format creation stack trace as string for error messages.
+   * Returns formatted stack trace if SD_GCC_FUNCTRACE is enabled, empty string otherwise.
+   */
+  std::string getCreationTraceAsString() const;
+  void printHostDevice(long offset);
+  static void memcpy(DataBuffer *dst, DataBuffer *src, sd::LongType startingOffset, sd::LongType dstOffset);
+  /**
+   * Print detailed buffer information including host and device content if available
+   * @param msg - Optional message to display
+   * @param offset - Starting offset for printing buffer contents
+   * @param limit - Maximum number of elements to print
+   */
+#ifndef __JAVACPP_HACK__
+  void printBufferDebug(const char* msg = nullptr, sd::LongType offset = 0, sd::LongType limit = 10);
+#endif
+};
+///// IMPLEMENTATION OF INLINE METHODS /////
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+T *DataBuffer::primaryAsT() {
+  return reinterpret_cast<T *>(_primaryBuffer);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+T *DataBuffer::specialAsT() {
+  return reinterpret_cast<T *>(_specialBuffer);
+}
+
+}  // namespace sd
+
+#endif  // DEV_TESTS_DATABUFFER_H

libnd4j/include/array/DataType.h

@@ -0,0 +1,55 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef ND4J_DATATYPE_H
+#define ND4J_DATATYPE_H
+
+namespace sd {
+enum DataType {
+  INHERIT = 0,
+  BOOL = 1,
+  FLOAT8 = 2,
+  HALF = 3,
+  HALF2 = 4,
+  FLOAT32 = 5,
+  DOUBLE = 6,
+  INT8 = 7,
+  INT16 = 8,
+  INT32 = 9,
+  INT64 = 10,
+  UINT8 = 11,
+  UINT16 = 12,
+  UINT32 = 13,
+  UINT64 = 14,
+  QINT8 = 15,
+  QINT16 = 16,
+  BFLOAT16 = 17,
+  UTF8 = 50,
+  UTF16 = 51,
+  UTF32 = 52,
+  ANY = 100,
+  AUTO = 200,
+  UNKNOWN = 255
+};
+}
+
+#endif

libnd4j/include/array/DataTypeConversions.h

@@ -0,0 +1,276 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 21.11.17.
+//
+
+#ifndef LIBND4J_DATATYPECONVERSIONS_H
+#define LIBND4J_DATATYPECONVERSIONS_H
+
+#include <array/DataType.h>
+#include <execution/Threads.h>
+#include <helpers/BitwiseUtils.h>
+#include <helpers/logger.h>
+#include <loops/type_conversions.h>
+#include <system/common.h>
+#include <system/op_boilerplate.h>
+#include <types/float16.h>
+
+namespace sd {
+template <typename T>
+class SD_LIB_EXPORT DataTypeConversions {
+ private:
+
+  template <typename T2>
+  static SD_INLINE void rconv(bool isBe, bool canKeep, T *buffer, LongType length, void *src) {
+    if (std::is_same<T, T2>::value && canKeep) {
+      // When the types match and we can keep the data as-is,
+      // copy the data from src to buffer.
+      ops::safe_copy(buffer, static_cast<const T*>(src), static_cast<size_t>(length));
+    } else {
+      // Otherwise, allocate a temporary array of type T2.
+      auto tmp = new T2[length];
+      ops::safe_copy(tmp, static_cast<const T2*>(src), static_cast<size_t>(length));
+
+      // If any conversion from T2 to T is required, perform it here.
+      // For example, if T can be constructed from T2, you might do:
+      // for (LongType i = 0; i < length; ++i) {
+      //     buffer[i] = static_cast<T>(tmp[i]);
+      // }
+
+      // Free the temporary storage.
+      delete[] tmp;
+    }
+  }
+
+ public:
+  // The convertType function, modified to use safe_copy instead of direct memcpy.
+  // (Assumes that T, DataType, ByteOrder, LongType, BitwiseUtils, DataTypeConversions,
+  //  TypeCast, PRAGMA_THREADS_FOR, samediff::Threads, sd_printf, THROW_EXCEPTION, etc.
+  //  are defined elsewhere in your code base.)
+  static SD_INLINE void convertType(void *vbuffer, void *src, DataType dataType, ByteOrder order, LongType length) {
+    auto buffer = reinterpret_cast<T *>(vbuffer);
+    bool isBe = BitwiseUtils::isBE();
+    bool canKeep = (isBe && order == BE) || (!isBe && order == LE);
+
+    switch (dataType) {
+#ifdef HAS_BOOL
+      case BOOL: {
+        DataTypeConversions<T>::template rconv<bool>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_UINT8
+      case UINT8: {
+        DataTypeConversions<T>::template rconv<uint8_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_INT8
+      case INT8: {
+        DataTypeConversions<T>::template rconv<int8_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_INT16
+      case INT16: {
+        DataTypeConversions<T>::template rconv<int16_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_INT32
+      case INT32: {
+        DataTypeConversions<T>::template rconv<int>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_LONG
+      case INT64: {
+        DataTypeConversions<T>::template rconv<LongType>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_FLOAT32
+      case FLOAT32: {
+        if (std::is_same<T, float>::value && canKeep) {
+          ops::safe_copy(buffer, static_cast<const float*>(src), static_cast<size_t>(length));
+        } else {
+          auto tmp = new float[length];
+          ops::safe_copy(tmp, static_cast<const float*>(src), static_cast<size_t>(length));
+
+#if __GNUC__ <= 4
+          if (!canKeep) {
+            for (sd::LongType e = 0; e < length; e++)
+              buffer[e] = BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          } else {
+            TypeCast::convertGeneric<float, T>(nullptr, tmp, length, buffer);
+          }
+#else
+          auto func = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++)
+              buffer[e] = canKeep
+                              ? static_cast<T>(tmp[e])
+                              : BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          };
+          samediff::Threads::parallel_for(func, 0, length);
+#endif
+          delete[] tmp;
+        }
+      } break;
+#endif
+
+#ifdef HAS_DOUBLE
+      case DOUBLE: {
+        if (std::is_same<T, double>::value && canKeep) {
+          ops::safe_copy(buffer, static_cast<const double*>(src), static_cast<size_t>(length));
+        } else {
+          auto tmp = new double[length];
+          ops::safe_copy(tmp, static_cast<const double*>(src), static_cast<size_t>(length));
+
+#if __GNUC__ <= 4
+          if (!canKeep) {
+            for (sd::LongType e = 0; e < length; e++)
+              buffer[e] = BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          } else {
+            TypeCast::convertGeneric<double, T>(nullptr, tmp, length, buffer);
+          }
+#else
+          auto func = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++)
+              buffer[e] = canKeep
+                              ? static_cast<T>(tmp[e])
+                              : BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          };
+          samediff::Threads::parallel_for(func, 0, length);
+#endif
+          delete[] tmp;
+        }
+      } break;
+#endif
+
+#ifdef HAS_FLOAT16
+      case HALF: {
+        if (std::is_same<T, float16>::value && canKeep) {
+          ops::safe_copy(buffer, static_cast<const float16*>(src), static_cast<size_t>(length));
+        } else {
+          auto tmp = new float16[length];
+          ops::safe_copy(tmp, static_cast<const float16*>(src), static_cast<size_t>(length));
+
+#if __GNUC__ <= 4
+          if (!canKeep) {
+            for (sd::LongType e = 0; e < length; e++)
+              buffer[e] = BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          } else {
+            TypeCast::convertGeneric<float16, T>(nullptr, tmp, length, buffer);
+          }
+#else
+          auto func = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++)
+              buffer[e] = canKeep
+                              ? static_cast<T>(tmp[e])
+                              : BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          };
+          samediff::Threads::parallel_for(func, 0, length);
+#endif
+          delete[] tmp;
+        }
+      } break;
+#endif
+
+#ifdef HAS_BFLOAT16
+      case BFLOAT16: {
+        if (std::is_same<T, bfloat16>::value && canKeep) {
+          ops::safe_copy(buffer, static_cast<const bfloat16*>(src), static_cast<size_t>(length));
+        } else {
+          auto tmp = new bfloat16[length];
+          ops::safe_copy(tmp, static_cast<const bfloat16*>(src), static_cast<size_t>(length));
+
+#if __GNUC__ <= 4
+          if (!canKeep) {
+            for (sd::LongType e = 0; e < length; e++)
+              buffer[e] = BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          } else {
+            TypeCast::convertGeneric<bfloat16, T>(nullptr, tmp, length, buffer);
+          }
+#else
+          auto func = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++)
+              buffer[e] = canKeep
+                              ? static_cast<T>(tmp[e])
+                              : BitwiseUtils::swap_bytes<T>(static_cast<T>(tmp[e]));
+          };
+          samediff::Threads::parallel_for(func, 0, length);
+#endif
+          delete[] tmp;
+        }
+      } break;
+#endif
+
+#ifdef HAS_UINT16
+      case UINT16: {
+        DataTypeConversions<T>::template rconv<uint16_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_UINT32
+      case UINT32: {
+        DataTypeConversions<T>::template rconv<uint32_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_UNSIGNEDLONG
+      case UINT64: {
+        DataTypeConversions<T>::template rconv<uint64_t>(isBe, canKeep, buffer, length, src);
+      } break;
+#endif
+
+#ifdef HAS_UTF8
+      case UTF8: {
+        // UTF8 handling would go here if needed
+        sd_print("UTF8 DataType conversion not implemented\n");
+        THROW_EXCEPTION("UTF8 DataType conversion not implemented");
+      } break;
+#endif
+
+#ifdef HAS_UTF16
+      case UTF16: {
+        // UTF16 handling would go here if needed
+        sd_print("UTF16 DataType conversion not implemented\n");
+        THROW_EXCEPTION("UTF16 DataType conversion not implemented");
+      } break;
+#endif
+
+#ifdef HAS_UTF32
+      case UTF32: {
+        // UTF32 handling would go here if needed
+        sd_print("UTF32 DataType conversion not implemented\n");
+        THROW_EXCEPTION("UTF32 DataType conversion not implemented");
+      } break;
+#endif
+
+      default: {
+        sd_printf("Unsupported DataType requested: [%i]\n", static_cast<int>(dataType));
+        THROW_EXCEPTION("Unsupported DataType");
+      }
+    }
+  }
+};
+}  // namespace sd
+
+#endif  // LIBND4J_DATATYPECONVERSIONS_H

libnd4j/include/array/DataTypeValidation.cpp

@@ -0,0 +1,394 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// Implementation of robust data type validation and error reporting
+//
+
+#include <array/DataType.h>
+#include <array/DataTypeValidation.h>
+#include <exceptions/allocation_exception.h>
+#include <system/op_enums.h>
+
+#include <sstream>
+
+#include "system/op_boilerplate.h"
+#include <helpers/logger.h>
+namespace sd {
+    
+    // Static member definitions
+    std::unordered_set<DataType> DataTypeValidation::compiledTypes_;
+    std::unordered_map<DataType, std::string> DataTypeValidation::typeNames_;
+    std::unordered_map<int, DataType> DataTypeValidation::enumMapping_;
+    bool DataTypeValidation::initialized_ = false;
+    
+    void DataTypeValidation::initialize() {
+        if (initialized_) return;
+        
+        populateTypeNames();
+        populateCompiledTypes();
+        
+        // Build enum ID to DataType mapping
+        for (const auto& pair : typeNames_) {
+            enumMapping_[static_cast<int>(pair.first)] = pair.first;
+        }
+        
+        initialized_ = true;
+    }
+    
+    void DataTypeValidation::initializeIfNeeded() {
+        if (!initialized_) {
+            initialize();
+        }
+    }
+    
+    void DataTypeValidation::populateTypeNames() {
+        if (initialized_) return;
+
+        // All possible data types defined in DataType enum
+        typeNames_[DataType::INHERIT] = "INHERIT";
+        typeNames_[DataType::BOOL] = "BOOL";
+        typeNames_[DataType::FLOAT8] = "FLOAT8";
+        typeNames_[DataType::HALF] = "HALF";
+        typeNames_[DataType::HALF2] = "HALF2";
+        typeNames_[DataType::FLOAT32] = "FLOAT32";
+        typeNames_[DataType::DOUBLE] = "DOUBLE";
+        typeNames_[DataType::INT8] = "INT8";
+        typeNames_[DataType::INT16] = "INT16";
+        typeNames_[DataType::INT32] = "INT32";
+        typeNames_[DataType::INT64] = "INT64";
+        typeNames_[DataType::UINT8] = "UINT8";
+        typeNames_[DataType::UINT16] = "UINT16";
+        typeNames_[DataType::UINT32] = "UINT32";
+        typeNames_[DataType::UINT64] = "UINT64";
+        typeNames_[DataType::QINT8] = "QINT8";
+        typeNames_[DataType::QINT16] = "QINT16";
+        typeNames_[DataType::BFLOAT16] = "BFLOAT16";
+        typeNames_[DataType::UTF8] = "UTF8";
+        typeNames_[DataType::UTF16] = "UTF16";
+        typeNames_[DataType::UTF32] = "UTF32";
+        typeNames_[DataType::UNKNOWN] = "UNKNOWN";
+    }
+    void DataTypeValidation::populateCompiledTypes() {
+        if (initialized_) return;
+
+        // Auto-register types based on compile-time flags
+        // This mirrors the logic from the CMake type system
+
+        // Boolean type
+#if defined(HAS_BOOL)
+        registerCompiledType(DataType::BOOL, "BOOL");
+#endif
+
+        // Floating point types - check both CMake-generated and alias defines
+#if defined(HAS_FLOAT) || defined(HAS_FLOAT32)
+        registerCompiledType(DataType::FLOAT32, "FLOAT32");
+#endif
+
+#if defined(HAS_DOUBLE) || defined(HAS_FLOAT64)
+        registerCompiledType(DataType::DOUBLE, "DOUBLE");
+#endif
+
+#if defined(HAS_FLOAT16) || defined(HAS_HALF)
+        registerCompiledType(DataType::HALF, "HALF");
+#endif
+
+#if defined(HAS_BFLOAT16) || defined(HAS_BFLOAT)
+        registerCompiledType(DataType::BFLOAT16, "BFLOAT16");
+#endif
+
+        // Integer types - check both normalized (_T) and alias forms
+#if defined(HAS_INT8_T) || defined(HAS_INT8)
+        registerCompiledType(DataType::INT8, "INT8");
+#endif
+
+#if defined(HAS_INT16_T) || defined(HAS_INT16)
+        registerCompiledType(DataType::INT16, "INT16");
+#endif
+
+#if defined(HAS_INT32_T) || defined(HAS_INT32) || defined(HAS_INT)
+        registerCompiledType(DataType::INT32, "INT32");
+#endif
+
+#if defined(HAS_INT64_T) || defined(HAS_INT64) || defined(HAS_LONG)
+        registerCompiledType(DataType::INT64, "INT64");
+#endif
+
+        // Unsigned integer types
+#if defined(HAS_UINT8_T) || defined(HAS_UINT8)
+        registerCompiledType(DataType::UINT8, "UINT8");
+#endif
+
+#if defined(HAS_UINT16_T) || defined(HAS_UINT16)
+        registerCompiledType(DataType::UINT16, "UINT16");
+#endif
+
+#if defined(HAS_UINT32_T) || defined(HAS_UINT32)
+        registerCompiledType(DataType::UINT32, "UINT32");
+#endif
+
+#if defined(HAS_UINT64_T) || defined(HAS_UINT64) || defined(HAS_UNSIGNEDLONG)
+        registerCompiledType(DataType::UINT64, "UINT64");
+#endif
+
+        // String types - only if string operations are enabled
+#if defined(SD_ENABLE_STRING_OPERATIONS)
+#if defined(HAS_STD_STRING) || defined(HAS_UTF8)
+        registerCompiledType(DataType::UTF8, "UTF8");
+#endif
+
+#if defined(HAS_STD_U16STRING) || defined(HAS_UTF16)
+        registerCompiledType(DataType::UTF16, "UTF16");
+#endif
+
+#if defined(HAS_STD_U32STRING) || defined(HAS_UTF32)
+        registerCompiledType(DataType::UTF32, "UTF32");
+#endif
+#endif
+
+    }
+    
+    void DataTypeValidation::registerCompiledType(DataType type, const std::string& name) {
+        compiledTypes_.insert(type);
+        typeNames_[type] = name;
+    }
+    
+    bool DataTypeValidation::isValidDataType(DataType dataType) {
+        initializeIfNeeded();
+        return typeNames_.find(dataType) != typeNames_.end();
+    }
+    
+    bool DataTypeValidation::isValidDataType(int dataTypeId) {
+        initializeIfNeeded();
+        return enumMapping_.find(dataTypeId) != enumMapping_.end();
+    }
+
+    std::vector<DataType> DataTypeValidation::getAvailableDataTypes() {
+        initializeIfNeeded();
+        std::vector<DataType> result;
+        result.reserve(compiledTypes_.size());
+        for (const auto& type : compiledTypes_) {
+            result.push_back(type);
+        }
+        return result;
+    }
+    
+    bool DataTypeValidation::isCompiledDataType(DataType dataType) {
+        initializeIfNeeded();
+        return compiledTypes_.find(dataType) != compiledTypes_.end();
+    }
+    
+    std::string DataTypeValidation::getDataTypeErrorMessage(int dataTypeId, const char* context) {
+        initializeIfNeeded();
+
+        std::ostringstream oss;
+
+        if (!isValidDataType(dataTypeId)) {
+            oss << "Invalid data type ID: " << dataTypeId << " (not a valid enum value)";
+        } else {
+            DataType dt = enumMapping_[dataTypeId];
+            const char* typeName = getDataTypeName(dt);
+
+            if (!isCompiledDataType(dt)) {
+              oss << "Data type " << typeName << " (ID: " << dataTypeId
+                  << ") is not available in this build";
+            } else {
+              oss << "Unexpected error with data type " << typeName
+                  << " (ID: " << dataTypeId << ") - type is valid and compiled";
+            }
+        }
+
+        if (context) {
+            oss << "\nContext: " << context;
+        }
+
+        // Add build configuration info with better diagnosis
+#if defined(SD_SELECTIVE_TYPES)
+        oss << "\nBuild configuration: Selective types enabled";
+
+        // List the define flags that would enable missing types
+        if (isValidDataType(dataTypeId)) {
+            DataType dt = enumMapping_[dataTypeId];
+            if (!isCompiledDataType(dt)) {
+              const char* typeName = getDataTypeName(dt);
+              oss << "\nTo enable this type, rebuild with the type included in SD_TYPES_LIST";
+              oss << "\nExample: cmake -DSD_TYPES_LIST=\"float32;int32;int64;" << typeName << "\" ..";
+            }
+        }
+#else
+        oss << "\nBuild configuration: All types enabled (SD_SELECTIVE_TYPES not defined)";
+        oss << "\nThis error suggests a build system configuration issue.";
+#endif
+
+        // Add available types
+        auto availableTypes = getAvailableDataTypes();
+        if (!availableTypes.empty()) {
+            oss << "\nAvailable types in this build (" << availableTypes.size() << "): ";
+            for (size_t i = 0; i < availableTypes.size(); ++i) {
+              if (i > 0) oss << ", ";
+              oss << getDataTypeName(availableTypes[i]);
+            }
+        } else {
+            oss << "\nCRITICAL: No data types are available in this build!";
+            oss << "\nThis indicates a serious build configuration problem.";
+
+            // Debug information to help diagnose the issue
+            oss << "\nDiagnostic information:";
+            oss << "\n  Total possible types: " << typeNames_.size();
+            oss << "\n  Compiled types: " << compiledTypes_.size();
+
+            // Show which defines are actually set
+            oss << "\n  Active defines:";
+#if defined(HAS_BOOL)
+            oss << " HAS_BOOL";
+#endif
+#if defined(HAS_FLOAT)
+            oss << " HAS_FLOAT";
+#endif
+#if defined(HAS_FLOAT32)
+            oss << " HAS_FLOAT32";
+#endif
+#if defined(HAS_DOUBLE)
+            oss << " HAS_DOUBLE";
+#endif
+#if defined(HAS_INT32_T)
+            oss << " HAS_INT32_T";
+#endif
+#if defined(HAS_INT32)
+            oss << " HAS_INT32";
+#endif
+#if defined(HAS_INT)
+            oss << " HAS_INT";
+#endif
+#if defined(HAS_INT64_T)
+            oss << " HAS_INT64_T";
+#endif
+#if defined(HAS_LONG)
+            oss << " HAS_LONG";
+#endif
+#if defined(SD_SELECTIVE_TYPES)
+            oss << " SD_SELECTIVE_TYPES";
+#endif
+
+            if (compiledTypes_.empty()) {
+              oss << "\n  No type defines matched - check CMake type generation";
+            }
+        }
+
+        // Add suggestions based on the situation
+        if (!isValidDataType(dataTypeId)) {
+            oss << "\nSuggestion: Check NDArray creation and data type initialization";
+        } else if (availableTypes.empty()) {
+            oss << "\nSuggestion: Check CMake build configuration and type define generation";
+            oss << "\nVerify that setup_type_definitions() was called in CMake";
+        } else {
+            oss << "\nSuggestion: Use one of the available types or rebuild with the required type enabled";
+        }
+
+        return oss.str();
+    }
+    
+    const char* DataTypeValidation::getDataTypeName(DataType dataType) {
+        initializeIfNeeded();
+        auto it = typeNames_.find(dataType);
+        return (it != typeNames_.end()) ? it->second.c_str() : "UNKNOWN";
+    }
+    
+    std::string DataTypeValidation::getBuildInfo() {
+        initializeIfNeeded();
+        std::ostringstream oss;
+
+        oss << "=== DATA TYPE VALIDATION BUILD INFO ===\n";
+
+#if defined(SD_SELECTIVE_TYPES)
+        oss << "Build type: Selective types (SD_SELECTIVE_TYPES defined)\n";
+#else
+        oss << "Build type: All types (SD_SELECTIVE_TYPES not defined)\n";
+#endif
+
+        oss << "Total possible types: " << typeNames_.size() << "\n";
+        oss << "Compiled types: " << compiledTypes_.size() << "\n";
+
+        if (!compiledTypes_.empty()) {
+            oss << "Available types: ";
+            auto types = getAvailableDataTypes();
+            for (size_t i = 0; i < types.size(); ++i) {
+              if (i > 0) oss << ", ";
+              oss << getDataTypeName(types[i]);
+            }
+            oss << "\n";
+        } else {
+            oss << "CRITICAL: No types compiled!\n";
+            oss << "This indicates a build configuration error.\n";
+
+            // Show active defines for debugging
+            oss << "Active defines:";
+#if defined(HAS_BOOL)
+            oss << " HAS_BOOL";
+#endif
+#if defined(HAS_FLOAT)
+            oss << " HAS_FLOAT";
+#endif
+#if defined(HAS_FLOAT32)
+            oss << " HAS_FLOAT32";
+#endif
+#if defined(HAS_DOUBLE)
+            oss << " HAS_DOUBLE";
+#endif
+#if defined(HAS_INT32_T)
+            oss << " HAS_INT32_T";
+#endif
+#if defined(HAS_INT32)
+            oss << " HAS_INT32";
+#endif
+#if defined(HAS_INT)
+            oss << " HAS_INT";
+#endif
+#if defined(HAS_INT64_T)
+            oss << " HAS_INT64_T";
+#endif
+#if defined(HAS_LONG)
+            oss << " HAS_LONG";
+#endif
+
+            if (compiledTypes_.empty()) {
+              oss << "\nNo defines matched - check CMake setup_type_definitions()";
+            }
+        }
+
+        oss << "==========================================\n";
+
+        return oss.str();
+    }
+    
+    // DataTypeValidator implementation
+    void DataTypeValidator::validateOrThrow() {
+        if (!DataTypeValidation::isValidDataType(dataType_)) {
+            std::string errorMsg = DataTypeValidation::getDataTypeErrorMessage(
+                static_cast<int>(dataType_), context_.c_str());
+            THROW_EXCEPTION(errorMsg.c_str());
+        }
+        
+        if (!DataTypeValidation::isCompiledDataType(dataType_)) {
+            std::string errorMsg = DataTypeValidation::getDataTypeErrorMessage(
+                dataType_, context_.c_str());
+            THROW_EXCEPTION(errorMsg.c_str());
+        }
+    }
+
+} // namespace sd

libnd4j/include/array/DataTypeValidation.h

@@ -0,0 +1,156 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 for robust data type validation and error reporting
+//
+
+#ifndef ND4J_DATATYPE_VALIDATION_H
+#define ND4J_DATATYPE_VALIDATION_H
+
+#include <array/DataType.h>
+#include <string>
+#include <vector>
+#include <unordered_set>
+#include <unordered_map>
+#include <system/common.h>
+
+namespace sd {
+    
+    /**
+     * Comprehensive data type validation and error reporting system
+     * 
+     * This class provides robust validation for data types, distinguishing between:
+     * - Invalid data type IDs (not in enum)
+     * - Valid but uncompiled data types (selective builds)
+     * - Runtime vs compile-time mismatches
+     */
+    class SD_LIB_EXPORT DataTypeValidation {
+    private:
+        static std::unordered_set<DataType> compiledTypes_;
+        static std::unordered_map<DataType, std::string> typeNames_;
+        static std::unordered_map<int, DataType> enumMapping_;
+        static bool initialized_;
+        
+        static void initializeIfNeeded();
+        static void populateTypeNames();
+        static void populateCompiledTypes();
+        
+    public:
+        /**
+         * Initialize the validation system (called automatically)
+         */
+        static void initialize();
+        
+        /**
+         * Check if a data type enum value is valid
+         * @param dataType The data type to check
+         * @return true if the data type exists in the enum
+         */
+        static bool isValidDataType(DataType dataType);
+        
+        /**
+         * Check if a numeric data type ID is valid
+         * @param dataTypeId The numeric data type ID to check
+         * @return true if the ID maps to a valid enum value
+         */
+        static bool isValidDataType(int dataTypeId);
+        
+        /**
+         * Check if a data type is compiled into this build
+         * @param dataType The data type to check
+         * @return true if the type is available in this binary
+         */
+        static bool isCompiledDataType(DataType dataType);
+        
+
+        
+        /**
+         * Get a descriptive error message for numeric data type ID
+         * @param dataTypeId The problematic data type ID
+         * @param context Context string (function name, etc.)
+         * @return Detailed error message with suggestions
+         */
+        static std::string getDataTypeErrorMessage(int dataTypeId, const char* context);
+        
+        /**
+         * Get list of available data types in this build
+         * @return Vector of compiled data types
+         */
+        static std::vector<DataType> getAvailableDataTypes();
+        
+        /**
+         * Get human-readable name for data type
+         * @param dataType The data type
+         * @return String name (e.g., "FLOAT32", "INT64")
+         */
+        static const char* getDataTypeName(DataType dataType);
+        
+        /**
+         * Get build configuration information
+         * @return String describing the build configuration and available types
+         */
+        static std::string getBuildInfo();
+        
+        /**
+         * Register a compiled data type (used internally by build system)
+         * @param type The data type to register
+         * @param name Human-readable name
+         */
+        static void registerCompiledType(DataType type, const std::string& name);
+    };
+    
+    /**
+     * RAII helper for data type validation in critical paths
+     * 
+     * Usage:
+     *   DataTypeValidator validator(myDataType, "MyFunction::myMethod()");
+     *   // Will throw if data type is invalid or not compiled
+     */
+    class SD_LIB_EXPORT DataTypeValidator {
+    private:
+        DataType dataType_;
+        std::string context_;
+        
+    public:
+        /**
+         * Constructor that validates the data type
+         * @param dt Data type to validate
+         * @param context Context for error messages
+         * @throws std::exception if validation fails
+         */
+        DataTypeValidator(DataType dt, const std::string& context) 
+            : dataType_(dt), context_(context) {
+            validateOrThrow();
+        }
+        
+        /**
+         * Perform validation and throw if invalid
+         * @throws std::exception with detailed error message
+         */
+        void validateOrThrow();
+        
+        /**
+         * Get the validated data type
+         * @return The data type (guaranteed to be valid and compiled)
+         */
+        DataType getDataType() const { return dataType_; }
+    };
+}
+
+#endif // ND4J_DATATYPE_VALIDATION_H

libnd4j/include/array/DeallocatorServiceLifecycleTracker.h

@@ -0,0 +1,169 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// DeallocatorServiceLifecycleTracker - Tracks DeallocatorService statistics
+// from the Java side for memory leak detection.
+// Only active when SD_GCC_FUNCTRACE is defined.
+//
+
+#ifndef LIBND4J_DEALLOCATORSERVICELIFECYCLETRACKER_H
+#define LIBND4J_DEALLOCATORSERVICELIFECYCLETRACKER_H
+
+#include <cstddef>
+#include <atomic>
+#include <system/common.h>
+
+namespace sd {
+
+namespace analysis {
+class ComprehensiveLeakAnalyzer;
+}
+
+namespace array {
+
+/**
+ * DeallocatorServiceLifecycleTracker - Singleton class for tracking DeallocatorService
+ * statistics from the Java side for memory leak detection.
+ *
+ * This receives snapshot data from Java's DeallocatorService and tracks
+ * memory management patterns.
+ */
+class DeallocatorServiceLifecycleTracker {
+    friend class sd::analysis::ComprehensiveLeakAnalyzer;
+
+public:
+    /**
+     * Get singleton instance
+     */
+    static DeallocatorServiceLifecycleTracker& getInstance() {
+        static DeallocatorServiceLifecycleTracker instance;
+        return instance;
+    }
+
+    /**
+     * Enable tracking
+     */
+    void enable() {
+        _enabled.store(true);
+    }
+
+    /**
+     * Disable tracking
+     */
+    void disable() {
+        _enabled.store(false);
+    }
+
+    /**
+     * Check if tracking is enabled
+     */
+    bool isEnabled() const {
+        return _enabled.load();
+    }
+
+    /**
+     * Record a snapshot from Java DeallocatorService
+     * @param totalAllocations Total allocations count
+     * @param totalDeallocations Total deallocations count
+     * @param totalBytesAllocated Total bytes allocated
+     * @param totalBytesDeallocated Total bytes deallocated
+     * @param peakLiveCount Peak live object count
+     * @param peakBytes Peak bytes in use
+     */
+    void recordSnapshot(LongType totalAllocations, LongType totalDeallocations,
+                       LongType totalBytesAllocated, LongType totalBytesDeallocated,
+                       LongType peakLiveCount, LongType peakBytes) {
+        if (!_enabled.load()) return;
+
+        _totalAllocations.store(totalAllocations);
+        _totalDeallocations.store(totalDeallocations);
+        _totalBytesAllocated.store(totalBytesAllocated);
+        _totalBytesDeallocated.store(totalBytesDeallocated);
+        _peakLiveCount.store(peakLiveCount);
+        _peakBytes.store(peakBytes);
+    }
+
+    /**
+     * Get current live count (allocations - deallocations)
+     */
+    LongType getCurrentLiveCount() const {
+        return _totalAllocations.load() - _totalDeallocations.load();
+    }
+
+    /**
+     * Get current bytes in use
+     */
+    LongType getCurrentBytesInUse() const {
+        return _totalBytesAllocated.load() - _totalBytesDeallocated.load();
+    }
+
+    /**
+     * Get total allocations
+     */
+    LongType getTotalAllocations() const {
+        return _totalAllocations.load();
+    }
+
+    /**
+     * Get total deallocations
+     */
+    LongType getTotalDeallocations() const {
+        return _totalDeallocations.load();
+    }
+
+    /**
+     * Get peak live count
+     */
+    LongType getPeakLiveCount() const {
+        return _peakLiveCount.load();
+    }
+
+    /**
+     * Get peak bytes
+     */
+    LongType getPeakBytes() const {
+        return _peakBytes.load();
+    }
+
+private:
+    DeallocatorServiceLifecycleTracker() : _enabled(false),
+        _totalAllocations(0), _totalDeallocations(0),
+        _totalBytesAllocated(0), _totalBytesDeallocated(0),
+        _peakLiveCount(0), _peakBytes(0) {}
+    ~DeallocatorServiceLifecycleTracker() = default;
+
+    // Disable copy and move
+    DeallocatorServiceLifecycleTracker(const DeallocatorServiceLifecycleTracker&) = delete;
+    DeallocatorServiceLifecycleTracker& operator=(const DeallocatorServiceLifecycleTracker&) = delete;
+    DeallocatorServiceLifecycleTracker(DeallocatorServiceLifecycleTracker&&) = delete;
+    DeallocatorServiceLifecycleTracker& operator=(DeallocatorServiceLifecycleTracker&&) = delete;
+
+    std::atomic<bool> _enabled;
+    std::atomic<LongType> _totalAllocations;
+    std::atomic<LongType> _totalDeallocations;
+    std::atomic<LongType> _totalBytesAllocated;
+    std::atomic<LongType> _totalBytesDeallocated;
+    std::atomic<LongType> _peakLiveCount;
+    std::atomic<LongType> _peakBytes;
+};
+
+} // namespace array
+} // namespace sd
+
+#endif // LIBND4J_DEALLOCATORSERVICELIFECYCLETRACKER_H

libnd4j/include/array/ExtraArguments.h

@@ -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
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+
+#ifndef DEV_TESTS_EXTRAARGUMENTS_H
+#define DEV_TESTS_EXTRAARGUMENTS_H
+
+#include <array/DataType.h>
+#include <stdlib.h>
+#include <system/common.h>
+
+#include <initializer_list>
+#include <vector>
+
+namespace sd {
+class SD_LIB_EXPORT ExtraArguments {
+ private:
+  std::vector<double> _fpArgs;
+  std::vector<LongType> _intArgs;
+
+  std::vector<Pointer> _pointers;
+
+  template <typename T>
+  void convertAndCopy(Pointer pointer, LongType offset);
+
+  void *allocate(size_t length, size_t elementSize);
+
+ public:
+  explicit ExtraArguments(std::initializer_list<double> arguments);
+  explicit ExtraArguments(std::initializer_list<LongType> arguments);
+
+  explicit ExtraArguments(const std::vector<double> &arguments);
+  explicit ExtraArguments(const std::vector<int> &arguments);
+  explicit ExtraArguments(const std::vector<LongType> &arguments);
+
+  explicit ExtraArguments();
+  ~ExtraArguments();
+
+  template <typename T>
+  void *argumentsAsT(LongType offset = 0);
+
+  void *argumentsAsT(DataType dataType, LongType offset = 0);
+
+  size_t length();
+};
+}  // namespace sd
+
+#endif  // DEV_TESTS_EXTRAARGUMENTS_H

libnd4j/include/array/InteropDataBuffer.h

@@ -0,0 +1,120 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/DataBuffer.h>
+#include <array/DataType.h>
+#include <system/common.h>
+
+#include <memory>
+
+#ifndef LIBND4J_INTEROPDATABUFFER_H
+#define LIBND4J_INTEROPDATABUFFER_H
+
+namespace sd {
+/**
+ * This class is a wrapper for DataBuffer, suitable for sharing DataBuffer between front-end and back-end languages
+ */
+class SD_LIB_EXPORT InteropDataBuffer {
+ private:
+  DataBuffer *_dataBuffer = nullptr;
+  bool owner;
+  DataType _dataType = DataType::UNKNOWN;
+ public:
+  size_t _cachedLenInBytes = 0;
+  void* _cachedPrimaryPtr = nullptr;  // Cached for deallocation tracking
+  void* _cachedSpecialPtr = nullptr;  // Cached for deallocation tracking
+  bool _closed = false;
+  bool isConstant = false;
+
+  InteropDataBuffer(InteropDataBuffer *dataBuffer, uint64_t length);
+  InteropDataBuffer(DataBuffer * databuffer);
+  InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth);
+  ~InteropDataBuffer() {
+    // Mark as closed FIRST to prevent any concurrent access
+    _closed = true;
+    // If we own the DataBuffer and it's not constant, clean it up
+    if (owner && !isConstant && _dataBuffer != nullptr) {
+      delete _dataBuffer;
+    }
+    // Always null out the pointer to prevent use-after-free
+    _dataBuffer = nullptr;
+  }
+
+#ifndef __JAVACPP_HACK__
+  DataBuffer * getDataBuffer() const;
+  DataBuffer * dataBuffer();
+  // In InteropDataBuffer class, add these public methods:
+  // Check if the underlying DataBuffer pointer is still valid
+  // by checking if it's non-null
+  bool hasValidDataBuffer() const { return _dataBuffer != nullptr; }
+
+  // Get the DataBuffer pointer directly without validation
+  DataBuffer* getDataBufferDirect() const { return _dataBuffer; }
+
+  // Mark the DataBuffer as invalid (set to null)
+  void invalidateDataBuffer() { _dataBuffer = nullptr; }
+  bool isOwner() const { return owner; }
+
+  // Safe accessor that doesn't touch other fields
+  DataBuffer* getDataBufferUnsafe() const {
+    return _dataBuffer;
+  }
+#endif
+
+  void printDbAllocationTrace();
+
+  void *primary() const;
+  void *special() const;
+
+  void markConstant(bool reallyConstant) {
+    isConstant = reallyConstant;
+    dataBuffer()->markConstant(reallyConstant);
+  }
+
+  void setPrimary(void *ptr, size_t length);
+  void setSpecial(void *ptr, size_t length);
+
+  void expand(size_t newlength);
+
+  int deviceId() const;
+  void setDeviceId(int deviceId);
+
+  //updates whether the buffer is the owner of its associated buffers or not.
+  void markOwner(bool owner);
+
+  int useCount() const;
+
+  static void registerSpecialUse(const std::vector<const InteropDataBuffer *> &writeList,
+                                 const std::vector<const InteropDataBuffer *> &readList);
+  static void prepareSpecialUse(const std::vector<const InteropDataBuffer *> &writeList,
+                                const std::vector<const InteropDataBuffer *> &readList,
+                                bool synchronizeWritables = false);
+
+  static void registerPrimaryUse(const std::vector<const InteropDataBuffer *> &writeList,
+                                 const std::vector<const InteropDataBuffer *> &readList);
+  static void preparePrimaryUse(const std::vector<const InteropDataBuffer *> &writeList,
+                                const std::vector<const InteropDataBuffer *> &readList,
+                                bool synchronizeWritables = false);
+};
+}  // namespace sd
+
+#endif  // LIBND4J_INTEROPDATABUFFER_H

libnd4j/include/array/NDArrayFactory.h

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// Created by raver119 on 2018-09-16.
+// @author Oleg Semeniv <oleg.semeniv@gmail.com>
+// @author Abdelrauf
+
+#ifndef DEV_TESTS_NDARRAYFACTORY_H
+#define DEV_TESTS_NDARRAYFACTORY_H
+#include <array/NDArray.h>
+
+#include <initializer_list>
+#include <vector>
+#include <execution/LaunchContext.h>
+
+#include <string>
+
+namespace sd {
+
+class SD_LIB_EXPORT NDArrayFactory {
+ private:
+  template <typename T>
+  static void memcpyFromVector(void *ptr, const std::vector<T> &vector);
+
+ public:
+  template <typename T>
+  static NDArray *empty_(LaunchContext *context = LaunchContext ::defaultContext());
+
+  static NDArray *empty_(DataType dataType, LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray empty(LaunchContext *context = LaunchContext ::defaultContext());
+
+  static NDArray empty(DataType dataType, LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *valueOf(const std::initializer_list<LongType> &shape, T value, char order = 'c',
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *valueOf(std::vector<LongType> &shape, T value, const char order = 'c',
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  static NDArray *valueOf(std::vector<LongType> &shape, NDArray&value, const char order = 'c',
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *linspace(T from, T to, LongType numElements);
+
+  static NDArray create(ShapeDescriptor *shapeDescriptor, LaunchContext *context = LaunchContext ::defaultContext());
+
+  static NDArray create(DataType dtype, LaunchContext *context = LaunchContext ::defaultContext());
+
+
+  template <typename T>
+  static NDArray *create_(const T value, LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray *create_(DataType dtype, LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray create(const T value, LaunchContext *context = LaunchContext ::defaultContext());
+  template <typename T>
+  static NDArray create(DataType type, const T scalar, LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *vector(LongType length, T startingValue = (T)0,
+                         LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *create_(const char order, std::vector<LongType> &shape,
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  static NDArray *create_(const char order, std::vector<LongType> &shape, DataType dataType,
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray *create_(char order, const std::vector<LongType> &shape, const std::vector<T> &data,
+                          LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray create(char order, const std::vector<LongType> &shape, const std::vector<T> &data,
+                        LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray create(char order, const std::vector<LongType> &shape,
+                        LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray create(char order, const std::vector<LongType> &shape, DataType dtype,
+                        LaunchContext *context = LaunchContext ::defaultContext());
+
+  template <typename T>
+  static NDArray create(const std::vector<T> &values, LaunchContext *context = LaunchContext ::defaultContext());
+
+#ifndef __JAVACPP_HACK__
+  // this method only available out of javacpp
+
+  template <typename T>
+  static NDArray create(T *buffer, char order, const std::initializer_list<LongType> &shape,
+                        LaunchContext *context = LaunchContext ::defaultContext());
+
+  /**
+   * This method creates NDArray from .npy file
+   * @param fileName
+   * @return
+   */
+  static NDArray fromNpyFile(const char *fileName);
+
+  #if defined(HAS_UTF8)
+  /**
+   * This factory create array from utf8 string
+   * @return NDArray default dataType UTF8
+   */
+  static NDArray string(const char *string, DataType dtype = UTF8,
+                        LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray *string_(const char *string, DataType dtype = UTF8,
+                          LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray *string_(const std::string &string, DataType dtype = UTF8,
+                          LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray string(const std::string &string, DataType dtype = UTF8,
+                        LaunchContext *context = LaunchContext::defaultContext());
+
+  static NDArray string(std::vector<LongType> &shape, const std::vector<const char *> &strings,
+                        DataType dataType = UTF8, LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::vector<std::string> &string,
+                        DataType dataType = UTF8, LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<const char *> &strings,
+                          DataType dataType = UTF8, LaunchContext *context = LaunchContext ::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<std::string> &string,
+                          DataType dataType = UTF8, LaunchContext *context = LaunchContext ::defaultContext());
+#endif
+
+#if defined(HAS_UTF16)
+  /**
+   * This factory create array from utf16 string
+   * @return NDArray default dataType UTF16
+   */
+  static NDArray string(const char16_t *u16string, DataType dtype = UTF16,
+                        LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(const char16_t *u16string, DataType dtype = UTF16,
+                          LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(const std::u16string &u16string, DataType dtype = UTF16,
+                          LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(const std::u16string &u16string, DataType dtype = UTF16,
+                        LaunchContext *context = LaunchContext::defaultContext());
+
+  /**
+   * This factory create array from vector of utf16 strings
+   * @return NDArray default dataType UTF16
+   */
+  static NDArray string(std::vector<LongType> &shape, const std::initializer_list<const char16_t *> &strings,
+                        DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::initializer_list<std::u16string> &string,
+                        DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::vector<const char16_t *> &strings,
+                        DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::vector<std::u16string> &string,
+                        DataType dtype = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape,
+                          const std::initializer_list<const char16_t *> &strings,
+                          DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::initializer_list<std::u16string> &string,
+                          DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<const char16_t *> &strings,
+                          DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<std::u16string> &string,
+                          DataType dataType = UTF16, LaunchContext *context = LaunchContext::defaultContext());
+#endif
+
+#if defined(HAS_UTF32)
+  /**
+   * This factory create array from utf32 string
+   * @return NDArray default dataType UTF32
+   */
+  static NDArray string(const char32_t *u32string, DataType dtype = UTF32,
+                        LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(const char32_t *u32string, DataType dtype = UTF32,
+                          LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(const std::u32string &u32string, DataType dtype = UTF32,
+                          LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(const std::u32string &u32string, DataType dtype = UTF32,
+                        LaunchContext *context = LaunchContext::defaultContext());
+
+  /**
+   * This factory create array from vector of utf32 strings
+   * @return NDArray default dataType UTF32
+   */
+  static NDArray string(std::vector<LongType> &shape, const std::initializer_list<const char32_t *> &strings,
+                        DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::initializer_list<std::u32string> &string,
+                        DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray string(std::vector<LongType> &shape, const std::vector<const char32_t *> &strings, DataType dtype,
+                        LaunchContext *context);
+  static NDArray string(std::vector<LongType> &shape, const std::vector<std::u32string> &string,
+                        DataType dtype = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape,
+                          const std::initializer_list<const char32_t *> &strings,
+                          DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::initializer_list<std::u32string> &string,
+                          DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<const char32_t *> &strings,
+                          DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+  static NDArray *string_(std::vector<LongType> &shape, const std::vector<std::u32string> &string,
+                          DataType dataType = UTF32, LaunchContext *context = LaunchContext::defaultContext());
+#endif
+
+#endif
+};
+}  // namespace sd
+
+#endif  // DEV_TESTS_NDARRAYFACTORY_H

libnd4j/include/array/NDArrayHelpers.hXX

@@ -0,0 +1,769 @@
+/*
+*  ******************************************************************************
+*  *
+*  *
+*  * 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
+*  *****************************************************************************
+*/
+
+/**
+* Streamlined implementation of NDArray operators
+ */
+#include <array/NDArray.h>
+#include <array/NDArrayFactory.h>
+#include <legacy/NativeOpExecutioner.h>
+#include <system/Environment.h>
+
+namespace sd {
+
+/////////////////////////////////////////////////////
+// Scalar-Array Operators - Helper Functions
+/////////////////////////////////////////////////////
+SD_LIB_EXPORT NDArray* operator+(NDArray& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator+(NDArray&& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator+(NDArray& arr1, NDArray&& arr2);
+SD_LIB_EXPORT NDArray* operator+(NDArray&& arr1, NDArray&& arr2);
+
+// Subtraction: NDArray - NDArray
+SD_LIB_EXPORT NDArray* operator-(NDArray& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator-(NDArray&& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator-(NDArray& arr1, NDArray&& arr2);
+SD_LIB_EXPORT NDArray* operator-(NDArray&& arr1, NDArray&& arr2);
+
+// Multiplication: NDArray * NDArray
+SD_LIB_EXPORT NDArray* operator*(NDArray& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator*(NDArray&& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator*(NDArray& arr1, NDArray&& arr2);
+SD_LIB_EXPORT NDArray* operator*(NDArray&& arr1, NDArray&& arr2);
+
+// Division: NDArray / NDArray
+SD_LIB_EXPORT NDArray* operator/(NDArray& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator/(NDArray&& arr1, NDArray& arr2);
+SD_LIB_EXPORT NDArray* operator/(NDArray& arr1, NDArray&& arr2);
+SD_LIB_EXPORT NDArray* operator/(NDArray&& arr1, NDArray&& arr2);
+
+
+template <typename T>
+NDArray* scalarArrayOpHelper(NDArray* arr, T scalar, sd::scalar::Ops op, bool inPlace) {
+  if (arr->isS()) {
+    std::string errorMessage = "scalarArrayOpHelper: you can't use this method on String array!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto tmp = NDArrayFactory::create(arr->dataType(), scalar, arr->getContext());
+
+  if (inPlace && !arr->isView()) {
+    NDArray::prepareSpecialUse({arr}, {arr, tmp});
+    NativeOpExecutioner::execScalar(arr->getContext(), op,
+                                    arr->buffer(), arr->shapeInfo(), arr->specialBuffer(), arr->specialShapeInfo(),
+                                    arr->buffer(), arr->shapeInfo(), arr->specialBuffer(), arr->specialShapeInfo(),
+                                    tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({arr}, {arr, tmp});
+    delete tmp;
+    return arr;
+  } else {
+    NDArray *result = new NDArray(arr->shapeInfo(), DataTypeUtils::pickPairwiseResultType(arr->dataType(), DataTypeUtils::fromT<T>()),
+                                  false, arr->getContext());
+    NDArray::prepareSpecialUse({result}, {arr, tmp});
+    NativeOpExecutioner::execScalar(arr->getContext(), op,
+                                    arr->buffer(), arr->shapeInfo(), arr->specialBuffer(), arr->specialShapeInfo(),
+                                    result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(),
+                                    tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {arr, tmp});
+    delete tmp;
+    return result;
+  }
+}
+
+/////////////////////////////////////////////////////
+// Array-Array Operators - Helper Function
+/////////////////////////////////////////////////////
+
+NDArray* arrayArrayOpHelper(NDArray* arr1, NDArray* arr2, sd::pairwise::Ops op, bool inPlace) {
+  if (arr1->isS() || arr2->isS()) {
+    std::string errorMessage = "arrayArrayOpHelper: you can't use this method on String arrays!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (!Environment::getInstance().isExperimentalBuild() && arr1->dataType() != arr2->dataType() &&
+      (arr1->dataType() != DataType::BOOL || arr2->dataType() != BOOL)) {
+    std::string errorMessage;
+    errorMessage += "arrayArrayOpHelper: Cannot operate on arrays of different types: ";
+    errorMessage += DataTypeUtils::asString(arr1->dataType());
+    errorMessage += " and ";
+    errorMessage += DataTypeUtils::asString(arr2->dataType());
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+  if (arr1->lengthOf() == arr2->lengthOf() && arr1->rankOf() == arr2->rankOf()) {
+    NDArray* result = nullptr;
+    const bool canUseArr1 = inPlace && !arr1->isView();
+    const bool canUseArr2 = inPlace && !arr2->isView();
+
+    if (canUseArr1)
+      result = arr1;
+    else if (canUseArr2)
+      result = arr2;
+    else
+      result = new NDArray(arr1->shapeInfo(),
+                           DataTypeUtils::pickPairwiseResultType(arr1->shapeInfo(), arr2->shapeInfo()),
+                           false, arr1->getContext());
+
+    NDArray::prepareSpecialUse({result}, {arr1, arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1->getContext(), op,
+        arr1->buffer(), arr1->shapeInfo(), arr1->specialBuffer(), arr1->specialShapeInfo(),
+        arr2->buffer(), arr2->shapeInfo(), arr2->specialBuffer(), arr2->specialShapeInfo(),
+        result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {arr1, arr2});
+
+    return result;
+  }
+
+  // Broadcast case
+  sd::BroadcastOpsTuple broadcastOp;
+  switch (op) {
+    case sd::pairwise::Add: broadcastOp = sd::BroadcastOpsTuple::Add(); break;
+    case sd::pairwise::Subtract: broadcastOp = sd::BroadcastOpsTuple::Subtract(); break;
+    case sd::pairwise::Multiply: broadcastOp = sd::BroadcastOpsTuple::Multiply(); break;
+    case sd::pairwise::Divide: broadcastOp = sd::BroadcastOpsTuple::Divide(); break;
+    default:
+      std::string errorMessage = "arrayArrayOpHelper: Unsupported pairwise operation for broadcasting";
+      THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  return arr1->applyTrueBroadcast(broadcastOp, arr2);
+}
+
+/////////////////////////////////////////////////////
+// Addition Operators (NDArray + scalar)
+/////////////////////////////////////////////////////
+
+template <typename T, typename>
+NDArray* operator+(NDArray& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Add, false);
+}
+
+template <typename T, typename>
+NDArray* operator+(NDArray&& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Add, true);
+}
+
+template <typename T, typename>
+NDArray* operator+(T scalar, NDArray& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Add, false);
+}
+
+template <typename T, typename>
+NDArray* operator+(T scalar, NDArray&& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Add, true);
+}
+
+#define INSTANTIATE_SCALAR_OP(OP, T) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(T), \
+    template SD_LIB_EXPORT NDArray* operator OP<GET_SECOND(T), typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<GET_SECOND(T)>::value>::type>(NDArray& arr, GET_SECOND(T) scalar); \
+    template SD_LIB_EXPORT NDArray* operator OP<GET_SECOND(T), typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<GET_SECOND(T)>::value>::type>(NDArray&& arr, GET_SECOND(T) scalar); \
+    template SD_LIB_EXPORT NDArray* operator OP<GET_SECOND(T), typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<GET_SECOND(T)>::value>::type>(GET_SECOND(T) scalar, NDArray& arr); \
+    template SD_LIB_EXPORT NDArray* operator OP<GET_SECOND(T), typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<GET_SECOND(T)>::value>::type>(GET_SECOND(T) scalar, NDArray&& arr); \
+))
+#define INSTANTIATE_ALL_SCALAR_OPS(T) \
+ INSTANTIATE_SCALAR_OP(+, T) \
+ INSTANTIATE_SCALAR_OP(-, T) \
+ INSTANTIATE_SCALAR_OP(*, T) \
+ INSTANTIATE_SCALAR_OP(/, T)
+
+ITERATE_LIST((SD_NUMERIC_TYPES), INSTANTIATE_ALL_SCALAR_OPS)
+
+/////////////////////////////////////////////////////
+// Subtraction Operators (NDArray - scalar)
+/////////////////////////////////////////////////////
+
+template <typename T, typename>
+NDArray* operator-(NDArray& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Subtract, false);
+}
+
+template <typename T, typename>
+NDArray* operator-(NDArray&& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Subtract, true);
+}
+
+template <typename T, typename>
+NDArray* operator-(T scalar, NDArray& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::ReverseSubtract, false);
+}
+
+template <typename T, typename>
+NDArray* operator-(T scalar, NDArray&& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::ReverseSubtract, true);
+}
+
+/////////////////////////////////////////////////////
+// Multiplication Operators (NDArray * scalar)
+/////////////////////////////////////////////////////
+
+template <typename T, typename>
+NDArray* operator*(NDArray& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Multiply, false);
+}
+
+template <typename T, typename>
+NDArray* operator*(NDArray&& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Multiply, true);
+}
+
+template <typename T, typename>
+NDArray* operator*(T scalar, NDArray& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Multiply, false);
+}
+
+template <typename T, typename>
+NDArray* operator*(T scalar, NDArray&& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Multiply, true);
+}
+
+/////////////////////////////////////////////////////
+// Division Operators (NDArray / scalar)
+/////////////////////////////////////////////////////
+
+template <typename T, typename>
+NDArray* operator/(NDArray& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Divide, false);
+}
+
+template <typename T, typename>
+NDArray* operator/(NDArray&& arr, T scalar) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::Divide, true);
+}
+
+template <typename T, typename>
+NDArray* operator/(T scalar, NDArray& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::ReverseDivide, false);
+}
+
+template <typename T, typename>
+NDArray* operator/(T scalar, NDArray&& arr) {
+  return scalarArrayOpHelper(&arr, scalar, sd::scalar::ReverseDivide, true);
+}
+
+/////////////////////////////////////////////////////
+// Array-Array Operators
+/////////////////////////////////////////////////////
+
+// Addition: NDArray + NDArray
+NDArray* operator+(NDArray& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Add, false);
+}
+
+NDArray* operator+(NDArray&& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Add, true);
+}
+
+NDArray* operator+(NDArray& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Add, true);
+}
+
+NDArray* operator+(NDArray&& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Add, true);
+}
+
+// Subtraction: NDArray - NDArray
+NDArray* operator-(NDArray& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Subtract, false);
+}
+
+NDArray* operator-(NDArray&& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Subtract, true);
+}
+
+NDArray* operator-(NDArray& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Subtract, true);
+}
+
+NDArray* operator-(NDArray&& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Subtract, true);
+}
+
+// Multiplication: NDArray * NDArray
+NDArray* operator*(NDArray& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Multiply, false);
+}
+
+NDArray* operator*(NDArray&& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Multiply, true);
+}
+
+NDArray* operator*(NDArray& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Multiply, true);
+}
+
+NDArray* operator*(NDArray&& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Multiply, true);
+}
+
+// Division: NDArray / NDArray
+NDArray* operator/(NDArray& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Divide, false);
+}
+
+NDArray* operator/(NDArray&& arr1, NDArray& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Divide, true);
+}
+
+NDArray* operator/(NDArray& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Divide, true);
+}
+
+NDArray* operator/(NDArray&& arr1, NDArray&& arr2) {
+  return arrayArrayOpHelper(&arr1, &arr2, sd::pairwise::Divide, true);
+}
+
+/////////////////////////////////////////////////////
+// Compound Assignment Operators
+/////////////////////////////////////////////////////
+
+// Direct implementation for compound assignment with scalar values
+template <typename T, typename>
+void NDArray::operator+=(const T scalar) {
+  if (isS()) {
+    std::string errorMessage = "NDArray::operator+=: you can't use this method on String array!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto tmp = NDArrayFactory::create(dataType(), scalar, getContext());
+
+  NDArray::prepareSpecialUse({this}, {this, tmp});
+  NativeOpExecutioner::execScalar(getContext(), sd::scalar::Add,
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+  NDArray::registerSpecialUse({this}, {this, tmp});
+  delete tmp;
+
+}
+
+template <typename T, typename>
+void NDArray::operator-=(const T scalar) {
+  if (isS()) {
+    std::string errorMessage = "NDArray::operator-=: you can't use this method on String array!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto tmp = NDArrayFactory::create(dataType(), scalar, getContext());
+
+  NDArray::prepareSpecialUse({this}, {this, tmp});
+  NativeOpExecutioner::execScalar(getContext(), sd::scalar::Subtract,
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+  NDArray::registerSpecialUse({this}, {this, tmp});
+  delete tmp;
+
+}
+
+template <typename T, typename>
+void NDArray::operator*=(const T scalar) {
+  if (isS()) {
+    std::string errorMessage = "NDArray::operator*=: you can't use this method on String array!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto tmp = NDArrayFactory::create(dataType(), scalar, getContext());
+
+  NDArray::prepareSpecialUse({this}, {this, tmp});
+  NativeOpExecutioner::execScalar(getContext(), sd::scalar::Multiply,
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+  NDArray::registerSpecialUse({this}, {this, tmp});
+  delete tmp;
+
+}
+
+template <typename T, typename>
+void NDArray::operator/=(const T scalar) {
+  if (isS()) {
+    std::string errorMessage = "NDArray::operator/=: you can't use this method on String array!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto tmp = NDArrayFactory::create(dataType(), scalar, getContext());
+
+  NDArray::prepareSpecialUse({this}, {this, tmp});
+  NativeOpExecutioner::execScalar(getContext(), sd::scalar::Divide,
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                  tmp->buffer(), tmp->shapeInfo(), tmp->specialBuffer(), tmp->specialShapeInfo(), nullptr);
+  NDArray::registerSpecialUse({this}, {this, tmp});
+  delete tmp;
+
+}
+
+// Define instantiations matching the header file declarations exactly
+#define INSTANTIATE_COMPOUND_ASSIGN_OP(OP, T) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(T), \
+    template SD_LIB_EXPORT void NDArray::operator OP<GET_SECOND(T), typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<GET_SECOND(T)>::value>::type>(const GET_SECOND(T) scalar); \
+))
+#define INSTANTIATE_ALL_COMPOUND_ASSIGN_OPS(T) \
+ INSTANTIATE_COMPOUND_ASSIGN_OP(+=, T) \
+ INSTANTIATE_COMPOUND_ASSIGN_OP(-=, T) \
+ INSTANTIATE_COMPOUND_ASSIGN_OP(*=, T) \
+ INSTANTIATE_COMPOUND_ASSIGN_OP(/=, T)
+
+ITERATE_LIST((SD_NUMERIC_TYPES), INSTANTIATE_ALL_COMPOUND_ASSIGN_OPS)
+
+// Helper for compound assignment operators with arrays
+void compoundAssignArray(NDArray& arr1, NDArray& arr2, sd::pairwise::Ops op, sd::BroadcastOpsTuple broadcastOp) {
+  if (arr1.isS() || arr2.isS()) {
+    std::string errorMessage = "compoundAssignArray: you can't use this method on String arrays!";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (!Environment::getInstance().isExperimentalBuild() && arr1.dataType() != arr2.dataType() &&
+      (arr1.dataType() != DataType::BOOL || arr2.dataType() != BOOL)) {
+    std::string errorMessage;
+    errorMessage += "compoundAssignArray: Cannot operate on arrays of different types: ";
+    errorMessage += DataTypeUtils::asString(arr1.dataType());
+    errorMessage += " and ";
+    errorMessage += DataTypeUtils::asString(arr2.dataType());
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (arr1.lengthOf() == arr2.lengthOf() && arr1.rankOf() == arr2.rankOf()) {
+    NDArray::prepareSpecialUse({&arr1}, {&arr1, &arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1.getContext(), op,
+        arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(), arr1.specialShapeInfo(),
+        arr2.buffer(), arr2.shapeInfo(), arr2.specialBuffer(), arr2.specialShapeInfo(),
+        arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(), arr1.specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({&arr1}, {&arr1, &arr2});
+  } else {
+    sd::LongType *bShape = nullptr;
+    if (!ShapeUtils::evalBroadcastShapeInfo(arr1.shapeInfo(), arr2.shapeInfo(), true, bShape, arr1.getContext()->getWorkspace())) {
+      std::string errorMessage = "compoundAssignArray: the shapes are not suitable for broadcast operation!";
+      THROW_EXCEPTION(errorMessage.c_str());
+    }
+
+    if (shape::equalsTypesAndShapesSoft(arr1.shapeInfo(), bShape)) {
+      arr1.applyTrueBroadcast(broadcastOp, &arr2, &arr1, false);
+    } else {
+      NDArray result(bShape, true, arr1.getContext());
+      arr1.applyTrueBroadcast(broadcastOp, &arr2, &result, false);
+      arr1 = std::move(result);
+    }
+  }
+}
+
+// Compound assignment with arrays
+void NDArray::operator+=(NDArray& other) {
+  compoundAssignArray(*this, other, sd::pairwise::Add, sd::BroadcastOpsTuple::Add());
+}
+
+void NDArray::operator+=(NDArray&& other) {
+  operator+=(other);
+}
+
+void NDArray::operator-=(NDArray& other) {
+  compoundAssignArray(*this, other, sd::pairwise::Subtract, sd::BroadcastOpsTuple::Subtract());
+}
+
+void NDArray::operator-=(NDArray&& other) {
+  operator-=(other);
+}
+
+void NDArray::operator*=(NDArray& other) {
+  compoundAssignArray(*this, other, sd::pairwise::Multiply, sd::BroadcastOpsTuple::Multiply());
+}
+
+void NDArray::operator*=(NDArray&& other) {
+  operator*=(other);
+}
+
+void NDArray::operator/=(NDArray& other) {
+  compoundAssignArray(*this, other, sd::pairwise::Divide, sd::BroadcastOpsTuple::Divide());
+}
+
+void NDArray::operator/=(NDArray&& other) {
+  operator/=(other);
+}
+
+/////////////////////////////////////////////////////
+// Pointer-Based Binary Operators (Forward Declarations)
+/////////////////////////////////////////////////////
+
+
+/////////////////////////////////////////////////////
+// Pointer-Scalar Operator Templates (Declarations)
+/////////////////////////////////////////////////////
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator+(NDArray* arr, T scalar);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator+(T scalar, NDArray* arr);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator-(NDArray* arr, T scalar);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator-(T scalar, NDArray* arr);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator*(NDArray* arr, T scalar);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator*(T scalar, NDArray* arr);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator/(NDArray* arr, T scalar);
+
+template <typename T, typename = typename std::enable_if<DataTypeUtils::scalarTypesForNDarray<T>::value>::type>
+SD_LIB_EXPORT NDArray* operator/(T scalar, NDArray* arr);
+
+// Unary negative operator
+NDArray NDArray::operator-() & {
+  NDArray result(shapeInfo(), false, getContext());
+
+  NDArray::prepareSpecialUse({&result}, {this});
+  NativeOpExecutioner::execTransformSame(getContext(), sd::transform::Neg,
+                                         buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                         result.buffer(), result.shapeInfo(), result.specialBuffer(),
+                                         result.specialShapeInfo(), nullptr, nullptr, nullptr);
+  NDArray::registerSpecialUse({&result}, {this});
+
+  return result;
+}
+
+NDArray NDArray::operator-() && {
+  NDArray::prepareSpecialUse({this}, {this});
+  NativeOpExecutioner::execTransformSame(getContext(), sd::transform::Neg,
+                                         buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                         buffer(), shapeInfo(), specialBuffer(), specialShapeInfo(),
+                                         nullptr, nullptr, nullptr);
+  NDArray::registerSpecialUse({this}, {this});
+
+  return std::move(*this);
+}
+
+template <typename T1, typename T2, typename>
+NDArray* operator-(T1 &&arr1, T2 &&arr2) {
+  if (arr1.isS() || arr2.isS())
+    THROW_EXCEPTION("operator-(T&& arr1, T&& arr2): you can't use this method on String arrays!");
+  if (!Environment::getInstance().isExperimentalBuild() && arr1.dataType() != arr2.dataType() &&
+      (arr1.dataType() != DataType::BOOL || arr2.dataType() != BOOL)) {
+    // Always throw exception - never use _exit(1) which kills the entire process
+    THROW_EXCEPTION(sd::datatype_exception::build("operator-(T&& arr1, T&& arr2): Cannot subtract different types",
+                                        arr1.dataType(), arr2.dataType()).what());
+  }
+
+  PointersManager pointersManager(arr1.getContext(), "operator-(T&& arr1, T&& arr2)");
+
+  if (arr1.lengthOf() == arr2.lengthOf() && arr1.rankOf() == arr2.rankOf()) {
+    const bool isArr1Rvalue = !std::is_reference<T1>::value && !arr1.isView();
+    const bool isArr2Rvalue = !std::is_reference<T2>::value && !arr2.isView();
+
+    NDArray *result = nullptr;
+    if (isArr1Rvalue)
+      result = const_cast<NDArray *>(&arr1);
+    else if (isArr2Rvalue)
+      result = const_cast<NDArray *>(&arr2);
+    else
+      result = new NDArray(arr1.shapeInfo(), DataTypeUtils::pickPairwiseResultType(arr1.shapeInfo(), arr2.shapeInfo()),
+                           false, arr1.getContext());
+
+    NDArray::prepareSpecialUse({result}, {&arr1, &arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1.getContext(), sd::pairwise::Subtract, arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(),
+        arr1.specialShapeInfo(), arr2.buffer(), arr2.shapeInfo(), arr2.specialBuffer(), arr2.specialShapeInfo(),
+        result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {&arr1, &arr2});
+
+    return result;
+  }
+
+  NDArray* arr1Addr = &arr1;
+  NDArray* arr2Addr = &arr2;
+  return std::forward<T1>(arr1).applyTrueBroadcast(sd::BroadcastOpsTuple::Subtract(), arr2Addr);
+}
+
+
+
+template SD_LIB_EXPORT NDArray* operator-<NDArray &, NDArray &, void>(NDArray &arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator-<NDArray &, NDArray, void>(NDArray &arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* operator-<NDArray, NDArray &, void>(NDArray &&arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator-<NDArray, NDArray, void>(NDArray &&arr1, NDArray &&arr2);
+
+
+template <typename T1, typename T2, typename>
+NDArray* operator*(T1 &&arr1, T2 &&arr2) {
+  if (arr1.isS() || arr2.isS()) {
+    THROW_EXCEPTION("operator*(T&& arr1, T&& arr2): you can't use this method on String arrays!");
+  }
+  if (!Environment::getInstance().isExperimentalBuild() && arr1.dataType() != arr2.dataType() &&
+      (arr1.dataType() != DataType::BOOL || arr2.dataType() != BOOL)) {
+    std::string errorMessage;
+    errorMessage += "operator*(T&& arr1, T&& arr2): Cannot multiply different types";
+    errorMessage += " arr1.dataType()=";
+    errorMessage += DataTypeUtils::asString(arr1.dataType());
+    errorMessage += " arr2.dataType()=";
+    errorMessage += DataTypeUtils::asString(arr2.dataType());
+    errorMessage += " arr1.shapeInfo()=";
+    errorMessage += ShapeUtils::shapeAsString(arr1.shapeInfo());
+    errorMessage += " arr2.shapeInfo()=";
+    errorMessage += ShapeUtils::shapeAsString(arr2.shapeInfo());
+    errorMessage += " arr1.ordering()=";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+  PointersManager pointersManager(arr1.getContext(), "operator*(T&& arr1, T&& arr2)");
+
+  if (arr1.lengthOf() == arr2.lengthOf() && arr1.rankOf() == arr2.rankOf()) {
+    const bool isArr1Rvalue = !std::is_reference<T1>::value && !arr1.isView();
+    const bool isArr2Rvalue = !std::is_reference<T2>::value && !arr2.isView();
+
+    NDArray *result = nullptr;
+    if (isArr1Rvalue) {
+      result = const_cast<NDArray *>(&arr1);
+    }    else if (isArr2Rvalue) {
+      result = const_cast<NDArray *>(&arr2);
+    } else {
+      result = new NDArray(arr1.shapeInfo(), DataTypeUtils::pickPairwiseResultType(arr1.shapeInfo(), arr2.shapeInfo()),
+                           false, arr1.getContext());
+    }
+    NDArray::prepareSpecialUse({result}, {&arr1, &arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1.getContext(), sd::pairwise::Multiply, arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(),
+        arr1.specialShapeInfo(), arr2.buffer(), arr2.shapeInfo(), arr2.specialBuffer(), arr2.specialShapeInfo(),
+        result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {&arr1, &arr2});
+
+    return result;
+  }
+
+  NDArray* arr1Addr = &arr1;
+  NDArray* arr2Addr = &arr2;
+  return std::forward<T1>(arr1).applyTrueBroadcast(sd::BroadcastOpsTuple::Multiply(), arr2Addr);
+}
+template SD_LIB_EXPORT NDArray* operator*<NDArray &, NDArray &, void>(NDArray &arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator*<NDArray &, NDArray, void>(NDArray &arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* operator*<NDArray, NDArray &, void>(NDArray &&arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator*<NDArray&&, NDArray &&, void>(NDArray &&arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* sd::operator*<NDArray, NDArray, void>(sd::NDArray&&, sd::NDArray&&);
+
+
+template <typename T1, typename T2, typename>
+NDArray* operator+(T1 &&arr1, T2 &&arr2) {
+  if (arr1.isS() || arr2.isS())
+    THROW_EXCEPTION("operator+(T&& arr1, T&& arr2): you can't use this method on String arrays!");
+  if (arr1.dataType() != arr2.dataType() &&
+      (arr1.dataType() != DataType::BOOL || arr2.dataType() != BOOL)) {
+    std::string errorMessage;
+    errorMessage += "operator+(T&& arr1, T&& arr2): Cannot multiply different types";
+    errorMessage += " arr1.dataType()=";
+    errorMessage += DataTypeUtils::asString(arr1.dataType());
+    errorMessage += " arr2.dataType()=";
+    errorMessage += DataTypeUtils::asString(arr2.dataType());
+    errorMessage += " arr1.shapeInfo()=";
+    errorMessage += ShapeUtils::shapeAsString(arr1.shapeInfo());
+    errorMessage += " arr2.shapeInfo()=";
+    errorMessage += ShapeUtils::shapeAsString(arr2.shapeInfo());
+    errorMessage += " arr1.ordering()=";
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+  PointersManager pointersManager(arr1.getContext(), "operator+(T&& arr1, T&& arr2)");
+
+  if (arr1.lengthOf() == arr2.lengthOf() && arr1.rankOf() == arr2.rankOf()) {
+    const bool isArr1Rvalue = !std::is_reference<T1>::value && !arr1.isView();
+    const bool isArr2Rvalue = !std::is_reference<T2>::value && !arr2.isView();
+
+    NDArray *result = nullptr;
+    if (isArr1Rvalue)
+      result = const_cast<NDArray *>(&arr1);
+    else if (isArr2Rvalue)
+      result = const_cast<NDArray *>(&arr2);
+    else
+      result = new NDArray(arr1.shapeInfo(), DataTypeUtils::pickPairwiseResultType(arr1.shapeInfo(), arr2.shapeInfo()),
+                           false, arr1.getContext());
+
+    NDArray::prepareSpecialUse({result}, {&arr1, &arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1.getContext(), sd::pairwise::Add, arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(),
+        arr1.specialShapeInfo(), arr2.buffer(), arr2.shapeInfo(), arr2.specialBuffer(), arr2.specialShapeInfo(),
+        result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {&arr1, &arr2});
+
+    return result;
+  }
+
+  NDArray* arr1Addr = &arr1;
+  NDArray* arr2Addr = &arr2;
+  return std::forward<T1>(arr1).applyTrueBroadcast(sd::BroadcastOpsTuple::Add(), arr2Addr);
+}
+template SD_LIB_EXPORT NDArray* operator+<NDArray &, NDArray &, void>(NDArray &arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator+<NDArray &, NDArray, void>(NDArray &arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* operator+<NDArray, NDArray &, void>(NDArray &&arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator+<NDArray, NDArray, void>(NDArray &&arr1, NDArray &&arr2);
+
+
+template <typename T1, typename T2, typename>
+NDArray* operator/(T1 &&arr1, T2 &&arr2) {
+  if (arr1.isS() || arr2.isS())
+    THROW_EXCEPTION("operator/(T&& arr1, T&& arr2): you can't use this method on String arrays!");
+  if (!Environment::getInstance().isExperimentalBuild() && arr1.dataType() != arr2.dataType() &&
+      (arr1.dataType() != DataType::BOOL || arr2.dataType() != BOOL)) {
+    // Always throw exception - never use _exit(1) which kills the entire process
+    THROW_EXCEPTION(sd::datatype_exception::build("operator/(T&& arr1, T&& arr2): Cannot divide different types",
+                                        arr1.dataType(), arr2.dataType()).what());
+  }
+
+  PointersManager pointersManager(arr1.getContext(), "operator/(T&& arr1, T&& arr2)");
+
+  if (arr1.lengthOf() == arr2.lengthOf() && arr1.rankOf() == arr2.rankOf()) {
+    const bool isArr1Rvalue = !std::is_reference<T1>::value && !arr1.isView();
+    const bool isArr2Rvalue = !std::is_reference<T2>::value && !arr2.isView();
+
+    NDArray *result = nullptr;
+    if (isArr1Rvalue)
+      result = const_cast<NDArray *>(&arr1);
+    else if (isArr2Rvalue)
+      result = const_cast<NDArray *>(&arr2);
+    else
+      result = new NDArray(arr1.shapeInfo(), DataTypeUtils::pickPairwiseResultType(arr1.shapeInfo(), arr2.shapeInfo()),
+                           false, arr1.getContext());
+
+    NDArray::prepareSpecialUse({result}, {&arr1, &arr2});
+    NativeOpExecutioner::execPairwiseTransform(
+        arr1.getContext(), sd::pairwise::Divide, arr1.buffer(), arr1.shapeInfo(), arr1.specialBuffer(),
+        arr1.specialShapeInfo(), arr2.buffer(), arr2.shapeInfo(), arr2.specialBuffer(), arr2.specialShapeInfo(),
+        result->buffer(), result->shapeInfo(), result->specialBuffer(), result->specialShapeInfo(), nullptr);
+    NDArray::registerSpecialUse({result}, {&arr1, &arr2});
+
+    return result;
+  }
+
+  NDArray* arr1Addr = &arr1;
+  NDArray* arr2Addr = &arr2;
+  return std::forward<T1>(arr1).applyTrueBroadcast(sd::BroadcastOpsTuple::Divide(), arr2Addr);
+}
+template SD_LIB_EXPORT NDArray* operator/<NDArray &, NDArray &, void>(NDArray &arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator/<NDArray &, NDArray, void>(NDArray &arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* operator/<NDArray, NDArray &, void>(NDArray &&arr1, NDArray &arr2);
+template SD_LIB_EXPORT NDArray* operator/<NDArray&&, NDArray &&, void>(NDArray &&arr1, NDArray &&arr2);
+template SD_LIB_EXPORT NDArray* operator/<NDArray, NDArray, void>(sd::NDArray&&, sd::NDArray&&);
+
+
+
+} // namespace sd

libnd4j/include/array/NDArrayLambda.hXX

@@ -0,0 +1,507 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+#ifndef CUDA_LAMBDA_HELPER
+#define CUDA_LAMBDA_HELPER
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <helpers/shape.h>
+#include <system/op_boilerplate.h>
+
+
+
+static sd::LongType SD_DEVICE  length(const sd::LongType *shapeInfo) { return shape::length(shapeInfo); }
+
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaKernel(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                   const sd::LongType *zShapeInfo, Lambda lambda);
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaIndexedKernel(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                          const sd::LongType *zShapeInfo, Lambda lambda);
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaIndexedPairwiseKernel(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                                  const sd::LongType *yShapeInfo, void *vz,
+                                                  const sd::LongType *zShapeInfo, Lambda lambda);
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaPairwiseKernel(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                           const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                           Lambda lambda);
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaPairwiseKernel(const void *scalarPtr, const void *vx, const sd::LongType *xShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                           Lambda lambda);
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaTriplewiseKernel(const void *vw, const sd::LongType *wShapeInfo, const void *vx,
+                                             const sd::LongType *xShapeInfo, const void *vy,
+                                             const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                             Lambda lambda);
+
+template <typename T>
+class LambdaHelper {
+ public:
+  template <typename Lambda>
+  SD_INLINE static void lambdaLauncher(cudaStream_t *stream, const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                       const sd::LongType *zShapeInfo, Lambda lambda) {
+    lambdaKernel<T, Lambda><<<256, 512, 1024, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, lambda);
+    auto err = cudaStreamSynchronize(*stream);
+    if (err != 0) THROW_EXCEPTION("NDArray::applyLambda execution failed");
+  }
+
+  template <typename Lambda>
+  SD_INLINE static void lambdaIndexedLauncher(cudaStream_t *stream, const void *vx, const sd::LongType *xShapeInfo,
+                                              void *vz, const sd::LongType *zShapeInfo, Lambda lambda) {
+    lambdaIndexedKernel<T, Lambda><<<256, 512, 1024, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, lambda);
+    auto err = cudaStreamSynchronize(*stream);
+    if (err != 0) THROW_EXCEPTION("NDArray::applyIndexedLambda execution failed");
+  }
+
+  template <typename Lambda>
+  SD_INLINE static void lambdaPairwiseLauncher(cudaStream_t *stream, const void *vx, const sd::LongType *xShapeInfo, bool otherIsScalar,
+                                               const void *vy, const sd::LongType *yShapeInfo, void *vz,
+                                               const sd::LongType *zShapeInfo, Lambda lambda) {
+    if (otherIsScalar) {
+      lambdaPairwiseKernel<T, Lambda><<<256, 512, 1024, *stream>>>(vy, vx, xShapeInfo, vz, zShapeInfo, lambda);
+    } else {
+      lambdaPairwiseKernel<T, Lambda>
+      <<<256, 512, 1024, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, lambda);
+    }
+    auto err = cudaStreamSynchronize(*stream);
+    if (err != 0) THROW_EXCEPTION("NDArray::applyPairwiseLambda execution failed");
+  }
+
+  template <typename Lambda>
+  SD_INLINE static void lambdaIndexedPairwiseLauncher(cudaStream_t *stream, const void *vx,
+                                                      const sd::LongType *xShapeInfo, const void *vy,
+                                                      const sd::LongType *yShapeInfo, void *vz,
+                                                      const sd::LongType *zShapeInfo, Lambda lambda) {
+    lambdaIndexedPairwiseKernel<T, Lambda>
+    <<<256, 512, 1024, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, lambda);
+    auto err = cudaStreamSynchronize(*stream);
+    if (err != 0) THROW_EXCEPTION("NDArray::applyIndexedPairwiseLambda execution failed");
+  }
+
+  template <typename Lambda>
+  SD_INLINE static void lambdaTriplewiseLauncher(cudaStream_t *stream, const void *vw, const sd::LongType *wShapeInfo,
+                                                 const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                                 const sd::LongType *yShapeInfo, void *vz,
+                                                 const sd::LongType *zShapeInfo, Lambda lambda) {
+    lambdaTriplewiseKernel<T, Lambda>
+    <<<256, 512, 1024, *stream>>>(vw, wShapeInfo, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, lambda);
+    auto err = cudaStreamSynchronize(*stream);
+    if (err != 0) THROW_EXCEPTION("NDArray::applyTriplewiseLambda execution failed");
+  }
+};
+
+////////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaKernel(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                   const sd::LongType *zShapeInfo, Lambda lambda) {
+  auto x = reinterpret_cast<const T *>(vx);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+
+  __syncthreads();
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e, xRank, xShape, xCoords);
+    COORDS2INDEX(xRank,xStride, xCoords, xOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(x[xOffset]);
+  }
+}
+////////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaIndexedKernel(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                          const sd::LongType *zShapeInfo, Lambda lambda) {
+  auto x = reinterpret_cast<const T *>(vx);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e, xRank, xShape, xCoords);
+    COORDS2INDEX(xRank,xStride, xCoords, xOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(e, x[xOffset]);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaIndexedPairwiseKernel(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                                  const sd::LongType *yShapeInfo, void *vz,
+                                                  const sd::LongType *zShapeInfo, Lambda lambda) {
+  auto x = reinterpret_cast<const T *>(vx);
+  auto y = reinterpret_cast<const T *>(vy);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  yRank;
+  __shared__ sd::LongType *yShape;
+  __shared__ sd::LongType *yStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    yShape = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+
+  __syncthreads();
+
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e, xRank,xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+    INDEX2COORDS(e, yRank, yShape, yCoords);
+    COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank,zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(e, x[xOffset], y[yOffset]);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaPairwiseKernel(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                           const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                           Lambda lambda) {
+  auto x = reinterpret_cast<const T *>(vx);
+  auto y = reinterpret_cast<const T *>(vy);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  yRank;
+  __shared__ sd::LongType *yShape;
+  __shared__ sd::LongType *yStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    yShape = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+
+  __syncthreads();
+
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e, xRank, xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+    INDEX2COORDS(e, yRank, yShape, yCoords);
+    COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank,zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(x[xOffset], y[yOffset]);
+  }
+}
+///////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaPairwiseKernel(const void *scalarPtr, const void *vx, const sd::LongType *xShapeInfo,
+                                           void *vz, const sd::LongType *zShapeInfo, Lambda lambda) {
+  auto x = reinterpret_cast<const T *>(vx);
+  auto y = reinterpret_cast<const T *>(scalarPtr);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto yVal = *y;
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  yRank;
+  __shared__ sd::LongType *yShape;
+  __shared__ sd::LongType *yStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    yShape = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+
+  __syncthreads();
+
+
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e,xRank,xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(x[xOffset], yVal);
+  }
+}
+////////////////////////////////////////////////////////////////////////
+template <typename T, typename Lambda>
+static SD_KERNEL void lambdaTriplewiseKernel(const void *vw, const sd::LongType *wShapeInfo, const void *vx,
+                                             const sd::LongType *xShapeInfo, const void *vy,
+                                             const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                             Lambda lambda) {
+  auto w = reinterpret_cast<const T *>(vw);
+  auto x = reinterpret_cast<const T *>(vx);
+  auto y = reinterpret_cast<const T *>(vy);
+  auto z = reinterpret_cast<T *>(vz);
+
+  auto zLength = length(zShapeInfo);
+
+  auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+  __shared__ sd::LongType  xRank;
+  __shared__ sd::LongType *xShape;
+  __shared__ sd::LongType *xStride;
+  __shared__ sd::LongType  yRank;
+  __shared__ sd::LongType *yShape;
+  __shared__ sd::LongType *yStride;
+  __shared__ sd::LongType  zRank;
+  __shared__ sd::LongType *zShape;
+  __shared__ sd::LongType *zStride;
+  __shared__ sd::LongType wRank;
+  __shared__ sd::LongType *wShape;
+  __shared__ sd::LongType *wStride;
+
+  if(threadIdx.x == 0) {
+    xRank = shape::rank(xShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    yShape = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+    wRank = shape::rank(wShapeInfo);
+    wShape = shape::shapeOf(wShapeInfo);
+    wStride = shape::stride(wShapeInfo);
+  }
+
+  __syncthreads();
+
+
+  for (sd::LongType e = tid; e < zLength; e += blockDim.x * gridDim.x) {
+    sd::LongType wCoords[SD_MAX_RANK];
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType wOffset;
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(e, wRank, wShape, wCoords);
+    COORDS2INDEX(wRank, wStride, wCoords, wOffset);
+    INDEX2COORDS(e, xRank,xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+    INDEX2COORDS(e, yRank, yStride, yCoords);
+    COORDS2INDEX(yRank,yStride, yCoords, yOffset);
+    INDEX2COORDS(e, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = lambda(w[wOffset], x[xOffset], y[yOffset]);
+  }
+}
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+template <typename Lambda>
+void NDArray::applyLambda(Lambda func, NDArray *target) {
+  auto dtype = this->dataType();
+
+  if (dtype != target->dataType()) THROW_EXCEPTION("NDArray::applyLambda X/Z data types must be the same");
+  prepareSpecialUse({&target}, {this});
+  BUILD_SINGLE_SELECTOR(
+      dtype, LambdaHelper,
+      ::lambdaLauncher(this->_context->getCudaStream(), this->specialBuffer(), this->specialShapeInfo(),
+                       target->specialBuffer(), target->specialShapeInfo(), func),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this});
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename Lambda>
+void NDArray::applyPairwiseLambda(NDArray *other, Lambda func, NDArray *target) {
+  auto dtype = this->dataType();
+
+  if (dtype != target->dataType() || dtype != other->dataType())
+    THROW_EXCEPTION("NDArray::applyPairwiseLambda X/Y/Z data types must be the same");
+  bool otherIsScalar = other->isScalar();
+  prepareSpecialUse({&target}, {this, &other});
+  BUILD_SINGLE_SELECTOR(
+      dtype, LambdaHelper,
+      ::lambdaPairwiseLauncher(this->_context->getCudaStream(), this->specialBuffer(), this->specialShapeInfo(), otherIsScalar,
+                               other->specialBuffer(), other->specialShapeInfo(), target->specialBuffer(),
+                               target->specialShapeInfo(), func),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this, &other});
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename Lambda>
+void NDArray::applyIndexedLambda(Lambda func, NDArray *target) {
+  auto dtype = this->dataType();
+  if (dtype != target->dataType())
+    THROW_EXCEPTION("NDArray::applyIndexedLambda X/Z data types must be the same");
+
+  prepareSpecialUse({&target}, {this});
+  BUILD_SINGLE_SELECTOR(
+      dtype, LambdaHelper,
+      ::lambdaIndexedLauncher(this->_context->getCudaStream(), this->specialBuffer(), this->specialShapeInfo(),
+                              target->specialBuffer(), target->specialShapeInfo(), func),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this});
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename Lambda>
+void NDArray::applyIndexedPairwiseLambda(NDArray *other, Lambda func, NDArray *target) {
+  auto dtype = this->dataType();
+  if (dtype != target->dataType() || dtype != other->dataType())
+    THROW_EXCEPTION("NDArray::applyIndexedPairwiseLambda X/Y/Z data types must be the same");
+
+  prepareSpecialUse({&target}, {this, &other});
+  BUILD_SINGLE_SELECTOR(
+      dtype, LambdaHelper,
+      ::lambdaIndexedPairwiseLauncher(this->_context->getCudaStream(), this->specialBuffer(), this->specialShapeInfo(),
+                                      other->specialBuffer(), other->specialShapeInfo(), target->specialBuffer(),
+                                      target->specialShapeInfo(), func),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this, &other});
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename Lambda>
+void NDArray::applyTriplewiseLambda(NDArray *second, NDArray *third, Lambda func, NDArray *target) {
+  auto dtype = this->dataType();
+
+  if (dtype != target->dataType() || dtype != second.dataType() || dtype != third.dataType())
+    THROW_EXCEPTION("NDArray::applyTriplewiseLambda X/Y/Z data types must be the same");
+
+  prepareSpecialUse({&target}, {this, &second, &third});
+  BUILD_SINGLE_SELECTOR(
+      dtype, LambdaHelper,
+      ::lambdaTriplewiseLauncher(this->_context->getCudaStream(), this->specialBuffer(), this->specialShapeInfo(),
+                                 second.specialBuffer(), second.specialShapeInfo(), third.specialBuffer(),
+                                 third.specialShapeInfo(), target->specialBuffer(), target->specialShapeInfo(), func),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this, &second, &third});
+}

libnd4j/include/array/NDArrayList.h

@@ -0,0 +1,102 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// This class describes collection of NDArrays
+//
+// @author raver119!gmail.com
+//
+
+#ifndef NDARRAY_LIST_H
+#define NDARRAY_LIST_H
+#include <array/NDArray.h>
+#include <memory/Workspace.h>
+#include <system/common.h>
+
+#include <atomic>
+#include <string>
+#include <unordered_map>
+
+namespace sd {
+class SD_LIB_EXPORT NDArrayList {
+ private:
+  // workspace where chunks belong to
+  // sd::memory::Workspace* _workspace = nullptr;
+  sd::LaunchContext *_context = sd::LaunchContext ::defaultContext();
+
+  // numeric and symbolic ids of this list
+  std::pair<int, int> _id;
+  std::string _name;
+
+  sd::DataType _dtype;
+
+  // stored chunks
+  SD_MAP_IMPL<int, sd::NDArray *> _chunks;
+
+  // just a counter, for stored elements
+  std::atomic<int> _elements;
+  std::atomic<int> _counter;
+
+  // reference shape
+  std::vector<sd::LongType> _shape;
+
+  // unstack axis
+  sd::LongType _axis = 0;
+
+  //
+  bool _expandable = false;
+
+  // maximum number of elements
+  int _height = 0;
+
+ public:
+  NDArrayList(int height, bool expandable = false);
+  ~NDArrayList();
+
+  sd::DataType dataType();
+
+  NDArray *remove(int idx);
+  NDArray *read(int idx);
+  NDArray *readRaw(int idx);
+  sd::Status write(int idx, NDArray *array);
+
+  NDArray *pick(std::initializer_list<LongType> indices);
+  NDArray *pick(std::vector<LongType> &indices);
+  bool isWritten(int index);
+
+  std::vector<sd::LongType> &shape();
+
+  NDArray *stack();
+  void unstack(NDArray *array, LongType axis);
+
+  std::pair<int, int> &id();
+  std::string &name();
+  // sd::memory::Workspace* workspace();
+  sd::LaunchContext *context();
+  NDArrayList *clone();
+
+  bool equals(NDArrayList &other);
+
+  int elements();
+  int height();
+
+  int counter();
+};
+}  // namespace sd
+
+#endif

libnd4j/include/array/PointerDeallocator.h

@@ -0,0 +1,42 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+//
+
+#ifndef SD_POINTERDEALLOCATOR_H_
+#define SD_POINTERDEALLOCATOR_H_
+
+#include <system/common.h>
+
+namespace sd {
+
+class SD_LIB_EXPORT PointerDeallocator {
+ public:
+  PointerDeallocator() = default;
+  virtual ~PointerDeallocator() = default;
+
+  virtual void release(void *ptr);
+};
+
+}  // namespace sd
+
+#endif  // SD_POINTERDEALLOCATOR_H_

libnd4j/include/array/PointerWrapper.h

@@ -0,0 +1,52 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+//
+
+#ifndef SD_ARRAY_POINTER_H_
+#define SD_ARRAY_POINTER_H_
+
+#include <array/PointerDeallocator.h>
+
+#include <memory>
+
+namespace sd {
+class SD_LIB_EXPORT PointerWrapper {
+ private:
+  void *_pointer = nullptr;
+  std::shared_ptr<PointerDeallocator> _deallocator;
+
+ public:
+  PointerWrapper(void *ptr, const std::shared_ptr<PointerDeallocator> &deallocator = {});
+  PointerWrapper() = default;
+  ~PointerWrapper();
+
+  void *pointer() const;
+
+  template <typename T>
+  T *pointerAsT() const {
+    return reinterpret_cast<T *>(pointer());
+  }
+};
+}  // namespace sd
+
+#endif  // SD_ARRAY_POINTER_H_

libnd4j/include/array/PrimaryPointerDeallocator.h

@@ -0,0 +1,40 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+//  @author raver119@gmail.com
+//
+
+#ifndef SD_PRIMARYPOINTERDEALLOCATOR_H_
+#define SD_PRIMARYPOINTERDEALLOCATOR_H_
+
+#include <array/PointerDeallocator.h>
+
+namespace sd {
+class SD_LIB_EXPORT PrimaryPointerDeallocator : public PointerDeallocator {
+ public:
+  PrimaryPointerDeallocator() = default;
+  virtual ~PrimaryPointerDeallocator() = default;
+
+  void release(void *ptr) override;
+};
+}  // namespace sd
+
+#endif  // SD_PRIMARYPOINTERDEALLOCATOR_H_

libnd4j/include/array/ResultSet.h

@@ -0,0 +1,78 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// This class is suited for execution results representation.
+//
+// PLESE NOTE: It will delete all stored NDArrays upon destructor call
+//
+// @author raver119@gmail.com
+//
+
+#ifndef LIBND4J_RESULTSET_H
+#define LIBND4J_RESULTSET_H
+#include <graph/generated/result_generated.h>
+#include <system/common.h>
+
+#include <vector>
+
+namespace sd {
+
+class NDArray;  // forward declaration of template class NDArray
+
+class SD_LIB_EXPORT ResultSet {
+ private:
+  std::vector<NDArray *> _content;
+  Status _status = Status::OK;
+  bool _removable = true;
+
+  void delContent();
+
+ public:
+  explicit ResultSet();
+
+#ifndef __JAVACPP_HACK__
+  ResultSet(const ::graph::FlatResult *result);
+#endif
+
+  ResultSet(const ResultSet &other) noexcept;
+
+  ResultSet &operator=(const ResultSet &other) noexcept;
+
+  // move constructor
+  ResultSet(ResultSet &&other) noexcept;
+
+  // move assignment operator
+  ResultSet &operator=(ResultSet &&other) noexcept;
+
+  ~ResultSet();
+
+  int size();
+  NDArray *at(const unsigned long idx) const;
+  NDArray *operator[](const unsigned long idx) const;
+  void push_back(NDArray *array);
+
+  Status status();
+  void setStatus(Status status);
+  void purge();
+  void setNonRemovable();
+  void printIndexedBuffers();
+};
+}  // namespace sd
+
+#endif  // LIBND4J_RESULTSET_H

libnd4j/include/array/ShapeCacheLifecycleTracker.h

@@ -0,0 +1,172 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// ShapeCacheLifecycleTracker - Tracks shape cache allocations and deallocations
+// for memory leak detection.
+// Only active when SD_GCC_FUNCTRACE is defined.
+//
+
+#ifndef LIBND4J_SHAPECACHELIFECYCLETRACKER_H
+#define LIBND4J_SHAPECACHELIFECYCLETRACKER_H
+
+#include <cstddef>
+#include <string>
+#include <atomic>
+#include <mutex>
+#include <ostream>
+#include <system/common.h>
+
+namespace sd {
+
+namespace analysis {
+class ComprehensiveLeakAnalyzer;
+}
+
+namespace array {
+
+/**
+ * Statistics structure for shape cache lifecycle tracking
+ */
+struct ShapeCacheStats {
+    size_t totalAllocations = 0;
+    size_t totalDeallocations = 0;
+    size_t currentLive = 0;
+    size_t peakLive = 0;
+    size_t totalBytesAllocated = 0;
+    size_t totalBytesDeallocated = 0;
+};
+
+/**
+ * ShapeCacheLifecycleTracker - Singleton class for tracking shape cache allocations
+ * and deallocations for memory leak detection.
+ *
+ * This is a stub implementation that provides the expected interface.
+ * When SD_GCC_FUNCTRACE is enabled, this tracks shape cache lifecycle events.
+ */
+class ShapeCacheLifecycleTracker {
+    friend class sd::analysis::ComprehensiveLeakAnalyzer;
+
+public:
+    /**
+     * Get singleton instance
+     */
+    static ShapeCacheLifecycleTracker& getInstance() {
+        static ShapeCacheLifecycleTracker instance;
+        return instance;
+    }
+
+    /**
+     * Enable or disable tracking
+     */
+    void setEnabled(bool enabled) {
+        _enabled.store(enabled);
+    }
+
+    /**
+     * Check if tracking is enabled
+     */
+    bool isEnabled() const {
+        return _enabled.load();
+    }
+
+    /**
+     * Record a shape cache allocation
+     * @param shapeInfo Pointer to the shape info being allocated
+     */
+    void recordAllocation(LongType* shapeInfo) {
+        if (!_enabled.load()) return;
+
+        std::lock_guard<std::mutex> lock(_mutex);
+        _stats.totalAllocations++;
+        _stats.currentLive++;
+        if (_stats.currentLive > _stats.peakLive) {
+            _stats.peakLive = _stats.currentLive;
+        }
+    }
+
+    /**
+     * Record a shape cache deallocation
+     * @param shapeInfo Pointer to the shape info being deallocated
+     */
+    void recordDeallocation(LongType* shapeInfo) {
+        if (!_enabled.load()) return;
+
+        std::lock_guard<std::mutex> lock(_mutex);
+        _stats.totalDeallocations++;
+        if (_stats.currentLive > 0) {
+            _stats.currentLive--;
+        }
+    }
+
+    /**
+     * Get statistics
+     */
+    ShapeCacheStats getStats() const {
+        return _stats;
+    }
+
+    /**
+     * Print statistics to output stream
+     */
+    void printStatistics(std::ostream& out) const {
+        out << "ShapeCache Statistics: allocations=" << _stats.totalAllocations
+            << ", deallocations=" << _stats.totalDeallocations
+            << ", live=" << _stats.currentLive << "\n";
+    }
+
+    /**
+     * Print current memory leaks to output stream
+     */
+    void printCurrentLeaks(std::ostream& out) const {
+        out << "ShapeCache Current Leaks: " << _stats.currentLive << " shapes\n";
+    }
+
+    /**
+     * Log shape info for a specific address (for crash debugging)
+     * @param address Address to look up
+     * @param out Output stream
+     * @return true if shape was found
+     */
+    bool logShapeForAddress(void* address, std::ostream& out) const {
+        std::lock_guard<std::mutex> lock(_mutex);
+        out << "ShapeCache address lookup for " << address << ": ";
+        out << "tracking " << (_enabled.load() ? "enabled" : "disabled");
+        out << ", live shapes: " << _stats.currentLive << "\n";
+        return false;
+    }
+
+private:
+    ShapeCacheLifecycleTracker() : _enabled(false) {}
+    ~ShapeCacheLifecycleTracker() = default;
+
+    // Disable copy and move
+    ShapeCacheLifecycleTracker(const ShapeCacheLifecycleTracker&) = delete;
+    ShapeCacheLifecycleTracker& operator=(const ShapeCacheLifecycleTracker&) = delete;
+    ShapeCacheLifecycleTracker(ShapeCacheLifecycleTracker&&) = delete;
+    ShapeCacheLifecycleTracker& operator=(ShapeCacheLifecycleTracker&&) = delete;
+
+    std::atomic<bool> _enabled;
+    mutable std::mutex _mutex;
+    ShapeCacheStats _stats;
+};
+
+} // namespace array
+} // namespace sd
+
+#endif // LIBND4J_SHAPECACHELIFECYCLETRACKER_H

libnd4j/include/array/ShapeDescriptor.h

@@ -0,0 +1,257 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//  @author AbdelRauf
+
+#ifndef DEV_TESTS_SHAPEDESCRIPTOR_H
+#define DEV_TESTS_SHAPEDESCRIPTOR_H
+#include <array/ArrayOptions.hXX>
+#include <array/DataType.h>
+#include <helpers/shape.h>
+#include <system/common.h>
+
+#include <initializer_list>
+#include <unordered_map>
+#include <vector>
+namespace sd {
+
+#define SHAPE_DESC_OK 0
+#define SHAPE_DESC_INCORRECT_STRIDES 1  // strides does not match shapes
+#define SHAPE_DESC_INCORRECT_EWS 2      // ews neither matches stride nor continuity
+#define SHAPE_DESC_INCORRECT_RANK 4     // rank > 32 or shape size and rank does not match
+#define SHAPE_DESC_INVALID_EMPTY 5     // rank > 32 or shape size and rank does not match
+
+class SD_LIB_EXPORT ShapeDescriptor {
+
+
+ private:
+  int _rank = 0;
+  LongType * _shape_strides = nullptr;
+  char _order = 'c';
+  DataType _dataType;
+  LongType _extraProperties = 0;
+  LongType _paddedAllocSize = 0;
+  LongType _offset = 0;
+
+ public:
+  bool ownsShapeStrides = false;
+  // Hash caching
+  mutable uint64_t _cached_hash;
+  mutable bool _hash_computed;
+#ifndef __JAVACPP_HACK__
+#if defined(SD_GCC_FUNCTRACE)
+  StackTrace st;
+  //stack trace when stored in cache.
+  StackTrace storeStackTrace;
+#endif
+  ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape, LongType extras);
+  ShapeDescriptor(const ShapeDescriptor &other);
+  ShapeDescriptor(const LongType *shapeInfo, bool validateDataType = true, bool overrideStrides = false);
+  explicit ShapeDescriptor(const LongType *shapeInfo, const DataType dtypeOverride, const bool overrideStrides);
+  explicit ShapeDescriptor(const LongType *shapeInfo, const LongType *dtypeOverride);
+  explicit ShapeDescriptor(const LongType *shapeInfo, const LongType *dtypeOverride,
+                           const LongType *orderOverride);
+  explicit ShapeDescriptor(const DataType type, const LongType length);
+  explicit ShapeDescriptor(const DataType type, const char order, const LongType *shape, const LongType rank);
+  explicit ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape);
+  explicit ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape,
+                           const std::vector<LongType> &strides);
+  explicit ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape,
+                           const std::vector<LongType> &strides, const LongType ews);
+  explicit ShapeDescriptor(const DataType type, const char order, const LongType *shape,
+                           const LongType *strides, const LongType rank, LongType extras);
+
+  ShapeDescriptor() = default;
+  ~ShapeDescriptor();
+#endif
+  int rank() const;
+
+  void invalidateHash() const {
+    _hash_computed = false;
+    _cached_hash = 0;
+  }
+
+  uint64_t getCachedHash() const {
+    return _cached_hash;
+  }
+  LongType arrLength() const;
+  LongType offset();
+  char order() const;
+  DataType dataType() const;
+  bool isEmpty() const;
+  sd::LongType * shape_strides();
+  const LongType *stridesPtr() const;
+  LongType extra() const {
+    return _extraProperties;
+  }
+
+
+  void collectStoreStackTrace();
+  void print() const;
+  // returns minimal allocation length
+  LongType allocLength() const;
+
+  // returns Status for the correctness
+  LongType validate() const;
+
+  // we use default copy assignment operator
+  // Modify assignment operator to reset hash cache:
+  ShapeDescriptor& operator=(const ShapeDescriptor& other) {
+    if (this != &other) {
+      // Existing cleanup code
+      if (_shape_strides != nullptr && ownsShapeStrides) {
+        delete[] _shape_strides;
+        _shape_strides = nullptr;
+      }
+
+      // Copy all basic members
+      _rank = other._rank;
+      _extraProperties = other._extraProperties;
+      _dataType = other._dataType;
+      _order = other._order;
+      _paddedAllocSize = other._paddedAllocSize;
+      _offset = other._offset;
+
+      // Reset hash cache
+      _cached_hash = 0;
+      _hash_computed = false;
+
+      // Handle shape_strides - make a deep copy if source has data
+      if (other._shape_strides != nullptr) {
+        const int size = (_rank < 1 ? 1 : _rank) * 2;
+        _shape_strides = new LongType[size];
+        std::memcpy(_shape_strides, other._shape_strides, size * sizeof(LongType));
+        ownsShapeStrides = true;
+      } else {
+        _shape_strides = nullptr;
+        ownsShapeStrides = false;
+      }
+    }
+    return *this;
+  }
+  // we use default move assignment operator
+  ShapeDescriptor &operator=(ShapeDescriptor &&other) noexcept = default;
+
+  // equal to operator
+  bool operator==(const ShapeDescriptor &other) const;
+
+  // less than operator
+  bool operator<(const ShapeDescriptor &other) const;
+
+  LongType *toShapeInfo() const;
+
+  const char * toString() {
+    std::string message;
+    message += " Rank:" ;
+    message += std::to_string(_rank);
+    message += " Shape and Strides:";
+    if(_shape_strides == nullptr) {
+      message += " Null";
+    } else {
+      for (int i = 0; i < _rank * 2; i++) {
+        message += " ";
+        message += std::to_string(_shape_strides[i]);
+      }
+
+    }
+    message += "Data type:";
+    message += std::to_string(_dataType);
+    message += " Order:";
+    message += std::to_string(_order);
+    message += " Extra Properties:";
+    message += std::to_string(_extraProperties);
+    message += " Padded Alloc Size: ";
+    message += std::to_string(_paddedAllocSize);
+    message += " Offset: ";
+    message += std::to_string(_offset);
+    //need this in order to avoid deallocation
+    std::string *ret = new std::string(message.c_str());
+    return ret->c_str();
+  }
+  static ShapeDescriptor * emptyDescriptor(const DataType type);
+  static ShapeDescriptor  * scalarDescriptor(const DataType type);
+  static ShapeDescriptor * vectorDescriptor(const LongType length, const DataType type);
+
+  // create Descriptor with padded buffer.
+  static ShapeDescriptor * paddedBufferDescriptor(const DataType type, const char order,
+                                                  const std::vector<LongType> &shape,
+                                                  const std::vector<LongType> &paddings);
+
+  static  const char *messageForShapeDescriptorError(const int errorCode) {
+    switch (errorCode) {
+      case SHAPE_DESC_OK:
+        return "OK";
+      case SHAPE_DESC_INCORRECT_STRIDES:
+        return "Incorrect strides";
+      case SHAPE_DESC_INCORRECT_EWS:
+        return "Incorrect ews";
+      case SHAPE_DESC_INCORRECT_RANK:
+        return "Incorrect rank";
+      case SHAPE_DESC_INVALID_EMPTY:
+        return "Invalid empty";
+      default:
+        return "Unknown error";
+    }
+  }
+  bool isScalar() const;
+
+  SD_INLINE void fillStrides() {
+    if(_rank == 0) {
+      return;
+    }
+
+    if(_shape_strides == nullptr) {
+      return;
+    }
+
+    // double checks if the _rank and _shape_strides are set correctly before filling strides
+    auto _shape = _shape_strides;
+    auto _strides = _shape_strides + _rank;
+    if (_rank > 0) {
+      if (_order == 'c')
+        shape::calcStrides(_shape, _rank, _strides);
+      else
+        shape::calcStridesFortran(_shape, _rank, _strides);
+
+    } else {
+      for (int i = 0; i < _rank; i++) {
+        _strides[i] = 0;
+      }
+    }
+
+
+  }
+
+};
+}  // namespace sd
+
+#ifndef __JAVACPP_HACK__
+
+namespace std {
+template <>
+class SD_LIB_EXPORT hash<sd::ShapeDescriptor> {
+ public:
+  size_t operator()(sd::ShapeDescriptor k) const;
+};
+}  // namespace std
+
+#endif
+
+#endif  // DEV_TESTS_SHAPEDESCRIPTOR_H

libnd4j/include/array/ShapeList.h

@@ -0,0 +1,59 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef LIBND4J_SHAPELIST_H
+#define LIBND4J_SHAPELIST_H
+#include <helpers/shape.h>
+#include <system/common.h>
+
+#include <vector>
+
+namespace sd {
+class SD_LIB_EXPORT ShapeList {
+ protected:
+
+  std::vector< sd::LongType *> _shapes;
+  bool _destroyed = false;
+  bool _autoremovable = false;
+  bool _workspace = false;
+
+ public:
+  ShapeList( sd::LongType *shape = nullptr);
+  ShapeList(const std::vector< sd::LongType *> &shapes, bool isWorkspace);
+  ShapeList(const std::vector< sd::LongType *> &shapes);
+
+  ~ShapeList();
+
+  void destroy();
+  int size() const;
+   sd::LongType *at(int idx);
+  void push_back( sd::LongType *shape);
+
+  /**
+   * PLEASE NOTE: This method should be called ONLY if shapes were generated at workspaces. Otherwise you'll get memory
+   * leak
+   */
+  void detach();
+};
+}  // namespace sd
+
+#endif  // LIBND4J_SHAPELIST_H

libnd4j/include/array/SpaceType.h

@@ -0,0 +1,34 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef ND4J_SPACE_TYPE_H
+#define ND4J_SPACE_TYPE_H
+
+namespace sd {
+enum SpaceType {
+  CONTINUOUS = 1,
+  COMPLEX = 2,
+  QUANTIZED = 3,
+};
+}
+
+#endif

libnd4j/include/array/SparseType.h

@@ -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
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+
+#ifndef LIBND4J_SPARSETYPE_H
+#define LIBND4J_SPARSETYPE_H
+
+namespace sd {
+enum SparseType {
+  CSR = 1,
+  CSC = 2,
+  COO = 3,
+  LIL = 4,
+};
+}
+
+#endif  // LIBND4J_SPARSETYPE_H

libnd4j/include/array/TADCacheLifecycleTracker.h

@@ -0,0 +1,261 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+//
+// TADCacheLifecycleTracker - Tracks TAD (Tensor Along Dimension) cache allocations
+// and deallocations for memory leak detection.
+// Only active when SD_GCC_FUNCTRACE is defined.
+//
+
+#ifndef LIBND4J_TADCACHELIFECYCLETRACKER_H
+#define LIBND4J_TADCACHELIFECYCLETRACKER_H
+
+#include <cstddef>
+#include <string>
+#include <vector>
+#include <atomic>
+#include <mutex>
+#include <map>
+#include <ostream>
+#include <fstream>
+#include <system/common.h>
+
+namespace sd {
+
+// Forward declaration
+class TadPack;
+
+namespace analysis {
+class ComprehensiveLeakAnalyzer;
+}
+
+namespace array {
+
+/**
+ * Statistics structure for TAD cache lifecycle tracking
+ */
+struct TADCacheStats {
+    size_t totalAllocations = 0;
+    size_t totalDeallocations = 0;
+    size_t currentLive = 0;
+    size_t peakLive = 0;
+    size_t totalBytesAllocated = 0;
+    size_t totalBytesDeallocated = 0;
+};
+
+/**
+ * TADCacheLifecycleTracker - Singleton class for tracking TAD cache allocations
+ * and deallocations for memory leak detection.
+ *
+ * This is a stub implementation that provides the expected interface.
+ * When SD_GCC_FUNCTRACE is enabled, this tracks TAD cache lifecycle events.
+ */
+class TADCacheLifecycleTracker {
+    friend class sd::analysis::ComprehensiveLeakAnalyzer;
+
+public:
+    /**
+     * Get singleton instance
+     */
+    static TADCacheLifecycleTracker& getInstance() {
+        static TADCacheLifecycleTracker instance;
+        return instance;
+    }
+
+    /**
+     * Enable or disable tracking
+     */
+    void setEnabled(bool enabled) {
+        _enabled.store(enabled);
+    }
+
+    /**
+     * Check if tracking is enabled
+     */
+    bool isEnabled() const {
+        return _enabled.load();
+    }
+
+    /**
+     * Record a TAD cache allocation
+     * @param tadPack Pointer to the TadPack being allocated
+     * @param numTads Number of TADs
+     * @param shapeInfoBytes Size of shape info in bytes
+     * @param offsetsBytes Size of offsets in bytes
+     * @param dimensions Dimensions vector
+     */
+    void recordAllocation(void* tadPack, LongType numTads,
+                         size_t shapeInfoBytes, size_t offsetsBytes,
+                         const std::vector<LongType>& dimensions) {
+        if (!_enabled.load()) return;
+
+        std::lock_guard<std::mutex> lock(_mutex);
+        _stats.totalAllocations++;
+        _stats.currentLive++;
+        size_t totalBytes = shapeInfoBytes + offsetsBytes;
+        _stats.totalBytesAllocated += totalBytes;
+        if (_stats.currentLive > _stats.peakLive) {
+            _stats.peakLive = _stats.currentLive;
+        }
+    }
+
+    /**
+     * Record a TAD cache deallocation
+     * @param tadPack Pointer to the TadPack being deallocated
+     */
+    void recordDeallocation(void* tadPack) {
+        if (!_enabled.load()) return;
+
+        std::lock_guard<std::mutex> lock(_mutex);
+        _stats.totalDeallocations++;
+        if (_stats.currentLive > 0) {
+            _stats.currentLive--;
+        }
+    }
+
+    /**
+     * Get statistics
+     */
+    TADCacheStats getStats() const {
+        return _stats;
+    }
+
+    /**
+     * Print statistics to output stream
+     */
+    void printStatistics(std::ostream& out) const {
+        out << "TADCache Statistics: allocations=" << _stats.totalAllocations
+            << ", deallocations=" << _stats.totalDeallocations
+            << ", live=" << _stats.currentLive << "\n";
+    }
+
+    /**
+     * Print current memory leaks to output stream
+     */
+    void printCurrentLeaks(std::ostream& out) const {
+        out << "TADCache Current Leaks: " << _stats.currentLive << " TadPacks\n";
+    }
+
+    /**
+     * Log TAD info for a specific address (for crash debugging)
+     * @param address Address to look up
+     * @param out Output stream
+     * @return true if TAD was found
+     */
+    bool logTADForAddress(void* address, std::ostream& out) const {
+        std::lock_guard<std::mutex> lock(_mutex);
+        out << "TADCache address lookup for " << address << ": ";
+        out << "tracking " << (_enabled.load() ? "enabled" : "disabled");
+        out << ", live TadPacks: " << _stats.currentLive << "\n";
+        return false;
+    }
+
+    /**
+     * Generate temporal leak report
+     * @param outputPath Path to write report
+     * @param windowCount Number of time windows
+     * @param windowDurationSec Duration of each window in seconds
+     */
+    void generateTemporalLeakReport(const std::string& outputPath, int windowCount, double windowDurationSec) const {
+        std::lock_guard<std::mutex> lock(_mutex);
+        std::ofstream out(outputPath);
+        if (out.is_open()) {
+            out << "=== TADCache Temporal Leak Report ===\n";
+            out << "Note: Temporal tracking requires timestamp storage - not yet implemented\n";
+            out << "Current Statistics:\n";
+            out << "  Total Allocations: " << _stats.totalAllocations << "\n";
+            out << "  Total Deallocations: " << _stats.totalDeallocations << "\n";
+            out << "  Current Live: " << _stats.currentLive << "\n";
+            out.close();
+        }
+    }
+
+    /**
+     * Capture a leak snapshot for later comparison
+     * @return Snapshot ID
+     */
+    LongType captureLeakSnapshot() {
+        std::lock_guard<std::mutex> lock(_mutex);
+        LongType id = _nextSnapshotId++;
+        _snapshots[id] = _stats;
+        return id;
+    }
+
+    /**
+     * Generate diff between two snapshots
+     * @param snapshot1 First snapshot ID
+     * @param snapshot2 Second snapshot ID
+     * @param outputPath Path to write diff report
+     */
+    void generateSnapshotDiff(LongType snapshot1, LongType snapshot2, const std::string& outputPath) const {
+        std::lock_guard<std::mutex> lock(_mutex);
+        std::ofstream out(outputPath);
+        if (out.is_open()) {
+            out << "=== TADCache Snapshot Diff Report ===\n";
+            out << "Comparing snapshot " << snapshot1 << " to snapshot " << snapshot2 << "\n\n";
+
+            auto it1 = _snapshots.find(snapshot1);
+            auto it2 = _snapshots.find(snapshot2);
+
+            if (it1 == _snapshots.end()) {
+                out << "ERROR: Snapshot " << snapshot1 << " not found\n";
+            } else if (it2 == _snapshots.end()) {
+                out << "ERROR: Snapshot " << snapshot2 << " not found\n";
+            } else {
+                const auto& s1 = it1->second;
+                const auto& s2 = it2->second;
+                out << "Allocations:   " << s1.totalAllocations << " -> " << s2.totalAllocations;
+                out << " (diff: " << (long long)(s2.totalAllocations - s1.totalAllocations) << ")\n";
+                out << "Deallocations: " << s1.totalDeallocations << " -> " << s2.totalDeallocations;
+                out << " (diff: " << (long long)(s2.totalDeallocations - s1.totalDeallocations) << ")\n";
+                out << "Live:          " << s1.currentLive << " -> " << s2.currentLive;
+                out << " (diff: " << (long long)(s2.currentLive - s1.currentLive) << ")\n";
+            }
+            out.close();
+        }
+    }
+
+    /**
+     * Clear all snapshots
+     */
+    void clearSnapshots() {
+        std::lock_guard<std::mutex> lock(_mutex);
+        _snapshots.clear();
+    }
+
+private:
+    TADCacheLifecycleTracker() : _enabled(false), _nextSnapshotId(1) {}
+    ~TADCacheLifecycleTracker() = default;
+
+    // Disable copy and move
+    TADCacheLifecycleTracker(const TADCacheLifecycleTracker&) = delete;
+    TADCacheLifecycleTracker& operator=(const TADCacheLifecycleTracker&) = delete;
+    TADCacheLifecycleTracker(TADCacheLifecycleTracker&&) = delete;
+    TADCacheLifecycleTracker& operator=(TADCacheLifecycleTracker&&) = delete;
+
+    std::atomic<bool> _enabled;
+    mutable std::mutex _mutex;
+    TADCacheStats _stats;
+    std::atomic<LongType> _nextSnapshotId;
+    std::map<LongType, TADCacheStats> _snapshots;
+};
+
+} // namespace array
+} // namespace sd
+
+#endif // LIBND4J_TADCACHELIFECYCLETRACKER_H

libnd4j/include/array/TadCalculator.h

@@ -0,0 +1,81 @@
+/* ******************************************************************************
+*
+* 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
+******************************************************************************/
+
+//
+// @author Adam Gibson
+//
+
+#ifndef DEV_TESTS_TADCALCULATOR_H
+#define DEV_TESTS_TADCALCULATOR_H
+
+#include <array/TadPack.h>
+#include <system/common.h>
+#include <helpers/ConstantHelper.h>
+#include <helpers/ConstantShapeHelper.h>
+#include <array/ConstantShapeBuffer.h>
+#include <array/ConstantOffsetsBuffer.h>
+#include <vector>
+#include <memory>
+
+namespace sd {
+
+/**
+* TadCalculator handles the computation of Tensor Along Dimension (TAD) information
+* including shapes and offsets for sub-arrays.
+ */
+class SD_LIB_EXPORT TadCalculator {
+ private:
+  LongType* _originalShape;       // Original shape info pointer
+  ConstantShapeBuffer *_tadShape;        // Calculated TAD shape buffer
+  ConstantOffsetsBuffer *_tadOffsets;    // Calculated TAD offsets buffer
+  LongType _numTads;                    // Number of TADs
+
+ public:
+  /**
+    * Constructor for TadCalculator
+    * @param originalShape Pointer to the original shape information
+   */
+  explicit TadCalculator(LongType* originalShape);
+  ~TadCalculator();
+
+  /**
+    * Creates a TAD pack for the given dimensions
+    * @param dimensions Vector of dimensions to calculate TADs for
+   */
+  void createTadPack(const std::vector<LongType>& dimensions);
+
+  /**
+    * Returns the calculated TAD shape buffer
+    * @return ConstantShapeBuffer containing TAD shape information
+   */
+  ConstantShapeBuffer *tadShape() const { return _tadShape; }
+
+  /**
+    * Returns the calculated TAD offsets buffer
+    * @return ConstantOffsetsBuffer containing TAD offset information
+   */
+  ConstantOffsetsBuffer *tadOffsets() const { return _tadOffsets; }
+
+  /**
+    * Returns the number of TADs calculated
+    * @return Number of TADs
+   */
+  LongType numberOfTads() const { return _numTads; }
+};
+
+} // namespace sd
+
+#endif // DEV_TESTS_TADCALCULATOR_H

libnd4j/include/array/TadDescriptor.h

@@ -0,0 +1,74 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef DEV_TESTS_TADDESCRIPTOR_H
+#define DEV_TESTS_TADDESCRIPTOR_H
+
+#include <array/ShapeDescriptor.h>
+
+namespace sd {
+class SD_LIB_EXPORT TadDescriptor {
+ private:
+  sd::LongType *_originalShape;
+
+  std::vector<LongType> _axis;
+
+  bool _unitiesInShape;
+
+ public:
+  explicit TadDescriptor(const LongType *originalShape, const LongType *dimensions, const LongType length,
+                         const bool keepUnitiesInShape = false);
+  ~TadDescriptor() = default;
+
+
+  // NCC has issues with copy constructors
+  TadDescriptor &operator=(const TadDescriptor &other) = default;
+  // we use default move assignment operator
+  TadDescriptor &operator=(TadDescriptor &&other) noexcept = default;
+  explicit TadDescriptor(const TadDescriptor &other);
+
+
+  // equal to operator
+  bool operator==(const TadDescriptor &other) const;
+
+  // less than operator
+  bool operator<(const TadDescriptor &other) const;
+
+  std::vector<LongType> &axis();
+  LongType *originalShape();
+  bool areUnitiesinShape() const;
+};
+}  // namespace sd
+
+#ifndef __JAVACPP_HACK__
+
+namespace std {
+template <>
+class SD_LIB_EXPORT hash<sd::TadDescriptor> {
+ public:
+  size_t operator()(const sd::TadDescriptor &k) const;
+};
+}  // namespace std
+
+#endif
+
+#endif  // DEV_TESTS_TADDESCRIPTOR_H

libnd4j/include/array/TadPack.h

@@ -0,0 +1,127 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef DEV_TESTS_TADPACK_H
+#define DEV_TESTS_TADPACK_H
+
+#include <array/ConstantOffsetsBuffer.h>
+#include <array/ConstantShapeBuffer.h>
+#include <system/common.h>
+
+#include <array/NDArray.h>
+#ifndef __JAVACPP_HACK__
+namespace sd {
+class SD_LIB_EXPORT TadPack {
+ private:
+  ConstantShapeBuffer *_tadShape;
+  ConstantOffsetsBuffer *_tadOffsets;
+  LongType _numTads = 0;
+  LongType _shapeInfoLength = 0;
+  LongType* _dimensions = nullptr;
+  LongType _dimensionsLength = 0;
+  size_t _packHash = 0;  // Cache the hash for quick comparison
+
+ public:
+  explicit TadPack( ConstantShapeBuffer *shapes,
+                    ConstantOffsetsBuffer *offets, LongType numTads,
+                   LongType* dimensions = nullptr, LongType dimLength = 0);
+  TadPack() = default;
+  ~TadPack() {};
+
+  LongType* primaryShapeInfo();
+  LongType* primaryOffsets();
+
+  LongType* specialShapeInfo();
+  LongType* specialOffsets();
+
+  LongType numberOfTads() const;
+  LongType shapeInfoLength();
+  /**
+   * Extracts an NDArray view for the given TAD index.
+   * @param input The input NDArray.
+   * @param tadIndex The index of the TAD to extract.
+   * @return A new NDArray view representing the TAD.
+   */
+  NDArray *extractTadView(NDArray* input, sd::LongType tadIndex) {
+    auto shapeInfo = primaryShapeInfo();
+    auto offsets = primaryOffsets();
+
+    auto tadOffset = offsets[tadIndex];
+    auto x = input->buffer();
+    NDArray *ret = new NDArray(x,shapeInfo,LaunchContext::defaultContext(),false,tadOffset);
+    return ret;
+  }
+
+  void computeHash() {
+    size_t hash = 17;
+
+    // Add dimensions to hash
+    for (LongType i = 0; i < _dimensionsLength; i++) {
+      hash = hash * 31 + static_cast<size_t>(_dimensions[i]);
+    }
+
+    // Add shape info to hash if available
+    LongType* primaryShape = primaryShapeInfo();
+    if (primaryShape) {
+      int rank = shape::rank(primaryShape);
+
+      // Add rank
+      hash = hash * 13 + static_cast<size_t>(rank);
+
+      // Add shape dimensions
+      LongType* shapeValues = shape::shapeOf(primaryShape);
+      for (int i = 0; i < rank; i++) {
+        hash = hash * 17 + static_cast<size_t>(shapeValues[i]);
+      }
+
+      // Add strides
+      LongType* strides = shape::stride(primaryShape);
+      for (int i = 0; i < rank; i++) {
+        hash = hash * 23 + static_cast<size_t>(strides[i]);
+      }
+
+      // Add order and data type
+      hash = hash * 29 + static_cast<size_t>(shape::order(primaryShape));
+      hash = hash * 37 + static_cast<size_t>(ArrayOptions::dataType(primaryShape));
+    }
+
+    // Add number of TADs
+    hash = hash * 41 + static_cast<size_t>(_numTads);
+
+    _packHash = hash;
+  }
+
+
+  /**
+   * These methods return either primary or special pointers depending on platform binaries were compiled for
+   * @return
+   */
+  LongType* platformShapeInfo();
+  LongType* platformOffsets();
+
+  void print(const char* msg);
+  bool operator==( TadPack& other);
+};
+}  // namespace sd
+
+#endif  // DEV_TESTS_TADPACK_H
+#endif

libnd4j/include/array/cpu/DataBuffer.cpp

@@ -0,0 +1,372 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <array/DataBuffer.h>
+#include <array/DataTypeUtils.h>
+#include <types/types.h>
+#include <system/type_boilerplate.h>
+
+
+namespace sd {
+void DataBuffer::expand(const uint64_t size) {
+  if (static_cast<LongType>(size) > _lenInBytes) {
+    // allocate new buffer
+    int8_t* newBuffer = nullptr;
+    ALLOCATE(newBuffer, _workspace, size, int8_t);
+
+    // copy data from existing buffer
+    std::memcpy(newBuffer, _primaryBuffer, _lenInBytes);
+
+    if (_isOwnerPrimary) {
+      RELEASE(reinterpret_cast<int8_t*>(_primaryBuffer), _workspace);
+    }
+
+    _primaryBuffer = newBuffer;
+    _lenInBytes = size;
+    _isOwnerPrimary = true;
+  }
+}
+
+template <typename T>
+void DataBuffer::printHostBufferContent(void* buffer, sd::LongType offset, sd::LongType length) {
+  T* typedBuffer = reinterpret_cast<T*>(buffer);
+
+  sd_printf("[ ", 0);
+  for (sd::LongType i = offset; i < offset + length; i++) {
+    // For numeric types, cast to double for consistent formatting
+    if (std::is_arithmetic<T>::value) {
+      sd_printf("%g ", (double)typedBuffer[i]);
+    } else {
+      // For non-numeric types, print as hex
+      sd_printf("0x%x ", *reinterpret_cast<int*>(&typedBuffer[i]));
+    }
+  }
+  sd_printf("]", 0);
+}
+BUILD_SINGLE_TEMPLATE(void DataBuffer::printHostBufferContent,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
+
+
+
+
+// DataBuffer implementation for .cpp file
+void DataBuffer::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
+  if (msg) sd_printf("%s:\n", msg);
+
+  // Print metadata
+  sd_printf("DataBuffer: DataType=%s, Length=%lld elements, DeviceId=%d\n",
+            DataTypeUtils::asString(_dataType).c_str(), (long long)getNumElements(), deviceId());
+
+  // Print host buffer content
+  if (_primaryBuffer != nullptr) {
+    sd_printf("Host buffer (@%p): ", _primaryBuffer);
+
+    sd::LongType len = getNumElements();
+    sd::LongType printLen = limit < 0 ? len : std::min(len - offset, limit);
+
+    // Print based on datatype
+    BUILD_SINGLE_SELECTOR(_dataType, printHostBufferContent,
+                          (_primaryBuffer, offset, printLen), SD_COMMON_TYPES);
+
+    if (offset + printLen < len) sd_printf("... ", 0);
+    sd_printf("\n", 0);
+  } else {
+    sd_printf("Host buffer: nullptr\n", 0);
+  }
+
+  // Print device info but not content (CPU version)
+  if (_specialBuffer != nullptr) {
+    sd_printf("Device buffer (@%p): [Not accessible from CPU]\n", _specialBuffer);
+  } else {
+    sd_printf("Device buffer: nullptr\n", 0);
+  }
+
+#if defined(SD_CUDA)
+  // Print sync state counters
+  sd_printf("Sync state: _counter=%lld, _writePrimary=%lld, _writeSpecial=%lld, _readPrimary=%lld, _readSpecial=%lld\n",
+            (long long)_counter.load(), (long long)_writePrimary.load(), (long long)_writeSpecial.load(),
+            (long long)_readPrimary.load(), (long long)_readSpecial.load());
+  sd_printf("isPrimaryActual=%d, isSpecialActual=%d\n", isPrimaryActual(), isSpecialActual());
+#endif
+}
+
+// Helper template to print host buffer content
+template <typename T>
+void printHostBufferContent(void* buffer, sd::LongType offset, sd::LongType length) {
+  T* typedBuffer = reinterpret_cast<T*>(buffer);
+
+  sd_printf("[ ", 0);
+  for (sd::LongType i = offset; i < offset + length; i++) {
+    // For numeric types, cast to double for consistent formatting
+    if (std::is_arithmetic<T>::value) {
+      sd_printf("%g ", (double)typedBuffer[i]);
+    } else {
+      // For non-numeric types, print as hex
+      sd_printf("0x%x ", *reinterpret_cast<int*>(&typedBuffer[i]));
+    }
+  }
+  sd_printf("]", 0);
+}
+
+void DataBuffer::printSpecialAllocationTraces() {
+  //no op on purpose
+}
+
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::allocateBuffers(const bool allocBoth) {  // always allocate primary buffer only (cpu case)
+  allocatePrimary();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyBufferFrom(const DataBuffer& other,
+                                size_t sizeToCopyinBytes,
+                                const sd::LongType offsetThis,
+                                const sd::LongType offsetOther) {
+  if(other._dataType != _dataType) {
+    THROW_EXCEPTION("DataBuffer::copyBufferFrom: data types of buffers are different");
+  }
+  if (sizeToCopyinBytes == 0) {
+    LongType otherBytes = other.getLenInBytes() - offsetOther;
+    LongType thisBytes = getLenInBytes() - offsetThis;
+    sizeToCopyinBytes = otherBytes < thisBytes ? otherBytes : thisBytes;
+  }
+  if (sizeToCopyinBytes == 0) return;
+  if(static_cast<LongType>(sizeToCopyinBytes) > other._lenInBytes - offsetOther) {
+    std::string errorMessage;
+    errorMessage = "DataBuffer::copyBufferFrom: size to copy is larger than source buffer ";
+    errorMessage += std::to_string(sizeToCopyinBytes);
+    errorMessage += " > ";
+    errorMessage += std::to_string(other._lenInBytes - offsetOther);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if(sizeToCopyinBytes > getLenInBytes() - offsetThis) {
+    std::string errorMessage;
+    errorMessage = "DataBuffer::copyBufferFrom: size to copy is larger than destination buffer ";
+    errorMessage += std::to_string(sizeToCopyinBytes);
+    errorMessage += " > ";
+    errorMessage += std::to_string(getLenInBytes() - offsetThis);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if (other._primaryBuffer != nullptr) {
+    auto sizeOfElement = DataTypeUtils::sizeOfElement(_dataType);
+    auto sizeOfOtherElement = DataTypeUtils::sizeOfElement(_dataType);
+    if(sizeOfElement != sizeOfOtherElement) {
+      THROW_EXCEPTION("DataBuffer::copyBufferFrom: size of elements in buffers are different");
+    }
+    std::memcpy(
+        static_cast<int8_t*>(_primaryBuffer) + offsetThis * sizeOfElement,
+        static_cast<const int8_t*>(other._primaryBuffer) + offsetOther * sizeOfOtherElement,
+        sizeToCopyinBytes);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyBufferFromHost(const void* hostBuffer, size_t sizeToCopyinBytes, const sd::LongType offsetThis,
+                                    const sd::LongType offsetHostBuffer) {
+  if (sizeToCopyinBytes == 0) sizeToCopyinBytes = getLenInBytes();
+  if (sizeToCopyinBytes == 0) return;
+
+  if (hostBuffer != nullptr)
+    std::memcpy(static_cast<int8_t*>(_primaryBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
+                static_cast<const int8_t*>(hostBuffer) + offsetHostBuffer * DataTypeUtils::sizeOfElement(_dataType),
+                sizeToCopyinBytes);
+}
+
+/////////////////////////
+
+template <typename T>
+void memcpyWithT(DataBuffer* dst, DataBuffer* src, sd::LongType startingOffset, sd::LongType dstOffset) {
+  if(src->getLenInBytes() != dst->getLenInBytes()) {
+    THROW_EXCEPTION("DataBuffer::memcpy: source and destination buffers have different length in bytes");
+  }
+
+  std::memcpy(dst->primaryAtOffset<T>(dstOffset), src->primaryAtOffset<T>(startingOffset), src->getLenInBytes());
+  dst->readPrimary();
+}
+
+void DataBuffer::memcpy(DataBuffer* dst, DataBuffer* src,
+                        sd::LongType startingOffset, sd::LongType dstOffset) {
+  BUILD_SINGLE_SELECTOR(dst->_dataType, memcpyWithT,(dst, src, startingOffset, dstOffset),
+                        SD_COMMON_TYPES);
+
+  dst->readPrimary();
+}
+
+
+
+
+////////////////////////////////////////////////////////////////////////
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::deleteSpecial() {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::syncToPrimary(const LaunchContext* context, const bool forceSync) {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setCountersToZero() {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyCounters(const DataBuffer& other) {}
+
+void DataBuffer::writePrimary() const {}
+void DataBuffer::writeSpecial() const {}
+void DataBuffer::readPrimary() const {}
+void DataBuffer::readSpecial() const {}
+bool DataBuffer::isPrimaryActual() const { return true; }
+bool DataBuffer::isSpecialActual() const { return false; }
+void DataBuffer::showBufferLimited() {}
+
+
+
+DataBuffer DataBuffer::dup() {
+  DataBuffer result;
+  result._dataType = _dataType;
+  result._lenInBytes = _lenInBytes;
+  result._primaryBuffer = _primaryBuffer;
+  result._specialBuffer = _specialBuffer;
+  result._isOwnerPrimary = _isOwnerPrimary;
+  result._isOwnerSpecial = _isOwnerSpecial;
+  result.allocateBuffers(true);
+  result.copyCounters(*this);
+  result.copyBufferFrom(*this);
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setSpecial(void* special, const bool isOwnerSpecial) {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setToZeroBuffers(const bool both) { memset(primary(), 0, getLenInBytes()); }
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::syncToSpecial(const bool forceSync) {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::allocateSpecial() {}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::migrate() {}
+
+
+template <typename T>
+void _printHostBuffer(DataBuffer* buffer, long offset) {
+  sd::LongType len = buffer->getNumElements();
+  auto buff = buffer->template primaryAsT<T>();
+
+
+  sd::LongType limit = len;
+  if (limit == -1 || limit >= static_cast<LongType>(buffer->getNumElements())) {
+    limit = buffer->getNumElements();
+  }
+
+  const char* msg = nullptr;
+  if (msg != nullptr) {
+    printf("%s: ", msg);
+  } else {
+    printf("[");
+  }
+
+  sd::DataType dataType = buffer->getDataType();
+  auto baseOffset = offset;
+  if (dataType == sd::DataType::DOUBLE || dataType == sd::DataType::FLOAT32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (e > offset) printf(", ");
+      if (dataType == sd::DataType::DOUBLE) {
+        double val = static_cast<double>(buff[e]);
+        printf("%.15f",val);
+      } else {
+        printf("%.15f", static_cast<float>(buff[e]));
+      }
+    }
+  } else if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64 ||
+             dataType == sd::DataType::INT32 || dataType == sd::DataType::UINT32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64) {
+        printf("%lld", static_cast<long long>(buff[e]));
+      } else {
+        printf("%d", static_cast<int>(buff[e]));
+      }
+
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  } else if (dataType == sd::DataType::BOOL) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (static_cast<bool>(buff[e])) {
+        printf("true");
+      } else {
+        printf("false");
+      }
+
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  } else if (dataType == sd::DataType::UTF8 || dataType == sd::DataType::UTF16 ||
+             dataType == sd::DataType::UTF32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      printf("\"%s\"", reinterpret_cast<const char*>(&buff[e]));
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  }
+
+  printf("]\n");
+  fflush(stdout);
+}
+void DataBuffer::printHostDevice(long offset) {
+  auto xType = getDataType();
+  BUILD_SINGLE_SELECTOR(xType, _printHostBuffer,(this,offset),SD_COMMON_TYPES);
+
+}
+
+
+void DataBuffer::showCounters(const char* msg1, const char* msg2) {
+
+}
+template <typename T>
+void* DataBuffer::primaryAtOffset(const LongType offset) {
+  if(_primaryBuffer == nullptr)
+    return nullptr;
+  T *type = reinterpret_cast<T*>(_primaryBuffer);
+  return reinterpret_cast<void *>(type + offset);
+}
+
+#define PRIMARYOFFSET(T) template void* DataBuffer::primaryAtOffset<GET_SECOND(T)>(const LongType offset);
+ITERATE_LIST((SD_COMMON_TYPES),PRIMARYOFFSET)
+
+template <typename T>
+void* DataBuffer::specialAtOffset(const LongType offset) {
+  if(_specialBuffer == nullptr)
+    return nullptr;
+  T *type = reinterpret_cast<T*>(_specialBuffer);
+  return reinterpret_cast<void *>(type + offset);
+}
+
+#define SPECIALOFFSET(T) template void* DataBuffer::specialAtOffset<GET_SECOND(T)>(const LongType offset);
+ITERATE_LIST((SD_COMMON_TYPES),SPECIALOFFSET)
+}  // namespace sd

libnd4j/include/array/cpu/NDArray.cpp

@@ -0,0 +1,541 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+#ifndef NDARRAY_CPP
+#define NDARRAY_CPP
+#include <array/NDArray.h>
+
+#include <helpers/ArrayUtils.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/ShapeUtils.h>
+#include <helpers/logger.h>
+
+#include <indexing/NDIndex.h>
+#include <loops/BroadcastPairwiseConverter.h>
+
+#include <loops/random.h>
+
+#include <ops/ops.h>
+
+#include <array/NDArray.hXX>
+
+#include <memory>
+#include <sstream>
+#include <stdexcept>
+#include <system/selective_rendering.h>
+namespace sd {
+
+
+// NDArray implementation for .cpp file
+void NDArray::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
+  if (msg) sd_printf("%s:\n", msg);
+  if(limit < 0) limit = lengthOf();
+
+  // Print array info
+  sd_printf("NDArray: Shape=[", 0);
+  for (int i = 0; i < rankOf(); i++) {
+    sd_printf("%lld", (long long)sizeAt(i));
+    if (i < rankOf() - 1) sd_printf(",", 0);
+  }
+  sd_printf("], DataType=%s, EWS=%lld, Order=%c\n",
+            DataTypeUtils::asString(dataType()).c_str(), (long long)ews(), ordering());
+
+  // Print buffer state
+  if (_buffer != nullptr) {
+    _buffer->printBufferDebug("Buffer contents", offset, limit);
+  } else {
+    sd_printf("Buffer is nullptr\n", 0);
+  }
+}
+////////////////////////////////////////////////////////////////////////
+
+void* NDArray::platformBuffer() { return buffer(); }
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+void NDArray::fillAsTriangular(const float val, int lower, int upper, NDArray& target, const char direction, const bool includeEdges) {
+  if (isS()) THROW_EXCEPTION("NDArray::fillArrayAsTriangular: you can't use this method on String array!");
+
+  if (!isSameShape(target) &&
+      !(rankOf() == 1 && target.rankOf() == 2 && sizeAt(0) == target.sizeAt(0) && sizeAt(0) == target.sizeAt(1)))
+    THROW_EXCEPTION("NDArray::fillArrayAsTriangular method: wrong shape of target array !");
+
+  const T value = static_cast<T>(val);
+  const auto x = reinterpret_cast<const T*>(buffer());
+  auto z = reinterpret_cast<T*>(target.buffer());
+
+  const int xRank = rankOf();
+  const int zRank = target.rankOf();
+
+  const auto zLen = target.lengthOf();
+
+  const bool areSameOffsets = shape::haveSameShapeAndStrides(shapeInfo(), target.shapeInfo());
+
+  sd::LongType *targetShape = shape::shapeOf(target.shapeInfo());
+  sd::LongType *targetStride = shape::stride(target.shapeInfo());
+  sd::LongType targetRank = target.rankOf();
+
+  sd::LongType *xShape = shape::shapeOf(shapeInfo());
+  sd::LongType *xStride = shape::stride(shapeInfo());
+  sd::LongType thisRank = this->rankOf();
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType coords[SD_MAX_RANK], temp;
+    sd::LongType vectorCoord[1] = {0};
+
+    bool notVectorScalar = targetRank == 2 && thisRank == 2;
+    bool thisNotVectorScalar = !shape::isScalar(this->shapeInfo()) && !shape::isVector(this->shapeInfo());
+    bool targetNotVectorScalar = !shape::isScalar(target.shapeInfo()) && !shape::isVector(target.shapeInfo());
+
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, targetRank,targetShape, coords);
+      sd::LongType row = targetNotVectorScalar ? coords[zRank - 2] : 0;
+      sd::LongType col = targetNotVectorScalar ? coords[zRank - 1] : 1;
+      sd::LongType zOffset, xOffset;
+
+      if (target.rankOf() < 2) {
+        COORDS2INDEX(targetRank, targetStride, vectorCoord, zOffset);
+      } else {
+        COORDS2INDEX(targetRank, targetStride, coords, zOffset);
+      }
+
+      if (!areSameOffsets && rankOf() < 2) {
+        COORDS2INDEX(thisRank, xStride, vectorCoord, xOffset);
+      } else if (areSameOffsets) {
+        xOffset = zOffset;
+      } else {
+        COORDS2INDEX(thisRank, xStride, coords, xOffset);
+      }
+
+      bool rowExclusive = this->rankOf() == target.rankOf();
+      bool colExclusive = this->rankOf() == target.rankOf();
+      auto lCompare = includeEdges ? row <= (col - lower) : row < (col - lower);
+      auto uCompare = includeEdges ? row >= (col - upper) : row > (col - upper);
+
+      if ((direction == 'u' && lCompare) || (direction == 'l' && uCompare)) {
+        z[zOffset] = value;
+      } else {
+        z[zOffset] = x[xOffset];
+      }
+
+      if (this != &target) {
+        if (xRank != zRank) {
+          temp = coords[0];
+          coords[0] = coords[1];
+        }
+
+        if (xRank != zRank)  // restore first coordinate
+          coords[0] = temp;
+      }
+
+      if (vectorCoord[0] == this->lengthOf() - 1) {
+        vectorCoord[0] = 0;
+      } else {
+        vectorCoord[0] = vectorCoord[0] + 1;
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, zLen);
+}
+BUILD_SINGLE_TEMPLATE( void NDArray::fillAsTriangular,
+                      (const float val, int lower, int upper, NDArray& target, const char direction,const bool includeEdges), SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::setIdentity() {
+  if (isS()) THROW_EXCEPTION("NDArray::setIdentity: you can't use this method on String array!");
+
+  this->nullify();
+
+  int rank = rankOf();
+  auto shape = shapeOf();
+  int minDim = SD_MAX_INT;
+  sd::LongType indices[SD_MAX_RANK];
+  for (int j = 0; j < rank; ++j) indices[j] = 1;
+
+  sd::LongType offset;
+  COORDS2INDEX(rank, shape::stride(shapeInfo()), indices, offset);
+
+  for (int i = 0; i < rank; ++i)
+    if (minDim > shape[i]) minDim = shape[i];
+
+  float v = 1.0f;
+
+  for (int i = 0; i < minDim; ++i) templatedSet<float,float>(buffer(), i * offset, this->dataType(), &v);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+static void templatedSwap(void* xBuffer, void* yBuffer, const sd::LongType* xShapeInfo, const sd::LongType* yShapeInfo,
+                          sd::LongType length) {
+  auto x = reinterpret_cast<T*>(xBuffer);
+  auto y = reinterpret_cast<T*>(yBuffer);
+
+  const bool isSameOrders = shape::order(xShapeInfo) == shape::order(xShapeInfo);
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *yShape = shape::shapeOf(yShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *yStride = shape::stride(yShapeInfo);
+
+  auto func = PRAGMA_THREADS_FOR {
+    if (isSameOrders) {
+      for (sd::LongType i = start; i < stop; i++) {
+        LongType xCoords[SD_MAX_RANK];
+        LongType yCoords[SD_MAX_RANK];
+        LongType xOffset;
+        LongType yOffset;
+
+        INDEX2COORDS(i, xRank, xShape, xCoords);
+        COORDS2INDEX(xRank, shape::stride(xShapeInfo), xCoords, xOffset);
+        INDEX2COORDS(i, yRank,yShape, yCoords);
+        COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+
+        sd::math::sd_swap(x[xOffset], y[yOffset]);
+      }
+    } else if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)) {
+      for (sd::LongType i = start; i < stop; i++) {
+        LongType coords[SD_MAX_RANK];
+        LongType ind;
+
+        INDEX2COORDS(i, xRank, xShape, coords);
+        COORDS2INDEX(xRank, xStride, coords, ind);
+
+        sd::math::sd_swap(x[ind], y[ind]);
+      }
+    } else {
+      for (sd::LongType i = start; i < stop; i++) {
+        LongType xCoords[SD_MAX_RANK];
+        LongType yCoords[SD_MAX_RANK];
+        LongType xInd;
+        LongType yInd;
+
+        INDEX2COORDS(i, xRank, xShape, xCoords);
+        COORDS2INDEX(xRank, xStride, xCoords, xInd);
+        INDEX2COORDS(i, yRank, yShape, yCoords);
+        COORDS2INDEX(yRank, yStride, yCoords, yInd);
+
+        sd::math::sd_swap(x[xInd], y[yInd]);
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, length);
+}
+BUILD_SINGLE_TEMPLATE( void templatedSwap,
+                      (void* xBuffer, void* yBuffer, const sd::LongType* xShapeInfo, const sd::LongType* yShapeInfo,
+                          sd::LongType length),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::swapUnsafe(NDArray& other) {
+  auto xType = this->dataType();
+
+  if (xType != other.dataType())
+    THROW_EXCEPTION("NDArray::swapUnsage method: both arrays must have the same data type");
+
+  if (buffer() == nullptr || other.buffer() == nullptr)
+    THROW_EXCEPTION("NDArray::swapUnsafe method: input array should not be empty!");
+
+  if (lengthOf() != other.lengthOf())
+    THROW_EXCEPTION("NDArray::swapUnsafe method: input arrays should have the same length!");
+
+  BUILD_SINGLE_SELECTOR(xType, templatedSwap,
+                        (buffer(), other.buffer(), shapeInfo(), other.shapeInfo(), this->lengthOf()), SD_COMMON_TYPES);
+}
+
+////////////////////////////////////////////////////////////////////////
+
+
+void NDArray::synchronize(const char* msg)  {
+  // no-op
+}
+
+void NDArray::syncToDevice()  {}
+void NDArray::syncToHost()  {}
+void NDArray::tickWriteHost()  {}
+void NDArray::tickWriteDevice()  {}
+void NDArray::tickReadHost()  {}
+void NDArray::tickReadDevice()  {}
+void NDArray::tickBothActual()  {}
+bool NDArray::isActualOnHostSide()  { return true; }
+bool NDArray::isActualOnDeviceSide()  { return true; }
+void NDArray::makeBothBuffersActual()  {}
+
+void NDArray::preparePrimaryUse(const std::vector<NDArray*>& writeList,
+                                const std::vector<NDArray*>& readList, bool synchronizeWritables) {
+  // no-op
+}
+void NDArray::registerPrimaryUse(const std::vector<NDArray*>& writeList,
+                                 const std::vector<NDArray*>& readList) {
+  // no-op
+}
+
+
+void NDArray::prepareSpecialUse(const std::vector<NDArray*>& writeList,
+                                const std::vector<NDArray*>& readList, bool synchronizeWritables) {
+  // no-op
+}
+void NDArray::registerSpecialUse(const std::vector<NDArray*>& writeList,
+                                 const std::vector<NDArray*>& readList) {
+  // no-op
+}
+
+void NDArray::syncShape()  {
+  // no-op
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void NDArray::printCurrentBuffer(const bool host, const char* msg, const int precision)  {}
+template void NDArray::printCurrentBuffer<int>(const bool host, const char* msg, const int precision) ;
+template void NDArray::printCurrentBuffer<float>(const bool host, const char* msg, const int precision) ;
+template void NDArray::printCurrentBuffer<double>(const bool host, const char* msg, const int precision);
+template void NDArray::printCurrentBuffer<sd::LongType>(const bool host, const char* msg, const int precision) ;
+
+////////////////////////////////////////////////////////////////////////
+void* NDArray::specialBuffer() {
+  if (_buffer == nullptr) {
+    THROW_EXCEPTION("NDArray::specialBuffer(): _buffer is nullptr - array not properly initialized");
+  }
+
+  void* specialBuf = _buffer->special();
+
+  // On CPU, special buffer is nullptr (only used for GPU/CUDA) - this is expected and normal
+  if (specialBuf == nullptr) {
+    return nullptr;
+  }
+
+  return static_cast<int8_t*>(specialBuf) + (_offset * sizeOfT());
+}
+
+
+//////////////////////////////////////////////////////////////////////////
+// change an array by repeating it the number of times given by reps.
+NDArray NDArray::tile(const std::vector<sd::LongType>& reps)  {
+  const int repsSize = reps.size();
+
+  sd::LongType product = 1;
+  for (const auto& item : reps) product *= item;
+  if (product == 0) THROW_EXCEPTION("NDArray::tile method: one of the elements in reps array is zero !");
+
+  int rankOld = rankOf();
+  int diff = rankOld - repsSize;
+  if (product == 1) {  // in this case 2 possibilities are present: just reshape or nothing to do
+    NDArray result(*this);
+    if (diff < 0) {  // reshape to higher dimension
+      std::vector<sd::LongType> shapeNew =
+          reps;  // there is requirement to have unities at first "diff" positions of new shape
+      memcpy(&shapeNew[-diff], result.shapeInfo() + 1,
+             rankOld * sizeof(sd::LongType));  // put old shape numbers at rest of positions
+      result.reshapei(ordering(), shapeNew);
+    }
+    return result;  // nothing to do, if diff >= 0 -> identity tile
+  }
+
+  // evaluate shapeInfo for resulting array
+  auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
+  // create new buffer, in any case the memory amount new buffer points to is bigger then those for old _buffer
+  DataBuffer * newBuff =
+      new DataBuffer(shape::length(newShapeInfo) * sizeOfT(), dataType(), getContext()->getWorkspace());
+  // assign new shape and new buffer to resulting array
+  NDArray result(newBuff,newShapeInfo , getContext());
+  // fill newBuff, loop through all elements of newBuff
+  // looping through _buffer goes automatically by means of getSubArrayIndex applying
+  const auto resultLen = result.lengthOf();
+  auto xType = this->dataType();
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto i = start; i < stop; i++) {
+      auto xOffset = result.getOffset(i);
+      auto yOffset = shape::subArrayOffset(i, newShapeInfo, shapeInfo());
+      BUILD_SINGLE_SELECTOR(xType, this->template templatedAssign,
+                            (result.buffer(), xOffset, this->buffer(), yOffset), SD_COMMON_TYPES);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, resultLen);
+  result.tickWriteHost();
+  return result;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// change an array by repeating it the number of times given by reps.
+void NDArray::tile(const std::vector<LongType>& reps, NDArray& target) {
+  // Validate the tile operation
+  auto repProd = shape::prodLong(reps.data(), reps.size());
+  if (repProd < 1)
+    THROW_EXCEPTION("NDArray::tile: reps can't contain 0s");
+
+  // Validate the target shape
+  auto correctShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
+  if (!shape::equalsSoft(correctShapeInfo, target.shapeInfo())) {
+    THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation!");
+  }
+
+  const auto targetLen = target.lengthOf();
+
+  // Safely calculate source array offset
+  for (LongType i = 0; i < targetLen; ++i) {
+    // Calculate target array offset
+    auto xOffset = target.getOffset(i);
+
+    // Calculate source coordinates based on target coordinates
+    LongType targetCoords[SD_MAX_RANK];
+    INDEX2COORDS(i, shape::rank(target.shapeInfo()), shape::shapeOf(target.shapeInfo()), targetCoords);
+
+    // Map target coordinates to source coordinates manually
+    LongType sourceCoords[SD_MAX_RANK];
+    for (int d = 0; d < shape::rank(shapeInfo()); d++) {
+      // Apply modulo for each dimension
+      sourceCoords[d] = targetCoords[d] % shape::sizeAt(shapeInfo(), d);
+    }
+
+    // Calculate source offset from source coordinates
+    LongType sourceOffset;
+    COORDS2INDEX(shape::rank(shapeInfo()), shape::stride(shapeInfo()), sourceCoords, sourceOffset);
+
+    auto targetDataType = target.dataType();
+    auto selfDType = dataType();
+    // Copy the value
+    BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
+                          (target.buffer(), xOffset, buffer(), sourceOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+void NDArray::tile(NDArray& target)  {
+  if (rankOf() > target.rankOf())
+    THROW_EXCEPTION(
+        "NDArray::tile method - rank of target array must be bigger or equal to the rank of this array !");
+
+  if (!ShapeUtils::areShapesBroadcastable(*this, target))
+    THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
+
+  // fill newBuff, loop through all elements of newBuff
+  // looping through _buffer goes automatically by means of getSubArrayIndex applying
+  const auto ews = target.ews();
+  const auto targetLen = target.lengthOf();
+  if (target.ordering() == 'c' && ews >= 1) {
+    for (sd::LongType i = 0; i < targetLen; ++i) {
+      auto yOffset = shape::subArrayOffset(i, target.shapeInfo(), shapeInfo());
+      auto targetDataType = target.dataType();
+      auto selfDType = dataType();
+      BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
+                            (target.buffer(), i * ews, buffer(), yOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
+    }
+  } else {
+    for (sd::LongType i = 0; i < targetLen; ++i) {
+      auto xOffset = target.getOffset(i);
+      auto yOffset = shape::subArrayOffset(i, target.shapeInfo(), shapeInfo());
+      auto targetDataType = target.dataType();
+      auto selfDType = dataType();
+      BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
+                            (target.buffer(), xOffset, buffer(), yOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+static void repeat_(NDArray& input, NDArray& output, const std::vector<LongType>& repeats, const LongType axis) {
+  const X* x = input.bufferAsT<X>();
+  Z* z = output.bufferAsT<Z>();
+
+  const sd::LongType rank = input.rankOf();     // xRank = zRank
+  const sd::LongType zLen = output.lengthOf();  // xLen <= zLen
+  const sd::LongType repSize = repeats.size();
+
+  sd::LongType outputRank = output.rankOf();
+  sd::LongType* outputShape = shape::shapeOf(output.shapeInfo());
+  sd::LongType* outputStride = shape::stride(output.shapeInfo());
+
+  sd::LongType inputRank = input.rankOf();
+  sd::LongType* inputShape = shape::shapeOf(input.shapeInfo());
+  sd::LongType* inputStride = shape::stride(input.shapeInfo());
+
+  // loop through input array
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType coords[SD_MAX_RANK], temp;
+
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, outputRank, outputShape, coords);
+      sd::LongType zOffset;
+      COORDS2INDEX(outputRank, outputStride, coords, zOffset);
+
+      temp = coords[axis];
+
+      if (repSize > 1) {
+        for (sd::LongType j = 0; j < repSize; ++j) {
+          coords[axis] -= repeats[j];
+          if (coords[axis] < 0) {
+            coords[axis] = j;
+            break;
+          }
+        }
+      } else
+        coords[axis] /= repeats[0];
+
+      sd::LongType xOffset;
+      COORDS2INDEX(inputRank,inputStride, coords, xOffset);
+
+      z[zOffset] = x[xOffset];
+
+      coords[axis] = temp;
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, zLen);
+}
+
+//////////////////////////////////////////////////////////////////////////
+// create new array by repeating it the number of times given by repeats
+NDArray NDArray::repeat(const int axis, const std::vector<LongType>& repeats)  {
+  NDArray *thisArr = const_cast<NDArray*>(this);
+  std::vector<sd::LongType> repeatShape = ShapeUtils::evalRepeatShape(axis, repeats, *thisArr);
+  NDArray output('c',repeatShape, dataType(), getContext());
+
+  BUILD_SINGLE_SELECTOR_TWICE(dataType(), repeat_, (*this, output, repeats, axis), SD_COMMON_TYPES);
+
+  return output;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// fill array by repeating it the number of times given by reps
+void NDArray::repeat(const int axis, const std::vector<LongType>& repeats, NDArray& target)  {
+  NDArray *thisArr = const_cast<NDArray*>(this);
+  if (!target.isSameShape(ShapeUtils::evalRepeatShape(axis, repeats, *thisArr)))
+    THROW_EXCEPTION(
+        "NDArray::repeat(const int axis, const std::vector<int>& repeats, NDArray& target) method: wrong shape of "
+        "target array!");
+  auto targetDataType = target.dataType();
+  auto selfDType = dataType();
+  BUILD_DOUBLE_SELECTOR(dataType(), target.dataType(), repeat_, (*this, target, repeats, axis), SD_COMMON_TYPES,
+                        SD_COMMON_TYPES);
+}
+
+//////////////////////////////////////////////////////////////////////////
+#ifndef __JAVACPP_HACK__
+#include "NDArrayLambda.hpp"
+
+#endif
+
+
+}  // namespace sd
+
+#endif

libnd4j/include/array/cpu/NDArray.macro

@@ -0,0 +1,149 @@
+################################################################################
+# 
+#
+# This program and the accompanying materials are made available under the
+# terms of the Apache License, Version 2.0 which is available at
+# https://www.apache.org/licenses/LICENSE-2.0.
+#
+# See the NOTICE file distributed with this work for additional
+# information regarding copyright ownership.
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+################################################################################
+
+#ifndef NDARRAY_MACRO
+#define NDARRAY_MACRO
+#include <op_boilerplate.h>
+
+//NDArray<T> *other, T *extraParams
+BUILD_CALL_1(template void NDArray<float>::template applyPairwiseTransform, float, (NDArray<float>* other, float* extraParams), PAIRWISE_TRANSFORM_OPS)
+BUILD_CALL_1(template void NDArray<float16>::applyPairwiseTransform, float16, (NDArray<float16>* other, float16* extraParams), PAIRWISE_TRANSFORM_OPS)
+BUILD_CALL_1(template void NDArray<double>::applyPairwiseTransform, double, (NDArray<double>* other, double* extraParams), PAIRWISE_TRANSFORM_OPS)
+
+// NDArray<T> *other, NDArray<T> *target, T *extraParams
+BUILD_CALL_1(template void sd::NDArray<float>::applyPairwiseTransform, float, (NDArray<float>* other, NDArray<float>* target, float* extraParams), PAIRWISE_TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::applyPairwiseTransform, float16, (NDArray<float16>* other, NDArray<float16>* target, float16* extraParams), PAIRWISE_TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyPairwiseTransform, double, (NDArray<double>* other, NDArray<double>* target, double* extraParams), PAIRWISE_TRANSFORM_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyScalar, float16, (NDArray<float16>& scalar, NDArray<float16>* target, float16 *extraParams) const, SCALAR_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::applyScalar, float16, (float16 scalar, NDArray<float16>* target, float16 *extraParams) const, SCALAR_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float>::applyScalar, float, (NDArray<float>& scalar, NDArray<float>* target, float *extraParams) const, SCALAR_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyScalar, float, (float scalar, NDArray<float>* target, float *extraParams) const, SCALAR_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<double>::applyScalar, double, (NDArray<double>& scalar, NDArray<double>* target, double *extraParams) const, SCALAR_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyScalar, double, (double scalar, NDArray<double>* target, double *extraParams) const, SCALAR_OPS)
+
+
+BUILD_CALL_1(template float16 sd::NDArray<float16>::reduceNumber, float16, (float16 *extraParams) const, REDUCE_OPS)
+BUILD_CALL_1(template float sd::NDArray<float>::reduceNumber, float, (float *extraParams) const, REDUCE_OPS)
+BUILD_CALL_1(template double sd::NDArray<double>::reduceNumber, double, (double *extraParams) const, REDUCE_OPS)
+
+BUILD_CALL_1(template LongType sd::NDArray<float16>::indexReduceNumber, float16, (float16 *extraParams), INDEX_REDUCE_OPS)
+BUILD_CALL_1(template LongType sd::NDArray<float>::indexReduceNumber, float, (float *extraParams), INDEX_REDUCE_OPS)
+BUILD_CALL_1(template LongType sd::NDArray<double>::indexReduceNumber, double, (double *extraParams), INDEX_REDUCE_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyBroadcast, float16, (std::initializer_list<int> list, const sd::NDArray<float16>* a, sd::NDArray<float16>* b, float16* c), BROADCAST_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyBroadcast, float, (std::initializer_list<int> list, const sd::NDArray<float>* a, sd::NDArray<float>* b, float* c), BROADCAST_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyBroadcast, double, (std::initializer_list<int> list, const sd::NDArray<double>* a, sd::NDArray<double>* b, double* c), BROADCAST_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyTrueBroadcast, float16,(const sd::NDArray<float16>* a, sd::NDArray<float16>* target, const bool checkTargetShape, float16* c) const, BROADCAST_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyTrueBroadcast,   float,  (const sd::NDArray<float>*   a, sd::NDArray<float>*   target, const bool checkTargetShape, float* c)   const, BROADCAST_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyTrueBroadcast,  double, (const sd::NDArray<double>*  a, sd::NDArray<double>*  target, const bool checkTargetShape, double* c)  const, BROADCAST_OPS)
+
+BUILD_CALL_1(template sd::NDArray<float16>* sd::NDArray<float16>::applyTrueBroadcast, float16, (const sd::NDArray<float16>* a, float16* c) const, BROADCAST_OPS)
+BUILD_CALL_1(template sd::NDArray<float>* sd::NDArray<float>::applyTrueBroadcast, float, (const sd::NDArray<float>* a, float* c) const, BROADCAST_OPS)
+BUILD_CALL_1(template sd::NDArray<double>* sd::NDArray<double>::applyTrueBroadcast, double, (const sd::NDArray<double>* a, double* c) const, BROADCAST_OPS)
+
+BUILD_CALL_1(template sd::NDArray<float16> sd::NDArray<float16>::applyTrueBroadcast, float16, (const sd::NDArray<float16>& a, float16* c) const, BROADCAST_OPS)
+BUILD_CALL_1(template sd::NDArray<float> sd::NDArray<float>::applyTrueBroadcast, float, (const sd::NDArray<float>& a, float* c) const, BROADCAST_OPS)
+BUILD_CALL_1(template sd::NDArray<double> sd::NDArray<double>::applyTrueBroadcast, double, (const sd::NDArray<double>& a, double* c) const, BROADCAST_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyTransform, float16, (NDArray<float16>* target, float16* extraParams), TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyTransform, float, (NDArray<float>* target, float* extraParams), TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyTransform, double, (NDArray<double>* target, double* extraParams), TRANSFORM_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyTransform, float16, (float16* extraParams), TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyTransform, float, (float* extraParams), TRANSFORM_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyTransform, double, (double* extraParams), TRANSFORM_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float16>::applyRandom, float16, (sd::random::RandomBuffer *buffer, NDArray<float16>* y, NDArray<float16>* z, float16* extraParams), RANDOM_OPS)
+BUILD_CALL_1(template void sd::NDArray<float>::applyRandom, float, (sd::random::RandomBuffer *buffer, NDArray<float>* y, NDArray<float>* z, float* extraParams), RANDOM_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::applyRandom, double, (sd::random::RandomBuffer *buffer, NDArray<double>* y, NDArray<double>* z, double* extraParams), RANDOM_OPS)
+
+BUILD_CALL_1(template NDArray<float16> sd::NDArray<float16>::transform, float16, (float16* extraParams) const, TRANSFORM_OPS)
+BUILD_CALL_1(template NDArray<float>   sd::NDArray<float>::transform, float, (float* extraParams) const, TRANSFORM_OPS)
+BUILD_CALL_1(template NDArray<double>  sd::NDArray<double>::transform, double, (double* extraParams) const, TRANSFORM_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template reduceAlongDimension, float, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template reduceAlongDimension, float16, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template reduceAlongDimension, double, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+
+BUILD_CALL_1(template NDArray<float> sd::NDArray<float>::template reduceAlongDims, float, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<float16> sd::NDArray<float16>::template reduceAlongDims, float16, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<double> sd::NDArray<double>::template reduceAlongDims, double, (const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template reduceAlongDimension, float, (const std::initializer_list<int>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template reduceAlongDimension, float16, (const std::initializer_list<int>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template reduceAlongDimension, double, (const std::initializer_list<int>& dimensions, const bool keepDims, const bool supportOldShapes) const, REDUCE_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float>::template reduceAlongDimension, float, (NDArray<float>* target, const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes,  float * extras) const, REDUCE_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::template reduceAlongDimension, float16, (NDArray<float16>* target, const std::vector<LongType>& dimensions, const bool keepDims, const bool supportOldShapes, float16 * extras) const, REDUCE_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::template reduceAlongDimension, double, (NDArray<double>* target, const std::vector<int>& dimension, const bool keepDims, const bool supportOldShapes, double * extras) const, REDUCE_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template varianceAlongDimension, float, (const bool biasCorrected, const std::initializer_list<int>& dimensions) const, SUMMARY_STATS_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template varianceAlongDimension, float16, (const bool biasCorrected, const std::initializer_list<int>& dimensions) const, SUMMARY_STATS_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template varianceAlongDimension, double, (const bool biasCorrected, const std::initializer_list<int>& dimensions) const, SUMMARY_STATS_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float>::template varianceAlongDimension, float, (const NDArray<float> *target, const bool biasCorrected, const std::initializer_list<int>& dimensions), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::template varianceAlongDimension, float16, (const NDArray<float16> *target,const bool biasCorrected, const std::initializer_list<int>& dimensions), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::template varianceAlongDimension, double, (const NDArray<double> *target, const bool biasCorrected, const std::initializer_list<int>& dimensions), SUMMARY_STATS_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float>::template varianceAlongDimension, float, (const NDArray<float> *target, const bool biasCorrected, const std::vector<LongType>& dimensions), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::template varianceAlongDimension, float16, (const NDArray<float16> *target,const bool biasCorrected, const std::vector<LongType>& dimensions), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::template varianceAlongDimension, double, (const NDArray<double> *target, const bool biasCorrected, const std::vector<LongType>& dimensions), SUMMARY_STATS_OPS)
+
+BUILD_CALL_1(template float sd::NDArray<float>::template varianceNumber, float, (bool biasCorrected), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template float16 sd::NDArray<float16>::template varianceNumber, float16, (bool biasCorrected), SUMMARY_STATS_OPS)
+BUILD_CALL_1(template double sd::NDArray<double>::template varianceNumber, double, (bool biasCorrected), SUMMARY_STATS_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template applyReduce3, float, (const NDArray<float>* other, const float* extraParams) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template applyReduce3, float16, (const NDArray<float16>* other, const float16* extraParams) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template applyReduce3, double, (const NDArray<double>* other, const double* extraParams) const, REDUCE3_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template applyReduce3, float, (const NDArray<float>* other, const std::vector<int> &dims, const float* extraParams) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template applyReduce3, float16, (const NDArray<float16>* other, const std::vector<int> &dims, const float16* extraParams) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template applyReduce3, double, (const NDArray<double>* other, const std::vector<int> &dims, const double* extraParams) const, REDUCE3_OPS)
+
+BUILD_CALL_1(template void sd::NDArray<float>::template applyIndexReduce, float, (const NDArray<float>* target, const std::vector<int> & alpha, const float* beta) const, INDEX_REDUCE_OPS)
+BUILD_CALL_1(template void sd::NDArray<float16>::template applyIndexReduce, float16, (const NDArray<float16>* target, const std::vector<int> & alpha, const float16* beta) const, INDEX_REDUCE_OPS)
+BUILD_CALL_1(template void sd::NDArray<double>::template applyIndexReduce, double, (const NDArray<double>* target, const std::vector<int> & alpha, const double* beta) const, INDEX_REDUCE_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template applyIndexReduce, float, (const std::vector<int> & alpha, const float* beta) const, INDEX_REDUCE_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template applyIndexReduce, float16, (const std::vector<int> & alpha, const float16* beta) const, INDEX_REDUCE_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template applyIndexReduce, double, (const std::vector<int> & alpha, const double* beta) const, INDEX_REDUCE_OPS)
+
+BUILD_CALL_1(template NDArray<float> *sd::NDArray<float>::template applyAllReduce3, float, (const sd::NDArray<float>* alpha, const std::vector<int> & beta, float const* gamma) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<float16> *sd::NDArray<float16>::template applyAllReduce3, float16, (const sd::NDArray<float16>* alpha, const std::vector<int> & beta, float16 const* gamma) const, REDUCE3_OPS)
+BUILD_CALL_1(template NDArray<double> *sd::NDArray<double>::template applyAllReduce3, double, (const sd::NDArray<double>* alpha, const std::vector<int> & beta, double const* gamma) const, REDUCE3_OPS)
+
+template NDArray<float>   mmul(const NDArray<float>&   left, const NDArray<float>& right);
+template NDArray<float16> mmul(const NDArray<float16>& left, const NDArray<float16>& right);
+template NDArray<double>  mmul(const NDArray<double>&  left, const NDArray<double>& right);
+
+// template NDArray<float>   operator-(const float,   const NDArray<float>&);
+// template NDArray<float16> operator-(const float16, const NDArray<float16>&);
+// template NDArray<double>  operator-(const double,  const NDArray<double>&);
+
+// template NDArray<float>   operator+(const float,   const NDArray<float>&);
+// template NDArray<float16> operator+(const float16, const NDArray<float16>&);
+// template NDArray<double>  operator+(const double,  const NDArray<double>&);
+
+
+#endif

libnd4j/include/array/cpu/NDArrayLambda.hpp

@@ -0,0 +1,559 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+template <typename T>
+SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
+                                                  std::function<T(T, T, T)>& func, NDArray* target) {
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyTriplewiseLambda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != second->dataType() || dataType() != third->dataType() || dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyTriplewiseLambda<T> method: bother four arrays (this, second, third, target) should have the "
+        "same type !");
+
+  if (this->lengthOf() != second->lengthOf() || this->lengthOf() != third->lengthOf() || !this->isSameShape(second) ||
+      !this->isSameShape(third)) {
+    std::string errorMessage;
+    errorMessage += "applyTriplewiseLambda requires all operands to have the same shape\n";
+    errorMessage += "this shape: " + ShapeUtils::shapeAsString(this->shapeInfo()) + "\n";
+    errorMessage += "second shape: " + ShapeUtils::shapeAsString(second->shapeInfo()) + "\n";
+    errorMessage += "third shape: " + ShapeUtils::shapeAsString(third->shapeInfo()) + "\n";
+    errorMessage += "target shape: " + ShapeUtils::shapeAsString(target->shapeInfo()) + "\n";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  auto f = this->bufferAsT<T>();
+  auto s = second->bufferAsT<T>();
+  auto t = third->bufferAsT<T>();
+  auto z = target->bufferAsT<T>();
+  if (f == z) {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto tOffset = this->getOffset(e);
+          auto uOffset = second->getOffset(e);
+          auto vOffset = third->getOffset(e);
+
+          f[tOffset] = func(f[tOffset], s[uOffset], t[vOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  } else {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto tOffset = this->getOffset(e);
+          auto uOffset = second->getOffset(e);
+          auto vOffset = third->getOffset(e);
+          auto zOffset = target->getOffset(e);
+
+          z[zOffset] = func(f[tOffset], s[uOffset], t[vOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  }
+
+}
+#if defined(HAS_DOUBLE)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
+                                                           std::function<double(double, double, double)>& func,
+                                                           NDArray* target);
+#endif
+#if defined(HAS_FLOAT32)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
+                                                           std::function<float(float, float, float)>& func,
+                                                           NDArray* target);
+#endif
+#if defined(HAS_FLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<float16(float16, float16, float16)>& func, NDArray* target);
+#endif
+#if defined(HAS_BFLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<bfloat16(bfloat16, bfloat16, bfloat16)>& func,
+    NDArray* target);
+#endif
+#if defined(HAS_INT64)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<sd::LongType(sd::LongType, sd::LongType, sd::LongType)>& func,
+    NDArray* target);
+    template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<long(long, long, long)>& func,
+    NDArray* target);
+#endif
+#if defined(HAS_INT32)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
+                                                           std::function<int(int, int, int)>& func,
+                                                           NDArray* target);
+#endif
+#if defined(HAS_INT16)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<int16_t(int16_t, int16_t, int16_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT8)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<char(char, char, char)>& func, NDArray* target);
+    template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<signed char(signed char, signed char, signed char)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT8)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<uint8_t(uint8_t, uint8_t, uint8_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT16)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<uint16_t(uint16_t, uint16_t, uint16_t)>& func,
+    NDArray* target);
+#endif
+#if defined(HAS_UINT32)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<uint32_t(uint32_t, uint32_t, uint32_t)>& func,
+    NDArray* target);
+#endif
+#if defined(HAS_UINT64)
+template SD_LIB_HIDDEN void NDArray::applyTriplewiseLambda(
+    NDArray* second, NDArray* third, std::function<uint64_t(uint64_t, uint64_t, uint64_t)>& func,
+    NDArray* target);
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other, std::function<T(T, T)>& func,
+                                                NDArray* target) {
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyPairwiseLambda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != other->dataType() || dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyPairwiseLambda<T> method: all three arrays (this, other, target) must have the same type !");
+
+  // scalar is broadcastable
+  if (this->lengthOf() != other->lengthOf() && !this->isScalar() && !other->isScalar()) {
+    THROW_EXCEPTION("applyPairwiseLambda requires both operands to have the same shape");
+  }
+
+  auto f = this->bufferAsT<T>();
+  auto s = other->bufferAsT<T>();
+  auto z = target->bufferAsT<T>();
+  auto isTargetOrderEws = !isView() && !target->isView() && this->ordering() == target->ordering() && (this->ews() == 1 && target->ews() == 1);
+  if (other->isScalar()) {
+    auto otherVal = s[other->getOffset(0)];
+    if (isTargetOrderEws) {
+      auto loop = PRAGMA_THREADS_FOR {
+          for (auto e = start; e < stop; e++) z[e] = func(f[e], otherVal);
+      };
+
+      samediff::Threads::parallel_for(loop, 0, _length);
+    } else {
+      if (f == z) {
+        auto loop = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++) {
+              auto xOffset = this->getOffset(e);
+              f[xOffset] = func(f[xOffset], otherVal);
+            }
+        };
+
+        samediff::Threads::parallel_for(loop, 0, _length);
+      } else {
+        auto loop = PRAGMA_THREADS_FOR {
+            for (auto e = start; e < stop; e++) {
+              auto xOffset = this->getOffset(e);
+              auto zOffset = target->getOffset(e);
+
+              z[zOffset] = func(f[xOffset], otherVal);
+            }
+        };
+
+        samediff::Threads::parallel_for(loop, 0, _length);
+      }
+    }
+  }
+
+  if (f == z && !this->isView() && !other->isView() && this->ordering() == other->ordering()) {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto xOffset = this->getOffset(e);
+          auto yOffset = other->getOffset(e);
+
+          f[xOffset] = func(f[xOffset], s[yOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  } else {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto xOffset = this->getOffset(e);
+          auto yOffset = other->getOffset(e);
+          auto zOffset = target->getOffset(e);
+
+          z[zOffset] = func(f[xOffset], s[yOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  }
+
+}
+
+#if defined(HAS_DOUBLE)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<double(double, double)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_FLOAT32)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<float(float, float)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_FLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<float16(float16, float16)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_BFLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<bfloat16(bfloat16, bfloat16)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_INT64)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(
+    NDArray* other, std::function<sd::LongType(sd::LongType, sd::LongType)>& func, NDArray* target);
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(
+    NDArray* other, std::function<long(long, long)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT32)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other, std::function<int(int, int)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_INT16)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<int16_t(int16_t, int16_t)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_UINT8)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<uint8_t(uint8_t, uint8_t)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_UINT16)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<uint16_t(uint16_t, uint16_t)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_UINT32)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<uint32_t(uint32_t, uint32_t)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_UNSIGNEDLONG)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<uint64_t(uint64_t, uint64_t)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_INT8)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<int8_t(int8_t, int8_t)>& func,
+                                                         NDArray* target);
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<SignedChar(SignedChar, SignedChar)>& func,
+                                                         NDArray* target);
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<char(char, char)>& func,
+                                                         NDArray* target);
+#endif
+#if defined(HAS_BOOL)
+template SD_LIB_HIDDEN void NDArray::applyPairwiseLambda(NDArray* other,
+                                                         std::function<bool(bool, bool)>& func, NDArray* target);
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_LIB_HIDDEN void NDArray::applyLambda(std::function<T(T)>& func, NDArray* target) {
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyLambda<T> method: wrong template parameter T, its type should be the same as type of this "
+        "array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION("NDArray::applyLambda<T> method: types of this and target array should match !");
+
+  auto f = this->bufferAsT<T>();
+  auto z = target->bufferAsT<T>();
+
+  if (f == z) {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e+= increment) {
+          auto xOffset = this->getOffset(e);
+          f[xOffset] = func(f[xOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length,1);
+  } else {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e+= increment) {
+          auto xOffset = this->getOffset(e);
+          auto zOffset = target->getOffset(e);
+          z[zOffset] = func(f[xOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length,1);
+  }
+
+}
+#if defined(HAS_DOUBLE)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<double(double)>& func, NDArray* target);
+#endif
+#if defined(HAS_FLOAT32)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<float(float)>& func, NDArray* target);
+#endif
+#if defined(HAS_FLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<float16(float16)>& func, NDArray* target);
+#endif
+#if defined(HAS_BFLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<bfloat16(bfloat16)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT64)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<sd::LongType(sd::LongType)>& func,
+                                                 NDArray* target);
+#endif
+#if defined(HAS_INT16)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<int16_t(int16_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT32)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<int32_t(int32_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT8)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<uint8_t(uint8_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT16)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<uint16_t(uint16_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT32)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<uint32_t(uint32_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UNSIGNEDLONG)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<uint64_t(uint64_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT8)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<int8_t(int8_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_BOOL)
+template SD_LIB_HIDDEN void NDArray::applyLambda(std::function<bool(bool)>& func, NDArray* target);
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<T(sd::LongType, T)>& func, NDArray* target) {
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedLambda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION("NDArray::applyIndexedLambda<T> method: types of this and target array should match !");
+
+  auto f = this->bufferAsT<T>();
+  auto z = target->bufferAsT<T>();
+
+  if (this->ordering() == target->ordering() && (this->ews() == 1 && target->ews() == 1)) {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) z[e] = func(e, f[e]);
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  } else {
+    if (f == z) {
+      auto loop = PRAGMA_THREADS_FOR {
+          for (auto e = start; e < stop; e++) {
+            auto xOffset = this->getOffset(e);
+
+            f[xOffset] = func(e, f[xOffset]);
+          }
+      };
+
+      samediff::Threads::parallel_for(loop, 0, _length);
+    } else {
+      auto loop = PRAGMA_THREADS_FOR {
+          for (auto e = start; e < stop; e++) {
+            auto xOffset = this->getOffset(e);
+            auto zOffset = target->getOffset(e);
+
+            z[zOffset] = func(e, f[xOffset]);
+          }
+      };
+
+      samediff::Threads::parallel_for(loop, 0, _length);
+    }
+  }
+}
+#if defined(HAS_DOUBLE)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<double(sd::LongType, double)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_FLOAT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<float(sd::LongType, float)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_FLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<float16(sd::LongType, float16)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_BFLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<bfloat16(sd::LongType, bfloat16)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_INT64)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(
+    std::function<sd::LongType(sd::LongType, sd::LongType)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<int(sd::LongType, int)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_INT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<int16_t(sd::LongType, int16_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_UINT8)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<uint8_t(sd::LongType, uint8_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_UINT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<uint16_t(sd::LongType, uint16_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_UINT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<uint32_t(sd::LongType, uint32_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_UNSIGNEDLONG)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<uint64_t(sd::LongType, uint64_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_INT8)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<int8_t(sd::LongType, int8_t)>& func,
+                                                        NDArray* target);
+#endif
+#if defined(HAS_BOOL)
+template SD_LIB_HIDDEN void NDArray::applyIndexedLambda(std::function<bool(sd::LongType, bool)>& func,
+                                                        NDArray* target);
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(NDArray* other, std::function<T(sd::LongType, T, T)>& func,
+                                                       NDArray* target) {
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedPairwiseLambda<T> method: wrong template parameter T, its type should be the same as "
+        "type of this array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedPairwiseLambda<T> method: types of this and target array should match !");
+  if (this->lengthOf() != other->lengthOf()) {
+    sd_printf("applyIndexedPairwiseLambda requires both operands to have the same shape\n", "");
+    THROW_EXCEPTION("Shapes mismatch");
+  }
+
+  auto f = this->bufferAsT<T>();
+  auto s = other->bufferAsT<T>();
+  auto z = target->bufferAsT<T>();
+  if (f == z) {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto xOffset = this->getOffset(e);
+          auto yOffset = other->getOffset(e);
+
+          f[xOffset] = func((sd::LongType)e, f[xOffset], s[yOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  } else {
+    auto loop = PRAGMA_THREADS_FOR {
+        for (auto e = start; e < stop; e++) {
+          auto xOffset = this->getOffset(e);
+          auto yOffset = other->getOffset(e);
+          auto zOffset = target->getOffset(e);
+
+          z[zOffset] = func((sd::LongType)e, f[xOffset], s[yOffset]);
+        }
+    };
+
+    samediff::Threads::parallel_for(loop, 0, _length);
+  }
+
+}
+#if defined(HAS_DOUBLE)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<double(sd::LongType, double, double)>& func, NDArray* target);
+#endif
+#if defined(HAS_FLOAT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<float(sd::LongType, float, float)>& func, NDArray* target);
+#endif
+#if defined(HAS_FLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<float16(sd::LongType, float16, float16)>& func, NDArray* target);
+#endif
+#if defined(HAS_BFLOAT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<bfloat16(sd::LongType, bfloat16, bfloat16)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT64)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<sd::LongType(sd::LongType, sd::LongType, sd::LongType)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(NDArray* other,
+                                                                std::function<int(sd::LongType, int, int)>& func,
+                                                                NDArray* target);
+#endif
+#if defined(HAS_INT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<int16_t(sd::LongType, int16_t, int16_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT8)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<uint8_t(sd::LongType, uint8_t, uint8_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT16)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<uint16_t(sd::LongType, uint16_t, uint16_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UINT32)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<uint32_t(sd::LongType, uint32_t, uint32_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_UNSIGNEDLONG)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<uint64_t(sd::LongType, uint64_t, uint64_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_INT8)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<int8_t(sd::LongType, int8_t, int8_t)>& func, NDArray* target);
+#endif
+#if defined(HAS_BOOL)
+template SD_LIB_HIDDEN void NDArray::applyIndexedPairwiseLambda(
+    NDArray* other, std::function<bool(sd::LongType, bool, bool)>& func, NDArray* target);
+#endif

libnd4j/include/array/cuda/CudaPointerDeallocator.cu

@@ -0,0 +1,47 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/CudaPointerDeallocator.h>
+
+namespace sd {
+
+void CudaPointerDeallocator::release(void *ptr) {
+  if (ptr == nullptr) return;
+
+  // Check if this is a valid device pointer before freeing
+  cudaPointerAttributes attributes;
+  cudaError_t result = cudaPointerGetAttributes(&attributes, ptr);
+
+  if (result == cudaSuccess) {
+    // Only free if it's a regular device pointer
+    // cudaMemoryTypeDevice is for regular allocations we can free
+    if (attributes.type == cudaMemoryTypeDevice) {
+      cudaFree(ptr);
+    }
+    // Don't free other types (like constant memory)
+  } else {
+    // Clear the error and don't try to free this pointer
+    cudaGetLastError(); // Clear the error state
+  }
+}
+}  // namespace sd

libnd4j/include/array/cuda/DataBuffer.cu

@@ -0,0 +1,608 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <array/DataTypeUtils.h>
+#include <exceptions/allocation_exception.h>
+#include <exceptions/cuda_exception.h>
+#include <execution/AffinityManager.h>
+#include <memory/MemoryCounter.h>
+#include <system/op_boilerplate.h>
+#include <system/type_boilerplate.h>
+
+#include "../DataBuffer.h"
+#include "helpers/DebugHelper.h"
+
+#if defined(SD_GCC_FUNCTRACE)
+#include <array/DataBufferLifecycleTracker.h>
+#endif
+
+namespace sd {
+void DataBuffer::expand(const uint64_t size) {
+  if (size > _lenInBytes) {
+    // allocate new buffer
+    int8_t* newBuffer = nullptr;
+    int8_t* newSpecialBuffer = nullptr;
+    ALLOCATE_SPECIAL(newSpecialBuffer, _workspace, size, int8_t);
+
+    // copy data from existing buffer
+    if (_primaryBuffer != nullptr) {
+      // there's non-zero chance that primary buffer doesn't exist yet
+      ALLOCATE(newBuffer, _workspace, size, int8_t);
+      std::memcpy(newBuffer, _primaryBuffer, _lenInBytes);
+
+      if (_isOwnerPrimary) {
+        auto ipb = reinterpret_cast<int8_t*>(_primaryBuffer);
+        RELEASE(ipb, _workspace);
+      }
+
+      _primaryBuffer = newBuffer;
+      _isOwnerPrimary = true;
+    }
+
+    cudaMemcpy(newSpecialBuffer, _specialBuffer, _lenInBytes, cudaMemcpyDeviceToDevice);
+
+    if (_isOwnerSpecial) {
+      auto isb = reinterpret_cast<int8_t*>(_specialBuffer);
+      RELEASE_SPECIAL(isb, _workspace);
+    }
+
+    _specialBuffer = newSpecialBuffer;
+    _lenInBytes = size;
+    _isOwnerSpecial = true;
+  }
+}
+
+DataBuffer DataBuffer::dup() {
+  DataBuffer result;
+  result._dataType = _dataType;
+  result._lenInBytes = _lenInBytes;
+  result._primaryBuffer = _primaryBuffer;
+  result._specialBuffer = _specialBuffer;
+  result._isOwnerPrimary = _isOwnerPrimary;
+  result._isOwnerSpecial = _isOwnerSpecial;
+  result.allocateBuffers(true);
+  result.copyCounters(*this);
+  result.copyBufferFrom(*this);
+  return result;
+}
+
+template <typename T>
+void* DataBuffer::primaryAtOffset(const LongType offset) {
+  if(_primaryBuffer == nullptr)
+    return nullptr;
+  T *type = reinterpret_cast<T*>(_primaryBuffer);
+  return reinterpret_cast<void *>(type + offset);
+}
+template <typename T>
+void* DataBuffer::specialAtOffset(const LongType offset) {
+  if(_specialBuffer == nullptr)
+    return nullptr;
+  T *type = reinterpret_cast<T*>(_specialBuffer);
+  return reinterpret_cast<void *>(type + offset);
+}
+
+#define PRIMARYOFFSET(T) template SD_LIB_EXPORT void* DataBuffer::primaryAtOffset<GET_SECOND(T)>(sd::LongType offset);
+ITERATE_LIST((SD_COMMON_TYPES),PRIMARYOFFSET)
+
+#define SPECIALOFFSET(T) template SD_LIB_EXPORT void* DataBuffer::specialAtOffset<GET_SECOND(T)>(sd::LongType offset);
+ITERATE_LIST((SD_COMMON_TYPES),SPECIALOFFSET)
+
+
+template <typename T>
+void _printHostBuffer(DataBuffer* buffer, long offset) {
+  sd::LongType len = buffer->getNumElements();
+  auto buff = buffer->template primaryAsT<T>();
+
+
+  sd::LongType limit = len;
+  if (limit == -1 || limit >= buffer->getNumElements()) {
+    limit = buffer->getNumElements();
+  }
+
+  const char* msg = nullptr;
+  if (msg != nullptr) {
+    printf("%s: ", msg);
+  } else {
+    printf("[");
+  }
+
+  sd::DataType dataType = buffer->getDataType();
+  auto baseOffset = offset;
+  if (dataType == sd::DataType::DOUBLE || dataType == sd::DataType::FLOAT32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (e > offset) printf(", ");
+      if (dataType == sd::DataType::DOUBLE) {
+        printf("%.15f", buff[e]);
+      } else {
+        printf("%.15f", static_cast<float>(buff[e]));
+      }
+    }
+  } else if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64 ||
+             dataType == sd::DataType::INT32 || dataType == sd::DataType::UINT32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64) {
+        printf("%lld", static_cast<long long>(buff[e]));
+      } else {
+        printf("%d", static_cast<int>(buff[e]));
+      }
+
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  } else if (dataType == sd::DataType::BOOL) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      if (static_cast<bool>(buff[e])) {
+        printf("true");
+      } else {
+        printf("false");
+      }
+
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  } else if (dataType == sd::DataType::UTF8 || dataType == sd::DataType::UTF16 ||
+             dataType == sd::DataType::UTF32) {
+    for (sd::LongType e = baseOffset; e < limit; e++) {
+      printf("\"%s\"", reinterpret_cast<const char*>(&buff[e]));
+      if (e < limit - 1) {
+        printf(", ");
+      }
+    }
+  }
+
+  printf("]\n");
+  fflush(stdout);
+}
+
+void DataBuffer::printHostDevice(long offset) {
+  THROW_EXCEPTION("");
+}
+
+void DataBuffer::printSpecialAllocationTraces() {
+  //no op on purpose
+}
+
+void DataBuffer::showBufferLimited() {
+
+}
+
+template <typename T>
+SD_KERNEL void printDeviceBufferKernel(void* buffer, sd::LongType offset, sd::LongType length) {
+  T* typedBuffer = reinterpret_cast<T*>(buffer);
+
+  if (threadIdx.x == 0 && blockIdx.x == 0) {
+    printf("[ ");
+    for (sd::LongType i = offset; i < offset + length; i++) {
+      // Cast to double for consistent formatting
+      printf("%g ", (double)typedBuffer[i]);
+    }
+    printf("]");
+  }
+}
+
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT  SD_KERNEL void printDeviceBufferKernel,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
+
+
+// Wrapper function to launch the kernel
+template <typename T>
+void launchPrintDeviceBufferKernel(void* buffer, sd::LongType offset, sd::LongType length) {
+  printDeviceBufferKernel<T><<<1, 1, 32*1024, *LaunchContext::defaultContext()->getCudaStream()>>>(
+      buffer, offset, length);
+  cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
+  sd::DebugHelper::checkErrorCode(LaunchContext::defaultContext()->getCudaStream(),
+                                  "printBufferDebug kernel failed");
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void launchPrintDeviceBufferKernel,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
+
+
+template <typename T>
+void DataBuffer::printHostBufferContent(void* buffer, sd::LongType offset, sd::LongType length) {
+  T* typedBuffer = reinterpret_cast<T*>(buffer);
+
+  sd_printf("[ ", 0);
+  for (sd::LongType i = offset; i < offset + length; i++) {
+    // For numeric types, cast to double for consistent formatting
+    if (std::is_arithmetic<T>::value) {
+      sd_printf("%g ", (double)typedBuffer[i]);
+    } else {
+      // For non-numeric types, print as hex
+      sd_printf("0x%x ", *reinterpret_cast<int*>(&typedBuffer[i]));
+    }
+  }
+  sd_printf("]", 0);
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void DataBuffer::printHostBufferContent,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
+
+
+// DataBuffer implementation for .cu file
+void DataBuffer::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
+  if (msg) sd_printf("%s:\n", msg);
+
+  // Print metadata
+  sd_printf("DataBuffer: DataType=%s, Length=%lld elements, DeviceId=%d\n",
+            DataTypeUtils::asString(_dataType).c_str(), (long long)getNumElements(), deviceId());
+
+  // Print host buffer content
+  if (_primaryBuffer != nullptr) {
+    sd_printf("Host buffer (@%p): ", _primaryBuffer);
+
+    sd::LongType len = getNumElements();
+    sd::LongType printLen = limit < 0 ? len : std::min(len - offset, limit);
+
+    // Print based on datatype
+    BUILD_SINGLE_SELECTOR(_dataType, printHostBufferContent,
+                          (_primaryBuffer, offset, printLen), SD_COMMON_TYPES);
+
+    if (offset + printLen < len) sd_printf("... ", 0);
+    sd_printf("\n", 0);
+  } else {
+    sd_printf("Host buffer: nullptr\n", 0);
+  }
+
+  // Print device buffer using kernel
+  if (_specialBuffer != nullptr) {
+    sd_printf("Device buffer (@%p): ", _specialBuffer);
+
+    sd::LongType len = getNumElements();
+    sd::LongType printLen = limit < 0 ? len : std::min(len - offset, limit);
+
+    // Launch kernel through wrapper function
+    BUILD_SINGLE_SELECTOR(_dataType, launchPrintDeviceBufferKernel,
+                          (_specialBuffer, offset, printLen), SD_COMMON_TYPES);
+
+    sd_printf("\n", 0);
+  } else {
+    sd_printf("Device buffer: nullptr\n", 0);
+  }
+
+  // Print sync state counters
+  sd_printf("Sync state: _counter=%lld, _writePrimary=%lld, _writeSpecial=%lld, _readPrimary=%lld, _readSpecial=%lld\n",
+            (long long)_counter.load(), (long long)_writePrimary.load(), (long long)_writeSpecial.load(),
+            (long long)_readPrimary.load(), (long long)_readSpecial.load());
+  sd_printf("isPrimaryActual=%d, isSpecialActual=%d\n", isPrimaryActual(), isSpecialActual());
+}
+
+
+
+void DataBuffer::showCounters(const char* msg1, const char* msg2) {
+  sd_debug("%s %s || primary %p special %p :: wP: %d wS: %d rP: %d rS: %d\n", msg1, msg2, _primaryBuffer,
+           _specialBuffer, (int)_writePrimary.load(), (int)_writeSpecial.load(), (int)_readPrimary.load(),
+           (int)_readSpecial.load());
+}
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::allocateSpecial() {
+  if (_specialBuffer != nullptr) {
+    return;
+  }
+
+  if (_lenInBytes == 0) {
+    std::string errorMessage;
+    errorMessage += "DataBuffer::allocateSpecial: ";
+    errorMessage += "Special buffer is already allocated";
+    errorMessage += " or length is 0";
+    errorMessage += "Length is: ";
+    errorMessage += std::to_string(getLenInBytes());
+    errorMessage += "Special buffer is nullptr : ";
+    errorMessage += std::to_string(_specialBuffer == nullptr);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+#if defined(SD_GCC_FUNCTRACE)
+  if(Environment::getInstance().isFuncTracePrintAllocate()) {
+    allocationStackTraceSpecial = new StackTrace();
+    allocationStackTraceSpecial->load_here();
+  }
+
+#endif
+
+  if (_specialBuffer == nullptr) {
+    auto deviceId = AffinityManager::currentDeviceId();
+
+    if (_workspace == nullptr) {
+      if (!memory::MemoryCounter::getInstance().validate(getLenInBytes())) {
+        std::string errorMessage;
+        errorMessage += "DataBuffer::allocateSpecial: ";
+        errorMessage += "Requested amount exceeds device limits";
+        errorMessage += "DeviceId: ";
+        errorMessage += std::to_string(deviceId);
+        errorMessage += "Device limit: ";
+        errorMessage += std::to_string(memory::MemoryCounter::getInstance().deviceLimit(deviceId));
+        errorMessage += "Requested amount: ";
+        errorMessage += std::to_string(getLenInBytes());
+        errorMessage += "Special buffer is nullptr : ";
+        errorMessage += std::to_string(_specialBuffer == nullptr);
+        THROW_EXCEPTION(errorMessage.c_str());
+      }
+    }
+
+    ALLOCATE_SPECIAL(_specialBuffer, _workspace, getLenInBytes(), int8_t);
+    _isOwnerSpecial = true;
+
+#if defined(SD_GCC_FUNCTRACE)
+    // Record SPECIAL (device) buffer allocation
+    array::DataBufferLifecycleTracker::getInstance().recordAllocation(
+        _specialBuffer, getLenInBytes(), getDataType(),
+        array::BufferType::SPECIAL, this, _workspace != nullptr);
+#endif
+
+    if (_workspace == nullptr) {
+      memory::MemoryCounter::getInstance().countIn(deviceId, getLenInBytes());
+      memory::MemoryCounter::getInstance().countIn(memory::MemoryType::DEVICE, getLenInBytes());
+
+    }
+  } else if(getLenInBytes() == 0) {
+    std::string errorMessage;
+    errorMessage += "DataBuffer::allocateSpecial: ";
+    errorMessage += "Special buffer is already allocated";
+    errorMessage += " or length is 0";
+    errorMessage += "Length is: ";
+    errorMessage += std::to_string(getLenInBytes());
+    errorMessage += "Special buffer is nullptr : ";
+    errorMessage += std::to_string(_specialBuffer == nullptr);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::syncToPrimary(const LaunchContext* context, const bool forceSync) {
+  if (isPrimaryActual() && !forceSync) {
+    return;
+  }
+
+  allocatePrimary();
+
+  auto res = cudaStreamSynchronize(*context->getCudaStream());
+  if (res != 0)  {
+    std::string errorMessage;
+    errorMessage += "DataBuffer::syncToPrimary: cudaStreamSynchronize failed: ";
+    errorMessage += std::to_string(getLenInBytes());
+    errorMessage += cudaGetErrorString(res);
+    errorMessage += "Special buffer is nullptr : ";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  res = cudaMemcpy(_primaryBuffer, _specialBuffer, getLenInBytes(), cudaMemcpyDeviceToHost);
+  if (res != 0) {
+        std::string errorMessage;
+        errorMessage += "DataBuffer::syncToPrimary: cudaMemcpy failed: ";
+        errorMessage += std::to_string(getLenInBytes());
+        errorMessage += cudaGetErrorString(res);
+        errorMessage += "Special buffer is nullptr : ";
+        errorMessage += std::to_string(_specialBuffer == nullptr);
+        THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  readPrimary();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::syncToSpecial(const bool forceSync) {
+  // in this case there's nothing to do here
+  if (_primaryBuffer == nullptr) return;
+
+  if (isSpecialActual() && !forceSync) {
+    return;
+  }
+
+  allocateSpecial();
+
+  auto res = cudaMemcpy(_specialBuffer, _primaryBuffer, getLenInBytes(), cudaMemcpyHostToDevice);
+  if (res != 0) {
+    std::string errorMessage;
+    errorMessage += "Failed to copy dataBuffer::syncToSpecial: ";
+    errorMessage += std::to_string(getLenInBytes());
+    errorMessage += cudaGetErrorString(res);
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+
+  readSpecial();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::deleteSpecial() {
+  if (_isOwnerSpecial && _specialBuffer != nullptr && getLenInBytes() != 0) {
+    auto p = reinterpret_cast<int8_t*>(_specialBuffer);
+#if defined(SD_GCC_FUNCTRACE)
+    // Record SPECIAL (device) buffer deallocation before releasing
+    array::DataBufferLifecycleTracker::getInstance().recordDeallocation(
+        _specialBuffer, array::BufferType::SPECIAL);
+#endif
+    RELEASE_SPECIAL(p, _workspace);
+    _specialBuffer = nullptr;
+    _isOwnerSpecial = false;
+
+    // count out towards DataBuffer device, only if we're not in workspace
+    if (_workspace == nullptr) {
+      sd::memory::MemoryCounter::getInstance().countOut(_deviceId, getLenInBytes());
+      sd::memory::MemoryCounter::getInstance().countOut(sd::memory::MemoryType::DEVICE, getLenInBytes());
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setCountersToZero() {
+  _counter.store(0L);
+  _writePrimary.store(0L);
+  _writeSpecial.store(0L);
+  _readPrimary.store(0L);
+  _readSpecial.store(0L);
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyCounters(const DataBuffer& other) {
+  _counter.store(other._counter);
+  _writePrimary.store(other._readSpecial);
+  _writeSpecial.store(other._readPrimary);
+  _readPrimary.store(other._writeSpecial);
+  _readSpecial.store(other._writePrimary);
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyBufferFrom(const DataBuffer& other, size_t sizeToCopyinBytes, const sd::LongType offsetThis,
+                                const sd::LongType offsetOther) {  // copies only to special buffer
+
+  if (other._primaryBuffer == nullptr && other._specialBuffer == nullptr) {
+    return;
+  }
+
+  if (sizeToCopyinBytes == 0) {
+    sizeToCopyinBytes = other.getLenInBytes();
+  }
+  if (sizeToCopyinBytes == 0) {
+    return;
+  }
+
+  if (other.isPrimaryActual()) {
+    auto res = cudaMemcpy(
+        static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
+        static_cast<const int8_t*>(other._primaryBuffer) + offsetOther * DataTypeUtils::sizeOfElement(other._dataType),
+        sizeToCopyinBytes, cudaMemcpyHostToDevice);
+    if (res != 0)
+      throw cuda_exception::build("DataBuffer::copyBufferFrom: cudaMemcpy_cudaMemcpyHostToDevice failed!", res);
+    other.readPrimary();
+  } else {
+    auto res = cudaMemcpy(
+        static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
+        static_cast<const int8_t*>(other._specialBuffer) + offsetOther * DataTypeUtils::sizeOfElement(other._dataType),
+        sizeToCopyinBytes, cudaMemcpyDeviceToDevice);
+    if (res != 0)
+      throw cuda_exception::build("DataBuffer::copyBufferFrom: cudaMemcpy_cudaMemcpyDeviceToDevice failed!", res);
+    other.readSpecial();
+  }
+
+  writeSpecial();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::copyBufferFromHost(const void* hostBuffer, size_t sizeToCopyinBytes, const sd::LongType offsetThis,
+                                    const sd::LongType offsetHostBuffer) {  // copies only to special buffer
+
+  if (hostBuffer == nullptr) return;
+
+  if (sizeToCopyinBytes == 0) sizeToCopyinBytes = getLenInBytes();
+  if (sizeToCopyinBytes == 0) return;
+
+  auto res =
+      cudaMemcpy(static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
+                 static_cast<const int8_t*>(hostBuffer) + offsetHostBuffer * DataTypeUtils::sizeOfElement(_dataType),
+                 sizeToCopyinBytes, cudaMemcpyHostToDevice);
+  if (res != 0)
+    throw cuda_exception::build("DataBuffer::copyBufferFromHost: cudaMemcpy_cudaMemcpyHostToDevice failed!", res);
+
+  writeSpecial();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setSpecial(void* special, const bool isOwnerSpecial) {
+  deleteSpecial();
+  _specialBuffer = special;
+  _isOwnerSpecial = isOwnerSpecial;
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::allocateBuffers(const bool allocBoth) {  // always allocate special buffer only (cuda case)
+  allocateSpecial();
+
+  if (allocBoth) allocatePrimary();
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setToZeroBuffers(const bool both) {
+  if(getLenInBytes() < 1 || special() == nullptr)
+    return;
+  cudaMemsetAsync(special(), 0, getLenInBytes(), *LaunchContext::defaultContext()->getCudaStream());
+  auto res = cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
+  if (res != 0) throw cuda_exception::build("DataBuffer::setToZeroBuffers: streamSync failed!", res);
+
+  writeSpecial();
+
+  if (both) {
+    memset(primary(), 0, getLenInBytes());
+    readPrimary();
+  }
+}
+
+
+
+/////////////////////////
+
+
+template <typename T>
+void memcpyWithT(DataBuffer* dst, DataBuffer* src, sd::LongType startingOffset, sd::LongType dstOffset) {
+  if (src->getLenInBytes() > dst->getLenInBytes())
+    THROW_EXCEPTION("DataBuffer::memcpy: Source data buffer is larger than destination");
+
+  int res = 0;
+  if (src->isSpecialActual()) {
+    res = cudaMemcpyAsync(dst->specialAtOffset<T>(dstOffset), src->specialAtOffset<T>(startingOffset), src->getLenInBytes(), cudaMemcpyDeviceToDevice,
+                          *LaunchContext::defaultContext()->getCudaStream());
+  } else if (src->isPrimaryActual()) {
+    res = cudaMemcpyAsync(dst->specialAtOffset<T>(dstOffset), src->specialAtOffset<T>(startingOffset), src->getLenInBytes(), cudaMemcpyHostToDevice,
+                          *LaunchContext::defaultContext()->getCudaStream());
+  }
+
+  if (res != 0) throw cuda_exception::build("DataBuffer::memcpy: cudaMemcpyAsync failed!", res);
+
+  res = cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
+  if (res != 0) throw cuda_exception::build("DataBuffer::memcpy: streamSync failed!", res);
+
+  dst->writeSpecial();
+}
+
+void DataBuffer::memcpy(DataBuffer* dst, DataBuffer* src,
+                        sd::LongType startingOffset, sd::LongType dstOffset) {
+  BUILD_SINGLE_TEMPLATE(memcpyWithT,(dst, src, startingOffset, dstOffset),
+                        SD_COMMON_TYPES);
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::migrate() {
+  memory::Workspace* newWorkspace = nullptr;
+  void* newBuffer;
+  ALLOCATE_SPECIAL(newBuffer, newWorkspace, getLenInBytes(), int8_t);
+  auto res = cudaMemcpy(newBuffer, _specialBuffer, getLenInBytes(), cudaMemcpyDeviceToDevice);
+  if (res != 0) throw cuda_exception::build("DataBuffer::migrate: cudaMemcpyAsync failed!", res);
+
+  if (_isOwnerSpecial) {
+    // now we're releasing original buffer
+    RELEASE_SPECIAL(_specialBuffer, _workspace);
+  }
+
+  _isOwnerSpecial = true;
+  _specialBuffer = newBuffer;
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::writePrimary() const { _writePrimary = ++_counter; }
+void DataBuffer::writeSpecial() const { _writeSpecial = ++_counter; }
+void DataBuffer::readPrimary() const { _readPrimary = ++_counter; }
+void DataBuffer::readSpecial() const { _readSpecial = ++_counter; }
+bool DataBuffer::isPrimaryActual() const {
+  return (_writePrimary.load() > _writeSpecial.load() || _readPrimary.load() > _writeSpecial.load());
+}
+bool DataBuffer::isSpecialActual() const {
+  return (_writeSpecial.load() > _writePrimary.load() || _readSpecial.load() > _writePrimary.load());
+}
+
+}  // namespace sd

libnd4j/include/array/cuda/NDArray.cu

@@ -0,0 +1,720 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+#ifndef NDARRAY_CPP
+#define NDARRAY_CPP
+#include <array/NDArray.h>
+#include <array/NDArrayFactory.h>
+#include <exceptions/cuda_exception.h>
+#include <exceptions/datatype_exception.h>
+#include <helpers/ArrayUtils.h>
+#include <helpers/ConstantShapeHelper.h>
+#include <helpers/MmulHelper.h>
+#include <helpers/PointersManager.h>
+#include <helpers/ShapeUtils.h>
+#include <helpers/logger.h>
+#include <helpers/threshold.h>
+#include <indexing/IndicesList.h>
+#include <indexing/NDIndex.h>
+#include <legacy/NativeOpExecutioner.h>
+#include <loops/broadcasting.h>
+#include <loops/pairwise_transform.h>
+#include <loops/random.h>
+#include <loops/special_kernels.h>
+#include <loops/transform_same.h>
+#include <memory/MemoryRegistrator.h>
+#include <memory/Workspace.h>
+#include <ops/ops.h>
+#include <ops/specials_cuda.h>
+
+#include <array/NDArray.hXX>
+#include <memory>
+#include <sstream>
+#include <stdexcept>
+#include <system/selective_rendering.h>
+#include "execution/cuda/LaunchDims.h"
+
+namespace sd {
+
+
+
+void* NDArray::platformBuffer() { return specialBuffer(); }
+
+
+
+void NDArray::syncToDevice()  {
+  auto currentDeviceId = AffinityManager::currentDeviceId();
+  if (currentDeviceId != _deviceId) {
+    // first of all we update shapeInfo
+    const_cast<NDArray*>(this)->setShapeInfo(this->shapeInfo());
+
+    // now we actually migrate data buffer
+    _buffer->migrate();
+  }
+
+  _buffer->syncToSpecial();
+}
+
+void NDArray::syncToHost()  { if(!isEmpty()) _buffer->syncToPrimary(getContext()); }
+void NDArray::tickWriteHost()  { if(!isEmpty()) _buffer->writePrimary(); }
+void NDArray::tickWriteDevice()  { if(!isEmpty()) _buffer->writeSpecial(); }
+void NDArray::tickReadHost()  { if(!isEmpty()) _buffer->readPrimary(); }
+void NDArray::tickReadDevice()  { if(!isEmpty()) _buffer->readSpecial(); }
+void NDArray::tickBothActual()  {
+  _buffer->writePrimary();
+  _buffer->readSpecial();
+}
+bool NDArray::isActualOnHostSide()  { return _buffer->isPrimaryActual(); }
+bool NDArray::isActualOnDeviceSide()  { return _buffer->isSpecialActual(); }
+void NDArray::makeBothBuffersActual()  {
+  if (!isActualOnHostSide()) syncToHost();
+  if (!isActualOnDeviceSide()) syncToDevice();
+}
+
+///////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL static void fillAsTriangularCuda(const void* vx, const LongType* xShapeInfo, void* vz,
+                                           const LongType* zShapeInfo, const T val, const int lower,
+                                           const int upper, char direction, bool includeEdges) {
+  const auto x = reinterpret_cast<const T*>(vx);
+  auto z = reinterpret_cast<T*>(vz);
+
+  __shared__ LongType zRank, xRank, areSameOffsets, *sharedMem;  // xRank == zRank always, except when xRank = 1, in this case zRank = 2
+  __shared__ LongType zLen, totalThreads;  // xLen == zLen, except when xRank = 1, in this case zLen = 2*xLen
+  __shared__ LongType *zShape;
+  __shared__ LongType *zStride;
+  __shared__ LongType *xShape;
+  __shared__ LongType *xStride;
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char shmem[];
+    sharedMem = reinterpret_cast<LongType*>(shmem);
+    areSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+    xRank = shape::rank(xShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+    zLen = shape::length(zShapeInfo);
+    totalThreads = gridDim.x * blockDim.x;
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+    xShape = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  auto coords = sharedMem + threadIdx.x * zRank;
+
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  bool dirU = direction == 'u';
+  bool dirL = direction == 'l';
+  for (LongType i = tid; i < zLen; i += totalThreads) {
+    INDEX2COORDS(i, zRank, zShape, coords);
+
+    LongType zOffset;
+    COORDS2INDEX(zRank, zStride, coords, zOffset);
+
+    auto row = coords[zRank - 2];
+    auto col = coords[zRank - 1];
+    auto lCompare = includeEdges ? row + lower <= col : row + lower < col;
+    auto uCompare = includeEdges ? row + upper >= col : row + upper > col;
+    if (dirU && lCompare || dirL && uCompare) {
+      z[zOffset] = val;
+    } else if (vx != vz) {  // when x and z are different arrays
+      if (xRank != zRank) coords[0] = coords[1];
+      LongType xOffset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+      z[zOffset] = x[xOffset];
+    }
+  }
+}
+///////////////////////////////////////////////////////////////////
+template <typename T>
+void NDArray::fillAsTriangular(const float val, int lower, int upper, NDArray& target, const char direction,
+                               const bool includeEdges) {
+  if (isS()) THROW_EXCEPTION("NDArray::fillAsTriangular: you can't use this method on String array!");
+
+  if (!isSameShape(target) &&
+      !(rankOf() == 1 && target.rankOf() == 2 && sizeAt(0) == target.sizeAt(0) && sizeAt(0) == target.sizeAt(1)))
+    throw std::string("NDArray::fillAsTriangular method: wrong shape of target array !");
+
+  const int threadsPerBlock = SD_MAX_NUM_THREADS / 4;
+  int len = target.isScalar() ? 1 : target.lengthOf();
+  const int blocksPerGrid = (len + threadsPerBlock - 1) / threadsPerBlock;
+  const int sharedMem = threadsPerBlock * sizeof(int) * target.rankOf() + 128;
+  dim3 launchDims = getFillTriLaunchDims(target.lengthOf(), target.rankOf());
+  PointersManager manager(getContext(), "NDArray::fillAsTriangular");
+
+  prepareSpecialUse({&target}, {this});
+  fillAsTriangularCuda<T><<<launchDims.y, launchDims.x, launchDims.z, *getContext()->getCudaStream()>>>(
+      platformBuffer(), specialShapeInfo(), target.platformBuffer(), target.specialShapeInfo(), static_cast<T>(val),
+      lower, upper, direction, includeEdges);
+  registerSpecialUse({&target}, {this});
+  sd::DebugHelper::checkGlobalErrorCode("fillTriangular  failed");
+
+  manager.synchronize();
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void NDArray::fillAsTriangular,
+                      (const float val, int lower, int upper, NDArray& target, const char direction,
+                          const bool includeEdges),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL static void identityMatrixCuda(void* vx, const LongType* xShapeInfo, const T val) {
+  auto x = reinterpret_cast<T*>(vx);
+
+  // Shared memory variables
+  __shared__ LongType rank;
+  __shared__ LongType len;
+  __shared__ LongType totalThreads;
+  __shared__ const LongType* shapePtr;
+  __shared__ const LongType* stridePtr;
+  __shared__ LongType* sharedMem;
+
+  // Initialize shared variables in thread 0
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char shmem[];
+    sharedMem = reinterpret_cast<LongType*>(shmem);
+
+    // Cache rank and length
+    rank = shape::rank(xShapeInfo);
+    len = shape::length(xShapeInfo);
+
+    // Cache pointers to shape and stride arrays
+    shapePtr = shape::shapeOf(xShapeInfo);
+    stridePtr = shape::stride(xShapeInfo);
+
+    // Calculate total number of threads
+    totalThreads = gridDim.x * blockDim.x;
+  }
+  __syncthreads();
+
+  // Each thread has its own coordinates array in shared memory
+  auto coords = sharedMem + threadIdx.x * rank;
+
+  // Calculate global thread ID
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  // Iterate over assigned elements
+  for (LongType i = tid; i < len; i += totalThreads) {
+    // Convert linear index to multi-dimensional coordinates using cached shape
+    INDEX2COORDS(i, rank, shapePtr, coords);
+
+    // Compute linear offset from coordinates using cached stride
+    LongType offset;
+    COORDS2INDEX(rank, stridePtr, coords, offset);
+
+    // Check if the current position is on the diagonal (row == col)
+    if (coords[rank - 2] == coords[rank - 1]) {  // Assuming 0-based indexing
+      x[offset] = val;
+    }
+    else {
+      x[offset] = static_cast<T>(0);
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////
+template <typename T>
+static void identityMatrixCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                                       const cudaStream_t* stream, void* vx, const LongType* xShapeInfo,
+                                       const float val) {
+  identityMatrixCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, static_cast<T>(val));
+  sd::DebugHelper::checkGlobalErrorCode("identityMatrix  failed");
+
+}
+BUILD_SINGLE_TEMPLATE( void identityMatrixCudaLauncher,
+                      (const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                          const cudaStream_t* stream, void* vx, const sd::LongType* xShapeInfo, const float val),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::setIdentity() {
+  if (isS()) THROW_EXCEPTION("NDArray::setIdentity: you can't use this method on String array!");
+
+  int len = isScalar() ? 1 : lengthOf();
+  dim3 launchDims = getIdentityLaunchDims(len, rankOf());
+
+  PointersManager manager(getContext(), "NDArray::setIdentity");
+
+  syncToDevice();
+  BUILD_SINGLE_SELECTOR(dataType(), identityMatrixCudaLauncher,
+                        (launchDims.y, launchDims.x,launchDims.z, getContext()->getCudaStream(), platformBuffer(),
+                            specialShapeInfo(), 1.f),
+                        SD_COMMON_TYPES);
+  tickWriteDevice();
+
+  manager.synchronize();
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+void NDArray::swapUnsafe(NDArray& other) {
+  auto xType = this->dataType();
+
+  if (xType != other.dataType())
+    THROW_EXCEPTION("NDArray::swapUnsage method: both arrays must have the same data type");
+
+  if (specialBuffer() == nullptr || other.specialBuffer() == nullptr)
+    THROW_EXCEPTION("NDArray::swapUnsafe method: input array should not be empty!");
+
+  if (lengthOf() != other.lengthOf())
+    THROW_EXCEPTION("NDArray::swapUnsafe method: input arrays should have the same length!");
+
+  PointersManager manager(getContext(), "NDArray::swapUnsafe");
+
+  prepareSpecialUse({&other, this}, {&other, this});
+  BUILD_SINGLE_SELECTOR(xType, templatedSwapUnsafe,
+                        (specialBuffer(), specialShapeInfo(), other.specialBuffer(), other.specialShapeInfo(),
+                            getContext()->getCudaStream()),
+                        SD_COMMON_TYPES);
+  registerSpecialUse({&other, this}, {&other, this});
+
+  manager.synchronize();
+}
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::synchronize(const char* msg)  {
+  auto res = cudaStreamSynchronize(*(getContext()->getCudaStream()));
+  if (res != 0) {
+    std::string message = msg + std::string(": synchronization failed !");
+    THROW_EXCEPTION(message.c_str());
+  }
+}
+
+
+// NDArray implementation for .cu file
+void NDArray::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
+  if (msg) sd_printf("%s:\n", msg);
+
+  if(limit < 0) limit = lengthOf();
+
+  // Print array info
+  sd_printf("NDArray: Shape=[", 0);
+  for (int i = 0; i < rankOf(); i++) {
+    sd_printf("%lld", (long long)sizeAt(i));
+    if (i < rankOf() - 1) sd_printf(",", 0);
+  }
+  sd_printf("], DataType=%s,  Order=%c\n",
+            DataTypeUtils::asString(dataType()).c_str(), ordering());
+
+#if defined(SD_GCC_FUNCTRACE)
+  printf("========================================================\n");
+  Printer p;
+  StackTrace st;
+  st.load_here();
+  p.print(st);
+  printf("========================================================\n");
+  fflush(stdout);
+#endif
+  // Print buffer state
+  if (_buffer != nullptr) {
+    _buffer->printBufferDebug("Buffer contents", offset, limit);
+  } else {
+    sd_printf("Buffer is nullptr\n", 0);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::prepareSpecialUse(const std::vector<NDArray*>& writeList,
+                                const std::vector<NDArray*>& readList, bool synchronizeWritables) {
+
+
+
+  for (const auto& a : readList)
+    if (a != nullptr) a->syncToDevice();
+
+  for (const auto& a : writeList) {
+    if (a != nullptr) {
+      a->getDataBuffer()->allocateSpecial();
+      if (synchronizeWritables) a->syncToDevice();
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::registerSpecialUse(const std::vector<NDArray*>& writeList,
+                                 const std::vector<NDArray*>& readList) {
+
+
+  for (const auto& p : readList)
+    if (p != nullptr) p->tickReadDevice();
+
+  for (const auto& p : writeList)
+    if (p != nullptr) p->tickWriteDevice();
+
+
+}
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::preparePrimaryUse(const std::vector<NDArray*>& writeList,
+                                const std::vector<NDArray*>& readList, bool synchronizeWritables) {
+  for (const auto& a : readList)
+    if (a != nullptr) a->syncToHost();
+
+  for (const auto& a : writeList) {
+    if (a != nullptr) {
+      a->getDataBuffer()->allocatePrimary();
+      if (synchronizeWritables) a->syncToHost();
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void NDArray::registerPrimaryUse(const std::vector<NDArray*>& writeList,
+                                 const std::vector<NDArray*>& readList) {
+  for (const auto& p : readList)
+    if (p != nullptr) p->tickReadHost();
+
+  for (const auto& p : writeList)
+    if (p != nullptr) p->tickWriteHost();
+
+}
+
+//////////////////////////////////////////////////////////////////////////
+void NDArray::syncShape()  {
+  cudaMemcpy(const_cast<LongType*>(specialShapeInfo()), shapeInfo(), shape::shapeInfoByteLength(shapeInfo()),
+             cudaMemcpyHostToDevice);
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+
+//////////////////////////////////////////////////////////////////////////
+// change an array by repeating it the number of times given by reps.
+NDArray NDArray::tile(const std::vector<LongType>& reps)  {
+  int dim = reps.size();
+  LongType product = 1;
+  for (const auto& item : reps) product *= item;
+
+  if (product < 1) THROW_EXCEPTION("NDArray::tile method: one of the elements in reps array is zero !");
+
+  int rankOld = rankOf();
+  int diff = rankOld - dim;
+  if (product == 1) {  // in this case 2 possibilities are present: just reshape or nothing to do
+    NDArray result(*this);
+    if (diff < 0) {                               // reshape to higher dimension
+      std::vector<LongType> shapeNew = reps;  // need to have unities at first "diff" positions of new shape
+      memcpy(&shapeNew[-diff], result.shapeInfo() + 1,
+             rankOld * sizeof(LongType));  // put old shape numbers at rest of positions
+      result.reshapei(ordering(), shapeNew);
+    }
+    return result;  // nothing to do, if diff >= 0 -> identity tile
+  }
+
+  // evaluate shapeInfo for resulting array
+  auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
+  // create new buffer, in any case the memory amount new buffer points to is bigger then those for old _buffer
+  DataBuffer *  newBuff = new DataBuffer(shape::length(newShapeInfo) * sizeOfT(),
+                                         dataType(), getContext()->getWorkspace(), true);
+  // assign new shape and new buffer to resulting array
+  NDArray result(newBuff,const_cast<sd::LongType *>(newShapeInfo) , getContext());
+  // fill newBuff, loop through all elements of newBuff
+  // looping through buffer() goes automatically by means of getSubArrayIndex applying
+  const auto resultLen = result.lengthOf();
+  auto xType = this->dataType();
+  auto stream = getContext()->getCudaStream();
+
+  prepareSpecialUse({&result}, {this});
+  BUILD_SINGLE_SELECTOR(xType, tileKernelH,
+                        (this->specialBuffer(), this->specialShapeInfo(), result.specialBuffer(),
+                            result.specialShapeInfo(), resultLen, stream),
+                        SD_COMMON_TYPES);
+  registerSpecialUse({&result}, {this});
+
+  return result;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// change an array by repeating it the number of times given by reps.
+void NDArray::tile(const std::vector<LongType>& reps, NDArray& target)  {
+  auto repProd = shape::prodLong(reps.data(), reps.size());
+  if (repProd < 1) THROW_EXCEPTION("NDArray::tile: reps can't contain 0s");
+
+  // evaluate true tile shapeInfo for comparison with target shapeInfo
+  auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
+  if (!shape::equalsSoft(newShapeInfo, target.shapeInfo())) {
+    THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
+  }
+
+  // fill newBuff, loop through all elements of newBuff
+  // looping through buffer() goes automatically by means of getSubArrayIndex applying
+  const int ews = target.ews();
+  const int targetLen = target.lengthOf();
+  auto stream = getContext()->getCudaStream();
+
+  prepareSpecialUse({&target}, {this});
+  BUILD_SINGLE_SELECTOR_TWICE(
+      target.dataType(), tileKernelHH,
+      (specialBuffer(), specialShapeInfo(), target.specialBuffer(), target.specialShapeInfo(), targetLen, stream),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this});
+}
+
+//////////////////////////////////////////////////////////////////////////
+void NDArray::tile(NDArray& target)  {
+  if (rankOf() > target.rankOf())
+    THROW_EXCEPTION(
+        "NDArray::tile method - rank of target array must be bigger or equal to the rank of this array !");
+
+  if (!ShapeUtils::areShapesBroadcastable(*this, target))
+    THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
+
+  // fill newBuff, loop through all elements of newBuff
+  // looping through getBuffer() goes automatically by means of getSubArrayIndex applying
+  const auto ews = target.ews();
+  const auto targetLen = target.lengthOf();
+  auto stream = getContext()->getCudaStream();
+
+  prepareSpecialUse({&target}, {this});
+  BUILD_SINGLE_SELECTOR_TWICE(
+      target.dataType(), tileKernelHH,
+      (specialBuffer(), specialShapeInfo(), target.specialBuffer(), target.specialShapeInfo(), targetLen, stream),
+      SD_COMMON_TYPES);
+  registerSpecialUse({&target}, {this});
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_KERNEL static void repeatCuda(const void* vx, const LongType* xShapeInfo, void* vz,
+                                 const LongType* zShapeInfo, const LongType* repeats, const LongType repSize,
+                                 const int axis) {
+  const X* x = reinterpret_cast<const X*>(vx);
+  Z* z = reinterpret_cast<Z*>(vz);
+
+  __shared__ LongType rank, *sharedMem;
+  __shared__ LongType zLen, totalThreads;  // xLen = zLen
+
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char shmem[];
+    sharedMem = reinterpret_cast<LongType*>(shmem);
+
+    rank = shape::rank(zShapeInfo);    // xRank = zRank
+    zLen = shape::length(zShapeInfo);  // xLen <= zLen
+
+    totalThreads = gridDim.x * blockDim.x;
+  }
+
+  __syncthreads();
+
+  auto coords = sharedMem + threadIdx.x * rank;
+
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  for (LongType i = tid; i < zLen; i += totalThreads) {
+    INDEX2COORDS(i, rank, shape::shapeOf(zShapeInfo), coords);
+
+    LongType zOffset;
+    COORDS2INDEX(rank, shape::stride(zShapeInfo), coords, zOffset);
+
+    if (repSize > 1) {
+      for (LongType j = 0; j < repSize; ++j) {
+        coords[axis] -= repeats[j];
+        if (coords[axis] < 0) {
+          coords[axis] = j;
+          break;
+        }
+      }
+    } else
+      coords[axis] /= repeats[0];
+
+    LongType xOffset;
+    COORDS2INDEX(rank, shape::stride(xShapeInfo), coords, xOffset);
+
+    z[zOffset] = x[xOffset];
+  }
+}
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+static void repeatCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                               const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, void* vz,
+                               const LongType* zShapeInfo, const LongType* repeats, const LongType repSize, const LongType axis) {
+  repeatCuda<X, Z>
+  <<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, repeats, repSize, axis);
+  DebugHelper::checkGlobalErrorCode("NDArray repeat cuda failed(...) failed");
+
+}
+BUILD_DOUBLE_TEMPLATE( void repeatCudaLauncher,
+                      (const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                          const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, void* vz,
+                          const sd::LongType* zShapeInfo, const sd::LongType* repeats, const sd::LongType repSize, const sd::LongType axis),
+                      SD_COMMON_TYPES, SD_COMMON_TYPES);
+
+//////////////////////////////////////////////////////////////////////////
+// create new array by repeating it the number of times given by repeats
+NDArray NDArray::repeat(const int axis, const std::vector<LongType>& repeats)  {
+  auto nonConst = const_cast<NDArray *>(this);
+  std::vector<sd::LongType> shape = ShapeUtils::evalRepeatShape(axis, repeats, *nonConst);
+  NDArray output('c',shape, dataType(), getContext());
+  dim3 launchDims = getRepeatLaunchDims(output.lengthOf(), output.rankOf());
+
+  PointersManager manager(getContext(), "NDArray::repeat(const int axis, const std::vector<int>& repeats)");
+
+  const LongType* reps = reinterpret_cast<LongType*>(manager.replicatePointer(repeats.data(), repeats.size() * sizeof(LongType)));
+
+  prepareSpecialUse({&output}, {this});
+  BUILD_SINGLE_SELECTOR_TWICE(
+      dataType(), repeatCudaLauncher,
+      (launchDims.y, launchDims.x, launchDims.z, getContext()->getCudaStream(), specialBuffer(), specialShapeInfo(),
+          output.specialBuffer(), output.specialShapeInfo(), reps, repeats.size(), axis),
+      SD_COMMON_TYPES);
+  prepareSpecialUse({&output}, {this});
+
+  manager.synchronize();
+
+  return output;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// fill array by repeating it the number of times given by repeats
+void NDArray::repeat(const int axis, const std::vector<LongType>& repeats, NDArray& target)  {
+  auto nonConst = const_cast<NDArray *>(this);
+  std::vector<sd::LongType> shape = ShapeUtils::evalRepeatShape(axis, repeats, *nonConst);
+
+  if (!target.isSameShape(shape))
+    THROW_EXCEPTION(
+        "NDArray::repeat(const int axis, const std::vector<int>& repeats, NDArray& target) method: wrong shape of "
+        "target array!");
+
+  dim3 launchDims = getRepeatLaunchDims(target.lengthOf(), target.rankOf());
+
+  PointersManager manager(getContext(), "NDArray::repeat(const int axis, const std::vector<int>& repeats)");
+
+  const LongType* reps = reinterpret_cast<LongType*>(manager.replicatePointer(repeats.data(), repeats.size() * sizeof(LongType)));
+  auto targetDataType = target.dataType();
+  auto selfDType = dataType();
+  prepareSpecialUse({&target}, {this});
+  BUILD_DOUBLE_SELECTOR(
+      dataType(), target.dataType(), repeatCudaLauncher,
+      (launchDims.y, launchDims.x, launchDims.z, getContext()->getCudaStream(), specialBuffer(), specialShapeInfo(),
+          target.specialBuffer(), target.specialShapeInfo(), reps, repeats.size(), axis),
+      SD_COMMON_TYPES, SD_COMMON_TYPES);
+  prepareSpecialUse({&target}, {this});
+
+  manager.synchronize();
+}
+
+////////////////////////////////////////////////////////////////////////
+void* NDArray::specialBuffer() {
+  if (_buffer == nullptr) {
+    THROW_EXCEPTION("NDArray::specialBuffer(): _buffer is nullptr - array not properly initialized");
+  }
+
+  void* specialBuf = _buffer->special();
+
+  if (specialBuf == nullptr) {
+    syncToDevice();
+    tickReadHost();
+    specialBuf = _buffer->special();
+    if (specialBuf == nullptr) {
+      THROW_EXCEPTION("NDArray::specialBuffer(): _buffer->special() returned nullptr even after syncToDevice - buffer not allocated");
+    }
+  }
+
+  // FIXME: this should be fixed once CUDA backend added
+  return static_cast<int8_t*>(specialBuf) + (offset() * sizeOfT());
+}
+
+
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+void NDArray::printCurrentBuffer(const bool host, const char* msg, const int precision)  {
+  if (!isScalar() && _length == 0) {
+    printf("NDArray::printActualBuffer: array length is zero !\n");
+    return;
+  }
+
+  if(isScalar()) {
+    if(host) {
+
+      if (msg) printf("%s", msg);
+
+      if (buffer() == nullptr ) {
+        printf("NDArray::printActualBuffer: host buffer is nullptr !\n");
+        return;
+      }
+
+      const T* buff = bufferAsT<T>();
+      if (msg) printf("%s", msg);
+      printf("%.*f\n", precision, (double)buff[getOffset(0)]);
+      return;
+    } else {
+      if (msg) printf("%s", msg);
+
+      if (specialBuffer() == nullptr) {
+        printf("NDArray::printSpecialBuffer: special buffer is nullptr !\n");
+        return;
+      }
+
+
+
+      const auto sizeOfBuffer = sizeOfT();
+
+      void* pHost = operator new(sizeOfBuffer);
+
+      cudaMemcpyAsync(pHost, specialBuffer(), sizeOfBuffer, cudaMemcpyDeviceToHost, *getContext()->getCudaStream());
+      cudaDeviceSynchronize();
+      cudaError_t cudaResult = cudaStreamSynchronize(*getContext()->getCudaStream());
+      auto cast = reinterpret_cast<T*>(pHost);
+      if (cudaResult != 0) THROW_EXCEPTION("NDArray::printSpecialBuffer: cudaStreamSynchronize failed!");
+      printf("%.*f\n", precision, (double)cast[0]);
+
+      return;
+    }
+
+  }
+
+
+
+  if (msg) printf("%s", msg);
+
+  if (host) {
+    if (buffer() == nullptr || _length == 0) {
+      printf("NDArray::printActualBuffer: host buffer is nullptr !\n");
+      return;
+    }
+
+    const T* buff = bufferAsT<T>();
+    for (LongType i = 0; i < _length; i++) printf("%.*f, ", precision, (double)buff[getOffset(i)]);
+    printf("\n");
+  } else {
+    if (specialBuffer() == nullptr) {
+      printf("NDArray::printSpecialBuffer: special buffer is nullptr !\n");
+      return;
+    }
+
+    const auto sizeOfBuffer = sizeOfT() * (getOffset(_length - 1) + 1);
+
+    void* pHost = operator new(sizeOfBuffer);
+
+    cudaMemcpyAsync(pHost, specialBuffer(), sizeOfBuffer, cudaMemcpyDeviceToHost, *getContext()->getCudaStream());
+
+    cudaError_t cudaResult = cudaStreamSynchronize(*getContext()->getCudaStream());
+    if (cudaResult != 0) THROW_EXCEPTION("NDArray::printSpecialBuffer: cudaStreamSynchronize failed!");
+
+    for (LongType i = 0; i < _length; i++)
+      printf("%.*f, ", precision, (double)reinterpret_cast<T*>(pHost)[getOffset(i)]);
+    printf("\n");
+
+    operator delete(pHost);
+  }
+}
+#define PRINT_BUFFER(T) template void NDArray::printCurrentBuffer<GET_SECOND(T)>(const bool host, const char* msg, const int precision);
+ITERATE_LIST((SD_COMMON_TYPES),PRINT_BUFFER)
+
+
+}  // end namespace sd
+#endif

libnd4j/include/array/cuda/NDArrayLambda.cu

@@ -0,0 +1,624 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/DataType.h>
+#include <array/DataTypeUtils.h>
+#include <array/NDArray.h>
+#include <exceptions/cuda_exception.h>
+#include <execution/ThreadPool.h>
+#include <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include "helpers/ShapeUtils.h"
+
+namespace sd {
+
+// ----------- Unary Lambda Operations ----------------
+
+template <typename T>
+SD_KERNEL void applyLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
+                                 void* vz, const sd::LongType* zShapeInfo,
+                                 void* vextraParams) {
+// Cast input and output pointers
+  auto x = reinterpret_cast<const T*>(vx);
+  auto z = reinterpret_cast<T*>(vz);
+  auto extraParams = reinterpret_cast<void*>(vextraParams);
+
+// Cache shape information for x buffer
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+// Cache shape information for z buffer
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length = shape::length(xShapeInfo);
+
+// Cache x shape information
+    xRank = shape::rank(xShapeInfo);
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+// Cache z shape information
+    zRank = shape::rank(zShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+// Apply the function using extraParams (this will be handled in the wrapper function)
+// For now, using a placeholder
+    z[zOffset] = x[xOffset]; // This will be replaced with the actual lambda function call
+  }
+}
+
+// ----------- Indexed Lambda Operations ----------------
+
+template <typename T>
+SD_KERNEL void applyIndexedLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
+                                        void* vz, const sd::LongType* zShapeInfo,
+                                        void* vextraParams) {
+// Cast input and output pointers
+  auto x = reinterpret_cast<const T*>(vx);
+  auto z = reinterpret_cast<T*>(vz);
+  auto extraParams = reinterpret_cast<void*>(vextraParams);
+
+// Cache shape information for x buffer
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+// Cache shape information for z buffer
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length = shape::length(xShapeInfo);
+
+// Cache x shape information
+    xRank = shape::rank(xShapeInfo);
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+// Cache z shape information
+    zRank = shape::rank(zShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+// Apply the indexed function - placeholder for actual lambda call
+    z[zOffset] = x[xOffset]; // This will be replaced with the actual indexed lambda function call
+  }
+}
+
+// ----------- Pairwise Lambda Operations ----------------
+
+template <typename T>
+SD_KERNEL void applyPairwiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
+                                         const void* vy, const sd::LongType* yShapeInfo,
+                                         void* vz, const sd::LongType* zShapeInfo,
+                                         void* vextraParams, bool isScalar) {
+// Cast input and output pointers
+  auto x = reinterpret_cast<const T*>(vx);
+  auto y = reinterpret_cast<const T*>(vy);
+  auto z = reinterpret_cast<T*>(vz);
+  auto extraParams = reinterpret_cast<void*>(vextraParams);
+
+// Cache shape information
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ sd::LongType yRank;
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ const sd::LongType* yStridePtr;
+  __shared__ const sd::LongType* zStridePtr;
+  __shared__ T scalarValue;
+  __shared__ sd::LongType yOffset0;
+
+  if (threadIdx.x == 0) {
+    length = shape::length(xShapeInfo);
+
+// Cache shape information
+    xRank = shape::rank(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    yShapePtr = shape::shapeOf(yShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+
+    xStridePtr = shape::stride(xShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+
+    if (isScalar) {
+      sd::LongType yCoords[SD_MAX_RANK];
+      for (int i = 0; i < yRank; i++) {
+        yCoords[i] = 0;
+      }
+      COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset0);
+      scalarValue = y[yOffset0];
+    }
+  }
+  __syncthreads();
+
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+    if (isScalar) {
+// Apply the pairwise function with scalar - placeholder
+      z[zOffset] = x[xOffset]; // Will be replaced with actual function call using scalarValue
+    } else {
+      INDEX2COORDS(i, yRank, yShapePtr, yCoords);
+      COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
+
+      // Apply the pairwise function - placeholder
+      z[zOffset] = x[xOffset]; // Will be replaced with actual function call using y[yOffset]
+    }
+  }
+}
+
+// ----------- Indexed Pairwise Lambda Operations ----------------
+
+template <typename T>
+SD_KERNEL void applyIndexedPairwiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
+                                                const void* vy, const sd::LongType* yShapeInfo,
+                                                void* vz, const sd::LongType* zShapeInfo,
+                                                void* vextraParams) {
+// Cast input and output pointers
+  auto x = reinterpret_cast<const T*>(vx);
+  auto y = reinterpret_cast<const T*>(vy);
+  auto z = reinterpret_cast<T*>(vz);
+  auto extraParams = reinterpret_cast<void*>(vextraParams);
+
+// Cache shape information
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ sd::LongType yRank;
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ const sd::LongType* yStridePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length = shape::length(xShapeInfo);
+
+// Cache shape information
+    xRank = shape::rank(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    yShapePtr = shape::shapeOf(yShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+
+    xStridePtr = shape::stride(xShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+    INDEX2COORDS(i, yRank, yShapePtr, yCoords);
+    COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
+    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+// Apply the indexed pairwise function - placeholder
+    z[zOffset] = x[xOffset]; // Will be replaced with actual function call
+  }
+}
+
+// ----------- Triplewise Lambda Operations ----------------
+
+template <typename T>
+SD_KERNEL void applyTriplewiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
+                                           const void* vy, const sd::LongType* yShapeInfo,
+                                           const void* vt, const sd::LongType* tShapeInfo,
+                                           void* vz, const sd::LongType* zShapeInfo,
+                                           void* vextraParams) {
+// Cast input and output pointers
+  auto x = reinterpret_cast<const T*>(vx);
+  auto y = reinterpret_cast<const T*>(vy);
+  auto t = reinterpret_cast<const T*>(vt);
+  auto z = reinterpret_cast<T*>(vz);
+  auto extraParams = reinterpret_cast<void*>(vextraParams);
+
+// Cache shape information
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ sd::LongType yRank;
+  __shared__ sd::LongType tRank;
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* tShapePtr;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ const sd::LongType* yStridePtr;
+  __shared__ const sd::LongType* tStridePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length = shape::length(xShapeInfo);
+
+// Cache shape information
+    xRank = shape::rank(xShapeInfo);
+    yRank = shape::rank(yShapeInfo);
+    tRank = shape::rank(tShapeInfo);
+    zRank = shape::rank(zShapeInfo);
+
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    yShapePtr = shape::shapeOf(yShapeInfo);
+    tShapePtr = shape::shapeOf(tShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+
+    xStridePtr = shape::stride(xShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+    tStridePtr = shape::stride(tShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType tCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType tOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+    INDEX2COORDS(i, yRank, yShapePtr, yCoords);
+    COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
+    INDEX2COORDS(i, tRank, tShapePtr, tCoords);
+    COORDS2INDEX(tRank, tStridePtr, tCoords, tOffset);
+    INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+    COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+// Apply the triplewise function - placeholder
+    z[zOffset] = x[xOffset]; // Will be replaced with actual function call
+  }
+}
+
+// ---------------------- Wrapper functions -----------------------
+
+// Helper class for CUDA Lambda operations
+template <typename T>
+class NDArrayLambdaCuda {
+ public:
+  static int constexpr LAMBDA_THREADS = 256;
+  static int constexpr LAMBDA_BLOCKS = 512;
+
+// Unary lambda wrapper
+  static void executeLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
+                            void* z, const sd::LongType* zShapeInfo, void* extraParams) {
+    if(stream == nullptr) {
+      THROW_EXCEPTION("executeLambda: Stream must not be nullptr!");
+    }
+    dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
+    applyLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x, xShapeInfo, z, zShapeInfo, extraParams);
+    sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeLambda failed");
+  }
+
+// Indexed lambda wrapper
+  static void executeIndexedLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
+                                   void* z, const sd::LongType* zShapeInfo, void* extraParams) {
+    if(stream == nullptr) {
+      THROW_EXCEPTION("executeIndexedLambda: Stream must not be nullptr!");
+    }
+    dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
+    applyIndexedLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x, xShapeInfo, z, zShapeInfo, extraParams);
+    sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeIndexedLambda failed");
+  }
+
+// Pairwise lambda wrapper
+  static void executePairwiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
+                                    const void* y, const sd::LongType* yShapeInfo,
+                                    void* z, const sd::LongType* zShapeInfo, void* extraParams, bool isScalar) {
+    dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
+    if(stream == nullptr) {
+      THROW_EXCEPTION("executePairwiseLambda: Stream must not be nullptr!");
+    }
+    applyPairwiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, isScalar);
+    sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executePairwiseLambda failed");
+  }
+
+// Indexed pairwise lambda wrapper
+  static void executeIndexedPairwiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
+                                           const void* y, const sd::LongType* yShapeInfo,
+                                           void* z, const sd::LongType* zShapeInfo, void* extraParams) {
+    dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
+    applyIndexedPairwiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+    sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeIndexedPairwiseLambda failed");
+  }
+
+// Triplewise lambda wrapper
+  static void executeTriplewiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
+                                      const void* y, const sd::LongType* yShapeInfo,
+                                      const void* t, const sd::LongType* tShapeInfo,
+                                      void* z, const sd::LongType* zShapeInfo, void* extraParams) {
+    if(stream == nullptr) {
+      THROW_EXCEPTION("executeTriplewiseLambda: Stream must not be nullptr!");
+    }
+    dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
+    applyTriplewiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x, xShapeInfo, y, yShapeInfo, t, tShapeInfo, z, zShapeInfo, extraParams);
+    sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeTriplewiseLambda failed");
+  }
+};
+
+
+
+// Implementation of the NDArray Lambda methods for CUDA
+template <typename T>
+SD_LIB_EXPORT void NDArray::applyLambda(std::function<T(T)>& func, NDArray* target) {
+// Validate types
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of this "
+        "array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION("NDArray::applyLambdaCuda<T> method: types of this and target array should match!");
+
+// Get device pointers and stream
+  auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
+  auto x = this->specialBuffer();
+  auto z = target->specialBuffer();
+  auto xShapeInfo = this->specialShapeInfo();
+  auto zShapeInfo = target->specialShapeInfo();
+
+// Create and set up extraParams
+  void* extraParams = nullptr; // This would hold the function pointer for the lambda
+
+// Execute the CUDA kernel
+  NDArrayLambdaCuda<T>::executeLambda(stream, x, xShapeInfo, z, zShapeInfo, extraParams);
+}
+
+template <typename T>
+SD_LIB_EXPORT void NDArray::applyIndexedLambda(std::function<T(sd::LongType, T)>& func, NDArray* target) {
+// Validate types
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION("NDArray::applyIndexedLambdaCuda<T> method: types of this and target array should match!");
+
+// Get device pointers and stream
+  auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
+  auto x = this->specialBuffer();
+  auto z = target->specialBuffer();
+  auto xShapeInfo = this->specialShapeInfo();
+  auto zShapeInfo = target->specialShapeInfo();
+
+// Create and set up extraParams
+  void* extraParams = nullptr; // This would hold the function pointer for the lambda
+
+// Execute the CUDA kernel
+  NDArrayLambdaCuda<T>::executeIndexedLambda(stream, x, xShapeInfo, z, zShapeInfo, extraParams);
+}
+
+template <typename T>
+SD_LIB_EXPORT void NDArray::applyPairwiseLambda(NDArray* other,  std::function<T(T, T)>& func,
+                                                NDArray* target) {
+// Validate types
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyPairwiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != other->dataType() || dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyPairwiseLambdaCuda<T> method: all three arrays (this, other, target) must have the same type!");
+
+// Check for scalar or same length
+  bool isScalar = other->isScalar();
+  if (this->lengthOf() != other->lengthOf() && !this->isScalar() && !isScalar) {
+    THROW_EXCEPTION("applyPairwiseLambdaCuda requires both operands to have the same shape or one to be a scalar");
+  }
+
+// Get device pointers and stream
+  auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
+  auto x = this->specialBuffer();
+  auto y = other->specialBuffer();
+  auto z = target->specialBuffer();
+  auto xShapeInfo = this->specialShapeInfo();
+  auto yShapeInfo = other->specialShapeInfo();
+  auto zShapeInfo = target->specialShapeInfo();
+
+// Create and set up extraParams
+  void* extraParams = nullptr; // This would hold the function pointer for the lambda
+
+// Execute the CUDA kernel
+  NDArrayLambdaCuda<T>::executePairwiseLambda(stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, isScalar);
+}
+
+template <typename T>
+SD_LIB_EXPORT void NDArray::applyIndexedPairwiseLambda(NDArray* other,  std::function<T(sd::LongType, T, T)>& func,
+                                                       NDArray* target) {
+// Validate types
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedPairwiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as "
+        "type of this array!");
+  if (dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyIndexedPairwiseLambdaCuda<T> method: types of this and target array should match!");
+  if (this->lengthOf() != other->lengthOf()) {
+    THROW_EXCEPTION("applyIndexedPairwiseLambdaCuda requires both operands to have the same shape");
+  }
+
+// Get device pointers and stream
+  auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
+  auto x = this->specialBuffer();
+  auto y = other->specialBuffer();
+  auto z = target->specialBuffer();
+  auto xShapeInfo = this->specialShapeInfo();
+  auto yShapeInfo = other->specialShapeInfo();
+  auto zShapeInfo = target->specialShapeInfo();
+
+// Create and set up extraParams
+  void* extraParams = nullptr; // This would hold the function pointer for the lambda
+
+// Execute the CUDA kernel
+  NDArrayLambdaCuda<T>::executeIndexedPairwiseLambda(stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+}
+
+template <typename T>
+SD_LIB_EXPORT void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
+                                                  std::function<T(T, T, T)>& func, NDArray* target) {
+// Validate types
+  if (dataType() != DataTypeUtils::fromT<T>())
+    THROW_EXCEPTION(
+        "NDArray::applyTriplewiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
+        "this array!");
+  if (dataType() != second->dataType() || dataType() != third->dataType() || dataType() != target->dataType())
+    THROW_EXCEPTION(
+        "NDArray::applyTriplewiseLambdaCuda<T> method: all four arrays (this, second, third, target) should have the "
+        "same type!");
+
+  if (this->lengthOf() != second->lengthOf() || this->lengthOf() != third->lengthOf() || !this->isSameShape(second) ||
+      !this->isSameShape(third)) {
+    std::string errorMessage;
+    errorMessage += "applyTriplewiseLambdaCuda requires all operands to have the same shape\n";
+    errorMessage += "this shape: " + ShapeUtils::shapeAsString(this->shapeInfo()) + "\n";
+    errorMessage += "second shape: " + ShapeUtils::shapeAsString(second->shapeInfo()) + "\n";
+    errorMessage += "third shape: " + ShapeUtils::shapeAsString(third->shapeInfo()) + "\n";
+    errorMessage += "target shape: " + ShapeUtils::shapeAsString(target->shapeInfo()) + "\n";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+// Get device pointers and stream
+  auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
+  auto x = this->specialBuffer();
+  auto y = second->specialBuffer();
+  auto t = third->specialBuffer();
+  auto z = target->specialBuffer();
+  auto xShapeInfo = this->specialShapeInfo();
+  auto yShapeInfo = second->specialShapeInfo();
+  auto tShapeInfo = third->specialShapeInfo();
+  auto zShapeInfo = target->specialShapeInfo();
+
+// Create and set up extraParams
+  void* extraParams = nullptr; // This would hold the function pointer for the lambda
+
+// Execute the CUDA kernel
+  NDArrayLambdaCuda<T>::executeTriplewiseLambda(stream, x, xShapeInfo, y, yShapeInfo, t, tShapeInfo,
+                                                z, zShapeInfo, extraParams);
+}
+
+
+#define INSTANTIATE_LAMBDA_METHODS(T) template SD_LIB_EXPORT void NDArray::applyLambda( std::function<GET_SECOND(T)(GET_SECOND(T))>& func, NDArray* target);
+ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS);
+
+#define INSTANTIATE_LAMBDA_METHODS_INDEXED(T) template SD_LIB_EXPORT void NDArray::applyIndexedLambda( std::function<GET_SECOND(T)(sd::LongType, GET_SECOND(T))>& func, NDArray* target);
+ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_INDEXED);
+
+#define INSTANTIATE_LAMBDA_METHODS_PAIRWISE(T) template SD_LIB_EXPORT void NDArray::applyPairwiseLambda(NDArray* other,  std::function<GET_SECOND(T)(GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
+ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_PAIRWISE);
+
+#define INSTANTIATE_LAMBDA_METHODS_INDEX_PAIR(T) template SD_LIB_EXPORT void NDArray::applyIndexedPairwiseLambda(NDArray* other,  std::function<GET_SECOND(T)(sd::LongType, GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
+ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_INDEX_PAIR);
+
+#define INSTANTIATE_LAMBDA_METHODS_TRIPLE(T) template void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,  std::function<GET_SECOND(T)(GET_SECOND(T), GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
+ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_TRIPLE);
+} // namespace sd

libnd4j/include/array/impl/ByteOrderUtils.cpp

@@ -0,0 +1,26 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 21.11.17.
+//
+#include <array/ByteOrderUtils.h>
+
+namespace sd {
+ByteOrder ByteOrderUtils::fromFlatByteOrder(graph::ByteOrder order) { return (ByteOrder)order; }
+}  // namespace sd

libnd4j/include/array/impl/ConstantDataBuffer.cpp

@@ -0,0 +1,72 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/ConstantDataBuffer.h>
+#include <array/DataTypeUtils.h>
+
+namespace sd {
+ConstantDataBuffer::ConstantDataBuffer(const std::shared_ptr<PointerWrapper>& primary, uint64_t numEelements,
+                                       DataType dtype)
+    : ConstantDataBuffer(primary, {}, numEelements, dtype) {
+  //
+}
+
+ConstantDataBuffer::ConstantDataBuffer(const std::shared_ptr<PointerWrapper>& primary,
+                                       const std::shared_ptr<PointerWrapper>& special, uint64_t numEelements,
+                                       DataType dtype)
+    : _primaryBuffer(primary), _specialBuffer(special), _length(numEelements) {
+  _sizeOf = DataTypeUtils::sizeOf(dtype);
+}
+
+void* ConstantDataBuffer::primary() const { return _primaryBuffer->pointer(); }
+
+void* ConstantDataBuffer::special() const { return _specialBuffer ? _specialBuffer->pointer() : nullptr; }
+
+uint8_t ConstantDataBuffer::sizeOf() const { return _sizeOf; }
+
+uint64_t ConstantDataBuffer::length() const { return _length; }
+
+ConstantDataBuffer::ConstantDataBuffer(const ConstantDataBuffer& other) {
+  _primaryBuffer = other._primaryBuffer;
+  _specialBuffer = other._specialBuffer;
+  _length = other._length;
+  _sizeOf = other._sizeOf;
+}
+
+template <typename T>
+T* ConstantDataBuffer::primaryAsT() const {
+  return reinterpret_cast<T*>(_primaryBuffer->pointer());
+}
+template SD_LIB_EXPORT float* ConstantDataBuffer::primaryAsT<float>() const;
+template SD_LIB_EXPORT double* ConstantDataBuffer::primaryAsT<double>() const;
+template SD_LIB_EXPORT int* ConstantDataBuffer::primaryAsT<int>() const;
+template SD_LIB_EXPORT LongType* ConstantDataBuffer::primaryAsT<LongType>() const;
+
+template <typename T>
+T* ConstantDataBuffer::specialAsT() const {
+  return reinterpret_cast<T*>(special());
+}
+template SD_LIB_EXPORT float* ConstantDataBuffer::specialAsT<float>() const;
+template SD_LIB_EXPORT double* ConstantDataBuffer::specialAsT<double>() const;
+template SD_LIB_EXPORT int* ConstantDataBuffer::specialAsT<int>() const;
+template SD_LIB_EXPORT LongType* ConstantDataBuffer::specialAsT<LongType>() const;
+
+}  // namespace sd

libnd4j/include/array/impl/ConstantDescriptor.cpp

@@ -0,0 +1,79 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/ConstantDescriptor.h>
+#include <array/DataTypeUtils.h>
+#include <helpers/ModularHasher.h>
+#include <stdexcept>
+
+namespace sd {
+ConstantDescriptor::ConstantDescriptor(double *values, int length) {
+  for (int e = 0; e < length; e++) _floatValues.emplace_back(values[e]);
+}
+
+ConstantDescriptor::ConstantDescriptor(LongType const *values, int length) {
+  for (int e = 0; e < length; e++) _integerValues.emplace_back(values[e]);
+}
+
+ConstantDescriptor::ConstantDescriptor(std::initializer_list<double> values) { _floatValues = values; }
+
+ConstantDescriptor::ConstantDescriptor(std::vector<LongType> &values) { _integerValues = values; }
+
+ConstantDescriptor::ConstantDescriptor(std::vector<double> &values) { _floatValues = values; }
+
+// equal to operator
+bool ConstantDescriptor::operator==(const ConstantDescriptor &other) const {
+  return std::tie(_floatValues, _integerValues) == std::tie(other._floatValues, other._integerValues);
+}
+
+// less than operator
+bool ConstantDescriptor::operator<(const ConstantDescriptor &other) const {
+  return std::tie(_floatValues, _integerValues) < std::tie(other._floatValues, other._integerValues);
+}
+
+bool ConstantDescriptor::isInteger() const { return !_integerValues.empty(); }
+
+bool ConstantDescriptor::isFloat() const { return !_floatValues.empty(); }
+
+const std::vector<LongType> &ConstantDescriptor::integerValues() const { return _integerValues; }
+
+const std::vector<double> &ConstantDescriptor::floatValues() const { return _floatValues; }
+
+LongType ConstantDescriptor::length() const {
+  return isInteger() ? _integerValues.size() : isFloat() ? _floatValues.size() : 0L;
+}
+}  // namespace sd
+
+namespace std {
+size_t hash<sd::ConstantDescriptor>::operator()(const sd::ConstantDescriptor &k) const {
+  uint64_t hash = sd::helpers::detail::ModularHasher::hash_scalar(k.isInteger());
+
+  // Hash the appropriate vector based on type
+  if (k.isInteger()) {
+    hash = sd::helpers::detail::ModularHasher::hash_vector(k.integerValues(), hash);
+  } else {
+    hash = sd::helpers::detail::ModularHasher::hash_vector(k.floatValues(), hash);
+  }
+
+  return hash;
+
+}
+}  // namespace std

libnd4j/include/array/impl/ConstantHolder.cpp

@@ -0,0 +1,66 @@
+/**
+* Copyright (c) 2019 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.
+ *
+ *  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 5/17/2019.
+//
+#include <array/ConstantHolder.h>
+#include <array/DataTypeUtils.h>
+#include <helpers/shape.h>
+
+namespace sd {
+ConstantHolder::ConstantHolder(const ConstantHolder& other) {
+  _buffers = other._buffers;
+  _deviceId = other._deviceId;
+}
+
+bool ConstantHolder::hasBuffer(DataType dataType) { return _buffers.count(dataType) > 0; }
+
+std::mutex* ConstantHolder::mutex() { return &_mutex; }
+
+template <typename T>
+bool ConstantHolder::hasBuffer() {
+  return hasBuffer(DataTypeUtils::fromT<T>());
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT bool ConstantHolder::hasBuffer, (void), SD_COMMON_TYPES);
+
+void ConstantHolder::addBuffer(ConstantDataBuffer& pointer, DataType dataType) { _buffers[dataType] = pointer; }
+
+template <typename T>
+void ConstantHolder::addBuffer(ConstantDataBuffer& pointer) {
+  addBuffer(pointer, DataTypeUtils::fromT<T>());
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void ConstantHolder::addBuffer, (ConstantDataBuffer & cb),
+                      SD_COMMON_TYPES);
+
+ConstantDataBuffer* ConstantHolder::getConstantDataBuffer(DataType dataType) {
+  if (!hasBuffer(dataType)) THROW_EXCEPTION("Requested dataType is absent in storage");
+
+  return &_buffers[dataType];
+}
+
+template <typename T>
+ConstantDataBuffer* ConstantHolder::getConstantDataBuffer() {
+  return getConstantDataBuffer(DataTypeUtils::fromT<T>());
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT ConstantDataBuffer* ConstantHolder::getConstantDataBuffer, (),
+                      SD_COMMON_TYPES);
+}  // namespace sd

libnd4j/include/array/impl/ConstantOffsetsBuffer.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 <array/ConstantOffsetsBuffer.h>
+
+namespace sd {
+ConstantOffsetsBuffer::ConstantOffsetsBuffer(const std::shared_ptr<PointerWrapper> &primary)
+    : ConstantOffsetsBuffer(primary, std::shared_ptr<PointerWrapper>(nullptr)) {
+  //
+}
+
+ConstantOffsetsBuffer::ConstantOffsetsBuffer(const std::shared_ptr<PointerWrapper> &primary,
+                                             const std::shared_ptr<PointerWrapper> &special) {
+  _magic = MAGIC_VALID;  // Mark as valid/constructed
+  _primaryOffsets = primary;
+  _specialOffsets = special;
+}
+
+ConstantOffsetsBuffer::ConstantOffsetsBuffer() {
+  _magic = MAGIC_VALID;  // Mark default-constructed objects as valid
+  // shared_ptr members are automatically initialized to nullptr
+}
+
+ConstantOffsetsBuffer::~ConstantOffsetsBuffer() {
+  _magic = 0xDEADBEEF;  // Mark as destroyed - helps detect use-after-free
+}
+
+LongType *ConstantOffsetsBuffer::primary() {
+  // Check magic number first to detect use-after-free (dangling pointers)
+  if (!isValid()) {
+    THROW_EXCEPTION("ConstantOffsetsBuffer::primary: Object has been destroyed! Use-after-free detected (magic number check failed).");
+  }
+  if (_primaryOffsets == nullptr || !_primaryOffsets) {
+    THROW_EXCEPTION("ConstantOffsetsBuffer::primary: _primaryOffsets is nullptr! Buffer was likely already destroyed.");
+  }
+  return reinterpret_cast<LongType *>(_primaryOffsets->pointer());
+}
+
+LongType *ConstantOffsetsBuffer::special() {
+  return _specialOffsets ? reinterpret_cast<LongType *>(_specialOffsets->pointer()) : nullptr;
+}
+
+LongType *ConstantOffsetsBuffer::platform() {
+#ifdef SD_CUDA
+  return special();
+#else
+  return primary();
+#endif  // CUDABLAS
+}
+
+}  // namespace sd

libnd4j/include/array/impl/ConstantShapeBuffer.cpp

@@ -0,0 +1,126 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/ConstantShapeBuffer.h>
+#include <sstream>
+
+namespace sd {
+ConstantShapeBuffer::ConstantShapeBuffer( PointerWrapper* primary)
+    : ConstantShapeBuffer(primary, nullptr) {
+#if defined(SD_GCC_FUNCTRACE)
+  st = backward::StackTrace();
+  st.load_here(32);
+#endif
+
+}
+ConstantShapeBuffer::ConstantShapeBuffer() : _magic(MAGIC_VALID), _refCount(1) {
+  _primaryShapeInfo = nullptr;
+  _specialShapeInfo = nullptr;
+}
+ConstantShapeBuffer::~ConstantShapeBuffer() {
+  // Clear magic number first to mark as invalid during destruction
+  _magic = 0;
+
+  if(_primaryShapeInfo != nullptr)
+    delete _primaryShapeInfo;
+  _primaryShapeInfo = nullptr;
+
+  if(_specialShapeInfo != nullptr)
+    delete _specialShapeInfo;
+  _specialShapeInfo = nullptr;
+}
+
+ConstantShapeBuffer::ConstantShapeBuffer( PointerWrapper* primary,
+                                          PointerWrapper* special) : _magic(MAGIC_VALID), _refCount(1) {
+  _primaryShapeInfo = primary;
+  _specialShapeInfo = special;
+#if defined(SD_GCC_FUNCTRACE)
+  st = backward::StackTrace();
+  st.load_here(32);
+#endif
+}
+
+LongType *ConstantShapeBuffer::primary()  {
+  if(_primaryShapeInfo != nullptr) {
+    return reinterpret_cast<LongType *>(_primaryShapeInfo->pointer());
+  }
+
+  return nullptr;
+}
+
+LongType *ConstantShapeBuffer::special()  {
+  if(_specialShapeInfo != nullptr) {
+    return reinterpret_cast<LongType *>(_specialShapeInfo->pointer());
+  }
+
+  return nullptr;
+}
+
+LongType *ConstantShapeBuffer::platform()  {
+#ifdef SD_CUDA
+  return special();
+#else
+  return primary();
+#endif  // CUDABLAS
+}
+
+std::string ConstantShapeBuffer::getStackTraceAsString() const {
+#if defined(SD_GCC_FUNCTRACE) && !defined(__JAVACPP_HACK__)
+  // Use backward::Printer to format the stack trace into a string
+  std::ostringstream oss;
+  backward::Printer p;
+  p.snippet = false;  // Don't include source code snippets
+  p.address = true;   // Include addresses
+  p.object = false;   // Don't include object file info
+  p.color_mode = backward::ColorMode::never;  // No ANSI colors in string
+
+  // Print to our string stream (we need to cast away const to use st)
+  // This is safe since print doesn't modify the StackTrace
+  backward::StackTrace& mutable_st = const_cast<backward::StackTrace&>(st);
+  p.print(mutable_st, oss);
+
+  return oss.str();
+#else
+  return "";  // Return empty string when functrace is not enabled
+#endif
+}
+
+void ConstantShapeBuffer::addRef() {
+  _refCount.fetch_add(1, std::memory_order_relaxed);
+}
+
+void ConstantShapeBuffer::release() {
+  // Decrement refcount - buffer stays in cache even when refcount reaches baseline
+  // The cache owns these buffers and is responsible for their lifecycle
+  // Don't delete when refcount reaches 1 - that means cache is the only owner
+  _refCount.fetch_sub(1, std::memory_order_acq_rel);
+
+  // NOTE: Buffers are deleted only when the cache itself is cleared/destroyed,
+  // not when temporary users release their references
+}
+
+int ConstantShapeBuffer::getRefCount() const {
+  return _refCount.load(std::memory_order_relaxed);
+}
+
+}  // namespace sd

libnd4j/include/array/impl/DataBuffer.cpp

@@ -0,0 +1,739 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <array/DataBuffer.h>
+#include <array/DataTypeUtils.h>
+#include <exceptions/allocation_exception.h>
+#include <execution/AffinityManager.h>
+#include <helpers/logger.h>
+#include <memory/MemoryCounter.h>
+#include <sstream>
+
+#if defined(SD_GCC_FUNCTRACE)
+#include <array/DataBufferLifecycleTracker.h>
+#endif
+
+namespace sd {
+///// IMPLEMENTATION OF COMMON METHODS /////
+
+////////////////////////////////////////////////////////////////////////
+// default constructor
+DataBuffer::DataBuffer() {
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer() default constructor\n");
+    fflush(stdout);
+  }
+  _primaryBuffer = nullptr;
+  _specialBuffer = nullptr;
+  _lenInBytes = 0;
+  _dataType = INT8;
+  _workspace = nullptr;
+  _isOwnerPrimary = false;
+  _isOwnerSpecial = false;
+  _deviceId = AffinityManager::currentDeviceId();
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+  setCountersToZero();
+}
+
+////////////////////////////////////////////////////////////////////////
+// copy constructor
+DataBuffer::DataBuffer(const DataBuffer& other) {
+  if(other._dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer(const DataBuffer& other) copy constructor\n");
+    fflush(stdout);
+  }
+  _lenInBytes = other._lenInBytes;
+  _dataType = other._dataType;
+  _workspace = other._workspace;
+#if defined(SD_GCC_FUNCTRACE)
+  // Don't share stack traces - they will be created fresh when we allocate
+  allocationStackTracePrimary = nullptr;
+  allocationStackTraceSpecial = nullptr;
+  creationStackTrace = nullptr;
+#endif
+  _primaryBuffer = other._primaryBuffer;
+  _specialBuffer = other._specialBuffer;
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+
+  _deviceId.store(other._deviceId.load());
+
+  setCountersToZero();
+
+  allocateBuffers();
+  copyBufferFrom(other);
+}
+
+////////////////////////////////////////////////////////////////////////
+DataBuffer::DataBuffer(void* primary, void* special, const size_t lenInBytes, const DataType dataType,
+                       const bool isOwnerPrimary, const bool isOwnerSpecial, memory::Workspace* workspace) {
+  if(dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf(
+        "DataBuffer::DataBuffer(void* primary, void* special, const size_t lenInBytes, const DataType dataType, const bool isOwnerPrimary, const bool isOwnerSpecial, memory::Workspace* workspace) constructor\n");
+    fflush(stdout);
+  }
+  _primaryBuffer = primary;
+  _specialBuffer = special;
+  _lenInBytes = lenInBytes;
+  _dataType = dataType;
+  _workspace = workspace;
+  _isOwnerPrimary = isOwnerPrimary;
+  _isOwnerSpecial = isOwnerSpecial;
+  _deviceId = AffinityManager::currentDeviceId();
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+  setCountersToZero();
+
+  if (primary != nullptr) {
+    readPrimary();
+  }
+  if (special != nullptr) {
+    readSpecial();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+DataBuffer::DataBuffer(void* primary, const size_t lenInBytes, const DataType dataType, const bool isOwnerPrimary,
+                       memory::Workspace* workspace)
+    : DataBuffer(primary, nullptr, lenInBytes, dataType, isOwnerPrimary, false, workspace) {
+  if(dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer(void* primary, const size_t lenInBytes, const DataType dataType, const bool isOwnerPrimary, memory::Workspace* workspace) constructor\n");
+    fflush(stdout);
+  }
+
+  if(primary != nullptr)
+    syncToSpecial(true);
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////
+// copies data from hostBuffer to own memory buffer
+DataBuffer::DataBuffer(const void* hostBuffer, const DataType dataType, const size_t lenInBytes,
+                       memory::Workspace* workspace) {
+  if(dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer(const void* hostBuffer, const DataType dataType, const size_t lenInBytes, memory::Workspace* workspace) constructor\n");
+    fflush(stdout);
+  }
+  if (hostBuffer == nullptr) {
+#if defined(SD_GCC_FUNCTRACE)
+    std::string traceInfo = getCreationTraceAsString();
+    std::string errorMsg = "DataBuffer constructor: can't be initialized with nullptr host buffer !";
+    if (!traceInfo.empty()) {
+      errorMsg += "\n\nDataBuffer allocation trace:\n" + traceInfo;
+    }
+    THROW_EXCEPTION(errorMsg.c_str());
+#else
+    THROW_EXCEPTION("DataBuffer constructor: can't be initialized with nullptr host buffer !");
+#endif
+  }
+  if (lenInBytes == 0) {
+#if defined(SD_GCC_FUNCTRACE)
+    std::string traceInfo = getCreationTraceAsString();
+    std::string errorMsg = "DataBuffer constructor: can't be initialized with zero length !";
+    if (!traceInfo.empty()) {
+      errorMsg += "\n\nDataBuffer allocation trace:\n" + traceInfo;
+    }
+    THROW_EXCEPTION(errorMsg.c_str());
+#else
+    THROW_EXCEPTION("DataBuffer constructor: can't be initialized with zero length !");
+#endif
+  }
+
+  _primaryBuffer = nullptr;
+  _specialBuffer = nullptr;
+  _lenInBytes = lenInBytes;
+  _dataType = dataType;
+  _workspace = workspace;
+
+  _deviceId = AffinityManager::currentDeviceId();
+
+  setCountersToZero();
+
+  allocateBuffers();
+
+  copyBufferFromHost(hostBuffer, lenInBytes);
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////
+DataBuffer::DataBuffer(const sd::LongType lenInBytes, const DataType dataType, memory::Workspace* workspace,
+                       const bool allocBoth) {
+
+  if(dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer(const size_t lenInBytes, const DataType dataType, memory::Workspace* workspace, const bool allocBoth) constructor\n");
+    fflush(stdout);
+  }
+
+
+
+  _dataType = dataType;
+  _workspace = workspace;
+  _lenInBytes = lenInBytes;
+
+  _primaryBuffer = nullptr;
+  _specialBuffer = nullptr;
+  _isOwnerPrimary = false;
+  _isOwnerSpecial = false;
+
+  _deviceId = AffinityManager::currentDeviceId();
+
+  setCountersToZero();
+
+  allocateBuffers(allocBoth);
+  writeSpecial();
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+
+}
+
+////////////////////////////////////////////////////////////////////////
+// move constructor
+DataBuffer::DataBuffer(DataBuffer&& other) {
+
+  if(other._dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer constructor: dataType is UNKNOWN !");
+  }
+
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::DataBuffer(DataBuffer&& other) move constructor\n");
+    fflush(stdout);
+  }
+  _primaryBuffer = other._primaryBuffer;
+  _specialBuffer = other._specialBuffer;
+  _lenInBytes = other._lenInBytes;
+  _dataType = other._dataType;
+  _workspace = other._workspace;
+  _isOwnerPrimary = other._isOwnerPrimary;
+  _isOwnerSpecial = other._isOwnerSpecial;
+  _deviceId.store(other._deviceId);
+
+  copyCounters(other);
+#if defined(SD_GCC_FUNCTRACE)
+  allocationStackTracePrimary = other.allocationStackTracePrimary;
+  allocationStackTraceSpecial = other.allocationStackTraceSpecial;
+  creationStackTrace = other.creationStackTrace;
+  // Transfer ownership - null out the source pointers to prevent double-free
+  other.allocationStackTracePrimary = nullptr;
+  other.allocationStackTraceSpecial = nullptr;
+  other.creationStackTrace = nullptr;
+#endif
+  other._primaryBuffer = other._specialBuffer = nullptr;
+  other.setAllocFlags(false, false);
+  other._lenInBytes = 0;
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////
+// assignment operator
+DataBuffer& DataBuffer::operator=(const DataBuffer& other) {
+  if(other._dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer assignment operator: dataType is UNKNOWN !");
+  }
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::operator=(const DataBuffer& other) assignment operator\n");
+    fflush(stdout);
+  }
+  if (this == &other) return *this;
+
+  deleteBuffers();
+
+  _lenInBytes = other._lenInBytes;
+  _dataType = other._dataType;
+  _workspace = other._workspace;
+
+  allocateBuffers();
+  copyBufferFrom(other);
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+  return *this;
+}
+
+////////////////////////////////////////////////////////////////////////
+// move assignment operator
+DataBuffer& DataBuffer::operator=(DataBuffer&& other) noexcept {
+  if(other._dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer move assignment operator: dataType is UNKNOWN !");
+  }
+
+  if(Environment::getInstance().isLogNativeNDArrayCreation()) {
+    printf("DataBuffer::operator=(DataBuffer&& other) move assignment operator\n");
+    fflush(stdout);
+  }
+  if (this == &other) return *this;
+
+  deleteBuffers();
+
+  _primaryBuffer = other._primaryBuffer;
+  _specialBuffer = other._specialBuffer;
+  _lenInBytes = other._lenInBytes;
+  _dataType = other._dataType;
+  _workspace = other._workspace;
+  _isOwnerPrimary = other._isOwnerPrimary;
+  _isOwnerSpecial = other._isOwnerSpecial;
+
+  copyCounters(other);
+
+#if defined(SD_GCC_FUNCTRACE)
+  allocationStackTracePrimary = other.allocationStackTracePrimary;
+  allocationStackTraceSpecial = other.allocationStackTraceSpecial;
+  creationStackTrace = other.creationStackTrace;
+  // Transfer ownership - null out the source pointers to prevent double-free
+  other.allocationStackTracePrimary = nullptr;
+  other.allocationStackTraceSpecial = nullptr;
+  other.creationStackTrace = nullptr;
+#endif
+
+  other._primaryBuffer = other._specialBuffer = nullptr;
+  other.setAllocFlags(false, false);
+  other._lenInBytes = 0;
+#if defined(SD_GCC_FUNCTRACE)
+  // - Stack trace capture via backward-cpp's backtrace() is NOT safe during early JVM initialization
+  // - The JVM's memory mappings and signal handlers aren't fully set up yet
+  // - This causes SIGSEGV crashes at addresses like 0x7f647edc2000 inside glibc internals
+  // - Session #953's try-catch doesn't work when C++ exceptions are disabled (common for performance)
+  // - DataBufferLifecycleTracker already captures stack traces separately for leak detection
+  // - The creationStackTrace was redundant and only used for constructor error messages
+  // - Solution: Leave creationStackTrace as nullptr (getCreationTraceAsString() handles this gracefully)
+  // - This eliminates crashes while preserving all leak detection functionality
+  creationStackTrace = nullptr;
+#endif
+  return *this;
+}
+
+
+void DataBuffer::markConstant(bool reallyConstant) {
+  isConstant = reallyConstant;
+}
+
+////////////////////////////////////////////////////////////////////////
+// Validation method following DirectShapeTrie pattern
+// Checks for use-after-free, corrupted pointers, and invalid state
+void DataBuffer::validateIntegrity() const {
+  // Check magic number first - if wrong, pointer is dangling/corrupted
+  if (_magicNumber != MAGIC_NUMBER) {
+    // Magic number doesn't match - this is a freed/corrupted DataBuffer!
+    std::stringstream ss;
+    ss << "DataBuffer integrity check FAILED!\n";
+    ss << "  Expected magic number: 0x" << std::hex << MAGIC_NUMBER << "\n";
+    ss << "  Actual magic number: 0x" << std::hex << _magicNumber << "\n";
+    ss << "  Likely causes:\n";
+    ss << "    1. Use-after-free: DataBuffer was deleted but pointer still used\n";
+    ss << "    2. Corrupted pointer: Pointer points to invalid memory\n";
+    ss << "    3. Uninitialized memory: DataBuffer was never properly constructed\n";
+    ss << "  This indicates a SERIOUS BUG in buffer lifecycle management!\n";
+    ss << "  Check where this DataBuffer pointer came from and ensure it's still valid.\n";
+    THROW_EXCEPTION(ss.str().c_str());
+  }
+
+  // Check if buffer has been closed
+  if (closed) {
+    std::stringstream ss;
+    ss << "DataBuffer integrity check FAILED!\n";
+    ss << "  Buffer has been closed (freed) but is still being accessed\n";
+    ss << "  Magic number is valid (0x" << std::hex << _magicNumber << ") but closed flag is true\n";
+    ss << "  This indicates use-after-close: buffer was explicitly closed but pointer retained\n";
+    THROW_EXCEPTION(ss.str().c_str());
+  }
+
+  // Sanity check data type
+  if (_dataType == DataType::UNKNOWN) {
+    std::stringstream ss;
+    ss << "DataBuffer integrity check FAILED!\n";
+    ss << "  DataType is UNKNOWN - buffer was not properly initialized\n";
+    THROW_EXCEPTION(ss.str().c_str());
+  }
+
+  // Sanity check length (negative or excessively large values indicate corruption)
+  if (_lenInBytes < 0 || _lenInBytes > (1LL << 40)) {  // 1TB limit
+    std::stringstream ss;
+    ss << "DataBuffer integrity check FAILED!\n";
+    ss << "  Length is invalid: " << _lenInBytes << " bytes\n";
+    ss << "  Valid range is 0 to " << (1LL << 40) << " bytes (1TB)\n";
+    ss << "  This indicates memory corruption\n";
+    THROW_EXCEPTION(ss.str().c_str());
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void* DataBuffer::primary() {
+  return _primaryBuffer;
+}
+
+////////////////////////////////////////////////////////////////////////
+void* DataBuffer::special() {
+  return _specialBuffer;
+}
+
+////////////////////////////////////////////////////////////////////////
+DataType DataBuffer::getDataType() { return _dataType; }
+
+////////////////////////////////////////////////////////////////////////
+size_t DataBuffer::getLenInBytes() const {
+  // Check if buffer has been closed/freed
+  if(closed) {
+    return 0;
+  }
+  //we need minimum 1 for scalars
+  if(_lenInBytes == 0) {
+   if(_dataType == DataType::UNKNOWN) {
+     THROW_EXCEPTION("DataBuffer getLenInBytes: dataType is UNKNOWN !");
+   }
+    return DataTypeUtils::sizeOfElement(_dataType);
+  }
+  return _lenInBytes;
+}
+size_t DataBuffer::getNumElements()   {
+  return _lenInBytes / DataTypeUtils::sizeOfElement(getDataType());
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::allocatePrimary() {
+#if defined(SD_GCC_FUNCTRACE)
+  // DataBufferLifecycleTracker already captures allocations for leak detection
+  if(allocationStackTracePrimary != nullptr) {
+    delete allocationStackTracePrimary;
+    allocationStackTracePrimary = nullptr;
+  }
+#endif
+  if (_primaryBuffer == nullptr) {
+    auto deviceId = AffinityManager::currentDeviceId();
+    // check if this allocation won't bring us above limit
+    if (_workspace == nullptr) {
+      if (Environment::getInstance().isCPU()) {
+        // on cpu backend we validate against device 0 for now
+        if (!memory::MemoryCounter::getInstance().validate(getLenInBytes()))
+          THROW_EXCEPTION(allocation_exception::build("Requested amount exceeds HOST device limits",
+                                            memory::MemoryCounter::getInstance().deviceLimit(deviceId),
+                                            getLenInBytes()).what());
+      } else {
+        // in heterogenuous mode we validate against device group
+        if (!memory::MemoryCounter::getInstance().validateGroup(memory::MemoryType::HOST, getLenInBytes()))
+          THROW_EXCEPTION(allocation_exception::build(
+              "Requested amount exceeds HOST group limits",
+              memory::MemoryCounter::getInstance().groupLimit(memory::MemoryType::HOST), getLenInBytes()).what());
+      }
+    }
+
+
+
+    ALLOCATE(_primaryBuffer, _workspace, getLenInBytes(), int8_t);
+    _isOwnerPrimary = true;
+
+    // count in towards current deviceId if we're not in workspace mode
+    if (_workspace == nullptr) {
+      if (Environment::getInstance().isCPU())  // we don't want this counter to be added to CUDA device
+        memory::MemoryCounter::getInstance().countIn(deviceId, getLenInBytes());
+
+      memory::MemoryCounter::getInstance().countIn(memory::MemoryType::HOST, getLenInBytes());
+    }
+
+#if defined(SD_GCC_FUNCTRACE)
+    // Record allocation in lifecycle tracker
+    array::DataBufferLifecycleTracker::getInstance().recordAllocation(
+        _primaryBuffer, getLenInBytes(), getDataType(),
+        array::BufferType::PRIMARY, this, _workspace != nullptr);
+#endif
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::setAllocFlags(const bool isOwnerPrimary, const bool isOwnerSpecial) {
+  _isOwnerPrimary = isOwnerPrimary;
+  _isOwnerSpecial = isOwnerSpecial;
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::deletePrimary() {
+#if defined(SD_GCC_FUNCTRACE)
+  printPrimaryAllocationStackTraces();
+
+#endif
+  if (_isOwnerPrimary && _primaryBuffer != nullptr) {
+    auto p = reinterpret_cast<int8_t*>(_primaryBuffer);
+
+    if(Environment::getInstance().isDeletePrimary()) {
+#if defined(SD_GCC_FUNCTRACE)
+      // Record deallocation before releasing memory
+      array::DataBufferLifecycleTracker::getInstance().recordDeallocation(
+          _primaryBuffer, array::BufferType::PRIMARY);
+#endif
+      RELEASE(p, _workspace);
+      _primaryBuffer = nullptr;
+    }
+
+    _isOwnerPrimary = false;
+
+    // count out towards DataBuffer device, only if we're not in workspace
+    if (_workspace == nullptr) {
+      if (Environment::getInstance().isCPU()) memory::MemoryCounter::getInstance().countOut(_deviceId, getLenInBytes());
+
+      memory::MemoryCounter::getInstance().countOut(memory::MemoryType::HOST, getLenInBytes());
+    }
+  }
+
+
+
+}
+
+void DataBuffer::printPrimaryAllocationStackTraces() {
+#if defined(SD_GCC_FUNCTRACE)
+
+#endif
+
+}
+
+////////////////////////////////////////////////////////////////////////
+void DataBuffer::deleteBuffers() {
+  if(isConstant || closed) {
+    return;
+  }
+
+  std::lock_guard<std::mutex> lock(_deleteMutex);
+  deletePrimary();
+  deleteSpecial();
+
+  // Clean up stack traces to prevent memory leak
+#if defined(SD_GCC_FUNCTRACE)
+  if(allocationStackTracePrimary != nullptr) {
+    delete allocationStackTracePrimary;
+    allocationStackTracePrimary = nullptr;
+  }
+  if(allocationStackTraceSpecial != nullptr) {
+    delete allocationStackTraceSpecial;
+    allocationStackTraceSpecial = nullptr;
+  }
+  if(creationStackTrace != nullptr) {
+    delete creationStackTrace;
+    creationStackTrace = nullptr;
+  }
+#endif
+
+  closed = true;
+  _lenInBytes = 0;
+}
+
+////////////////////////////////////////////////////////////////////////
+DataBuffer::~DataBuffer() {
+  // Clear magic number to detect use-after-free
+  // If anyone tries to use this buffer after destruction, validateIntegrity() will catch it
+  _magicNumber = 0xDEADBEEF;
+  deleteBuffers();
+}
+
+
+void DataBuffer::setPrimaryBuffer(void* buffer, size_t length) {
+  std::lock_guard<std::mutex> lock(_deleteMutex);
+#if defined(SD_GCC_FUNCTRACE)
+  // DataBufferLifecycleTracker already captures allocations for leak detection
+  if(allocationStackTracePrimary != nullptr) {
+    delete allocationStackTracePrimary;
+    allocationStackTracePrimary = nullptr;
+  }
+#endif
+  _primaryBuffer = buffer;
+  _isOwnerPrimary = false;
+  _lenInBytes = length * DataTypeUtils::sizeOf(_dataType);
+}
+
+void DataBuffer::setSpecialBuffer(void* buffer, size_t length) {
+  std::lock_guard<std::mutex> lock(_deleteMutex);
+#if defined(SD_GCC_FUNCTRACE)
+  // DataBufferLifecycleTracker already captures allocations for leak detection
+  if(allocationStackTraceSpecial != nullptr) {
+    delete allocationStackTraceSpecial;
+    allocationStackTraceSpecial = nullptr;
+  }
+#endif
+  this->setSpecial(buffer, false);
+  _lenInBytes = length * DataTypeUtils::sizeOf(_dataType);
+}
+
+void DataBuffer::setDataType(DataType dataType) {
+  if(dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("DataBuffer setDataType: dataType is UNKNOWN !");
+  }
+  _dataType = dataType;
+}
+
+void DataBuffer::printAllocationTrace() {
+  if(closed) {
+    printf("DataBuffer::printAllocationTrace() - buffer is closed\n");
+    fflush(stdout);
+  }
+#if defined(SD_GCC_FUNCTRACE)
+  //print whether each stack trace is null or not:
+  Printer p;
+  if(allocationStackTracePrimary != nullptr) {
+    p.print(*allocationStackTracePrimary);
+  }
+  if(allocationStackTraceSpecial != nullptr) {
+    p.print(*allocationStackTraceSpecial);
+  }
+  if(creationStackTrace != nullptr) {
+    p.print(*creationStackTrace);
+  }
+#endif
+}
+
+std::string DataBuffer::getCreationTraceAsString() const {
+#if defined(SD_GCC_FUNCTRACE)
+  if (creationStackTrace == nullptr || creationStackTrace->size() == 0) {
+    return "";
+  }
+
+  std::ostringstream oss;
+  backward::TraceResolver resolver;
+  resolver.load_stacktrace(*creationStackTrace);
+
+  for (size_t i = 0; i < creationStackTrace->size(); ++i) {
+    const backward::ResolvedTrace &trace = resolver.resolve((*creationStackTrace)[i]);
+
+    // Format: #frame function_name at source_file:line
+    oss << "#" << i << " ";
+
+    if (!trace.object_function.empty()) {
+      oss << trace.object_function;
+    } else {
+      oss << "???";
+    }
+
+    if (!trace.source.filename.empty()) {
+      oss << " at " << trace.source.filename;
+      if (trace.source.line > 0) {
+        oss << ":" << trace.source.line;
+      }
+    }
+
+    oss << "\n";
+  }
+
+  return oss.str();
+#else
+  return "";
+#endif
+}
+
+int DataBuffer::deviceId() const { return _deviceId.load(); }
+
+void DataBuffer::close() { this->deleteBuffers(); }
+
+void DataBuffer::setDeviceId(int deviceId) { _deviceId = deviceId; }
+}  // namespace sd

libnd4j/include/array/impl/DataTypeUtils.cpp

@@ -0,0 +1,39 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/DataType.h>
+#include <array/DataTypeUtils.h>
+#include <types/float16.h>
+#include <system/selective_rendering.h>
+
+namespace sd {
+DataType DataTypeUtils::fromInt(int val) { return (DataType)val; }
+
+DataType DataTypeUtils::fromFlatDataType(graph::DType dtype) { return (DataType)dtype; }
+
+int DataTypeUtils::asInt(DataType type) { return static_cast<int>(type); }
+
+/**
+ * Check if a triple of data types is enabled for compilation in selective rendering
+ */
+static bool isCompiledTypeTriple(DataType type1, DataType type2, DataType type3);
+
+}  // namespace sd

libnd4j/include/array/impl/ExtraArguments.cpp

@@ -0,0 +1,130 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/DataType.h>
+#include <array/DataTypeUtils.h>
+#include <array/ExtraArguments.h>
+#include <types/types.h>
+
+#include <stdexcept>
+
+#ifdef SD_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+namespace sd {
+ExtraArguments::ExtraArguments(std::initializer_list<double> arguments) { _fpArgs = arguments; }
+
+ExtraArguments::ExtraArguments(std::initializer_list<LongType> arguments) { _intArgs = arguments; }
+
+ExtraArguments::ExtraArguments(const std::vector<double> &arguments) { _fpArgs = arguments; }
+
+ExtraArguments::ExtraArguments(const std::vector<LongType> &arguments) { _intArgs = arguments; }
+
+ExtraArguments::ExtraArguments(const std::vector<int> &arguments) {
+  for (const auto &v : arguments) _intArgs.emplace_back(static_cast<LongType>(v));
+}
+
+ExtraArguments::ExtraArguments() {
+  // no-op
+}
+
+ExtraArguments::~ExtraArguments() {
+  for (auto p : _pointers) {
+#ifdef SD_CUDA
+    cudaFree(p);
+#else  // CPU branch
+    delete reinterpret_cast<int8_t *>(p);
+#endif
+  }
+}
+
+template <typename T>
+void ExtraArguments::convertAndCopy(Pointer pointer, LongType offset) {
+  auto length = this->length();
+  auto target = reinterpret_cast<T *>(pointer);
+#ifdef SD_CUDA
+  target = new T[length];
+#endif
+
+  if (!_fpArgs.empty()) {
+    for (size_t e = offset; e < _fpArgs.size(); e++) {
+      target[e] = static_cast<T>(_fpArgs[e]);
+    }
+  } else if (_intArgs.empty()) {
+    for (size_t e = offset; e < _intArgs.size(); e++) {
+      target[e] = static_cast<T>(_intArgs[e]);
+    }
+  }
+
+#ifdef SD_CUDA
+  cudaMemcpy(pointer, target, length * DataTypeUtils::sizeOf(DataTypeUtils::fromT<T>()), cudaMemcpyHostToDevice);
+  delete[] target;
+#endif
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void ExtraArguments::convertAndCopy,
+                      (sd::Pointer pointer, sd::LongType offset), SD_COMMON_TYPES);
+
+void *ExtraArguments::allocate(size_t length, size_t elementSize) {
+#ifdef SD_CUDA
+  Pointer ptr;
+  auto res = cudaMalloc(reinterpret_cast<void **>(&ptr), length * elementSize);
+  if (res != 0) THROW_EXCEPTION("Can't allocate CUDA memory");
+#else  // CPU branch
+  auto ptr = new int8_t[length * elementSize];
+  if (!ptr) THROW_EXCEPTION("Can't allocate memory");
+#endif
+
+  return ptr;
+}
+
+size_t ExtraArguments::length() {
+  if (!_fpArgs.empty())
+    return _fpArgs.size();
+  else if (!_intArgs.empty())
+    return _intArgs.size();
+  else
+    return 0;
+}
+
+template <typename T>
+void *ExtraArguments::argumentsAsT(LongType offset) {
+  return argumentsAsT(DataTypeUtils::fromT<T>(), offset);
+}
+BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void *ExtraArguments::argumentsAsT, (sd::LongType offset),
+                      SD_COMMON_TYPES);
+
+void *ExtraArguments::argumentsAsT(DataType dataType, LongType offset) {
+  if (_fpArgs.empty() && _intArgs.empty()) return nullptr;
+
+  // we allocate pointer
+  auto ptr = allocate(length() - offset, DataTypeUtils::sizeOf(dataType));
+
+  // fill it with data
+  BUILD_SINGLE_SELECTOR(dataType, convertAndCopy, (ptr, offset), SD_COMMON_TYPES);
+
+  // store it internally for future release
+  _pointers.emplace_back(ptr);
+
+  return ptr;
+}
+}  // namespace sd

libnd4j/include/array/impl/InteropDataBuffer.cpp

@@ -0,0 +1,231 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/DataTypeUtils.h>
+#include <array/InteropDataBuffer.h>
+#include <execution/AffinityManager.h>
+#include <helpers/logger.h>
+
+namespace sd {
+InteropDataBuffer::InteropDataBuffer(InteropDataBuffer* dataBuffer, uint64_t length) {
+ if(dataBuffer == nullptr) {
+        THROW_EXCEPTION("InteropDataBuffer::InteropDataBuffer(InteropDataBuffer& dataBuffer, uint64_t length, uint64_t offset) - dataBuffer is nullptr");
+ }
+  if(dataBuffer->_dataBuffer->getDataType() == DataType::UNKNOWN)
+    THROW_EXCEPTION("InteropDataBuffer::InteropDataBuffer(InteropDataBuffer& dataBuffer, uint64_t length, uint64_t offset) - dataBuffer has unknown data type");
+  _dataBuffer = dataBuffer->dataBuffer();
+  _dataType = dataBuffer->_dataType;
+
+  _cachedLenInBytes = length * DataTypeUtils::sizeOf(_dataType);
+  // Cache pointers for deallocation tracking
+  if (_dataBuffer != nullptr) {
+    _cachedPrimaryPtr = _dataBuffer->primary();
+    _cachedSpecialPtr = _dataBuffer->special();
+  }
+
+  owner = false;
+}
+
+InteropDataBuffer::InteropDataBuffer(DataBuffer * databuffer) {
+  _dataBuffer = databuffer;
+  _dataType = databuffer->getDataType();
+  if(_dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION(
+        "InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) - data type is unknown");
+  }
+  // Cache the size to avoid accessing freed memory later
+  _cachedLenInBytes = databuffer != nullptr ? databuffer->getLenInBytes() : 0;
+  // Cache pointers for deallocation tracking
+  if (databuffer != nullptr) {
+    _cachedPrimaryPtr = databuffer->primary();
+    _cachedSpecialPtr = databuffer->special();
+  }
+  // When wrapping an existing DataBuffer, we don't own it by default
+  owner = false;
+}
+
+InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) {
+  if(dtype == DataType::UNKNOWN) {
+    THROW_EXCEPTION(
+        "InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) - data type is unknown");
+  }
+
+  _cachedLenInBytes = lenInBytes;
+
+  if (lenInBytes == 0) {
+    _dataBuffer = nullptr;
+    this->_dataType = dtype;
+
+  } else {
+    //note this should be size in bytes hence why we multiply the number of elements by the size of the data type
+    _dataBuffer = new DataBuffer(lenInBytes, dtype, nullptr, allocateBoth);
+    this->_dataType = dtype;
+    this->markOwner(true);
+    // Cache pointers for deallocation tracking
+    _cachedPrimaryPtr = _dataBuffer->primary();
+    _cachedSpecialPtr = _dataBuffer->special();
+  }
+}
+
+
+void InteropDataBuffer::printDbAllocationTrace() {
+  if(_dataBuffer == nullptr)
+    return;
+  _dataBuffer->printAllocationTrace();
+}
+
+void InteropDataBuffer::markOwner(bool owner) {
+  this->owner = owner;
+  if(_dataBuffer != nullptr && !_closed) {
+    this->_dataBuffer->_isOwnerPrimary = owner;
+    this->_dataBuffer->_isOwnerSpecial = owner;
+  }
+}
+
+DataBuffer * InteropDataBuffer::getDataBuffer() const {
+  //this can effect size of calculations among others
+  if(_dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("All interop buffers must have a known data type.");
+  }
+  // Don't access _dataBuffer if it's been closed/freed
+  if(_dataBuffer != nullptr && !_closed && _dataBuffer->_dataType == DataType::UNKNOWN) {
+    _dataBuffer->_dataType = _dataType;
+  }
+  // Return nullptr if closed to prevent use-after-free
+  return _closed ? nullptr : _dataBuffer;
+}
+
+DataBuffer * InteropDataBuffer::dataBuffer() {
+  if(_dataBuffer == nullptr) {
+    return nullptr;
+  }
+  return _dataBuffer;
+}
+
+
+
+void* InteropDataBuffer::primary() const {
+  if(_dataBuffer == nullptr || _dataBuffer->primary() == nullptr) {
+    return nullptr;
+  }
+  return reinterpret_cast<int8_t*>(_dataBuffer->primary());
+}
+
+void* InteropDataBuffer::special() const {
+  if(_dataBuffer == nullptr)
+    return nullptr;
+  if(_dataBuffer->special() == nullptr) {
+    return nullptr;
+  }
+  return reinterpret_cast<int8_t*>(_dataBuffer->special());
+}
+
+void InteropDataBuffer::setSpecial(void* ptr, size_t length) {
+  if(_dataBuffer == nullptr)
+    THROW_EXCEPTION("InteropDataBuffer::setSpecial() - _dataBuffer is nullptr");
+  if(_closed)
+    return;  // Silently ignore if buffer was already closed
+  _dataBuffer->setSpecialBuffer(ptr, length);
+  // Update cached pointer
+  _cachedSpecialPtr = ptr;
+}
+
+void InteropDataBuffer::setPrimary(void* ptr, size_t length) {
+  if(_dataBuffer == nullptr)
+    THROW_EXCEPTION("InteropDataBuffer::setPrimary() - _dataBuffer is nullptr");
+  if(_closed)
+    return;  // Silently ignore if buffer was already closed
+  _dataBuffer->setPrimaryBuffer(ptr, length);
+  // Update cached pointer
+  _cachedPrimaryPtr = ptr;
+}
+
+void InteropDataBuffer::setDeviceId(int deviceId) {
+  if(_dataBuffer == nullptr || _closed)
+    return;
+  _dataBuffer->setDeviceId(deviceId);
+}
+int InteropDataBuffer::deviceId() const {
+  if(_dataBuffer == nullptr || _closed)
+    return 0;
+  return _dataBuffer->deviceId();
+}
+
+int InteropDataBuffer::useCount() const {
+  return 1;
+}
+
+void InteropDataBuffer::registerSpecialUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                           const std::vector<const InteropDataBuffer*>& readList) {
+  for (const auto& v : writeList) {
+    if (v == nullptr) continue;
+
+    v->getDataBuffer()->writeSpecial();
+  }
+}
+
+void InteropDataBuffer::prepareSpecialUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                          const std::vector<const InteropDataBuffer*>& readList,
+                                          bool synchronizeWritables) {
+  auto currentDeviceId = AffinityManager::currentDeviceId();
+  for (const auto& v : readList) {
+    if (v == nullptr || v->_closed) continue;
+
+    auto db = v->getDataBuffer();
+    if(db == nullptr) continue;
+
+    if (db->deviceId() != currentDeviceId) db->migrate();
+    db->syncToSpecial();
+  }
+
+  // we don't tick write list, only ensure the same device affinity
+  for (const auto& v : writeList) {
+    if (v == nullptr || v->_closed) continue;
+
+    auto db = v->getDataBuffer();
+    if(db == nullptr) continue;
+
+    // special case for legacy ops - views can be updated on host side, thus original array can be not updated
+    if (!db->isSpecialActual()) db->syncToSpecial();
+
+    if (db->deviceId() != currentDeviceId) db->migrate();
+  }
+}
+
+void InteropDataBuffer::registerPrimaryUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                           const std::vector<const InteropDataBuffer*>& readList) {
+
+}
+
+void InteropDataBuffer::preparePrimaryUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                          const std::vector<const InteropDataBuffer*>& readList,
+                                          bool synchronizeWritables) {
+
+}
+
+void InteropDataBuffer::expand(size_t newlength) {
+  if(_dataBuffer == nullptr || _closed)
+    return;  // Cannot expand a closed or null buffer
+  _dataBuffer->expand(newlength * DataTypeUtils::sizeOf(_dataBuffer->getDataType()));
+}
+
+
+}  // namespace sd

libnd4j/include/array/impl/NDArrayFactory.cpp

@@ -0,0 +1,744 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 GS <sgazeos@gmail.com> on 2018-12-20.
+// @author Oleg Semeniv <oleg.semeniv@gmail.com>
+//
+#include <array/NDArrayFactory.h>
+#include <exceptions/cuda_exception.h>
+#include <graph/GraphExecutioner.h>
+#include <helpers/ConstantHelper.h>
+#include <helpers/ConstantShapeHelper.h>
+#include <helpers/LoopsCoordsHelper.h>
+#include <helpers/ShapeUtils.h>
+#include <helpers/StringUtils.h>
+#include <legacy/NativeOps.h>
+
+#include <type_traits>
+
+namespace sd {
+
+SD_LIB_EXPORT NDArray* NDArrayFactory::create(ShapeDescriptor *shapeDescriptor, LaunchContext* context) {
+  auto status = shapeDescriptor->validate();
+  if (status != SHAPE_DESC_OK) {
+    THROW_EXCEPTION("NDArrayFactory::create: invalid ShapeDescriptor ");
+  }
+  LongType allocSize = shapeDescriptor->allocLength() * DataTypeUtils::sizeOfElement(shapeDescriptor->dataType());
+  DataBuffer *  buffer =
+      new DataBuffer(allocSize, shapeDescriptor->dataType(), context->getWorkspace());
+  NDArray *result = new NDArray(buffer, shapeDescriptor, context);
+  result->nullify();
+  return result;
+}
+
+
+////////////////////////////////////////////////////////////////////////
+template <>
+SD_LIB_EXPORT NDArray* NDArrayFactory::create<bool>(const char order, const std::vector<LongType>& shape,
+                                                   const std::vector<bool>& data, LaunchContext* context) {
+  if ((int)shape.size() > SD_MAX_RANK)
+    THROW_EXCEPTION("NDArrayFactory::create: rank of NDArray can't exceed 32 !");
+
+  ShapeDescriptor *descriptor = new ShapeDescriptor(BOOL, order, shape);
+
+  if (static_cast<size_t>(descriptor->arrLength()) != data.size()) {
+    sd_printf("NDArrayFactory::create: data size [%i] doesn't match shape length [%lld]\n", data.size(),
+              descriptor->arrLength());
+    THROW_EXCEPTION("NDArrayFactory::create: data size doesn't match shape");
+  }
+
+  bool* hostBuffer = nullptr;
+  ALLOCATE(hostBuffer, context->getWorkspace(), data.size(), bool);
+  std::copy(data.begin(), data.end(), hostBuffer);
+
+  DataBuffer * buffer = new DataBuffer(hostBuffer, data.size() * sizeof(bool), BOOL, true, context->getWorkspace());
+
+  NDArray *result = new NDArray(buffer, descriptor, context);
+  delete descriptor;
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::create(const char order,
+                                const std::vector<LongType>& shape,
+                                const std::vector<T>& data,
+                               LaunchContext* context) {
+  if (shape.size() > SD_MAX_RANK)
+    THROW_EXCEPTION("NDArrayFactory::create: rank of NDArray can't exceed 32 !");
+  ShapeDescriptor *descriptor = new ShapeDescriptor(DataTypeUtils::fromT<T>(), order, shape);
+
+  //scalars can be created with zero length
+  if (descriptor->arrLength() != 0 && data.size() != 1 && static_cast<size_t>(descriptor->arrLength()) != data.size()) {
+    delete descriptor;
+    sd_printf("NDArrayFactory::create: data size [%i] doesn't match shape length [%lld]\n", data.size(),
+              descriptor->arrLength());
+    THROW_EXCEPTION("NDArrayFactory::create: data size doesn't match shape");
+  }
+
+  T *hostData = nullptr;
+  ALLOCATE(hostData, context->getWorkspace(), data.size(), T);
+  std::copy(data.begin(), data.end(), hostData);
+
+  //note here we use data.size() to work around the scalar case. If the shape is zero but the data is actually length 1 we need this reflected
+  //to create a correct length data buffer
+  auto dtypeString = DataTypeUtils::asString(descriptor->dataType());
+  DataBuffer *  buffer = new DataBuffer(
+      hostData, DataTypeUtils::fromT<T>(), data.size() * sizeof(T), context->getWorkspace());
+
+  NDArray *result = new NDArray(buffer, descriptor, context);
+  delete descriptor;
+  return result;
+}
+
+// Update the instantiation macro to use the expanded type pattern
+#define TMPL_INSTANTIATE_CREATE_A_TYPE(TYPE) \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create<TYPE>(const char order, const std::vector<sd::LongType>& shape, \
+                                                      const std::vector<TYPE>& data, sd::LaunchContext* context);
+
+#define TMPL_INSTANTIATE_CREATE_A(T) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(T), \
+    CONCAT(EXPAND_TYPE_APPLY_, GET_SECOND(T))(TMPL_INSTANTIATE_CREATE_A_TYPE) \
+))
+
+ITERATE_LIST((SD_NUMERIC_TYPES), TMPL_INSTANTIATE_CREATE_A)
+
+#undef TMPL_INSTANTIATE_CREATE_A_TYPE
+#undef TMPL_INSTANTIATE_CREATE_A
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::create_(const char order, std::vector<LongType>& shape, LaunchContext* context) {
+  return create_(order, shape, DataTypeUtils::fromT<T>(), context);
+}
+BUILD_SINGLE_TEMPLATE(NDArray* NDArrayFactory::create_,
+                      (const char order,  std::vector<sd::LongType>& shape, sd::LaunchContext* context),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+void NDArrayFactory::memcpyFromVector(void* ptr, const std::vector<T>& vector) {
+  memcpy(ptr, vector.data(), vector.size() * sizeof(T));
+}
+
+template <>
+void SD_LIB_EXPORT NDArrayFactory::memcpyFromVector(void* ptr, const std::vector<bool>& vector) {
+  auto p = reinterpret_cast<bool*>(ptr);
+  for (size_t e = 0; e < vector.size(); e++) p[e] = vector[e];
+}
+
+
+#define TMPL_INSTANTIATE_MEMCPY(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT void NDArrayFactory::memcpyFromVector<GET_SECOND(TYPE)>(void* ptr, const std::vector<GET_SECOND(TYPE)>& vector); \
+))
+
+ITERATE_LIST((SD_NUMERIC_TYPES), TMPL_INSTANTIATE_MEMCPY)
+#undef TMPL_INSTANTIATE_MEMCPY
+
+#ifndef __JAVACPP_HACK__
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::valueOf(const std::initializer_list<LongType>& shape, const T value, const char order,
+                                 LaunchContext* context) {
+  std::vector<sd::LongType> shape2 = shape;
+  return valueOf(shape2, value, order);
+}
+
+#define TMPL_INSTANTIATE_VALUEOF_A(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::valueOf<GET_SECOND(TYPE)>(const std::initializer_list<sd::LongType>& shape, \
+                                                        const GET_SECOND(TYPE) value, const char order, \
+                                                        sd::LaunchContext* context); \
+))
+ITERATE_LIST((SD_NUMERIC_TYPES), TMPL_INSTANTIATE_VALUEOF_A)
+
+#undef TMPL_INSTANTIATE_VALUEOF_A
+
+
+
+#endif
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::create_(const T scalar, LaunchContext* context) {
+  sd::LongType  size = DataTypeUtils::sizeOfElement(DataTypeUtils::fromT<T>());
+  DataBuffer *  buffer =
+      new DataBuffer(size,
+                     DataTypeUtils::fromT<T>(),
+                     context->getWorkspace(),
+                     true);
+
+  auto desc = ShapeBuilders::createScalarShapeInfo(DataTypeUtils::fromT<T>());
+  auto constDesc = ConstantShapeHelper::getInstance().bufferForShapeInfo(desc);
+  auto recast = const_cast<LongType*>(constDesc->primary());
+  NDArray* res = new NDArray(buffer, recast, context);
+  res->p<T>(0,scalar);
+
+  res->tickWriteHost();
+  res->syncToDevice();
+
+  delete[] desc;  // Free allocated shape info
+
+  return res;
+}
+
+#define TMPL_INSTANTIATE_CREATE_C(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create_<GET_SECOND(TYPE)>(const GET_SECOND(TYPE) scalar, sd::LaunchContext* context); \
+))
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_CREATE_C)
+
+#undef TMPL_INSTANTIATE_CREATE_C
+
+NDArray* NDArrayFactory::create(DataType dtype, LaunchContext *context) {
+  return create(dtype,0, context);
+}
+
+template <typename T>
+NDArray* NDArrayFactory::create(DataType type, const T scalar, LaunchContext* context) {
+  if (type == DataTypeUtils::fromT<T>()) return NDArrayFactory::create(scalar, context);
+
+  NDArray *res = new NDArray(type, context);
+  res->p(0, scalar);
+  res->syncToDevice();
+
+  return res;
+}
+
+#define TMPL_INSTANTIATE_CREATE_D(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create<GET_SECOND(TYPE)>(DataType type, const GET_SECOND(TYPE) scalar, sd::LaunchContext* context); \
+))
+
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_CREATE_D)
+
+#undef TMPL_INSTANTIATE_CREATE_D
+
+template <typename T>
+NDArray* NDArrayFactory::create(const T scalar, LaunchContext* context) {
+  DataBuffer *  buffer =
+      new DataBuffer(1 * sizeof(T), DataTypeUtils::fromT<T>(), context->getWorkspace(), true);
+
+  auto desc = ShapeDescriptor::scalarDescriptor(DataTypeUtils::fromT<T>());
+  NDArray *res = new NDArray(buffer,desc , context);
+  res->bufferAsT<T>()[0] = scalar;
+
+  res->tickWriteHost();
+  res->syncToDevice();
+
+  return res;
+}
+
+#define TMPL_INSTANTIATE_CREATE_E(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create<GET_SECOND(TYPE)>(const GET_SECOND(TYPE) scalar, sd::LaunchContext* context); \
+))
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_CREATE_E)
+#undef TMPL_INSTANTIATE_CREATE_E
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::create_(const char order, const std::vector<LongType>& shape, const std::vector<T>& data,
+                                 LaunchContext* context) {
+  return NDArrayFactory::create<T>(order, shape, data, context);
+}
+
+#define TMPL_INSTANTIATE_CREATE_F(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create_<GET_SECOND(TYPE)>(const char order, const std::vector<sd::LongType>& shape, \
+                                                        const std::vector<GET_SECOND(TYPE)>& data, sd::LaunchContext* context); \
+))
+
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_CREATE_F)
+
+
+#undef TMPL_INSTANTIATE_CREATE_F
+
+////////////////////////////////////////////////////////////////////////
+template <>
+SD_LIB_EXPORT NDArray* NDArrayFactory::valueOf(std::vector<LongType>& shape, NDArray* value, const char order,
+                                               LaunchContext* context) {
+  auto result = create_(order, shape, value->dataType(), context);
+  result->assign(value);
+  return result;
+}
+
+template <>
+SD_LIB_EXPORT NDArray* NDArrayFactory::valueOf(std::vector<LongType>& shape, NDArray& value, const char order,
+                                               LaunchContext* context) {
+  auto result = create_(order, shape, value.dataType(), context);
+  result->assign(&value);
+  return result;
+}
+
+template <typename T>
+NDArray* NDArrayFactory::valueOf(std::vector<LongType>& shape,  T value, const char order,
+                                 LaunchContext* context) {
+  auto result = create_(order, shape, DataTypeUtils::fromT<T>());
+  result->assign(value);
+  return result;
+}
+
+// Replace TMPL_INSTANTIATE_VALUEOF
+#define TMPL_INSTANTIATE_VALUEOF(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* \
+    NDArrayFactory::valueOf<GET_SECOND(TYPE)>(std::vector<sd::LongType>& shape,  GET_SECOND(TYPE) value, \
+                                                        const char order, sd::LaunchContext* context); \
+))
+
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_VALUEOF)
+
+
+#undef TMPL_INSTANTIATE_VALUEOF
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::linspace(const T from, const T to, const LongType numElements) {
+  NDArray* result = NDArrayFactory::vector<T>(numElements);
+  // TO DO: linspace should be executed on DEVICE, but only CPU version implemnted!
+  for (LongType e = 0; e < numElements; e++) {
+    T step = (T)e / ((T)numElements - (T)1);
+    result->p<T>(e, (from * ((T)1 - step) + step * to));
+  }
+  result->syncToDevice();
+
+  return result;
+}
+
+#define TMPL_INSTANTIATE_LINSPACE(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::linspace<GET_SECOND(TYPE)>(const GET_SECOND(TYPE) from, const GET_SECOND(TYPE) to, \
+                                                         const sd::LongType numElements); \
+))
+ITERATE_LIST((SD_NUMERIC_TYPES), TMPL_INSTANTIATE_LINSPACE)
+
+
+
+#undef TMPL_INSTANTIATE_LINSPACE
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::vector(LongType length,  T value, LaunchContext* context) {
+  DataBuffer *  buffer =
+      new DataBuffer(length * sizeof(T), DataTypeUtils::fromT<T>(), context->getWorkspace(), true);
+  auto desc = ShapeBuilders::createVectorShapeInfo(DataTypeUtils::fromT<T>(),length);
+  auto constDesc = ConstantShapeHelper::getInstance().bufferForShapeInfo(desc);
+  auto recast = const_cast<LongType*>(constDesc->primary());
+  auto res = new NDArray(buffer, recast, context);
+  if (value == (T)0.0f)
+    res->nullify();
+  else
+    res->assign(value);
+
+  delete[] desc;  // Free allocated shape info
+
+  return res;
+}
+
+#define TMPL_INSTANTIATE_VECTOR(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::vector<GET_SECOND(TYPE)>(sd::LongType length, const GET_SECOND(TYPE) startingValue, \
+                                                       sd::LaunchContext* context); \
+))
+
+ITERATE_LIST((SD_COMMON_TYPES), TMPL_INSTANTIATE_VECTOR)
+
+
+
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray *NDArrayFactory::create(const char order, const std::vector<LongType>& shape, LaunchContext* context) {
+  return create(order, shape, DataTypeUtils::fromT<T>(), context);
+}
+BUILD_SINGLE_TEMPLATE(NDArray *NDArrayFactory::create,
+                      (const char order, const std::vector<sd::LongType>& shape, sd::LaunchContext* context),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::create(const char order, const std::vector<LongType>& shape, DataType dtype,
+                               LaunchContext* context) {
+  if ((int)shape.size() > SD_MAX_RANK)
+    THROW_EXCEPTION("NDArrayFactory::create: rank of NDArray can't exceed 32");
+
+
+  ShapeDescriptor *descriptor = new ShapeDescriptor(dtype, order, shape);
+
+  DataBuffer *  buffer = new DataBuffer(
+      descriptor->arrLength() * DataTypeUtils::sizeOfElement(dtype), dtype, context->getWorkspace());
+
+  NDArray *result = new NDArray(buffer, descriptor, context);
+  delete descriptor;
+  result->nullify();
+
+  return result;
+}
+
+NDArray* NDArrayFactory::create_(DataType dtype, LaunchContext* context) {
+  auto result = create(dtype, context);
+  return result;
+}
+
+template <typename T>
+static NDArray *create(DataType type, const std::vector<LongType>& shape, LaunchContext* context) {
+  auto buffer = new DataBuffer(DataTypeUtils::sizeOfElement(type) * shape::prodLong(shape.data(),shape.size()), type, context->getWorkspace());
+  auto desc = ShapeBuilders::createShapeInfo(type,'c',shape);
+  auto cachedDesc = ConstantShapeHelper::getInstance().bufferForShapeInfo(desc);
+  NDArray *result = new NDArray(buffer, cachedDesc->primary(), context);
+  delete[] desc;
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray *NDArrayFactory::create(const std::vector<T>& values, LaunchContext* context) {
+  DataBuffer *  buffer =
+      new DataBuffer(values.size() * sizeof(T), DataTypeUtils::fromT<T>(), context->getWorkspace(), true);
+
+  auto desc = ShapeDescriptor::vectorDescriptor(values.size(), DataTypeUtils::fromT<T>());
+  NDArray *res = new NDArray(buffer, desc, context);
+  memcpyFromVector<T>(res->buffer(), values);
+
+  res->tickWriteHost();
+  res->syncToDevice();
+
+  return res;
+}
+
+#define TMPL_INSTANTIATE_CREATE_G(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray* NDArrayFactory::create<GET_SECOND(TYPE)>(const std::vector<GET_SECOND(TYPE)>& values, sd::LaunchContext* context); \
+))
+
+ITERATE_LIST((SD_NUMERIC_TYPES), TMPL_INSTANTIATE_CREATE_G)
+
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray* NDArrayFactory::empty_(LaunchContext* context) {
+  auto shapeInfo = ShapeBuilders::createScalarShapeInfo(DataTypeUtils::fromT<T>(), context->getWorkspace());
+  ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  auto result = new NDArray(nullptr, shapeInfo, context, false, 0);
+
+  RELEASE(shapeInfo, context->getWorkspace());
+
+  return result;
+}
+BUILD_SINGLE_TEMPLATE(NDArray* NDArrayFactory::empty_, (sd::LaunchContext * context),
+                      SD_COMMON_TYPES);
+
+NDArray* NDArrayFactory::empty_(DataType dataType, LaunchContext* context) {
+  if (context == nullptr) context = LaunchContext ::defaultContext();
+
+  auto shapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, context->getWorkspace());
+  ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  auto result = new NDArray(nullptr, shapeInfo, context, false, 0);
+
+  RELEASE(shapeInfo, context->getWorkspace());
+
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray *NDArrayFactory::empty(LaunchContext* context) {
+  return empty(DataTypeUtils::fromT<T>(), context);
+}
+BUILD_SINGLE_TEMPLATE(NDArray *NDArrayFactory::empty, (sd::LaunchContext * context),
+                      SD_COMMON_TYPES);
+
+////////////////////////////////////////////////////////////////////////
+SD_LIB_EXPORT NDArray* NDArrayFactory::empty(DataType dataType, LaunchContext* context) {
+  auto shapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, context->getWorkspace());
+  ArrayOptions::setPropertyBit(shapeInfo, ARRAY_EMPTY);
+  NDArray *result= new NDArray(nullptr, shapeInfo, context, false, 0);
+
+  RELEASE(shapeInfo, context->getWorkspace());
+
+  return result;
+}
+
+////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::valueOf(std::vector<LongType>& shape, NDArray& value, const char order,
+                                 LaunchContext* context) {
+  auto res = create_(order, shape, value.dataType(), context);
+  res->assign(&value);
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::create_(const char order, std::vector<LongType>& shape, DataType dataType,
+                                 LaunchContext* context) {
+  return new NDArray(order, shape, dataType, context);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+NDArray *NDArrayFactory::create(T* buffer, const char order, const std::initializer_list<LongType>& shape,
+                               LaunchContext* context) {
+  if ((int)shape.size() > SD_MAX_RANK)
+    THROW_EXCEPTION("NDArrayFactory::create: Rank of NDArray can't exceed 32");
+
+  std::vector<LongType> shp(shape);
+  ShapeDescriptor *descriptor = new ShapeDescriptor(DataTypeUtils::fromT<T>(), order, shp);
+
+  DataBuffer *  pBuffer = new DataBuffer(
+      buffer, descriptor->arrLength() * sizeof(T), descriptor->dataType(), false, context->getWorkspace());
+
+  NDArray *result = new NDArray(pBuffer, descriptor, context);
+  delete descriptor;
+  return result;
+}
+
+// Replace TMPL_INSTANTIATE_CREATE_H
+#define TMPL_INSTANTIATE_CREATE_H(TYPE) \
+EVAL(SD_IF_SINGLE_ALIAS_COMPILED_DECL( \
+    GET_FIRST(TYPE), \
+    template SD_LIB_EXPORT NDArray *NDArrayFactory::create<GET_SECOND(TYPE)>(GET_SECOND(TYPE)* buffer, const char order, \
+                                                      const std::initializer_list<sd::LongType>& shape,  \
+                                                      sd::LaunchContext* context); \
+))
+ITERATE_LIST((SD_COMMON_TYPES),TMPL_INSTANTIATE_CREATE_H)
+
+
+
+
+#if defined(HAS_UTF16)
+/////////////////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(const char16_t* u16string, DataType dtype, LaunchContext* context) {
+  return new NDArray(u16string, dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const char16_t* u16string, DataType dtype, LaunchContext* context) {
+  return string_(std::u16string(u16string), dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const std::u16string& u16string, DataType dtype, LaunchContext* context) {
+  auto res = new NDArray(u16string, dtype, context);
+  return res;
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string(const std::u16string& u16string, DataType dtype, LaunchContext* context) {
+  return new NDArray(u16string, dtype, context);
+}
+#endif
+#if defined(HAS_UTF32)
+#if defined(HAS_UTF32)
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(const char32_t* u32string, DataType dtype, LaunchContext* context) {
+  return new NDArray(u32string, dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const char32_t* u32string, DataType dtype, LaunchContext* context) {
+  return string_(std::u32string(u32string), dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const std::u32string& u32string, DataType dtype, LaunchContext* context) {
+  auto res = new NDArray(u32string, dtype, context);
+  return res;
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray * NDArrayFactory::string(const std::u32string& u32string, DataType dtype, LaunchContext* context) {
+  return new NDArray(u32string, dtype, context);
+}
+#endif
+#endif
+#if defined(HAS_UTF8)
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(const char* str, DataType dtype, LaunchContext* context) {
+  return new NDArray(str, dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const char* str, DataType dtype, LaunchContext* context) {
+  return string_(std::string(str), dtype, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(const std::string& str, DataType dtype, LaunchContext* context) {
+  auto res = new NDArray(str, dtype, context);
+  return res;
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(const std::string& str, DataType dtype, LaunchContext* context) {
+  return new NDArray(str, dtype, context);
+}
+#endif
+#if defined(HAS_UTF8)
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<const char*>& strings,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, strings, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<const char*>& strings,
+                                 DataType dataType, LaunchContext* context) {
+  std::vector<std::string> vec(strings.size());
+  int cnt = 0;
+  for (auto s : strings) vec[cnt++] = std::string(s);
+
+  return string_(shape, vec, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<std::string>& string,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, string, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<std::string>& string,
+                                 DataType dataType, LaunchContext* context) {
+  auto res = new NDArray(shape, string, dataType, context);
+  return res;
+}
+#endif
+#if defined(HAS_UTF16)
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape,
+                               const std::initializer_list<const char16_t*>& strings, DataType dataType,
+                               LaunchContext* context) {
+  return new NDArray(shape, std::vector<const char16_t*>(strings), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<const char16_t*>& strings,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, strings, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string(std::vector<LongType>& shape,
+                               const std::initializer_list<std::u16string>& string,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, std::vector<std::u16string>(string), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape,
+                                 const std::initializer_list<const char16_t*>& strings, DataType dataType,
+                                 LaunchContext* context) {
+  return string_(shape, std::vector<const char16_t*>(strings), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<const char16_t*>& strings,
+                                 DataType dataType, LaunchContext* context) {
+  std::vector<std::u16string> vec(strings.size());
+  int cnt = 0;
+  for (auto s : strings) vec[cnt++] = std::u16string(s);
+
+  return string_(shape, vec, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape,
+                                 const std::initializer_list<std::u16string>& string, DataType dataType,
+                                 LaunchContext* context) {
+  return string_(shape, std::vector<std::u16string>(string), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<std::u16string>& string,
+                                 DataType dataType, LaunchContext* context) {
+  auto res = new NDArray(shape, string, dataType, context);
+  return res;
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray * NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<std::u16string>& string,
+                               DataType dtype, LaunchContext* context) {
+  return new NDArray(shape, string, dtype, context);
+}
+#endif
+#if defined(HAS_UTF32)
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape,
+                               const std::initializer_list<const char32_t*>& strings, DataType dataType,
+                               LaunchContext* context) {
+  return new NDArray(shape, std::vector<const char32_t*>(strings), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray * NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<const char32_t*>& strings,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, strings, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray *NDArrayFactory::string(std::vector<LongType>& shape,
+                               const std::initializer_list<std::u32string>& string,
+                               DataType dataType, LaunchContext* context) {
+  return new NDArray(shape, std::vector<std::u32string>(string), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape,
+                                 const std::initializer_list<const char32_t*>& strings, DataType dataType,
+                                 LaunchContext* context) {
+  return string_(shape, std::vector<const char32_t*>(strings), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<const char32_t*>& strings,
+                                 DataType dataType, LaunchContext* context) {
+  std::vector<std::u32string> vec(strings.size());
+  int cnt = 0;
+  for (auto s : strings) vec[cnt++] = std::u32string(s);
+  return string_(shape, vec, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape,
+                                 const std::initializer_list<std::u32string>& string, DataType dataType,
+                                 LaunchContext* context) {
+  return string_(shape, std::vector<std::u32string>(string), dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray* NDArrayFactory::string_(std::vector<LongType>& shape, const std::vector<std::u32string>& string,
+                                 DataType dataType, LaunchContext* context) {
+  // Default constructor was deleted to prevent uninitialized NDArrays (nullptr _shapeInfo)
+  return new NDArray(shape, string, dataType, context);
+}
+/////////////////////////////////////////////////////////////////////////
+NDArray * NDArrayFactory::string(std::vector<LongType>& shape, const std::vector<std::u32string>& string,
+                               DataType dtype, LaunchContext* context) {
+  return new NDArray(shape, string, dtype, context);
+}
+#endif
+
+NDArray NDArrayFactory::fromNpyFile(const char* fileName) {
+  auto size = getFileSize(fileName);
+  if (size < 0) THROW_EXCEPTION("File doesn't exit");
+
+  auto pNPY = reinterpret_cast<char*>(numpyFromFile(std::string(fileName)));
+
+  auto nBuffer = reinterpret_cast<void*>(dataPointForNumpy(pNPY));
+  auto shape = reinterpret_cast<LongType*>(shapeBufferForNumpy(pNPY));
+
+  auto length = shape::length(shape);
+  int8_t* buffer = nullptr;
+  memory::Workspace* workspace = nullptr;
+  auto byteLen = length * DataTypeUtils::sizeOfElement(ArrayOptions::dataType(shape));
+
+  ALLOCATE(buffer, workspace, byteLen, int8_t);
+  memcpy(buffer, nBuffer, byteLen);
+
+  free(pNPY);
+
+  return NDArray(buffer, shape, LaunchContext::defaultContext(), true, 0);
+}
+}  // namespace sd

libnd4j/include/array/impl/NDArrayList.cpp

@@ -0,0 +1,278 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/NDArrayList.h>
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/stack.h>
+
+#include <iterator>
+#if NOT_EXCLUDED(OP_stack)
+namespace sd {
+NDArrayList::NDArrayList(int height, bool expandable) {
+  _expandable = expandable;
+  _elements.store(0);
+  _counter.store(0);
+  _id.first = 0;
+  _id.second = 0;
+  _height = height;
+   sd_debug("\nCreating NDArrayList\n","");
+}
+
+NDArrayList::~NDArrayList() {
+  sd_debug("\nDeleting NDArrayList: [%i]\n", _chunks.size());
+  for (auto const& v : _chunks) delete v.second;
+
+  _chunks.clear();
+}
+
+NDArray* NDArrayList::read(int idx) { return new NDArray(readRaw(idx)->dup()); }
+
+sd::DataType NDArrayList::dataType() { return _dtype; }
+
+NDArray* NDArrayList::readRaw(int idx) {
+  if (_chunks.count(idx) < 1) {
+    sd_debug("Non-existent chunk requested: [%i]\n", idx);
+    THROW_EXCEPTION("Bad index");
+  }
+
+  return _chunks[idx];
+}
+
+
+NDArray* NDArrayList::remove(int idx) {
+  if(!isWritten(idx)) {
+    sd_debug("Non-existent chunk requested: [%i]\n", idx);
+    THROW_EXCEPTION("Bad index");
+  }
+
+  delete _chunks[idx];
+
+  _elements--;
+  return new NDArray(readRaw(idx)->dup());
+}
+
+
+sd::Status NDArrayList::write(int idx, NDArray* array) {
+  if (_chunks.count(idx) == 0)
+    _elements++;
+  else {
+    delete _chunks[idx];
+  }
+
+  // we store reference shape on first write
+  if (_chunks.empty()) {
+    _dtype = array->dataType();
+
+    if (_shape.empty()) {
+      // adding leading 1 to shape
+      _shape.emplace_back(1);
+      for (int e = 0; e < array->rankOf(); e++) _shape.emplace_back(array->sizeAt(e));
+    } else {
+      // if shape is inferred (say, from split_list)
+      if (static_cast<size_t>(array->rankOf()) == _shape.size()) {
+        // skipping first dim
+        for (size_t e = 1; e < _shape.size(); e++) {
+          if (_shape[e] != array->sizeAt(e))
+            return Logger::logStatusMsg(Status::BAD_INPUT,
+                                        "NDArrayList: all arrays must have same size along inner dimensions");
+        }
+      } else if (static_cast<size_t>(array->rankOf()) == _shape.size() - 1) {
+        // case like 2d _shape, and 1D rows
+        for (size_t e = 1; e < _shape.size(); e++)
+          if (_shape[e] != array->sizeAt(e - 1))
+            return Logger::logStatusMsg(Status::BAD_INPUT,
+                                        "NDArrayList: all arrays must have same size along inner dimensions");
+      } else
+        return Logger::logStatusMsg(Status::BAD_INPUT,
+                                    "NDArrayList: all arrays must have same size along inner dimensions");
+    }
+  } else {
+    if (array->dataType() != _dtype)
+      return Logger::logStatusMsg(Status::BAD_INPUT, "NDArrayList: all arrays must have same data type");
+
+    // if shape is inferred (say, from split_list)
+    if (static_cast<size_t>(array->rankOf()) == _shape.size()) {
+      // skipping first dim
+      for (size_t e = 1; e < _shape.size(); e++) {
+        if (_shape[e] != array->sizeAt(e))
+          return Logger::logStatusMsg(Status::BAD_INPUT,
+                                      "NDArrayList: all arrays must have same size along inner dimensions");
+      }
+    } else if (static_cast<size_t>(array->rankOf()) == _shape.size() - 1) {
+      // case like 2d _shape, and 1D rows
+      for (size_t e = 1; e < _shape.size(); e++)
+        if (_shape[e] != array->sizeAt(e - 1))
+          return Logger::logStatusMsg(Status::BAD_INPUT,
+                                      "NDArrayList: all arrays must have same size along inner dimensions");
+    } else
+      return Logger::logStatusMsg(Status::BAD_INPUT,
+                                  "NDArrayList: all arrays must have same size along inner dimensions");
+  }
+
+
+  // storing reference
+  _chunks[idx] = array;
+
+  return Status::OK;
+}
+
+std::vector<sd::LongType>& NDArrayList::shape() { return _shape; }
+
+int NDArrayList::counter() { return _counter++; }
+
+void NDArrayList::unstack(NDArray* array, LongType axis) {
+  _axis = axis;
+  std::vector<sd::LongType> args({axis});
+  auto newAxis = ShapeUtils::evalDimsToExclude(array->rankOf(),1, args.data());
+  auto result = array->allTensorsAlongDimension(*newAxis);
+  for (sd::LongType e = 0; e < result.size(); e++) {
+    auto chunk = result.at(e);
+    write(e, new NDArray(chunk->dup(array->ordering())));
+  }
+
+  delete newAxis;
+}
+
+NDArray* NDArrayList::stack() {
+  int numElements = _elements.load();
+  if(numElements < 1) {
+    return  new NDArray(NDArrayFactory::empty<double>());
+
+  }
+  std::vector<NDArray*> inputs(numElements);
+  for (int e = 0; e < numElements; e++) {
+    if(!_chunks[e]->isEmpty())
+      _chunks[e]->syncToDevice();
+    inputs[e] = _chunks[e];
+  }
+
+  if(inputs[0] == nullptr) {
+    THROW_EXCEPTION("First input element was a null ptr!");
+  }
+
+  auto inShapeInfo = inputs[0]->shapeInfo();
+  int rank = shape::rank(inShapeInfo);
+  NDArray* array = nullptr;
+
+  if (shape::isEmptyConst(inShapeInfo)) {
+    switch (rank) {
+      case 0: {
+        if (numElements == 1) {
+          std::vector<sd::LongType> shape = {0};
+          array = new NDArray(inputs[0]->ordering(), shape, ArrayOptions::dataType(inShapeInfo), inputs[0]->getContext());
+        } else {
+          std::vector<sd::LongType> shape =  {(sd::LongType)numElements, 0};
+          array = new NDArray('c', shape, ArrayOptions::dataType(inShapeInfo),
+                              inputs[0]->getContext());
+        }
+      }
+    }
+  } else {
+
+    std::vector<sd::LongType> outShape(inShapeInfo + 1, inShapeInfo + 1 + rank);
+    outShape.insert(outShape.begin(), (sd::LongType)numElements);
+    array =
+        new NDArray(shape::order(inShapeInfo), outShape, ArrayOptions::dataType(inShapeInfo), inputs[0]->getContext());
+  }
+
+  ops::helpers::stack(inputs[0]->getContext(), inputs, *array, 0);
+
+  return array;
+}
+
+std::pair<int, int>& NDArrayList::id() { return _id; }
+
+std::string& NDArrayList::name() { return _name; }
+
+sd::LaunchContext* NDArrayList::context() { return _context; }
+
+int NDArrayList::elements() { return _elements.load(); }
+
+int NDArrayList::height() {
+  return (int)_chunks.size();
+}
+
+bool NDArrayList::isWritten(int index) {
+  if (_chunks.count(index) > 0)
+    return true;
+  else
+    return false;
+}
+
+NDArray* NDArrayList::pick(std::initializer_list<LongType> indices) {
+  std::vector<LongType> idcs(indices);
+  return pick(idcs);
+}
+
+NDArray* NDArrayList::pick(std::vector<LongType>& indices) {
+  std::vector<sd::LongType> shape(_shape);
+
+  shape[_axis] = indices.size();
+  // do we have to enforce C order here?
+  auto array = new NDArray('c', shape, _chunks[0]->dataType(), _context);
+  const sd::LongType *axis2 = const_cast<sd::LongType *>(&_axis);
+  std::vector<sd::LongType> *axis = ShapeUtils::evalDimsToExclude(shape.size(),1, axis2);
+  auto tads = array->allTensorsAlongDimension(*axis);
+  int indicesSize = indices.size();
+
+  if (tads.size() != indicesSize) THROW_EXCEPTION("Number of TADs should match number of indices");
+
+  for (int e = 0; e < indicesSize; e++) tads.at(e)->assign(_chunks[indices[e]]);
+
+  delete axis;
+  return array;
+}
+
+NDArrayList* NDArrayList::clone() {
+  auto list = new NDArrayList(_height, _expandable);
+  list->_axis = _axis;
+  list->_id.first = _id.first;
+  list->_id.second = _id.second;
+  list->_name = _name;
+  list->_elements.store(_elements.load());
+
+  for (auto const& v : _chunks) {
+    list->_chunks[v.first] = new NDArray(v.second->dup());
+  }
+
+  return list;
+}
+
+bool NDArrayList::equals(NDArrayList& other) {
+  if (_axis != other._axis) return false;
+
+  if (_chunks.size() != other._chunks.size()) return false;
+
+  for (auto const& v : _chunks) {
+    if (other._chunks.count(v.first) == 0) return false;
+
+    auto arrThis = _chunks[v.first];
+    auto arrThat = other._chunks[v.first];
+
+    if (!arrThis->equalsTo(arrThat)) return false;
+  }
+
+  return true;
+}
+}  // namespace sd
+#endif

libnd4j/include/array/impl/PointerDeallocator.cpp

@@ -0,0 +1,32 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/PointerDeallocator.h>
+
+namespace sd {
+
+void PointerDeallocator::release(void *ptr) {
+  // noop
+}
+
+}  // namespace sd

libnd4j/include/array/impl/PointerWrapper.cpp

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/PointerWrapper.h>
+
+namespace sd {
+PointerWrapper::PointerWrapper(void *ptr, const std::shared_ptr<PointerDeallocator> &deallocator)
+    : _pointer(ptr), _deallocator(deallocator) {
+  //
+}
+
+PointerWrapper::~PointerWrapper() {
+  if (_deallocator.get() != nullptr) _deallocator->release(_pointer);
+}
+
+void *PointerWrapper::pointer() const { return _pointer; }
+
+}  // namespace sd

libnd4j/include/array/impl/PrimaryPointerDeallocator.cpp

@@ -0,0 +1,44 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 <array/PrimaryPointerDeallocator.h>
+
+#if defined(SD_GCC_FUNCTRACE)
+#include <array/ShapeCacheLifecycleTracker.h>
+#endif
+
+namespace sd {
+
+void PrimaryPointerDeallocator::release(void *ptr) {
+#if defined(SD_GCC_FUNCTRACE)
+  // Track shape cache deallocation before freeing
+  sd::array::ShapeCacheLifecycleTracker::getInstance().recordDeallocation(
+      reinterpret_cast<LongType*>(ptr));
+#endif
+
+  // Root cause of SIGSEGV crashes: shape buffers are allocated with new[] but were being
+  // freed with delete (single-object), causing heap corruption and undefined behavior.
+  delete[] reinterpret_cast<int8_t *>(ptr);
+}
+
+}  // namespace sd

libnd4j/include/array/impl/ResultSet.cpp

@@ -0,0 +1,146 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/ResultSet.h>
+#include <graph/FlatUtils.h>
+
+namespace sd {
+ResultSet::ResultSet() {
+  //
+}
+
+ResultSet::ResultSet(const ::graph::FlatResult* result) {
+  for (size_t e = 0; e < result->variables()->size(); e++) {
+    auto var = result->variables()->Get(e);
+
+    NDArray* array;
+
+    if (var->ndarray() != nullptr) {
+      array = graph::FlatUtils::fromFlatArray(var->ndarray());
+    } else if (var->shape() != nullptr) {
+      std::vector<LongType> shapeInfo;
+      for (size_t i = 0; i < var->shape()->size(); i++) {
+        shapeInfo.emplace_back(var->shape()->Get(i));
+      }
+
+      // we just create empty array here
+      int s0 = shapeInfo.at(0);
+
+      std::vector<LongType> shape;
+      for (int i = 0; i < s0; i++) {
+        shape.emplace_back(shapeInfo.at(i + 1));
+      }
+
+      array =
+          new NDArray((char)shapeInfo.at(shapeInfo.size() - 1), shape, DataTypeUtils::fromFlatDataType(var->dtype()));
+    } else {
+      sd_printf("Either shape or NDArray should be defined in FlatResult variable\n", "");
+      THROW_EXCEPTION("Empty variable");
+    }
+
+    _content.push_back(array);
+  }
+}
+
+ResultSet::ResultSet(const ResultSet& other) noexcept {
+  for (const auto v : other._content) _content.emplace_back(v);
+
+  _status = other._status;
+  _removable = false;
+}
+
+////////////////////////////////////////////////////////////////////////
+// move constructor
+ResultSet::ResultSet(ResultSet&& other) noexcept {
+  _content = std::move(other._content);
+  _status = other._status;
+  _removable = other._removable;
+  other._removable = false;
+}
+
+////////////////////////////////////////////////////////////////////////
+// move assignment operator
+ResultSet& ResultSet::operator=(ResultSet&& other) noexcept {
+  if (this == &other) return *this;
+
+  delContent();
+
+  _content = std::move(other._content);
+
+  _status = other._status;
+  _removable = other._removable;
+  other._removable = false;
+
+  return *this;
+}
+
+ResultSet& ResultSet::operator=(const ResultSet& other) noexcept {
+  if (this == &other) return *this;
+
+  delContent();
+
+  for (const auto v : other._content) _content.push_back(v);
+
+  _status = other._status;
+  _removable = false;
+
+  return *this;
+}
+
+void ResultSet::delContent() {
+  if (_removable) {
+    std::vector<DataBuffer *> deleted;
+    for (auto v : _content) {
+      auto buffer = v->dataBuffer();
+      deleted.push_back(buffer);
+      if (!v->isView() && v->shapeInfo() != nullptr) {
+        delete v;
+      }
+    }
+  }
+}
+
+ResultSet::~ResultSet() { delContent(); }
+
+void ResultSet::printIndexedBuffers() {
+  for (size_t e = 0; e < _content.size(); e++) {
+    auto array = _content.at(e);
+    auto strVal = "Array e: " + std::to_string(e) + " is: ";
+    array->printIndexedBuffer(strVal.c_str());
+  }
+
+}
+void ResultSet::setNonRemovable() { _removable = false; }
+
+int ResultSet::size() { return (int)_content.size(); }
+
+NDArray* ResultSet::at(const unsigned long idx) const { return _content.at(idx); }
+
+NDArray* ResultSet::operator[](const unsigned long idx) const { return _content[idx]; }
+
+void ResultSet::push_back(NDArray* array) { _content.emplace_back(array); }
+
+Status ResultSet::status() { return _status; }
+
+void ResultSet::setStatus(Status status) { _status = status; }
+
+void ResultSet::purge() { _content.clear(); }
+}  // namespace sd

libnd4j/include/array/impl/ShapeDescriptor.cpp

@@ -0,0 +1,873 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//  @author AbdelRauf
+
+#include <array/ShapeDescriptor.h>
+#include <helpers/ShapeBuilders.h>
+#include <helpers/shape.h>
+#include <helpers/ModularHasher.h>
+#include "helpers/ShapeUtils.h"
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////
+// equal to operator
+bool ShapeDescriptor::operator==(const ShapeDescriptor& other) const {
+  // First check scalar values to fail fast
+  if (_rank != other._rank ||
+      _dataType != other._dataType ||
+      _order != other._order ||
+      _extraProperties != other._extraProperties ||
+      _offset != other._offset) {
+    return false;
+  }
+
+  // Handle null pointers
+  if (_shape_strides == nullptr && other._shape_strides == nullptr) {
+    return true;
+  }
+  if (_shape_strides == nullptr || other._shape_strides == nullptr) {
+    return false;
+  }
+
+  // Compare shape and strides
+  const int total_length = (_rank < 1 ? 1 : _rank) * 2;
+  return memcmp(_shape_strides, other._shape_strides, total_length * sizeof(LongType)) == 0;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// less than operator
+bool ShapeDescriptor::operator<(const ShapeDescriptor &other) const {
+  return std::tie(_extraProperties, _rank, _dataType, _order, _shape_strides) <
+         std::tie(other._extraProperties, other._rank, other._dataType,  other._order, other._shape_strides);
+}
+
+LongType *ShapeDescriptor::toShapeInfo() const {
+  // for empty array use original
+  return ShapeBuilders::createShapeInfoFrom(const_cast<ShapeDescriptor *>(this));
+}
+
+ShapeDescriptor::~ShapeDescriptor() {
+  // no-op
+  if(_shape_strides != nullptr && this->ownsShapeStrides) {
+    delete[] _shape_strides;
+    _shape_strides = nullptr;
+  }
+
+}
+
+LongType ShapeDescriptor::offset() {
+  return _offset;
+}
+
+ShapeDescriptor::ShapeDescriptor(const DataType type, const char order, const LongType *shape, const LongType rank)
+    : _rank(rank), _order(order), _dataType(type) {
+  int rank2 = rank < 1 ? 1 : rank;
+  _shape_strides = new LongType[2 * rank2];
+  this->ownsShapeStrides = true;
+  if(order != 'c' && order != 'f') {
+    std::string errorMessage;
+    errorMessage += "Invalid ordering from shape buffer";
+    errorMessage += std::to_string(order);
+    delete[] _shape_strides;
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    delete[] _shape_strides;
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+  auto _shape = _shape_strides;
+  for (int i = 0; i < rank2; i++) {
+    _shape[i] = shape[i];
+  }
+
+  _extraProperties = ArrayOptions::flagForDataType(type);
+
+  fillStrides();
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+
+
+}
+
+ShapeDescriptor::ShapeDescriptor(const DataType type, const char order, const LongType *shape,
+                                 const LongType *strides, const LongType rank, LongType extras = -1) {
+  if(shape == nullptr)
+    THROW_EXCEPTION("ShapeDescriptor constructor: Shape can not be null!");
+  _shape_strides = nullptr;
+  ownsShapeStrides = false;
+  if(type  == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  ownsShapeStrides = true;
+  //note this used to operate directly on the vector buffer
+  //it now does manual copies with more checks.
+  //this is to handle the 0 length case.
+  if(rank < 1) {
+    _shape_strides = new LongType[2 * rank];
+    _dataType = type;
+    _order = order;
+    _rank = rank;
+    _extraProperties = extras;
+  } else {
+    _shape_strides = new LongType [2 * rank];
+    _dataType = type;
+    _order = order;
+    _rank = rank;
+    _extraProperties = extras;
+    auto _shape = _shape_strides;
+    auto _strides = _shape_strides + rank;
+    for (int e = 0; e < rank; e++) {
+      _shape[e] = shape[e];
+      if(rank > 1 && shape[e] == 0 && !ArrayOptions::hasPropertyBitSet(_extraProperties, ARRAY_EMPTY)) {
+        _extraProperties = ArrayOptions::setPropertyBitForFlagsValue(_extraProperties, ARRAY_EMPTY);
+      }
+      if(strides != nullptr)
+        _strides[e] = strides[e];
+    }
+
+    if(strides == nullptr)
+      fillStrides();
+  }
+
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+
+ShapeDescriptor::ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape)
+    : _order(order), _dataType(type) {
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+  _rank = shape.size();
+  _extraProperties = ArrayOptions::defaultFlag();
+  _extraProperties = ArrayOptions::setDataTypeValue(_extraProperties, type);
+  int rank2 = shape.size() < 1 ? 1 : shape.size();
+  _shape_strides = new LongType [2 * rank2];
+  this->ownsShapeStrides = true;
+  if(_rank > 0) {
+    auto _shape = _shape_strides;
+    for (int i = 0; i < _rank; i++) {
+      _shape[i] = shape[i];
+      if(shape[i] == 0 && !ArrayOptions::hasPropertyBitSet(_extraProperties, ARRAY_EMPTY)) {
+        _extraProperties = ArrayOptions::setPropertyBitForFlagsValue(_extraProperties, ARRAY_EMPTY);
+      }
+    }
+    fillStrides();
+  }
+
+  _order = order;
+  if(_order != 'c' && _order != 'f') {
+    std::string errorMessage;
+    errorMessage += "Invalid ordering from shape buffer";
+    errorMessage += std::to_string(_order);
+    delete[] _shape_strides;
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    delete[] _shape_strides;
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+
+
+//////////////////////////////////////////////////////////////////////////
+ShapeDescriptor::ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape,
+                                 const std::vector<LongType> &strides, const LongType ews)
+    : ShapeDescriptor(type, order, shape, strides) {
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+ShapeDescriptor::ShapeDescriptor(const DataType type, const LongType length)
+    : _rank(1),  _order('c'), _dataType(type), _extraProperties(0) {
+  _shape_strides = new LongType [2];
+  _shape_strides[0] = length;
+  _shape_strides[1] = 1;  //{shape, stride}
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+ShapeDescriptor::ShapeDescriptor(const LongType *shapeInfo, bool validateDataType, bool overrideStrides) {
+  if(shapeInfo == nullptr) {
+    THROW_EXCEPTION("ShapeDescriptor constructor: Shape info cannot be null!");
+  }
+
+
+  sd::LongType rankVal = shape::rank(shapeInfo);
+  if(rankVal < 0 || rankVal > SD_MAX_RANK) {
+    std::string errorMessage;
+    errorMessage += "Shape descriptor created with invalid rank: ";
+    errorMessage += std::to_string(rankVal);
+    errorMessage += ". Valid range is 0 to ";
+    errorMessage += std::to_string(SD_MAX_RANK);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  if(rankVal == 0) {
+    //detect when the shape buffer values are unset.
+    auto len = shape::shapeInfoLength(rankVal);
+    //min number of values in a shape info buffer
+    bool allZero = true;
+    for(int i = 0; i < len; i++) {
+      if(shapeInfo[i] != 0) {
+        allZero = false;
+        break;
+      }
+    }
+
+    if(allZero) {
+      THROW_EXCEPTION("Found shape buffer with all zero values. Values likely unset.");
+    }
+  }
+
+  _shape_strides = nullptr;
+  _order = shape::order(shapeInfo);
+  this->ownsShapeStrides = true;
+  if(_order != 'c' && _order != 'f') {
+    std::string errorMessage;
+    errorMessage += "Invalid ordering from shape buffer";
+    errorMessage += std::to_string(_order);
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+
+  _rank = static_cast<sd::LongType >(rankVal);
+  _extraProperties = shape::extra(shapeInfo);
+  _dataType = ArrayOptions::dataType(shapeInfo);
+
+  if(_rank > 0 && shape::isEmptyConst(shapeInfo)) {
+    _shape_strides = new LongType[2 * _rank];
+    auto _strides = _shape_strides + _rank;
+    auto shapePtr = shape::shapeOf(shapeInfo);
+    for (LongType e = 0; e < _rank; e++) {
+      _shape_strides[e] = shapePtr[e];
+      _strides[e] = 0;
+    }
+
+  }
+
+  else if (_rank > 0 && !shape::isEmptyConst(shapeInfo)) {
+    _shape_strides = new LongType[2 * rankVal];
+    auto _strides = _shape_strides + _rank;
+    auto shapePtr = shape::shapeOf(shapeInfo);
+    auto stridePtr = shape::stride(shapeInfo);
+    if(overrideStrides) {
+      LongType  *stridesNew = shape::order(shapeInfo) == 'c' ? shape::calcStrides(shapePtr, rankVal) : shape::calcStridesFortran(shapePtr, rankVal);
+      for (LongType e = 0; e < _rank; e++) {
+        _shape_strides[e] = shapePtr[e];
+        _shape_strides[e + _rank] = stridesNew[e];
+      }
+
+      delete[] stridesNew;
+
+    } else {
+      for (LongType e = 0; e < _rank; e++) {
+        _shape_strides[e] = shapePtr[e];
+        _shape_strides[e + _rank] = stridePtr[e];
+      }
+    }
+
+
+
+    //validate construction of the shape descriptor. This is to prevent flag regressions when modifying
+    //_extraProperties.
+    //ensure that we only validate this for array size > 1
+    if(!ArrayOptions::hasPropertyBitSet(_extraProperties, ARRAY_EMPTY) && this->arrLength() > 1) {
+      for(int i = 0; i < _rank; i++) {
+        if(_strides[i] == 0 && shapePtr[i] != 1) {
+          std::string errorMessage;
+          errorMessage += "Shape descriptor:";
+          errorMessage += toString();
+          errorMessage += "Array set as  not empty but stride is not 0. Index is ";
+          errorMessage += std::to_string(i);
+          errorMessage += " Stride is ";
+          errorMessage += std::to_string(_strides[i]);
+          //append the full _shape_strides data
+          errorMessage += " _shape_strides is ";
+          for(int j = 0; j < _rank * 2; j++) {
+            errorMessage += std::to_string(_shape_strides[j]);
+            if(j < _rank * 2 - 1) {
+              errorMessage += ", ";
+            }
+          }
+
+          THROW_EXCEPTION(errorMessage.c_str());
+        }
+      }
+    } else if(this->arrLength() > 1) {
+      for(int i = 0; i < _rank; i++) {
+        if(_strides[i] != 0) {
+          std::string errorMessage;
+          errorMessage += "Array set as not empty but stride is 0. Index is";
+          errorMessage += std::to_string(i);
+          THROW_EXCEPTION(errorMessage.c_str());
+        }
+      }
+    }
+
+  } else if(!shape::isEmptyConst(shapeInfo)) {  // Handle scalar case
+    _shape_strides = new LongType [2]; // Since we're setting shape and stride
+    _shape_strides[0] = 0;    // Shape for scalar
+    _shape_strides[1] = 1;    // Stride for scalar
+  } else {
+    _shape_strides = new LongType[2];
+    _shape_strides[0] = 0;
+    _shape_strides[1] = 0;
+  }
+  _order = shape::order(shapeInfo);
+  _dataType = ArrayOptions::dataType(shapeInfo);
+  if(validateDataType && _dataType  == UNKNOWN) {
+    std::string errorMessage;
+    errorMessage += "Shape descriptor created with invalid data type ";
+    errorMessage += DataTypeUtils::asString(_dataType);
+    errorMessage += " extra properties for data type was ";
+    errorMessage += DataTypeUtils::asString(ArrayOptions::dataTypeValue(_extraProperties));
+    errorMessage += " Underlying extra value  was ";
+    errorMessage += std::to_string(_extraProperties);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+
+}
+
+ShapeDescriptor::ShapeDescriptor(const LongType *shapeInfo, const DataType dtypeOverride, const bool overrideStrides)
+    : ShapeDescriptor(shapeInfo, false, overrideStrides) {
+  if(dtypeOverride == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  _dataType = dtypeOverride;
+  _order = shape::order(shapeInfo);
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+  //data type has already been set by another constructor. We need to update the _extraProperties
+  //to reflect the new data type. This is effectively a cast.
+  _extraProperties = ArrayOptions::propertyWithoutDataTypeValue(_extraProperties);
+  _extraProperties = ArrayOptions::setDataTypeValue(_extraProperties, dtypeOverride);
+
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+ShapeDescriptor::ShapeDescriptor(const LongType *shapeInfo, const LongType *dtypeOverride)
+    : ShapeDescriptor(shapeInfo, ArrayOptions::dataType(dtypeOverride), false) {
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+ShapeDescriptor::ShapeDescriptor(const LongType *shapeInfo, const LongType *dtypeOverride,
+                                 const LongType *orderOverride)
+    : ShapeDescriptor(shapeInfo, ArrayOptions::dataType(dtypeOverride), false) {
+  _order = shape::order(orderOverride);
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+int ShapeDescriptor::rank() const { return _rank; }
+
+
+LongType ShapeDescriptor::arrLength() const {
+  if(_shape_strides== nullptr) {
+    return 0;
+  }
+  // when _ews == 1 allocation length is also array length
+  LongType len = 1;
+  for (int i = 0; i < _rank; i++) len *= _shape_strides[i];
+
+  return len;
+}
+
+void ShapeDescriptor::print() const {
+  printf("ShapeDescriptor: [");
+  for (int i = 0; i < _rank; i++) {
+    printf("%lld", _shape_strides[i]);
+    if (i < _rank - 1) printf(", ");
+  }
+  printf("], [");
+  for (int i = _rank; i < 2 * _rank; i++) {
+    printf("%lld", _shape_strides[i]);
+    if (i < 2 * _rank - 1) printf(", ");
+  }
+  printf("], %c, %s, %lld\n", _order, DataTypeUtils::asString(_dataType).c_str(), _extraProperties);
+}
+
+LongType ShapeDescriptor::allocLength() const {
+  if (_paddedAllocSize > 0) return _paddedAllocSize;
+  auto _shape = _shape_strides;
+  auto _strides = _shape_strides + _rank;
+  int rank2 = _rank < 1 ? 1 : _rank;
+
+  LongType len = 1;
+  if (_rank > 1) {
+    // calculate using max stride
+    int ind = _order == 'c' ? 0 : rank2 - 1;
+    return _shape[ind] * _strides[ind];
+  }
+  for (int i = 0; i < rank2; i++) {
+    len += (_shape[i] - 1) * _strides[i];
+  }
+  return len;
+}
+
+void ShapeDescriptor::collectStoreStackTrace() {
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+  this->storeStackTrace = backward::StackTrace();
+#ifdef __cpp_exceptions
+  try {
+    this->storeStackTrace.load_here(32);
+  } catch (...) {
+    // Stack trace capture failed - storeStackTrace will remain empty (size() == 0)
+  }
+#else
+  this->storeStackTrace.load_here(32);
+#endif
+#endif
+}
+
+LongType ShapeDescriptor::validate() const {
+  auto status = SHAPE_DESC_OK;
+  bool is_continous = true;
+  //exclude scalars on purpose here
+  if (_rank > 0  || _rank > SD_MAX_RANK) status |= SHAPE_DESC_INCORRECT_RANK;
+  auto _shape = _shape_strides;
+  auto _strides = _shape_strides + _rank;
+  if(_order != 'c' && _order != 'f') {
+    THROW_EXCEPTION("Invalid ordering from shape buffer");
+  }
+
+  bool hasZero = false;
+  for (int i = 0; i < _rank; i++) {
+    if (_shape[i] == 0) {
+      hasZero = true;
+      break;
+    }
+  }
+  //this check isn't correct for vectors
+  if (_rank > 0 && !shape::isVector(_shape_strides,2) && !hasZero) {
+    if (_order == 'c') {
+      for (int j = _rank - 2; j >= 0; j--) {
+        LongType currentStride = _strides[j];
+        LongType allowedStride = _strides[j + 1] * _shape[j + 1];
+        if (currentStride < allowedStride) {
+          status = status | SHAPE_DESC_INCORRECT_STRIDES;
+          break;
+        }
+        is_continous = is_continous & (currentStride == allowedStride);
+      }
+    } else {
+      for (int j = 1; j < _rank; j++) {
+        LongType currentStride = _strides[j];
+        LongType allowedStride = _strides[j - 1] * _shape[j - 1];
+        if (currentStride < allowedStride) {
+          status = status | SHAPE_DESC_INCORRECT_STRIDES;
+          break;
+        }
+        is_continous = is_continous & (currentStride == allowedStride);
+      }
+    }
+
+    int index = (_order == 'c') ? _rank - 1 : 0;
+    auto correctEws = is_continous ? _strides[index] : 0;
+  }
+
+  if(isEmpty()) {
+    for(int i = 0; i < _rank; i++) {
+      if(_strides[i] != 0) {
+        std::string errorMessage;
+        errorMessage += "Array set as empty but stride is not 0. Index is ";
+        errorMessage += std::to_string(i);
+        errorMessage += " Stride is ";
+        errorMessage += std::to_string(_strides[i]);
+        THROW_EXCEPTION(errorMessage.c_str());
+        break;
+      }
+    }
+  }
+
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  }
+
+  return status;
+}
+
+char ShapeDescriptor::order() const { return _order; }
+
+DataType ShapeDescriptor::dataType() const {
+  if(!DataTypeUtils::validDataType(_dataType)) {
+    std::string errorMessage;
+    errorMessage += "Shape descriptor created with invalid data type";
+    errorMessage += DataTypeUtils::asString(_dataType);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+  return _dataType;
+}
+
+bool ShapeDescriptor::isEmpty() const { return (_extraProperties & ARRAY_EMPTY) == ARRAY_EMPTY; }
+bool ShapeDescriptor::isScalar() const { return (!isEmpty() && rank() == 0) || (rank() == 1 && arrLength() == 1); }
+
+sd::LongType * ShapeDescriptor::shape_strides() { return _shape_strides; }
+
+const LongType *ShapeDescriptor::stridesPtr() const {
+  return _shape_strides == nullptr ? nullptr :  _shape_strides + _rank;
+}
+
+ShapeDescriptor::ShapeDescriptor(const ShapeDescriptor &other) {
+  _rank = other._rank;
+  _extraProperties = other._extraProperties;
+  if(other._dataType == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  _dataType = other._dataType;
+  _order = other._order;
+  _shape_strides = other._shape_strides;
+  this->ownsShapeStrides = false;
+  _paddedAllocSize = other._paddedAllocSize;
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+}
+
+//////////////////////////////////////////////////////////////////////////
+ShapeDescriptor::ShapeDescriptor(const DataType type, const char order, const std::vector<LongType> &shape,
+                                 const std::vector<LongType> &strides)
+    : _order(order), _dataType(type) {
+  _rank = shape.size();
+  int rank2 = _rank < 1 ? 1 : _rank;
+
+  _shape_strides = new LongType [2 * rank2];
+  this->ownsShapeStrides = true;
+#if defined(SD_GCC_FUNCTRACE)
+  // - backward-cpp's backtrace() is NOT safe during very early JVM initialization
+#ifdef __cpp_exceptions
+  try {
+    this->st.load_here();
+  } catch (...) {
+    // Stack trace capture failed - st will remain empty (size() == 0)
+  }
+#else
+  this->st.load_here();
+#endif
+#endif
+  auto _shape = _shape_strides;
+  auto _strides = _shape_strides + rank2;
+  if (!shape.empty() && strides.size() != shape.size() ) {
+    for (int i = 0; i < rank2; i++) {
+      _shape[i] = shape[i];
+    }
+    fillStrides();
+  } else {
+    for (int i = 0; i < rank2; i++) {
+      _shape[i] = shape[i];
+      _strides[i] = strides[i];
+    }
+  }
+}
+
+ShapeDescriptor  * ShapeDescriptor::emptyDescriptor(const DataType type) {
+  ShapeDescriptor *descriptor = new ShapeDescriptor();
+  if(type == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  descriptor->_dataType = type;
+  descriptor->_extraProperties = ARRAY_EMPTY | ArrayOptions::flagForDataType(type);
+  descriptor->_rank = 0;
+  descriptor->_order = 'c';
+  descriptor->ownsShapeStrides = true;
+  descriptor->_shape_strides = new LongType [1];
+  descriptor->_shape_strides[0] = 0;
+  return descriptor;
+}
+
+ShapeDescriptor * ShapeDescriptor::scalarDescriptor(const DataType type) {
+  ShapeDescriptor *descriptor = new ShapeDescriptor();
+  if(type  == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+  descriptor->_dataType = type;
+  descriptor->_extraProperties = ArrayOptions::flagForDataType(type);
+  descriptor->_rank = 0;
+  descriptor->_order = 'c';
+  descriptor->ownsShapeStrides = true;
+  descriptor->_shape_strides = new LongType [2];
+  descriptor->_shape_strides[0] = 0;
+  descriptor->_shape_strides[1] = 1;
+  descriptor->_offset = 0;
+  return descriptor;
+}
+
+ShapeDescriptor * ShapeDescriptor::vectorDescriptor(const LongType length, const DataType type) {
+  ShapeDescriptor *descriptor = new ShapeDescriptor();
+  if(type  == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+
+  descriptor->_dataType = type;
+  descriptor->_shape_strides = new LongType [2];
+  descriptor->_shape_strides[0] = length;
+  descriptor->_shape_strides[1] = 0;
+  descriptor->ownsShapeStrides = true;
+  descriptor->_offset = 0;
+  if (length > 0) {
+    descriptor->_shape_strides[1] = 1;
+    descriptor->_extraProperties = ArrayOptions::flagForDataType(type);
+  }
+  else {
+    descriptor->_shape_strides[1] = 0;
+    descriptor->_extraProperties = ARRAY_EMPTY;
+    descriptor->_extraProperties = ArrayOptions::setDataTypeValue(descriptor->_extraProperties, type);
+  }
+
+  descriptor->_order = 'c';
+  descriptor->_rank = 1;
+
+  return descriptor;
+}
+
+ShapeDescriptor  * ShapeDescriptor::paddedBufferDescriptor(const DataType type, const char order,
+                                                           const std::vector<LongType> &shape,
+                                                           const std::vector<LongType> &paddings) {
+  ShapeDescriptor *descriptor = new ShapeDescriptor();
+  if(type  == UNKNOWN)
+    THROW_EXCEPTION("Shape descriptor created with invalid data type");
+
+  descriptor->_dataType = type;
+  descriptor->_order = order;
+  descriptor->_rank = shape.size();
+  descriptor->_extraProperties = ArrayOptions::flagForDataType(type);
+  descriptor->ownsShapeStrides = true;
+  if (descriptor->_rank < 1) {
+    return descriptor;
+  }
+
+  int rank2 = descriptor->_rank < 1 ? 1 : descriptor->_rank;
+
+  descriptor->_shape_strides = new LongType [2 * rank2];
+  auto _shape = descriptor->_shape_strides;
+  auto _strides = descriptor->_shape_strides + rank2;
+  for (size_t i = 0; i < shape.size(); i++) {
+    _shape[i] = shape[i];
+  }
+  // calculate strides with paddings
+  int min_rank = static_cast<size_t>(descriptor->_rank) > paddings.size() ? paddings.size() : rank2;
+  bool is_continous = true;
+  if (order == 'c') {
+    _strides[rank2 - 1] = 1L;
+    for (int j = descriptor->_rank - 2; j >= 0; j--) {
+      LongType pad = (j + 1 < min_rank) ? paddings[j + 1] : 0;
+      _strides[j] = _strides[j + 1] * (_shape[j + 1] + pad);
+      descriptor->_extraProperties = descriptor->_extraProperties | (_shape[j + 1] == 0);
+      if (pad != 0) is_continous = false;
+    }
+    if (!is_continous && descriptor->_rank > 0) {
+      LongType size_pad = paddings.size() > 0 ? paddings[0] : 0;
+      // alloc size should be supplied manually as we dont have place to store it
+      descriptor->_paddedAllocSize = _strides[0] * (_shape[0] + size_pad);
+    }
+  } else {
+    _strides[0] = 1L;
+    for (int j = 1; j < rank2; j++) {
+      LongType pad = (j - 1 < min_rank) ? paddings[j - 1] : 0;
+      _strides[j] = _strides[j - 1] * (_shape[j - 1] + pad);
+      descriptor->_extraProperties = descriptor->_extraProperties | (_shape[j - 1] == 0);
+      if (pad != 0) is_continous = false;
+    }
+    if (!is_continous && descriptor->_rank > 0) {
+      LongType size_pad = paddings.size() >= static_cast<size_t>(descriptor->_rank) ? paddings[descriptor->_rank - 1] : 0;
+      // alloc size should be supplied manually as we dont have place to store it
+      descriptor->_paddedAllocSize = _strides[descriptor->_rank - 1] * (_shape[descriptor->_rank - 1] + size_pad);
+    }
+  }
+
+  if (!is_continous) descriptor->_extraProperties |= ARRAY_HAS_PADDED_BUFFER;
+  return descriptor;
+}
+
+}  // namespace sd
+
+namespace std {
+size_t hash<sd::ShapeDescriptor>::operator()(sd::ShapeDescriptor k) const {
+  // Check cache first
+  if (k._hash_computed) {
+    return k._cached_hash;
+  }
+  // Check cache first
+  if (k._hash_computed) {
+    return k._cached_hash;
+  }
+
+  using namespace sd::helpers::detail;
+
+  // Combine order and dataType into initial hash
+  uint64_t hash = ModularHasher::combine_hashes({
+      static_cast<uint64_t>(k.order()),
+      static_cast<uint64_t>(k.dataType()) << 8
+  });
+
+  // Hash the shape and strides array
+  sd::LongType* shape_strides = k.shape_strides();
+  const int stop = k.rank() * 2;
+
+  if (shape_strides != nullptr) {
+    hash = DataChunkHasher<sd::LongType>::hash_data(
+        shape_strides,
+        stop,
+        hash
+    );
+  }
+
+  // Cache the computed hash
+  k._cached_hash = hash;
+  k._hash_computed = true;
+
+  return hash;
+}
+}  // namespace std

libnd4j/include/array/impl/ShapeList.cpp

@@ -0,0 +1,91 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <array/ShapeList.h>
+
+namespace sd {
+
+
+ShapeList::ShapeList(LongType* shape) {
+  if (shape != nullptr) push_back(shape);
+}
+
+ShapeList::~ShapeList() {
+  if (_autoremovable) destroy();
+}
+
+ShapeList::ShapeList(const std::vector< LongType*>& shapes, bool isWorkspace)
+
+{
+  for (size_t i = 0; i < shapes.size(); i++) {
+    push_back(shapes[i]);
+  }
+  _workspace = isWorkspace;
+}
+
+ShapeList::ShapeList(const std::vector< LongType*>& shapes) {
+  _shapes = shapes;
+}
+
+void ShapeList::destroy() {
+
+  if (_destroyed) return;
+  if (!_workspace) {
+    for (int i = 0; i < size(); i++){
+      if (_shapes[i] != nullptr) delete[] _shapes[i];
+    }
+  }
+  _destroyed = true;
+}
+
+int ShapeList::size() const {
+  return (int)_shapes.size();
+}
+
+LongType* ShapeList::at(int idx) {
+  if(idx < 0) {
+    idx += _shapes.size();
+  }
+
+  if (size() <= idx) {
+    std::string errorMessage;
+    errorMessage += "Can't find requested variable by index: ";
+    errorMessage += std::to_string(idx);
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  return _shapes[idx];
+}
+
+void ShapeList::push_back( LongType* shape) {
+  _shapes.push_back(shape);
+}
+
+void ShapeList::detach() {
+  for (int e = 0; e < size(); e++) {
+    _shapes[e] = shape::detachShape(_shapes[e]);
+  }
+
+  _autoremovable = true;
+  _workspace = false;
+}
+}  // namespace sd

libnd4j/include/array/impl/TadDescriptor.cpp

@@ -0,0 +1,79 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 "../TadDescriptor.h"
+
+#include <algorithm>
+#include <helpers/ModularHasher.h>
+namespace sd {
+TadDescriptor::TadDescriptor(const TadDescriptor &other) {
+  _originalShape = other._originalShape;
+  _axis = other._axis;
+  _unitiesInShape = other._unitiesInShape;
+}
+
+TadDescriptor::TadDescriptor(const LongType *originalShape, const LongType *dimensions, const LongType length,
+                             const bool keepUnitiesInShape) {
+
+  _axis.resize(length);
+  for (LongType e = 0; e < length; e++) {
+    _axis[e] = dimensions[e];
+  }
+
+  if (length > 1) std::sort(_axis.begin(), _axis.end());
+
+  _originalShape = const_cast<sd::LongType *>(originalShape);
+  _unitiesInShape = keepUnitiesInShape;
+}
+
+bool TadDescriptor::operator==(const TadDescriptor &other) const {
+  return std::tie(_originalShape, _axis, _unitiesInShape) ==
+         std::tie(other._originalShape, other._axis, other._unitiesInShape);
+}
+
+bool TadDescriptor::operator<(const TadDescriptor &other) const {
+  return std::tie(_originalShape, _axis, _unitiesInShape) <
+         std::tie(other._originalShape, other._axis, other._unitiesInShape);
+}
+
+std::vector<LongType> &TadDescriptor::axis() { return _axis; }
+
+LongType *TadDescriptor::originalShape() { return _originalShape; }
+
+bool TadDescriptor::areUnitiesinShape() const { return _unitiesInShape; }
+}  // namespace sd
+
+namespace std {
+size_t hash<sd::TadDescriptor>::operator()(const sd::TadDescriptor &k) const {
+  using namespace sd::helpers::detail;
+
+  // Start with initial hash from unities flag
+  uint64_t hash = ModularHasher::hash_scalar(k.areUnitiesinShape());
+  // Hash the axis vector
+  auto& axes = const_cast<sd::TadDescriptor&>(k).axis();
+  if (!axes.empty()) {
+    hash = ModularHasher::hash_vector(axes, hash);
+  }
+
+  return hash;
+}
+}  // namespace std

libnd4j/include/array/impl/TadPack.cpp

@@ -0,0 +1,252 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 "../TadPack.h"
+
+#include <helpers/shape.h>
+#include <system/Environment.h>
+#include <sstream>
+
+#if defined(SD_GCC_FUNCTRACE)
+#include <array/TADCacheLifecycleTracker.h>
+#endif
+
+namespace sd {
+TadPack::TadPack( ConstantShapeBuffer *shapes,
+                  ConstantOffsetsBuffer *offets, LongType numTads,
+                 LongType* dimensions, LongType dimLength)
+    : _tadShape(shapes),
+      _tadOffsets(offets) {
+  _numTads = numTads;
+  _dimensionsLength = dimLength;
+  if(dimensions != nullptr) {
+    _dimensions = new LongType[dimLength];
+    for(int i = 0; i < dimLength; i++) {
+      _dimensions[i] = dimensions[i];
+    }
+  }
+
+  computeHash();
+
+#if defined(SD_GCC_FUNCTRACE) && !defined(__JAVACPP_HACK__)
+  // Track TAD cache allocation - only capture stack trace if tracking is enabled
+  // Stack trace capture is expensive, so skip it when tracking is disabled
+  auto& tracker = sd::array::TADCacheLifecycleTracker::getInstance();
+  if (tracker.isEnabled()) {
+    // Capture stack trace for this TadPack allocation
+    _stackTrace = backward::StackTrace();
+    _stackTrace.load_here(32);
+  }
+
+  // Always record allocation for basic counting (cheap operation)
+  size_t shape_info_bytes = 0;
+  size_t offsets_bytes = 0;
+
+  if (_tadShape != nullptr && _tadShape->primary() != nullptr) {
+    shape_info_bytes = shape::shapeInfoByteLength(_tadShape->primary());
+  }
+
+  if (_tadOffsets != nullptr && _tadOffsets->primary() != nullptr) {
+    offsets_bytes = _numTads * sizeof(LongType);
+  }
+
+  std::vector<LongType> dims;
+  if (_dimensions != nullptr) {
+    dims.assign(_dimensions, _dimensions + _dimensionsLength);
+  }
+
+  tracker.recordAllocation(this, _numTads, shape_info_bytes, offsets_bytes, dims);
+#endif
+
+}
+
+TadPack::~TadPack() {
+#if defined(SD_GCC_FUNCTRACE) && !defined(__JAVACPP_HACK__)
+  // Track TAD cache deallocation before cleanup.
+  // Guard must match constructor guard for proper allocation/deallocation pairing.
+  sd::array::TADCacheLifecycleTracker::getInstance().recordDeallocation(this);
+#endif
+
+  // Clean up dimensions array that was allocated in constructor
+  if (_dimensions != nullptr) {
+    delete[] _dimensions;
+    _dimensions = nullptr;
+  }
+
+  // Clean up TAD offsets buffer if we own it
+  // This is owned when transferred via releaseOffsets() from TadCalculator
+  if (_tadOffsets != nullptr) {
+    delete _tadOffsets;
+    _tadOffsets = nullptr;
+  }
+
+  // DON'T delete _tadShape - it comes from ConstantShapeHelper cache
+}
+
+LongType* TadPack::primaryShapeInfo() {
+  if(_tadShape->primary() == nullptr)
+    THROW_EXCEPTION("TadPack::primaryShapeInfo: primary shape info is nullptr!");
+  return _tadShape->primary();
+}
+
+LongType* TadPack::primaryOffsets() {
+  return _tadOffsets->primary();
+}
+
+LongType* TadPack::specialShapeInfo() { return _tadShape->special(); }
+
+LongType* TadPack::specialOffsets() { return _tadOffsets->special(); }
+
+LongType TadPack::numberOfTads() const { return _numTads; }
+
+LongType* TadPack::platformShapeInfo() {
+  return Environment::getInstance().isCPU() ? primaryShapeInfo() : specialShapeInfo();
+}
+
+LongType* TadPack::platformOffsets() {
+  return Environment::getInstance().isCPU() ? primaryOffsets() : specialOffsets();
+}
+
+
+void TadPack::print(const char* msg) {
+  printf("---------------------------\n");
+  printf("%s: ", msg);
+  printf("Offsets:\n");
+  for (int e = 0; e < _numTads; e++) {
+    printf("%lld, ", _tadOffsets->primary()[e]);
+  }
+  printf("\n");
+
+  printf("Dimensions:\n");
+  if (_dimensions == nullptr || _dimensionsLength == 0) {
+    printf("none\n");
+  } else {
+    for (int i = 0; i < _dimensionsLength; i++) {
+      printf("%lld, ", _dimensions[i]);
+    }
+    printf("\n");
+  }
+
+  printf("tad pack shape info:");
+  shape::printShapeInfo(_tadShape->primary());
+  printf("\n");
+  printf("number of tads: %lld\n", _numTads);
+  printf("shape info length: %lld\n", _shapeInfoLength);
+  printf("---------------------------\n");
+}
+
+LongType TadPack::shapeInfoLength() { return shape::shapeInfoLength(primaryShapeInfo()); }
+bool TadPack::operator==( TadPack& other)  {
+  // Compare number of TADs
+  if (_numTads != other._numTads)
+    return false;
+
+  // Compare shape information
+  LongType* thisShape = primaryShapeInfo();
+  LongType* otherShape = other.primaryShapeInfo();
+
+  // Check for null shape info
+  if ((thisShape == nullptr) != (otherShape == nullptr))
+    return false;
+
+  if (thisShape != nullptr) {
+    // Compare rank
+    const int thisRank = shape::rank(thisShape);
+    const int otherRank = shape::rank(otherShape);
+    if (thisRank != otherRank)
+      return false;
+
+    // Compare shape order
+    if (shape::order(thisShape) != shape::order(otherShape))
+      return false;
+
+    // Compare data type
+    if (ArrayOptions::dataType(thisShape) != ArrayOptions::dataType(otherShape))
+      return false;
+
+    // Compare shape dimensions
+    for (int i = 0; i < thisRank; i++) {
+      if (shape::shapeOf(thisShape)[i] != shape::shapeOf(otherShape)[i])
+        return false;
+    }
+
+    // Compare shape strides
+    for (int i = 0; i < thisRank; i++) {
+      if (shape::stride(thisShape)[i] != shape::stride(otherShape)[i])
+        return false;
+    }
+  }
+
+  // Compare dimensions array
+  if ((_dimensions == nullptr) != (other._dimensions == nullptr))
+    return false;
+
+  if (_dimensions != nullptr) {
+    if (_dimensionsLength != other._dimensionsLength)
+      return false;
+
+    for (LongType i = 0; i < _dimensionsLength; i++) {
+      if (_dimensions[i] != other._dimensions[i])
+        return false;
+    }
+  }
+
+  // Compare offsets
+  LongType* thisOffsets = primaryOffsets();
+  LongType* otherOffsets = other.primaryOffsets();
+
+  // Check for null offsets
+  if ((thisOffsets == nullptr) != (otherOffsets == nullptr))
+    return false;
+
+  if (thisOffsets != nullptr) {
+    for (LongType i = 0; i < _numTads; i++) {
+      if (thisOffsets[i] != otherOffsets[i])
+        return false;
+    }
+  }
+
+  return true;
+}
+
+std::string TadPack::getStackTraceAsString() const {
+#if defined(SD_GCC_FUNCTRACE) && !defined(__JAVACPP_HACK__)
+  // Use backward::Printer to format the stack trace into a string
+  std::ostringstream oss;
+  backward::Printer p;
+  p.snippet = false;  // Don't include source code snippets
+  p.address = true;   // Include addresses
+  p.object = false;   // Don't include object file info
+  p.color_mode = backward::ColorMode::never;  // No ANSI colors in string
+
+  // Print to our string stream (we need to cast away const to use _stackTrace)
+  // This is safe since print doesn't modify the StackTrace
+  backward::StackTrace& mutable_st = const_cast<backward::StackTrace&>(_stackTrace);
+  p.print(mutable_st, oss);
+
+  return oss.str();
+#else
+  return "";  // Return empty string when functrace is not enabled
+#endif
+}
+
+
+}  // namespace sd