chore: import upstream snapshot with attribution

libnd4j/include/loops/BroadcastPairwiseConverter.h

@@ -0,0 +1,141 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com), created on 09.04.2019
+//
+
+#ifndef DEV_TESTS_BROADCASTPAIRWISECONVERTER_H
+#define DEV_TESTS_BROADCASTPAIRWISECONVERTER_H
+#include <system/op_boilerplate.h>
+
+#include <stdexcept>
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////
+inline pairwise::Ops fromBroadcastToPairwise(broadcast::Ops op) {
+  switch (op) {
+    case broadcast::Add:
+      return pairwise::Add;
+    case broadcast::Subtract:
+      return pairwise::Subtract;
+    case broadcast::Multiply:
+      return pairwise::Multiply;
+    case broadcast::Divide:
+      return pairwise::Divide;
+    case broadcast::ReverseDivide:
+      return pairwise::ReverseDivide;
+    case broadcast::ReverseSubtract:
+      return pairwise::ReverseSubtract;
+    case broadcast::CopyPws:
+      return pairwise::CopyPws;
+    case broadcast::Pow:
+      return pairwise::Pow;
+    case broadcast::MinPairwise:
+      return pairwise::MinPairwise;
+    case broadcast::MaxPairwise:
+      return pairwise::MaxPairwise;
+    case broadcast::AMinPairwise:
+      return pairwise::AMinPairwise;
+    case broadcast::AMaxPairwise:
+      return pairwise::AMaxPairwise;
+    case broadcast::SquaredSubtract:
+      return pairwise::SquaredSubtract;
+    case broadcast::FloorMod:
+      return pairwise::FloorMod;
+    case broadcast::FloorDiv:
+      return pairwise::FloorDiv;
+    case broadcast::ReverseMod:
+      return pairwise::ReverseMod;
+    case broadcast::SafeDivide:
+      return pairwise::SafeDivide;
+    case broadcast::Mod:
+      return pairwise::Mod;
+    case broadcast::TruncateDiv:
+      return pairwise::TruncateDiv;
+    case broadcast::Atan2:
+      return pairwise::Atan2;
+    case broadcast::LogicalOr:
+      return pairwise::LogicalOr;
+    case broadcast::LogicalXor:
+      return pairwise::LogicalXor;
+    case broadcast::LogicalNot:
+      return pairwise::LogicalNot;
+    case broadcast::LogicalAnd:
+      return pairwise::LogicalAnd;
+    case broadcast::PowDerivative:
+      return pairwise::PowDerivative;
+    default:
+      THROW_EXCEPTION("fromBroadcastToPairwise: Not convertible operation");
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+inline pairwise::BoolOps fromBroadcastToPairwiseBool(broadcast::BoolOps op) {
+  switch (op) {
+    case broadcast::EqualTo:
+      return pairwise::EqualTo;
+    case broadcast::GreaterThan:
+      return pairwise::GreaterThan;
+    case broadcast::LessThan:
+      return pairwise::LessThan;
+    case broadcast::Epsilon:
+      return pairwise::Epsilon;
+    case broadcast::GreaterThanOrEqual:
+      return pairwise::GreaterThanOrEqual;
+    case broadcast::LessThanOrEqual:
+      return pairwise::LessThanOrEqual;
+    case broadcast::NotEqualTo:
+      return pairwise::NotEqualTo;
+    case broadcast::And:
+      return pairwise::And;
+    case broadcast::Or:
+      return pairwise::Or;
+    case broadcast::Xor:
+      return pairwise::Xor;
+    case broadcast::Not:
+      return pairwise::Not;
+    default:
+      THROW_EXCEPTION("fromBroadcastToPairwiseBool: Not convertible operation");
+  }
+}
+
+inline pairwise::IntOps fromBroadcastToPairwiseInt(broadcast::IntOps op) {
+  switch (op) {
+    case broadcast::IntOps::IntAnd:
+      return pairwise::IntOps::IntAnd;
+    case broadcast::IntOps::IntOr:
+      return pairwise::IntOps::IntOr;
+    case broadcast::IntOps::IntXor:
+      return pairwise::IntOps::IntXor;
+    case broadcast::IntOps::ShiftLeft:
+      return pairwise::IntOps::ShiftLeft;
+    case broadcast::IntOps::ShiftRight:
+      return pairwise::IntOps::ShiftRight;
+    case broadcast::IntOps::CyclicShiftLeft:
+      return pairwise::IntOps::CyclicShiftLeft;
+    case broadcast::IntOps::CyclicShiftRight:
+      return pairwise::IntOps::CyclicShiftRight;
+    default:
+      THROW_EXCEPTION("fromBroadcastToPairwiseInt: Not convertible operation");
+  }
+}
+}  // namespace sd
+
+#endif  // DEV_TESTS_BROADCASTPAIRWISECONVERTER_H

libnd4j/include/loops/BroadcastScalarConverter.h

@@ -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
+ */
+#ifndef DEV_TESTS_BROADCASTSCALARCONVERTER_H
+#define DEV_TESTS_BROADCASTSCALARCONVERTER_H
+#include <system/op_boilerplate.h>
+#include <system/op_enums.h>
+
+#include <stdexcept>
+
+namespace sd {
+inline bool isConvertibleToScalar(broadcast::Ops op) {
+  int opNum = (int)op;
+
+  if (opNum <= 17) return true;
+
+  return false;
+}
+
+inline scalar::Ops convertToScalar(broadcast::Ops op) {
+  switch (op) {
+    case broadcast::Add:
+      return scalar::Add;
+    case broadcast::Subtract:
+      return scalar::Subtract;
+    case broadcast::Multiply:
+      return scalar::Multiply;
+    case broadcast::Divide:
+      return scalar::Divide;
+    case broadcast::ReverseDivide:
+      return scalar::ReverseDivide;
+    case broadcast::ReverseSubtract:
+      return scalar::ReverseSubtract;
+    case broadcast::CopyPws:
+      return scalar::CopyPws;
+    case broadcast::Pow:
+      return scalar::Pow;
+    case broadcast::MinPairwise:
+      return scalar::MinPairwise;
+    case broadcast::MaxPairwise:
+      return scalar::MaxPairwise;
+    case broadcast::AMinPairwise:
+      return scalar::AMinPairwise;
+    case broadcast::AMaxPairwise:
+      return scalar::AMaxPairwise;
+    case broadcast::SquaredSubtract:
+      return scalar::SquaredSubtract;
+    default:
+      THROW_EXCEPTION("Not convertible operation");
+  }
+}
+}  // namespace sd
+
+#endif  // DEV_TESTS_BROADCASTSCALARCONVERTER_H

libnd4j/include/loops/ReduceType.h

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+/**
+ * @author raver119@gmail.com
+ */
+
+#ifndef DEV_TESTS_REDUCETYPE_H
+#define DEV_TESTS_REDUCETYPE_H
+
+namespace functions {
+enum ReduceType { SUM, PRODUCT, MAX, MIN, ASUM, AMAX, AMIN };
+}
+
+#endif  // DEV_TESTS_REDUCETYPE_H

libnd4j/include/loops/broadcasting.h

@@ -0,0 +1,164 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+/*
+ * broadcasting.h
+ *
+ *  Created on: Dec 28, 2015
+ *      Author: agibsonccc
+ */
+
+#ifndef BROADCASTING_H_
+#define BROADCASTING_H_
+#include <helpers/DebugHelper.h>
+#include <helpers/shape.h>
+#include <math/templatemath.h>
+#include <ops/ops.h>
+#include <system/op_boilerplate.h>
+
+#ifdef __JNI__
+#include <jni.h>
+#endif
+#include <helpers/LoopKind.h>
+
+#include <loops/legacy_ops.h>
+
+namespace functions {
+namespace broadcast {
+
+/**
+ * Broadcast operation
+ * for broadcasting a smaller tensor
+ * along long a bigger one.
+ */
+template <typename X, typename Y, typename Z>
+class Broadcast {
+ public:
+#ifdef SD_CUDA
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                                      sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                      const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                      const sd::LongType *tadOffsetsZ);
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                            const sd::LongType *xShapeInfo, const void *y,
+                                            const sd::LongType *yShapeInfo, void *result,
+                                            const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                            sd::LongType dimensionLength,
+                                            const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                            const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                            const sd::LongType *xShapeInfo, const void *y,
+                                            const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                    const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                                    void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                    sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                    const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                    const sd::LongType *tadOffsetsZ);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                    const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                                    void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpType>
+  static SD_DEVICE void transformInverseCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                             const sd::LongType *yShapeInfo, void *result,
+                                             const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                             sd::LongType dimensionLength,
+                                             const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                             const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateInverseBroadcast(
+      dim3 launchDims, cudaStream_t *stream, const void *x, const sd::LongType *xShapeInfo, const void *y,
+      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+      sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+      const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  static SD_HOST void execInverseBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                           const sd::LongType *xShapeInfo, const void *y,
+                                           const sd::LongType *yShapeInfo, void *result,
+                                           const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                           sd::LongType dimensionLength,
+                                           const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                           const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+#else
+
+  static void execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                          const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                          sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo,
+                          const sd::LongType *tadOffset, const sd::LongType *tadShapeInfoZ,
+                          const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                   sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                   const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, sd::LoopKind::Kind loopKind,sd::LongType start, sd::LongType stop);
+
+  /**
+   * CPU execution
+   * @param x the input
+   * @param xShapeInfo the x shape information
+   * @param y the y data
+   * @param yShapeInfo the y shape information
+   * @param result the result
+   * @param resultShapeInfo the result shape information
+   * @param dimension the dimension to broadcast along long
+   * @param dimensionLength the length of the dimension buffer
+   */
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                   sd::LongType dimensionLength,
+                   const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset, const sd::LongType *tadShapeInfoZ,
+                   const sd::LongType *tadOffsetZ, sd::LoopKind::Kind loopKind,
+                   sd::LongType start, sd::LongType stop);
+
+  template <typename OpType>
+  static void execInverse(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                          void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength,
+                          const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                          const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, sd::LongType start,
+                          sd::LongType stop);
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *z, const sd::LongType *zShapeInfo);
+
+#endif
+};
+}  // namespace broadcast
+}  // namespace functions
+
+#endif /* BROADCASTING_H_ */

libnd4j/include/loops/broadcasting_bool.h

@@ -0,0 +1,162 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+/*
+ * broadcasting.h
+ *
+ *  Created on: Dec 28, 2015
+ *      Author: agibsonccc
+ */
+
+#ifndef BROADCASTING_BOOL_H_
+#define BROADCASTING_BOOL_H_
+#include <helpers/DebugHelper.h>
+#include <math/templatemath.h>
+#include <ops/ops.h>
+#include <system/op_boilerplate.h>
+
+#ifdef __JNI__
+#include <jni.h>
+#endif
+
+#include <loops/legacy_ops.h>
+
+namespace functions {
+namespace broadcast {
+
+/**
+ * Broadcast operation
+ * for broadcasting a smaller tensor
+ * along long a bigger one.
+ */
+template <typename X, typename Z>
+class BroadcastBool {
+ public:
+#ifdef __CUDACC__
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                                      void *extraParams,sd::LongType *dimension,sd::LongType dimensionLength,
+                                      const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                      const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                      void *extraParams);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, cudaStream_t *stream, void const *x,
+                                            sd::LongType const *xShapeInfo, void const *y,
+                                            sd::LongType const *yShapeInfo, void *result,
+                                            sd::LongType const *resultShapeInfo, void *extraParams,
+                                            sd::LongType *dimension,sd::LongType dimensionLength, sd::LongType const *tadOnlyShapeInfo,
+                                            sd::LongType const *tadOffsets, sd::LongType const *tadOnlyShapeInfoZ,
+                                            sd::LongType const *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, 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);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, void const *x,
+                                    sd::LongType const *xShapeInfo, void const *y, sd::LongType const *yShapeInfo,
+                                    void *result, sd::LongType const *resultShapeInfo, void *extraParams,
+                                    sd::LongType *dimension, sd::LongType dimensionLength, sd::LongType const *tadOnlyShapeInfo,
+                                    sd::LongType const *tadOffsets, sd::LongType const *tadOnlyShapeInfoZ,
+                                    sd::LongType const *tadOffsetsZ);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, const int opNum, const void *x,
+                                    const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                                    void *z, const sd::LongType *zShapeInfo, void *extraParams);
+
+  template <typename OpType>
+  static SD_DEVICE void transformInverseCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                             const sd::LongType *yShapeInfo, void *result,
+                                             const sd::LongType *resultShapeInfo, void *extraParams,
+                                             sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                             const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                             const sd::LongType *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateInverseBroadcast(
+      dim3 launchDims, cudaStream_t *stream, const void *x, const sd::LongType *xShapeInfo, const void *y,
+      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo, void *extraParams,
+      sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+      const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  static SD_HOST void execInverseBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                           const sd::LongType *xShapeInfo, const void *y,
+                                           const sd::LongType *yShapeInfo, void *result,
+                                           const sd::LongType *resultShapeInfo, void *extraParams,
+                                           sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                           const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                           const sd::LongType *tadOffsetsZ);
+
+#else
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo, void *extraParams,
+                   sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                   const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo, void *extraParams);
+
+  static void execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                          const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                          void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo,
+                          const sd::LongType *tadOffset, const sd::LongType *tadShapeInfoZ,
+                          const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  /**
+   * CPU execution
+   * @param x the input
+   * @param xShapeInfo the x shape information
+   * @param y the y data
+   * @param yShapeInfo the y shape information
+   * @param result the result
+   * @param resultShapeInfo the result shape information
+   * @param dimension the dimension to broadcast along long
+   * @param dimensionLength the length of the dimension buffer
+   */
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *result, const sd::LongType *resultShapeInfo, void *extraParams, sd::LongType *dimension,
+                   sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                   const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, uint64_t start, uint64_t stop);
+
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *z, const sd::LongType *zShapeInfo, void *extraParams);
+
+  template <typename OpType>
+  static void execInverse(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                          void *result, const sd::LongType *resultShapeInfo, void *extraParams,
+                          sd::LongType *dimension,
+                          sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                          const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, uint64_t start,
+                          uint64_t stop);
+#endif
+};
+}  // namespace broadcast
+}  // namespace functions
+
+#endif /* BROADCASTING_H_ */

libnd4j/include/loops/broadcasting_int.h

@@ -0,0 +1,160 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+ ******************************************************************************/
+
+/*
+ * broadcasting.h
+ *
+ *  Created on: Dec 28, 2015
+ *      Author: agibsonccc
+ */
+
+#ifndef BROADCASTING_INT_H_
+#define BROADCASTING_INT_H_
+#include <helpers/DebugHelper.h>
+#include <math/templatemath.h>
+#include <ops/ops.h>
+#include <system/op_boilerplate.h>
+
+#ifdef __JNI__
+#include <jni.h>
+#endif
+
+#include <loops/legacy_ops.h>
+
+namespace functions {
+namespace broadcast {
+
+/**
+ * Broadcast operation
+ * for broadcasting a smaller tensor
+ * along long a bigger one.
+ */
+template <typename X>
+class BroadcastInt {
+ public:
+#ifdef __CUDACC__
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                                      sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                      const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                      const sd::LongType *tadOffsetsZ);
+
+  template <typename OpType>
+  static SD_DEVICE void transformCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                            const sd::LongType *xShapeInfo, const void *y,
+                                            const sd::LongType *yShapeInfo, void *result,
+                                            const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                            sd::LongType dimensionLength,
+                                            const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                            const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateBroadcast(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                            const sd::LongType *xShapeInfo, const void *y,
+                                            const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                    const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                                    void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                    sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo,
+                                    const sd::LongType *tadOffsets, const sd::LongType *tadOnlyShapeInfoZ,
+                                    const sd::LongType *tadOffsetsZ);
+
+  static SD_HOST void execBroadcast(dim3 launchDims, cudaStream_t *stream, const int opNum, const void *x,
+                                    const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                                    void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpType>
+  static SD_DEVICE void transformInverseCuda(const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                             const sd::LongType *yShapeInfo, void *result,
+                                             const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                                             sd::LongType dimensionLength,
+                                             const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                             const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  template <typename OpClass>
+  static SD_HOST void intermediateInverseBroadcast(
+      dim3 launchDims, cudaStream_t *stream, const void *x, const sd::LongType *xShapeInfo, const void *y,
+      const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+      sd::LongType dimensionLength, const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+      const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+  static SD_HOST void execInverseBroadcast(dim3 launchDims, cudaStream_t *stream, int opNum, const void *x,
+                                           const sd::LongType *xShapeInfo, const void *y,
+                                           const sd::LongType *yShapeInfo, void *result,
+                                           const sd::LongType *resultShapeInfo, sd::LongType*dimension,
+                                           sd::LongType dimensionLength,
+                                           const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                                           const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ);
+
+#else
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                   sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                   const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  static void exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                   const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo);
+
+  static void execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                          const sd::LongType *yShapeInfo, void *result, const sd::LongType *resultShapeInfo,
+                          sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo,
+                          const sd::LongType *tadOffset, const sd::LongType *tadShapeInfoZ,
+                          const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  /**
+   * CPU execution
+   * @param x the input
+   * @param xShapeInfo the x shape information
+   * @param y the y data
+   * @param yShapeInfo the y shape information
+   * @param result the result
+   * @param resultShapeInfo the result shape information
+   * @param dimension the dimension to broadcast along long
+   * @param dimensionLength the length of the dimension buffer
+   */
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                   sd::LongType dimensionLength,
+                   const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset, const sd::LongType *tadShapeInfoZ,
+                   const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+
+  template <typename OpType>
+  static void exec(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                   void *z, const sd::LongType *zShapeInfo);
+
+  template <typename OpType>
+  static void execInverse(const void *x, const sd::LongType *xShapeInfo, const void *y, const sd::LongType *yShapeInfo,
+                          void *result, const sd::LongType *resultShapeInfo, sd::LongType *dimension,
+                          sd::LongType dimensionLength,
+                          const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset,
+                          const sd::LongType *tadShapeInfoZ, const sd::LongType *tadOffsetZ, sd::LongType start, sd::LongType stop);
+#endif
+};
+}  // namespace broadcast
+}  // namespace functions
+
+#endif /* BROADCASTING_H_ */

libnd4j/include/loops/cpu/broadcasting.hpp

@@ -0,0 +1,924 @@
+/* ******************************************************************************
+*
+*
+* 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 <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/LoopKind.h>
+#include <helpers/ShapeUtils.h>
+#include <loops/broadcasting.h>
+#include <loops/legacy_ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <cstdio>
+
+using namespace simdOps;
+
+namespace functions {
+namespace broadcast {
+
+template <typename X, typename Y, typename Z>
+void Broadcast<X, Y, Z>::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                     const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                     sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                     const sd::LongType *xTadOffset, const sd::LongType *zTadShapeInfo,
+                                     const sd::LongType *zTadOffset, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(execInverse,
+                        PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
+                               xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
+                        BROADCAST_OPS);
+}
+
+template <typename X, typename Y, typename Z>
+void Broadcast<X, Y, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                              const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                              sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                              const sd::LongType *xTadOffset,
+                              const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset,
+                              sd::LoopKind::Kind loopKind, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(exec,
+                        PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
+                               xTadOffset, zTadShapeInfo, zTadOffset, loopKind, start, stop),
+                        BROADCAST_OPS);
+}
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void Broadcast<X, Y, Z>::exec(const void* vx, const sd::LongType* xShapeInfo,
+                              const void* vy, const sd::LongType* yShapeInfo,
+                              void* vz, const sd::LongType* zShapeInfo,
+                              sd::LongType* dimension, sd::LongType dimensionLength,
+                              const sd::LongType* xTadShapeInfo, const sd::LongType* xTadOffset,
+                              const sd::LongType* zTadShapeInfo, const sd::LongType* zTadOffset,
+                              sd::LoopKind::Kind loopKind, sd::LongType start, sd::LongType stop) {
+
+  auto x = reinterpret_cast<const X*>(vx);
+  auto y = reinterpret_cast<const Y*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+
+  // Get rank information
+  const int xRank = shape::rank(xShapeInfo);
+  const int yRank = shape::rank(yShapeInfo);
+  const int zRank = shape::rank(zShapeInfo);
+  const int xTadRank = xTadShapeInfo ? shape::rank(xTadShapeInfo) : xRank;
+  const int zTadRank = zTadShapeInfo ? shape::rank(zTadShapeInfo) : zRank;
+
+  // Get shape information
+  const sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  const sd::LongType* yShape = shape::shapeOf(yShapeInfo);
+  const sd::LongType* zShape = shape::shapeOf(zShapeInfo);
+  const sd::LongType* xTadShape = shape::shapeOf(xTadShapeInfo);
+  const sd::LongType* zTadShape = shape::shapeOf(zTadShapeInfo);
+
+  // Get stride information
+  const sd::LongType* xStrides = shape::stride(xShapeInfo);
+  const sd::LongType* yStrides = shape::stride(yShapeInfo);
+  const sd::LongType* zStrides = shape::stride(zShapeInfo);
+  const sd::LongType* xTadStrides = shape::stride(xTadShapeInfo);
+  const sd::LongType* zTadStrides = shape::stride(zTadShapeInfo);
+
+  // Classify array types
+  // For X array or X TAD
+  bool isXScalar = xTadRank == 0 || (xTadRank == 1 && xTadShape[0] == 1);
+  bool isXVector = (xTadRank == 1 && xTadShape[0] > 1) ||
+                   (xTadRank == 2 && (xTadShape[0] == 1 || xTadShape[1] == 1));
+  bool isXRowVector = (xTadRank == 1 && xTadShape[0] > 1) ||
+                      (xTadRank == 2 && xTadShape[0] == 1 && xTadShape[1] > 1);
+  bool isXColumnVector = (xTadRank == 2 && xTadShape[0] > 1 && xTadShape[1] == 1);
+
+  // For Y array
+  bool isYScalar = yRank == 0 || (yRank == 1 && yShape[0] == 1);
+  bool isYVector = (yRank == 1 && yShape[0] > 1) ||
+                   (yRank == 2 && (yShape[0] == 1 || yShape[1] == 1));
+  bool isYRowVector = (yRank == 1 && yShape[0] > 1) ||
+                      (yRank == 2 && yShape[0] == 1 && yShape[1] > 1);
+  bool isYColumnVector = (yRank == 2 && yShape[0] > 1 && yShape[1] == 1);
+
+  // For Z array or Z TAD
+  bool isZScalar = zTadRank == 0 || (zTadRank == 1 && zTadShape[0] == 1);
+  bool isZVector = (zTadRank == 1 && zTadShape[0] > 1) ||
+                   (zTadRank == 2 && (zTadShape[0] == 1 || zTadShape[1] == 1));
+  bool isZRowVector = (zTadRank == 1 && zTadShape[0] > 1) ||
+                      (zTadRank == 2 && zTadShape[0] == 1 && zTadShape[1] > 1);
+  bool isZColumnVector = (zTadRank == 2 && zTadShape[0] > 1 && zTadShape[1] == 1);
+
+  // Handle scalar broadcasting as a special case first
+  if (isYScalar) {
+    // Scalar broadcast - apply same value to every element
+    const Y scalarY = y[0];
+    sd::LongType length = shape::length(xTadShapeInfo ? xTadShapeInfo : xShapeInfo);
+
+    if (xTadShapeInfo && zTadShapeInfo) {
+      // TAD case
+      for (auto i = start; i < stop; i++) {
+        auto oX = x + xTadOffset[i];
+        auto oZ = z + zTadOffset[i];
+
+        // Handle different X and Z shapes
+        if (isXVector && isZVector) {
+          sd::LongType len = shape::length(xTadShapeInfo);
+          PRAGMA_OMP_SIMD
+          for (sd::LongType f = 0; f < len; f++) {
+            sd::LongType xOffset = f * xTadStrides[xTadRank-1];
+            sd::LongType zOffset = f * zTadStrides[zTadRank-1];
+            oZ[zOffset] = OpType::op(oX[xOffset], scalarY);
+          }
+        } else {
+          // General case
+          for (sd::LongType f = 0; f < length; f++) {
+            // Calculate proper offsets for current position
+            sd::LongType xCoord[SD_MAX_RANK], zCoord[SD_MAX_RANK];
+            sd::LongType xOffset, zOffset;
+
+            INDEX2COORDS(f, xTadRank, xTadShape, xCoord);
+            INDEX2COORDS(f, zTadRank, zTadShape, zCoord);
+
+            COORDS2INDEX(xTadRank, xTadStrides, xCoord, xOffset);
+            COORDS2INDEX(zTadRank, zTadStrides, zCoord, zOffset);
+
+            oZ[zOffset] = OpType::op(oX[xOffset], scalarY);
+          }
+        }
+      }
+    } else {
+      // Non-TAD case
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < length; f++)
+        z[f] = OpType::op(x[f], scalarY);
+    }
+  }
+    // Handle 2D broadcasting
+  else if (loopKind == sd::LoopKind::BROADCAST_2D) {
+    // Determine shapes for broadcasting
+    sd::LongType nRows = zTadShape[0];
+    sd::LongType nCols = zTadRank > 1 ? zTadShape[1] : shape::length(zTadShapeInfo);
+
+    // Special vector broadcasting cases
+    if (isYVector && (isXRowVector || isXColumnVector || isXVector)) {
+      // Vector to vector broadcasting
+      if (isYRowVector && (isXRowVector || isXVector)) {
+        // Row vector to row vector
+        for (auto i = start; i < stop; i++) {
+          auto baseX = x + xTadOffset[i];
+          auto baseZ = z + zTadOffset[i];
+
+          sd::LongType xStride = xTadRank > 1 ? xTadStrides[xTadRank - 1] : xTadStrides[0];
+          sd::LongType yStride = yRank == 1 ? yStrides[0] : yStrides[1];
+          sd::LongType zStride = zTadRank ? zTadStrides[zTadRank - 1] : zTadStrides[0];
+
+
+          PRAGMA_OMP_SIMD
+          for (sd::LongType i1 = 0; i1 < nCols; i1++) {
+            auto rX = baseX + i1 * xStride;
+            auto rY = y + i1 * yStride;
+            auto rZ = baseZ + i1 * zStride;
+
+            *rZ = OpType::op(*rX, *rY);
+          }
+        }
+      }
+      else if (isYColumnVector && (isXColumnVector || isXVector)) {
+        // Column vector to column vector
+        // Row vector to row vector
+        for (auto i = start; i < stop; i++) {
+          auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
+          auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+          // Calculate correct strides based on shape and rank
+          sd::LongType xStride;
+          if (xTadRank == 1) {
+            xStride = xTadStrides[0];
+          } else { // xTadRank == 2
+            xStride = xTadStrides[0]; // For 2D column vector, use row stride
+          }
+
+          sd::LongType yStride;
+          if (yRank == 1) {
+            yStride = yStrides[0];
+          } else { // yRank == 2
+            yStride = yStrides[0]; // For 2D column vector, use row stride
+          }
+
+          sd::LongType zStride;
+          if (zTadRank == 1) {
+            zStride = zTadStrides[0];
+          } else { // zTadRank == 2
+            zStride = zTadStrides[0]; // For 2D column vector, use row stride
+          }
+
+          // Verify dimensions match
+          sd::LongType xLen = isXColumnVector ? (xTadRank == 2 ? xTadShape[0] : xTadShape[0]) : xTadShape[0];
+          sd::LongType yLen = yRank == 2 ? yShape[0] : yShape[0];
+          printf("xLen: %lld; yLen: %lld nRows %lld,xStride %lld,yStride %lld, zStride %lld\n", xLen, yLen,nRows,xStride,yStride,zStride);
+          PRAGMA_OMP_SIMD
+          for (sd::LongType i1 = 0; i1 < nRows; i1++) {
+            auto rX = baseX + i1 * xStride;
+            auto rY = y + i1 * yStride;
+            auto rZ = baseZ + i1 * zStride;
+
+            *rZ = OpType::op(*rX, *rY);
+          }
+        }
+      }
+      else if (isYColumnVector && isXRowVector) {
+        // Column vector to row vector (outer product)
+        for (auto i = start; i < stop; i++) {
+          auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
+          auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+          for (sd::LongType i0 = 0; i0 < nRows; i0++) {
+            auto colValue = y[i0 * yStrides[0]];
+
+            PRAGMA_OMP_SIMD
+            for (sd::LongType i1 = 0; i1 < nCols; i1++) {
+              auto rX = baseX + i1 * xTadStrides[xTadRank-1];
+              auto rZ = baseZ + i0 * zTadStrides[0] + i1 * zTadStrides[1];
+
+              *rZ = OpType::op(*rX, colValue);
+            }
+          }
+        }
+      }
+      else if (isYRowVector && isXColumnVector) {
+        // Row vector to column vector (outer product)
+        for (auto i = start; i < stop; i++) {
+          printf("4 2d tad: %lld\n", i);
+          fflush(stdout);
+          auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
+          auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+          for (sd::LongType i0 = 0; i0 < nRows; i0++) {
+            auto xValue = baseX[i0 * xTadStrides[0]];
+
+            PRAGMA_OMP_SIMD
+            for (sd::LongType i1 = 0; i1 < nCols; i1++) {
+              auto rY = y + i1 * (yRank == 1 ? yStrides[0] : yStrides[1]);
+              auto rZ = baseZ + i0 * zTadStrides[0] + i1 * zTadStrides[1];
+
+              *rZ = OpType::op(xValue, *rY);
+            }
+          }
+        }
+      }
+    }
+      // Matrix with vector broadcasting
+    else if ((isXRowVector && isYRowVector) || (isXColumnVector && isYColumnVector)) {
+      // Matching vectors - element-wise operation
+      for (auto i = start; i < stop; i++) {
+        auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
+        auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+        sd::LongType vecLength = isXRowVector ? nCols : nRows;
+        sd::LongType xStride = isXRowVector ? xTadStrides[xTadRank-1] : xTadStrides[0];
+        sd::LongType yStride = isYRowVector ? (yRank == 1 ? yStrides[0] : yStrides[1]) : yStrides[0];
+        sd::LongType zStride = isZRowVector ? zTadStrides[zTadRank-1] : zTadStrides[0];
+
+        PRAGMA_OMP_SIMD
+        for (sd::LongType i1 = 0; i1 < vecLength; i1++) {
+          auto rX = baseX + i1 * xStride;
+          auto rY = y + i1 * yStride;
+          auto rZ = baseZ + i1 * zStride;
+
+          *rZ = OpType::op(*rX, *rY);
+        }
+      }
+    }
+      // Matrix with row vector broadcasting
+    else if (isYRowVector && xTadRank == 2 && zTadRank == 2 &&
+             xTadShape[1] == (yRank == 1 ? yShape[0] : yShape[1])) {
+      // Broadcasting row vector (each element applied to a column)
+      for (auto i0 = start; i0 < stop; i0++) {
+        auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
+        auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
+
+        for (sd::LongType i1 = 0; i1 < nRows; i1++) {
+          for (sd::LongType i2 = 0; i2 < nCols; i2++) {
+            // Get element from X at current position
+            auto xOffset = i1 * xTadStrides[0] + i2 * xTadStrides[1];
+            // Get element from Y row vector based on column index only
+            auto yOffset = i2 * (yRank == 1 ? yStrides[0] : yStrides[1]);
+            // Get destination element in Z at current position
+            auto zOffset = i1 * zTadStrides[0] + i2 * zTadStrides[1];
+
+            // Apply operation
+            baseZ[zOffset] = OpType::op(baseX[xOffset], y[yOffset]);
+          }
+        }
+      }
+    }
+      // Matrix with column vector broadcasting
+    else if (isYColumnVector) {
+      // Broadcasting column vector (each element applied to a row)
+      for (auto i0 = start; i0 < stop; i0++) {
+        auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
+        auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
+
+        for (sd::LongType i1 = 0; i1 < nRows; i1++) {
+          // Get element from column vector based on row index
+          auto rY = y + i1 * yStrides[0];
+
+          PRAGMA_OMP_SIMD
+          for (sd::LongType i2 = 0; i2 < nCols; i2++) {
+            auto rX = baseX + i1 * xTadStrides[0] + i2 * xTadStrides[1];
+            auto rZ = baseZ + i1 * zTadStrides[0] + i2 * zTadStrides[1];
+
+            *rZ = OpType::op(*rX, *rY);
+          }
+        }
+      }
+    }
+      // Standard 2D broadcasting
+    else {
+      for (auto i0 = start; i0 < stop; i0++) {
+        auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
+        auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
+
+        for (sd::LongType i1 = 0; i1 < nRows; i1++) {
+          PRAGMA_OMP_SIMD
+          for (sd::LongType i2 = 0; i2 < nCols; i2++) {
+            auto rX = baseX + i1 * xTadStrides[0] + i2 * xTadStrides[1];
+            auto rY = y + i1 * yStrides[0] + i2 * yStrides[1];
+            auto rZ = baseZ + i1 * zTadStrides[0] + i2 * zTadStrides[1];
+
+            *rZ = OpType::op(*rX, *rY);
+          }
+        }
+      }
+    }
+  }
+    // Handle remaining loop kinds
+  else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_X) {
+    sd::LongType tadLength = shape::length(xTadShapeInfo);
+    for (auto i = start; i < stop; i++) {
+      auto oY = y + (i * tadLength);
+      auto oZ = z + (i * tadLength);
+      const auto oX = x[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++)
+        oZ[f] = OpType::op(oX, oY[f]);
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_Y) {
+    sd::LongType tadLength = shape::length(xTadShapeInfo);
+    for (auto i = start; i < stop; i++) {
+      auto oX = x + (i * tadLength);
+      auto oZ = z + (i * tadLength);
+      const auto oY = y[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++)
+        oZ[f] = OpType::op(oX[f], oY);
+    }
+  }
+    // Handle 3D broadcasting (generalized like 2D case)
+  else if (loopKind == sd::LoopKind::BROADCAST_3D) {
+    // Get TAD info
+    const sd::LongType tadRank = xTadShapeInfo ? shape::rank(xTadShapeInfo) : 3;
+    const sd::LongType* tadShape = xTadShapeInfo ? shape::shapeOf(xTadShapeInfo) : xShape;
+    const sd::LongType* tadStride = xTadShapeInfo ? shape::stride(xTadShapeInfo) : xStrides;
+    const sd::LongType tadLength = xTadShapeInfo ? shape::length(xTadShapeInfo) : shape::length(xShapeInfo);
+    if (isYVector) {
+      // Vector broadcasting
+      const sd::LongType yLength = yRank == 1 ? yShape[0] : (yShape[0] == 1 ? yShape[1] : yShape[0]);
+      const sd::LongType yStride = yRank == 1 ? yStrides[0] : (yShape[0] == 1 ? yStrides[1] : yStrides[0]);
+
+      // Determine which dimension this vector should be broadcast along
+      // For a 3D TAD, check if vector length matches any dimension
+      if (tadRank == 3) {
+        if (yLength == tadShape[2]) {
+          // Broadcast along last dimension
+          for (auto i = start; i < stop; i++) {
+            auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+            auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+            for (sd::LongType j = 0; j < tadLength; j++) {
+              // Calculate TAD coords
+              sd::LongType coords[SD_MAX_RANK];
+              INDEX2COORDS(j, tadRank, tadShape, coords);
+
+              // Get offsets
+              sd::LongType xOffset, zOffset;
+              COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+              COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+              // Get Y index - use the last dimension (coords[2])
+              sd::LongType yOffset = coords[2] * yStride;
+
+              // Apply operation
+              oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+            }
+          }
+        }
+        else if (yLength == tadShape[1]) {
+          // Broadcast along middle dimension
+          PRAGMA_OMP_SIMD
+          for (auto i = start; i < stop; i++) {
+            auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+            auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+            for (sd::LongType j = 0; j < tadLength; j++) {
+              // Calculate TAD coords
+              sd::LongType coords[SD_MAX_RANK];
+              INDEX2COORDS(j, tadRank, tadShape, coords);
+
+              // Get offsets
+              sd::LongType xOffset, zOffset;
+              COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+              COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+              // Get Y index - use the middle dimension (coords[1])
+              sd::LongType yOffset = coords[1] * yStride;
+
+              // Apply operation
+              oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+            }
+          }
+        }
+        else if (yLength == tadShape[0]) {
+          // Broadcast along first dimension
+          PRAGMA_OMP_SIMD
+          for (auto i = start; i < stop; i++) {
+            auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+            auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+            for (sd::LongType j = 0; j < tadLength; j++) {
+              // Calculate TAD coords
+              sd::LongType coords[SD_MAX_RANK];
+              INDEX2COORDS(j, tadRank, tadShape, coords);
+
+              // Get offsets
+              sd::LongType xOffset, zOffset;
+              COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+              COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+              // Get Y index - use the first dimension (coords[0])
+              sd::LongType yOffset = coords[0] * yStride;
+
+              // Apply operation
+              oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+            }
+          }
+        }
+        else {
+          // Default broadcasting behavior - broadcast along the last dimension
+          PRAGMA_OMP_SIMD
+          for (auto i = start; i < stop; i++) {
+            auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+            auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+            for (sd::LongType j = 0; j < tadLength; j++) {
+              // Calculate TAD coords
+              sd::LongType coords[SD_MAX_RANK];
+              INDEX2COORDS(j, tadRank, tadShape, coords);
+
+              // Get offsets
+              sd::LongType xOffset, zOffset;
+              COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+              COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+              // Get Y index with wrapping/broadcasting
+              sd::LongType yOffset = (coords[2] % yLength) * yStride;
+
+              // Apply operation
+              oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+            }
+          }
+        }
+      }
+      else {
+        // Handle lower rank TADs (1D or 2D)
+        PRAGMA_OMP_SIMD
+        for (auto i = start; i < stop; i++) {
+          auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+          auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+          for (sd::LongType j = 0; j < tadLength; j++) {
+            // Calculate TAD coords
+            sd::LongType coords[SD_MAX_RANK];
+            INDEX2COORDS(j, tadRank, tadShape, coords);
+
+            // Get offsets
+            sd::LongType xOffset, zOffset;
+            COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+            COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+            // Get Y index - for lower ranks, broadcast along the last available dimension
+            sd::LongType lastCoord = tadRank > 0 ? coords[tadRank - 1] : 0;
+            sd::LongType yOffset = (lastCoord % yLength) * yStride;
+
+            // Apply operation
+            oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+          }
+        }
+      }
+    }
+    else if (yRank == 2) {
+      // Y is a 2D matrix - determine which dimensions it aligns with
+      for (auto i = start; i < stop; i++) {
+        auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+        auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+        PRAGMA_OMP_SIMD
+        for (sd::LongType j = 0; j < tadLength; j++) {
+          // Calculate TAD coords
+          sd::LongType coords[SD_MAX_RANK];
+          INDEX2COORDS(j, tadRank, tadShape, coords);
+
+          // Get offsets
+          sd::LongType xOffset, zOffset;
+          COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+          COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+          // Calculate Y offset based on dimension matching
+          sd::LongType yOffset;
+
+          // Default behavior for different 2D matrix broadcasting patterns
+          if (tadRank == 3) {
+            if (yShape[0] == tadShape[0] && yShape[1] == tadShape[2]) {
+              // Y is aligned with dimensions 0 and 2
+              yOffset = coords[0] * yStrides[0] + coords[2] * yStrides[1];
+            }
+            else if (yShape[0] == tadShape[0] && yShape[1] == tadShape[1]) {
+              // Y is aligned with dimensions 0 and 1
+              yOffset = coords[0] * yStrides[0] + coords[1] * yStrides[1];
+            }
+            else if (yShape[0] == tadShape[1] && yShape[1] == tadShape[2]) {
+              // Y is aligned with dimensions 1 and 2
+              yOffset = coords[1] * yStrides[0] + coords[2] * yStrides[1];
+            }
+            else {
+              // Default: broadcast Y to match the last two dimensions with modulo
+              yOffset = (coords[1] % yShape[0]) * yStrides[0] + (coords[2] % yShape[1]) * yStrides[1];
+            }
+          }
+          else if (tadRank == 2) {
+            // Direct mapping for 2D TAD and 2D Y
+            yOffset = (coords[0] % yShape[0]) * yStrides[0] + (coords[1] % yShape[1]) * yStrides[1];
+          }
+          else {
+            // For 1D TAD, map to the first dimension of Y
+            yOffset = (coords[0] % yShape[0]) * yStrides[0];
+          }
+
+          // Apply operation
+          oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+        }
+      }
+    }
+    else if (yRank == 3) {
+      // Y is a 3D tensor
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+        auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+        for (sd::LongType j = 0; j < tadLength; j++) {
+          // Calculate TAD coords
+          sd::LongType coords[SD_MAX_RANK];
+          INDEX2COORDS(j, tadRank, tadShape, coords);
+
+          // Get offsets
+          sd::LongType xOffset, zOffset;
+          COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+          COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+          // Calculate Y offset with modulo for broadcasting if needed
+          sd::LongType yCoords[3] = {0, 0, 0};
+
+          // Map coordinates appropriately based on ranks
+          if (tadRank == 3) {
+            yCoords[0] = coords[0] % yShape[0];
+            yCoords[1] = coords[1] % yShape[1];
+            yCoords[2] = coords[2] % yShape[2];
+          }
+          else if (tadRank == 2) {
+            // Map 2D to last 2 dimensions of 3D
+            yCoords[0] = 0;  // First dimension is broadcasted
+            yCoords[1] = coords[0] % yShape[1];
+            yCoords[2] = coords[1] % yShape[2];
+          }
+          else {
+            // Map 1D to last dimension of 3D
+            yCoords[0] = 0;
+            yCoords[1] = 0;
+            yCoords[2] = coords[0] % yShape[2];
+          }
+
+          sd::LongType yOffset = yCoords[0] * yStrides[0] + yCoords[1] * yStrides[1] + yCoords[2] * yStrides[2];
+
+          // Apply operation
+          oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+        }
+      }
+    }
+    else {
+      // General case for other ranks of Y
+      for (auto i = start; i < stop; i++) {
+        auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
+        auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
+
+        for (sd::LongType j = 0; j < tadLength; j++) {
+          // Calculate TAD coords
+          sd::LongType coords[SD_MAX_RANK];
+          INDEX2COORDS(j, tadRank, tadShape, coords);
+
+          // Get offsets
+          sd::LongType xOffset, zOffset;
+          COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+          COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
+
+          // Calculate Y offset based on rank
+          sd::LongType yOffset = 0;
+
+          // Map coordinates to Y based on rank
+          for (int d = 0; d < tadRank && d < yRank; d++) {
+            yOffset += (coords[d] % yShape[d]) * yStrides[d];
+          }
+
+          // Apply operation
+          oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
+        }
+      }
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_4D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+    const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+
+    for (auto i = start; i < stop; i++) {
+      uint64_t i0 = i / nSize1;
+      uint64_t i1 = i % nSize1;
+
+      for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+        PRAGMA_OMP_SIMD
+        for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+          auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2 + xStrides[3] * i3);
+          auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2 + yStrides[3] * i3);
+          auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2 + zStrides[3] * i3);
+
+          *rZ = OpType::op(*rX, *rY);
+        }
+      }
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_5D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+    const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+    const sd::LongType nSize4 = shape::sizeAt(zShapeInfo, 4);
+
+    for (auto i = start; i < stop; i++) {
+      uint32_t i0 = i / nSize1;
+      uint32_t i1 = i % nSize1;
+
+      for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+        for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+          PRAGMA_OMP_SIMD
+          for (sd::LongType i4 = 0; i4 < nSize4; i4++) {
+            auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2 + xStrides[3] * i3 + xStrides[4] * i4);
+            auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2 + yStrides[3] * i3 + yStrides[4] * i4);
+            auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2 + zStrides[3] * i3 + zStrides[4] * i4);
+
+            *rZ = OpType::op(*rX, *rY);
+          }
+        }
+      }
+    }
+  }
+  else {
+    // Default case for other ranks - general purpose implementation
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+
+    for (auto i = start; i < stop; i++) {
+      // Calculate independent coordinates for each array
+      INDEX2COORDS(i, xRank, xShape, xCoords);
+      INDEX2COORDS(i, yRank, yShape, yCoords);
+      INDEX2COORDS(i, zRank, zShape, zCoords);
+
+      // Calculate offsets based on each array's coordinates and strides
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStrides, xCoords, xOffset);
+      COORDS2INDEX(yRank, yStrides, yCoords, yOffset);
+      COORDS2INDEX(zRank, zStrides, zCoords, zOffset);
+
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+    }
+  }
+}
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void Broadcast<X, Y, Z>::execInverse(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                     const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                     sd::LongType *dimension, sd::LongType dimensionLength,
+                                     const sd::LongType *yTadShapeInfo, const sd::LongType *yTadOffset,
+                                     const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset,
+                                     sd::LongType start, sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const Y *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+
+  // Handle TAD setup
+  auto yTadShapeShapeInfo = yTadShapeInfo;
+  auto tadOffsets = yTadOffset;
+
+  // When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
+  std::shared_ptr<sd::TadPack> tadPack = nullptr;
+
+  if (yTadShapeInfo == nullptr || tadOffsets == nullptr) {
+    tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dimension,
+                                                                    dimensionLength);
+    yTadShapeShapeInfo = tadPack->primaryShapeInfo();
+    tadOffsets = tadPack->primaryOffsets();
+  }
+
+  if (zTadShapeInfo == nullptr) {
+    zTadShapeInfo = yTadShapeShapeInfo;
+    zTadOffset = tadOffsets;
+  }
+
+  // Get shape information
+  const auto xRank = shape::rank(xShapeInfo);
+  const auto yTadRank = shape::rank(yTadShapeShapeInfo);
+  const auto zTadRank = shape::rank(zTadShapeInfo);
+
+  const auto xStrides = shape::stride(xShapeInfo);
+  const auto yTadStrides = shape::stride(yTadShapeShapeInfo);
+  const auto zTadStrides = shape::stride(zTadShapeInfo);
+
+  const auto xShape = shape::shapeOf(xShapeInfo);
+  const auto yTadShape = shape::shapeOf(yTadShapeShapeInfo);
+  const auto zTadShape = shape::shapeOf(zTadShapeInfo);
+
+  const sd::LongType tadLength = shape::length(yTadShapeShapeInfo);
+
+  if (yTadRank <= 3) {
+    // Optimized path for lower ranks
+    for (auto i = start; i < stop; i++) {
+      auto oZ = z + zTadOffset[i];
+      auto oY = y + tadOffsets[i];
+
+      if (yTadRank == 1) {
+        PRAGMA_OMP_SIMD
+        for (sd::LongType j = 0; j < tadLength; j++) {
+          oZ[j * zTadStrides[0]] = OpType::op(x[j * xStrides[0]], oY[j * yTadStrides[0]]);
+        }
+      }
+      else if (yTadRank == 2) {
+        const sd::LongType dim0 = yTadShape[0];
+        const sd::LongType dim1 = yTadShape[1];
+
+        for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+          PRAGMA_OMP_SIMD
+          for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+            const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1];
+            const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1];
+            const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1];
+            oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+          }
+        }
+      }
+      else { // rank 3
+        const sd::LongType dim0 = yTadShape[0];
+        const sd::LongType dim1 = yTadShape[1];
+        const sd::LongType dim2 = yTadShape[2];
+
+        for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+          for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+            PRAGMA_OMP_SIMD
+            for (sd::LongType j2 = 0; j2 < dim2; j2++) {
+              const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1] + j2 * xStrides[2];
+              const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1] + j2 * yTadStrides[2];
+              const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1] + j2 * zTadStrides[2];
+              oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+            }
+          }
+        }
+      }
+    }
+  }
+  else {
+    // Use macros for higher ranks
+    for (auto i = start; i < stop; i++) {
+      auto oZ = z + zTadOffset[i];
+      auto oY = y + tadOffsets[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(f, yTadRank, yTadShape, coords);
+
+        sd::LongType xOffset, yOffset, zOffset;
+        COORDS2INDEX(xRank, xStrides, coords, xOffset);
+        COORDS2INDEX(yTadRank, yTadStrides, coords, yOffset);
+        COORDS2INDEX(zTadRank, zTadStrides, coords, zOffset);
+
+        oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+      }
+    }
+  }
+}
+
+template <typename X, typename Y, typename Z>
+void Broadcast<X, Y, Z>::exec(const int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                              const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_TTT(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo), BROADCAST_OPS);
+}
+
+template <typename X, typename Y, typename Z, typename OpType>
+static void execDefault(const X *x, const sd::LongType *xShapeInfo, const Y *y, const sd::LongType *yShapeInfo, Z *z,
+                        const sd::LongType *zShapeInfo) {
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+
+  // C-style arrays CANNOT be captured by value in lambdas - they decay to pointers
+  // that point to stack memory. std::array CAN be captured by value, ensuring each
+  // parallel thread gets its own copy of the data with guaranteed lifetime.
+  std::array<sd::LongType, SD_MAX_RANK> xShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> zShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> xStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> zStrideLocal;
+
+  // Copy actual data from shapeInfo into std::arrays
+  std::memcpy(xShapeLocal.data(), shape::shapeOf(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(yShapeLocal.data(), shape::shapeOf(yShapeInfo), yRank * sizeof(sd::LongType));
+  std::memcpy(zShapeLocal.data(), shape::shapeOf(zShapeInfo), zRank * sizeof(sd::LongType));
+  std::memcpy(xStrideLocal.data(), shape::stride(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(yStrideLocal.data(), shape::stride(yShapeInfo), yRank * sizeof(sd::LongType));
+  std::memcpy(zStrideLocal.data(), shape::stride(zShapeInfo), zRank * sizeof(sd::LongType));
+
+  // Capture std::arrays by value - C++ will copy the entire array contents into the lambda's closure.
+  // This ensures each parallel thread has its own copy of the data with no dangling pointers.
+  auto func = [x, y, z, xRank, yRank, zRank, xShapeLocal, yShapeLocal, zShapeLocal, xStrideLocal, yStrideLocal, zStrideLocal](
+      sd::LongType thread_id, sd::LongType start, sd::LongType stop, sd::LongType increment) -> void {
+    for (auto i = start; i < stop; ++i) {
+      sd::LongType zCoords[SD_MAX_RANK];
+      sd::LongType xCoords[SD_MAX_RANK];
+      sd::LongType yCoords[SD_MAX_RANK];
+
+      // Convert linear index to coordinates based on Z (output) shape
+      INDEX2COORDS(i, zRank, zShapeLocal.data(), zCoords);
+
+      // Broadcast Z coordinates to X and Y shapes
+      // For broadcasting, we map Z coords to X and Y coords using modulo for smaller dimensions
+      // When a dimension is 1 in X or Y but larger in Z, we use index 0 (broadcast)
+      for (sd::LongType d = 0; d < xRank; d++) {
+        xCoords[d] = xShapeLocal[d] == 1 ? 0 : (zCoords[d] % xShapeLocal[d]);
+      }
+      for (sd::LongType d = 0; d < yRank; d++) {
+        yCoords[d] = yShapeLocal[d] == 1 ? 0 : (zCoords[d] % yShapeLocal[d]);
+      }
+
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStrideLocal.data(), xCoords, xOffset);
+      COORDS2INDEX(yRank, yStrideLocal.data(), yCoords, yOffset);
+      COORDS2INDEX(zRank, zStrideLocal.data(), zCoords, zOffset);
+
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), shape::length(zShapeInfo));
+}
+
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void Broadcast<X, Y, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                              const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  const Y *y = reinterpret_cast<const Y *>(vy);
+  Z *z = reinterpret_cast<Z *>(vz);
+
+  const int rank = shape::rank(zShapeInfo);  // xRank = yRank = zRank
+
+  switch (rank) {
+    default:
+      execDefault<X, Y, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  }
+}
+
+
+}  // namespace broadcast
+}  // namespace functions

libnd4j/include/loops/cpu/broadcasting_bool.hpp

@@ -0,0 +1,672 @@
+/* ******************************************************************************
+*
+*
+* 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 <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/LoopKind.h>
+#include <loops/broadcasting_bool.h>
+#include <loops/legacy_ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace broadcast {
+
+template <typename X, typename Z>
+void BroadcastBool<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                               const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                               void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+                               const sd::LongType *xTadShapeInfo, const sd::LongType *xTadOffset,
+                               const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, sd::LongType start,
+                               sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(exec,
+                       PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, dimension, dimensionLength,
+                              xTadShapeInfo, xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
+                       BROADCAST_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+void BroadcastBool<X, Z>::exec(const int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                               const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                               void *extraParams) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams), BROADCAST_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+void BroadcastBool<X, Z>::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                      const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                      void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+                                      const sd::LongType *xTadShapeInfo, const sd::LongType *xTadOffset,
+                                      const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, sd::LongType start,
+                                      sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(execInverse,
+                       PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, dimension, dimensionLength,
+                              xTadShapeInfo, xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
+                       BROADCAST_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void BroadcastBool<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                               const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                               void *vextraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+                               const sd::LongType *xTadShapeInfo, const sd::LongType *xTadOffset,
+                               const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, uint64_t start,
+                               uint64_t stop) {
+  auto x = reinterpret_cast<const X*>(vx);
+  auto y = reinterpret_cast<const X*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+  sd::LoopKind::Kind loopKind= sd::LoopKind::deduceKindOfLoopXYZ(xTadShapeInfo, yShapeInfo, zTadShapeInfo);
+
+  if (loopKind == sd::LoopKind::BROADCAST_SCALAR_X) {
+    sd::LongType tadLength = shape::length(xTadShapeInfo);
+    for (auto i = start; i < stop; i++) {
+      auto oY = y + (i * tadLength);
+      auto oZ = z + (i * tadLength);
+      const auto oX = x[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++)
+        oZ[f] = OpType::op(oX, oY[f],extraParams);
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_Y) {
+    sd::LongType tadLength = shape::length(xTadShapeInfo);
+    for (auto i = start; i < stop; i++) {
+      auto oX = x + (i * tadLength);
+      auto oZ = z + (i * tadLength);
+      const auto oY = y[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++)
+        oZ[f] = OpType::op(oX[f], oY,extraParams);
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_2D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType* xStrides = shape::stride(xTadShapeInfo);
+    const sd::LongType* yStrides = shape::stride(yShapeInfo);
+    const sd::LongType* zStrides = shape::stride(zTadShapeInfo);
+
+    for (auto i0 = start; i0 < stop; i0++) {
+      auto baseX = x + xTadOffset[i0];
+      auto baseZ = z + zTadOffset[i0];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i1 = 0; i1 < nSize1; i1++) {
+        auto rX = baseX + xStrides[1] * i1;
+        auto rY = y + yStrides[1] * i1;
+        auto rZ = baseZ + zStrides[1] * i1;
+
+        *rZ = OpType::op(*rX, *rY,extraParams);
+      }
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_3D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+
+    const sd::LongType* xStrides = shape::stride(xShapeInfo);
+    const sd::LongType* yStrides = shape::stride(yShapeInfo);
+    const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+    for (auto i0 = start; i0 < stop; i0++) {
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i1 = 0; i1 < nSize1; i1++) {
+        for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+          auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2);
+          auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2);
+          auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2);
+
+          *rZ = OpType::op(*rX, *rY,extraParams);
+        }
+      }
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_4D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+    const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+
+    const sd::LongType* xStrides = shape::stride(xShapeInfo);
+    const sd::LongType* yStrides = shape::stride(yShapeInfo);
+    const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+    for (auto i = start; i < stop; i++) {
+      uint64_t i0 = i / nSize1;
+      uint64_t i1 = i % nSize1;
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+        for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+          auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 +
+                         xStrides[2] * i2 + xStrides[3] * i3);
+          auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 +
+                         yStrides[2] * i2 + yStrides[3] * i3);
+          auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 +
+                         zStrides[2] * i2 + zStrides[3] * i3);
+
+          *rZ = OpType::op(*rX, *rY,extraParams);
+        }
+      }
+    }
+  }
+  else if (loopKind == sd::LoopKind::BROADCAST_5D) {
+    const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+    const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+    const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+    const sd::LongType nSize4 = shape::sizeAt(zShapeInfo, 4);
+
+    const sd::LongType* xStrides = shape::stride(xShapeInfo);
+    const sd::LongType* yStrides = shape::stride(yShapeInfo);
+    const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+    for (auto i = start; i < stop; i++) {
+      uint32_t i0 = i / nSize1;
+      uint32_t i1 = i % nSize1;
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+        for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+          for (sd::LongType i4 = 0; i4 < nSize4; i4++) {
+            auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 +
+                           xStrides[2] * i2 + xStrides[3] * i3 + xStrides[4] * i4);
+            auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 +
+                           yStrides[2] * i2 + yStrides[3] * i3 + yStrides[4] * i4);
+            auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 +
+                           zStrides[2] * i2 + zStrides[3] * i3 + zStrides[4] * i4);
+
+            *rZ = OpType::op(*rX, *rY,extraParams);
+          }
+        }
+      }
+    }
+  }
+  else {
+    // Default case for other ranks
+    const int xRank = shape::rank(xShapeInfo);
+    const int yRank = shape::rank(yShapeInfo);
+    const int zRank = shape::rank(zShapeInfo);
+
+    const sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+    const sd::LongType* yShape = shape::shapeOf(yShapeInfo);
+    const sd::LongType* zShape = shape::shapeOf(zShapeInfo);
+    const sd::LongType* xStrides = shape::stride(xShapeInfo);
+    const sd::LongType* yStrides = shape::stride(yShapeInfo);
+    const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+
+    for (auto i = start; i < stop; i++) {
+      // Calculate independent coordinates for each array
+      INDEX2COORDS(i, xRank, xShape, xCoords);
+      INDEX2COORDS(i, yRank, yShape, yCoords);
+      INDEX2COORDS(i, zRank, zShape, zCoords);
+
+      // Calculate offsets based on each array's coordinates and strides
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStrides, xCoords, xOffset);
+      COORDS2INDEX(yRank, yStrides, yCoords, yOffset);
+      COORDS2INDEX(zRank, zStrides, zCoords, zOffset);
+
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset],extraParams);
+    }
+  }
+}
+
+
+template <typename X, typename Z>
+template <typename OpType>
+void BroadcastBool<X, Z>::execInverse(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                      const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                      void *vextraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+                                      const sd::LongType *yTadShapeInfo, const sd::LongType *yTadOffset,
+                                      const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, uint64_t start,
+                                      uint64_t stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+  // Handle TAD setup
+  auto yTadShapeShapeInfo = yTadShapeInfo;
+  auto tadOffsets = yTadOffset;
+
+  // When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
+  std::shared_ptr<sd::TadPack> tadPack = nullptr;
+
+  if (yTadShapeInfo == nullptr || tadOffsets == nullptr) {
+    tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dimension,
+                                                                    dimensionLength);
+    yTadShapeShapeInfo = tadPack->primaryShapeInfo();
+    tadOffsets = tadPack->primaryOffsets();
+  }
+
+  if (zTadShapeInfo == nullptr) {
+    zTadShapeInfo = yTadShapeShapeInfo;
+    zTadOffset = tadOffsets;
+  }
+
+  // Get shape information
+  const auto xRank = shape::rank(xShapeInfo);
+  const auto yTadRank = shape::rank(yTadShapeShapeInfo);
+  const auto zTadRank = shape::rank(zTadShapeInfo);
+
+  const auto xStrides = shape::stride(xShapeInfo);
+  const auto yTadStrides = shape::stride(yTadShapeShapeInfo);
+  const auto zTadStrides = shape::stride(zTadShapeInfo);
+
+  const auto xShape = shape::shapeOf(xShapeInfo);
+  const auto yTadShape = shape::shapeOf(yTadShapeShapeInfo);
+  const auto zTadShape = shape::shapeOf(zTadShapeInfo);
+
+  const sd::LongType tadLength = shape::length(yTadShapeShapeInfo);
+
+  if (yTadRank <= 3) {
+    // Optimized path for lower ranks
+    for (auto i = start; i < stop; i++) {
+      auto oZ = z + zTadOffset[i];
+      auto oY = y + tadOffsets[i];
+
+      if (yTadRank == 1) {
+        PRAGMA_OMP_SIMD
+        for (sd::LongType j = 0; j < tadLength; j++) {
+          oZ[j * zTadStrides[0]] = OpType::op(x[j * xStrides[0]], oY[j * yTadStrides[0]],extraParams);
+        }
+      }
+      else if (yTadRank == 2) {
+        const sd::LongType dim0 = yTadShape[0];
+        const sd::LongType dim1 = yTadShape[1];
+
+        for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+          PRAGMA_OMP_SIMD
+          for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+            const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1];
+            const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1];
+            const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1];
+            oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset],extraParams);
+          }
+        }
+      }
+      else { // rank 3
+        const sd::LongType dim0 = yTadShape[0];
+        const sd::LongType dim1 = yTadShape[1];
+        const sd::LongType dim2 = yTadShape[2];
+
+        for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+          for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+            PRAGMA_OMP_SIMD
+            for (sd::LongType j2 = 0; j2 < dim2; j2++) {
+              const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1] + j2 * xStrides[2];
+              const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1] + j2 * yTadStrides[2];
+              const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1] + j2 * zTadStrides[2];
+              oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset],extraParams);
+            }
+          }
+        }
+      }
+    }
+  }
+  else {
+    // Use macros for higher ranks
+    for (auto i = start; i < stop; i++) {
+      auto oZ = z + zTadOffset[i];
+      auto oY = y + tadOffsets[i];
+
+      PRAGMA_OMP_SIMD
+      for (sd::LongType f = 0; f < tadLength; f++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(f, yTadRank, yTadShape, coords);
+
+        sd::LongType xOffset, yOffset, zOffset;
+        COORDS2INDEX(xRank, xStrides, coords, xOffset);
+        COORDS2INDEX(yTadRank, yTadStrides, coords, yOffset);
+        COORDS2INDEX(zTadRank, zTadStrides, coords, zOffset);
+
+        oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset],extraParams);
+      }
+    }
+  }
+}
+
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execRank1(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                      const sd::LongType *zShapeInfo, X *extraParams) {
+  // Cache shape-related values
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, static_cast<sd::LongType>(0));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, static_cast<sd::LongType>(0));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, static_cast<sd::LongType>(0));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR {
+    if (zStrd0 == 1 && xStrd0 == 1 && yStrd0 == 0) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(x[i0], *y, extraParams);
+    } else if (zStrd0 == 1 && xStrd0 == 0 && yStrd0 == 1) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(*x, y[i0], extraParams);
+    } else if (zStrd0 == 1 && xStrd0 == 1 && yStrd0 == 1) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(x[i0], y[i0], extraParams);
+    } else {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0 * zStrd0] = OpType::op(x[i0 * xStrd0], y[i0 * yStrd0], extraParams);
+    }
+  };
+  samediff::Threads::parallel_tad(func, static_cast<sd::LongType>(0), zAxis0);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execRank2(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                      const sd::LongType *zShapeInfo, X *extraParams) {
+
+  // Cache shape-related values
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(1));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(1));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(1));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto i0 = start; i0 < stop; ++i0) {
+      auto x0 = x + i0 * xStrd0;
+      auto y0 = y + i0 * yStrd0;
+      auto z0 = z + i0 * zStrd0;
+
+      if (zStrd1 == 1 && xStrd1 == 1 && yStrd1 == 0)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(x0[i1], *y0, extraParams);
+      else if (zStrd1 == 1 && xStrd1 == 0 && yStrd1 == 1)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(*x0, y0[i1], extraParams);
+      else if (zStrd1 == 1 && xStrd1 == 1 && yStrd1 == 1)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(x0[i1], y0[i1], extraParams);
+      else
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1)
+          z0[i1 * zStrd1] = OpType::op(x0[i1 * xStrd1], y0[i1 * yStrd1], extraParams);
+    }
+  };
+
+  samediff::Threads::parallel_tad(func, static_cast<sd::LongType>(0), zAxis0);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execRank3(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                      const sd::LongType *zShapeInfo, X *extraParams) {
+  // Cache shape-related values
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(2));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(2));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(2));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(2));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, static_cast<sd::LongType>(1));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, static_cast<sd::LongType>(1));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, static_cast<sd::LongType>(1));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis2 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd2 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR_2D {
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        auto x1 = x + i0 * xStrd0 + i1 * xStrd1;
+        auto y1 = y + i0 * yStrd0 + i1 * yStrd1;
+        auto z1 = z + i0 * zStrd0 + i1 * zStrd1;
+
+        if (zStrd2 == 1 && xStrd2 == 1 && yStrd2 == 0)
+          for (sd::LongType i2 = 0; i2 < zAxis2; ++i2) z1[i2] = OpType::op(x1[i2], *y1, extraParams);
+        else if (zStrd2 == 1 && xStrd2 == 0 && yStrd2 == 1)
+          for (sd::LongType i2 = 0; i2 < zAxis2; ++i2) z1[i2] = OpType::op(*x1, y1[i2], extraParams);
+        else if (zStrd2 == 1 && xStrd2 == 1 && yStrd2 == 1)
+          for (sd::LongType i2 = 0; i2 < zAxis2; ++i2) z1[i2] = OpType::op(x1[i2], y1[i2], extraParams);
+        else
+          for (sd::LongType i2 = 0; i2 < zAxis2; ++i2)
+            z1[i2 * zStrd2] = OpType::op(x1[i2 * xStrd2], y1[i2 * yStrd2], extraParams);
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zAxis0, static_cast<sd::LongType>(1), static_cast<sd::LongType>(0), zAxis1, static_cast<sd::LongType>(1));
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execRank4(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                      const sd::LongType *zShapeInfo, X *extraParams) {
+  // Cache shape-related values
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+
+  sd::LongType zAxis2 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType yStrd2 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis3 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd3 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd3 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd3 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR_3D {
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        for (auto i2 = start_z; i2 < stop_z; ++i2) {
+          auto x2 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2;
+          auto y2 = y + i0 * yStrd0 + i1 * yStrd1 + i2 * yStrd2;
+          auto z2 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2;
+
+          if (zStrd3 == 1 && xStrd3 == 1 && yStrd3 == 0)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(x2[i3], *y2, extraParams);
+          else if (zStrd3 == 1 && xStrd3 == 0 && yStrd3 == 1)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(*x2, y2[i3], extraParams);
+          else if (zStrd3 == 1 && xStrd3 == 1 && yStrd3 == 1)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(x2[i3], y2[i3], extraParams);
+          else
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3)
+              z2[i3 * zStrd3] = OpType::op(x2[i3 * xStrd3], y2[i3 * yStrd3], extraParams);
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zAxis0, static_cast<sd::LongType>(1),
+                                  static_cast<sd::LongType>(0), zAxis1, static_cast<sd::LongType>(1),
+                                  static_cast<sd::LongType>(0), zAxis2, static_cast<sd::LongType>(1));
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execRank5(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                      const sd::LongType *zShapeInfo, X *extraParams) {
+  // Cache shape-related values
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+
+  sd::LongType zAxis2 = shape::sizeAt(zShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType xStrd2 = shape::strideAt(xShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType yStrd2 = shape::strideAt(yShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType zStrd2 = shape::strideAt(zShapeInfo, static_cast<sd::LongType>(2));
+
+  sd::LongType zAxis3 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType xStrd3 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType yStrd3 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType zStrd3 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis4 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd4 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd4 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd4 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR_3D {
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        for (auto i2 = start_z; i2 < stop_z; ++i2) {
+          for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) {
+            auto x3 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3;
+            auto y3 = y + i0 * yStrd0 + i1 * yStrd1 + i2 * yStrd2 + i3 * yStrd3;
+            auto z3 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2 + i3 * zStrd3;
+
+            if (zStrd4 == 1 && xStrd4 == 1 && yStrd4 == 0)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(x3[i4], *y3, extraParams);
+            else if (zStrd4 == 1 && xStrd4 == 0 && yStrd4 == 1)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(*x3, y3[i4], extraParams);
+            else if (zStrd4 == 1 && xStrd4 == 1 && yStrd4 == 1)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(x3[i4], y3[i4], extraParams);
+            else
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4)
+                z3[i4 * zStrd4] = OpType::op(x3[i4 * xStrd4], y3[i4 * yStrd4], extraParams);
+          }
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zAxis0, static_cast<sd::LongType>(1),
+                                  static_cast<sd::LongType>(0), zAxis1, static_cast<sd::LongType>(1),
+                                  static_cast<sd::LongType>(0), zAxis2, static_cast<sd::LongType>(1));
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+static void execDefault(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, Z *z,
+                        const sd::LongType *zShapeInfo, X *extraParams) {
+  const bool xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+  const bool yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
+
+  // Cache shape-related values
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+
+  // C-style arrays CANNOT be captured by value in lambdas - they decay to pointers
+  // that point to stack memory. std::array CAN be captured by value, ensuring each
+  // parallel thread gets its own copy of the data with guaranteed lifetime.
+  std::array<sd::LongType, SD_MAX_RANK> zShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> zStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> xShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> xStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yStrideLocal;
+
+  std::memcpy(zShapeLocal.data(), shape::shapeOf(zShapeInfo), zRank * sizeof(sd::LongType));
+  std::memcpy(zStrideLocal.data(), shape::stride(zShapeInfo), zRank * sizeof(sd::LongType));
+  std::memcpy(xShapeLocal.data(), shape::shapeOf(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(xStrideLocal.data(), shape::stride(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(yShapeLocal.data(), shape::shapeOf(yShapeInfo), yRank * sizeof(sd::LongType));
+  std::memcpy(yStrideLocal.data(), shape::stride(yShapeInfo), yRank * sizeof(sd::LongType));
+
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType xOffset, yOffset, zOffset;
+
+    for (auto i = start; i < stop; ++i) {
+      // Convert linear index to coordinates based on Z (output) shape
+      INDEX2COORDS(i, zRank, zShapeLocal.data(), zCoords);
+      COORDS2INDEX(zRank, zStrideLocal.data(), zCoords, zOffset);
+
+      if (xzSameOffsets) {
+        xOffset = zOffset;
+      } else {
+        // Broadcast Z coordinates to X shape
+        for (sd::LongType d = 0; d < xRank; d++) {
+          xCoords[d] = xShapeLocal[d] == 1 ? 0 : (zCoords[d] % xShapeLocal[d]);
+        }
+        COORDS2INDEX(xRank, xStrideLocal.data(), xCoords, xOffset);
+      }
+      if (yzSameOffsets) {
+        yOffset = zOffset;
+      } else {
+        // Broadcast Z coordinates to Y shape
+        for (sd::LongType d = 0; d < yRank; d++) {
+          yCoords[d] = yShapeLocal[d] == 1 ? 0 : (zCoords[d] % yShapeLocal[d]);
+        }
+        COORDS2INDEX(yRank, yStrideLocal.data(), yCoords, yOffset);
+      }
+
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), shape::length(zShapeInfo));
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void BroadcastBool<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                               const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                               void *vextraParams) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  const X *y = reinterpret_cast<const X *>(vy);
+  Z *z = reinterpret_cast<Z *>(vz);
+
+  X *extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const int rank = shape::rank(zShapeInfo);  // xRank = yRank = zRank
+
+  switch (rank) {
+    case 1:
+      execRank1<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+      break;
+    case 2:
+      execRank2<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+      break;
+    case 3:
+      execRank3<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+      break;
+    case 4:
+      execRank4<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+      break;
+    case 5:
+      execRank5<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+      break;
+    default:
+      execDefault<X, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
+  }
+}
+
+} // namespace broadcast
+} // namespace functions

libnd4j/include/loops/cpu/broadcasting_int.hpp

@@ -0,0 +1,650 @@
+/* ******************************************************************************
+*
+*
+* 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 <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/LoopKind.h>
+#include <loops/broadcasting_int.h>
+#include <loops/legacy_ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace broadcast {
+
+template <typename X>
+void BroadcastInt<X>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                          sd::LongType*dimension,
+                          sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo, const sd::LongType *xTadOffset,
+                          const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, sd::LongType start,
+                          sd::LongType stop) {
+ DISPATCH_BY_OPNUM_T(exec,
+                     PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
+                            xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
+                     BROADCAST_INT_OPS);
+}
+
+template <typename X>
+void BroadcastInt<X>::exec(const int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo) {
+ DISPATCH_BY_OPNUM_T(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo), BROADCAST_INT_OPS);
+}
+
+template <typename X>
+void BroadcastInt<X>::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                 const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                 sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                 const sd::LongType *xTadOffset, const sd::LongType *zTadShapeInfo,
+                                 const sd::LongType *zTadOffset, sd::LongType start, sd::LongType stop) {
+ DISPATCH_BY_OPNUM_T(execInverse,
+                     PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
+                            xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
+                     BROADCAST_INT_OPS);
+}
+
+template <typename X>
+template <typename OpType>
+void BroadcastInt<X>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                           const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                           sd::LongType *dimension,
+                           sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo, const sd::LongType *xTadOffset,
+                           const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset, sd::LongType start,
+                           sd::LongType stop) {
+ auto x = reinterpret_cast<const X*>(vx);
+ auto y = reinterpret_cast<const X*>(vy);
+ auto z = reinterpret_cast<X*>(vz);
+ sd::LoopKind::Kind loopKind= sd::LoopKind::deduceKindOfLoopXYZ(xTadShapeInfo, yShapeInfo, zTadShapeInfo);
+
+ if (loopKind == sd::LoopKind::BROADCAST_SCALAR_X) {
+   sd::LongType tadLength = shape::length(xTadShapeInfo);
+   for (auto i = start; i < stop; i++) {
+     auto oY = y + (i * tadLength);
+     auto oZ = z + (i * tadLength);
+     const auto oX = x[i];
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType f = 0; f < tadLength; f++)
+       oZ[f] = OpType::op(oX, oY[f]);
+   }
+ }
+ else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_Y) {
+   sd::LongType tadLength = shape::length(xTadShapeInfo);
+   for (auto i = start; i < stop; i++) {
+     auto oX = x + (i * tadLength);
+     auto oZ = z + (i * tadLength);
+     const auto oY = y[i];
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType f = 0; f < tadLength; f++)
+       oZ[f] = OpType::op(oX[f], oY);
+   }
+ }
+ else if (loopKind == sd::LoopKind::BROADCAST_2D) {
+   const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+   const sd::LongType* xStrides = shape::stride(xTadShapeInfo);
+   const sd::LongType* yStrides = shape::stride(yShapeInfo);
+   const sd::LongType* zStrides = shape::stride(zTadShapeInfo);
+
+   for (auto i0 = start; i0 < stop; i0++) {
+     auto baseX = x + xTadOffset[i0];
+     auto baseZ = z + zTadOffset[i0];
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType i1 = 0; i1 < nSize1; i1++) {
+       auto rX = baseX + xStrides[1] * i1;
+       auto rY = y + yStrides[1] * i1;
+       auto rZ = baseZ + zStrides[1] * i1;
+
+       *rZ = OpType::op(*rX, *rY);
+     }
+   }
+ }
+ else if (loopKind == sd::LoopKind::BROADCAST_3D) {
+   const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+   const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+
+   const sd::LongType* xStrides = shape::stride(xShapeInfo);
+   const sd::LongType* yStrides = shape::stride(yShapeInfo);
+   const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+   for (auto i0 = start; i0 < stop; i0++) {
+     PRAGMA_OMP_SIMD
+     for (sd::LongType i1 = 0; i1 < nSize1; i1++) {
+       for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+         auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2);
+         auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2);
+         auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2);
+
+         *rZ = OpType::op(*rX, *rY);
+       }
+     }
+   }
+ }
+ else if (loopKind == sd::LoopKind::BROADCAST_4D) {
+   const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+   const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+   const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+
+   const sd::LongType* xStrides = shape::stride(xShapeInfo);
+   const sd::LongType* yStrides = shape::stride(yShapeInfo);
+   const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+   for (auto i = start; i < stop; i++) {
+     uint64_t i0 = i / nSize1;
+     uint64_t i1 = i % nSize1;
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+       for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+         auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 +
+                        xStrides[2] * i2 + xStrides[3] * i3);
+         auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 +
+                        yStrides[2] * i2 + yStrides[3] * i3);
+         auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 +
+                        zStrides[2] * i2 + zStrides[3] * i3);
+
+         *rZ = OpType::op(*rX, *rY);
+       }
+     }
+   }
+ }
+ else if (loopKind == sd::LoopKind::BROADCAST_5D) {
+   const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
+   const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
+   const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
+   const sd::LongType nSize4 = shape::sizeAt(zShapeInfo, 4);
+
+   const sd::LongType* xStrides = shape::stride(xShapeInfo);
+   const sd::LongType* yStrides = shape::stride(yShapeInfo);
+   const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+   for (auto i = start; i < stop; i++) {
+     uint32_t i0 = i / nSize1;
+     uint32_t i1 = i % nSize1;
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
+       for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
+         for (sd::LongType i4 = 0; i4 < nSize4; i4++) {
+           auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 +
+                          xStrides[2] * i2 + xStrides[3] * i3 + xStrides[4] * i4);
+           auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 +
+                          yStrides[2] * i2 + yStrides[3] * i3 + yStrides[4] * i4);
+           auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 +
+                          zStrides[2] * i2 + zStrides[3] * i3 + zStrides[4] * i4);
+
+           *rZ = OpType::op(*rX, *rY);
+         }
+       }
+     }
+   }
+ }
+ else {
+   // Default case for other ranks
+   const int xRank = shape::rank(xShapeInfo);
+   const int yRank = shape::rank(yShapeInfo);
+   const int zRank = shape::rank(zShapeInfo);
+
+   const sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+   const sd::LongType* yShape = shape::shapeOf(yShapeInfo);
+   const sd::LongType* zShape = shape::shapeOf(zShapeInfo);
+   const sd::LongType* xStrides = shape::stride(xShapeInfo);
+   const sd::LongType* yStrides = shape::stride(yShapeInfo);
+   const sd::LongType* zStrides = shape::stride(zShapeInfo);
+
+   sd::LongType xCoords[SD_MAX_RANK];
+   sd::LongType yCoords[SD_MAX_RANK];
+   sd::LongType zCoords[SD_MAX_RANK];
+
+   for (auto i = start; i < stop; i++) {
+     // Calculate independent coordinates for each array
+     INDEX2COORDS(i, xRank, xShape, xCoords);
+     INDEX2COORDS(i, yRank, yShape, yCoords);
+     INDEX2COORDS(i, zRank, zShape, zCoords);
+
+     // Calculate offsets based on each array's coordinates and strides
+     sd::LongType xOffset, yOffset, zOffset;
+     COORDS2INDEX(xRank, xStrides, xCoords, xOffset);
+     COORDS2INDEX(yRank, yStrides, yCoords, yOffset);
+     COORDS2INDEX(zRank, zStrides, zCoords, zOffset);
+
+     z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+   }
+ }
+}
+
+template <typename X>
+template <typename OpType>
+void BroadcastInt<X>::execInverse(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                  const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                  sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *yTadShapeInfo,
+                                  const sd::LongType *yTadOffset, const sd::LongType *zTadShapeInfo,
+                                  const sd::LongType *zTadOffset, sd::LongType start, sd::LongType stop) {
+ auto x = reinterpret_cast<const X *>(vx);
+ auto y = reinterpret_cast<const X *>(vy);
+ auto z = reinterpret_cast<X *>(vz);
+
+ // Handle TAD setup
+ auto yTadShapeShapeInfo = yTadShapeInfo;
+ auto tadOffsets = yTadOffset;
+
+ // When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
+ std::shared_ptr<sd::TadPack> tadPack = nullptr;
+
+ if (yTadShapeInfo == nullptr || tadOffsets == nullptr) {
+   tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dimension,
+                                                                   dimensionLength);
+   yTadShapeShapeInfo = tadPack->primaryShapeInfo();
+   tadOffsets = tadPack->primaryOffsets();
+ }
+
+ if (zTadShapeInfo == nullptr) {
+   zTadShapeInfo = yTadShapeShapeInfo;
+   zTadOffset = tadOffsets;
+ }
+
+ // Get shape information
+ const auto xRank = shape::rank(xShapeInfo);
+ const auto yTadRank = shape::rank(yTadShapeShapeInfo);
+ const auto zTadRank = shape::rank(zTadShapeInfo);
+
+ const auto xStrides = shape::stride(xShapeInfo);
+ const auto yTadStrides = shape::stride(yTadShapeShapeInfo);
+ const auto zTadStrides = shape::stride(zTadShapeInfo);
+
+ const auto xShape = shape::shapeOf(xShapeInfo);
+ const auto yTadShape = shape::shapeOf(yTadShapeShapeInfo);
+ const auto zTadShape = shape::shapeOf(zTadShapeInfo);
+
+ const sd::LongType tadLength = shape::length(yTadShapeShapeInfo);
+
+ if (yTadRank <= 3) {
+   // Optimized path for lower ranks
+   for (auto i = start; i < stop; i++) {
+     auto oZ = z + zTadOffset[i];
+     auto oY = y + tadOffsets[i];
+
+     if (yTadRank == 1) {
+       PRAGMA_OMP_SIMD
+       for (sd::LongType j = 0; j < tadLength; j++) {
+         oZ[j * zTadStrides[0]] = OpType::op(x[j * xStrides[0]], oY[j * yTadStrides[0]]);
+       }
+     }
+     else if (yTadRank == 2) {
+       const sd::LongType dim0 = yTadShape[0];
+       const sd::LongType dim1 = yTadShape[1];
+
+       for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+         PRAGMA_OMP_SIMD
+         for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+           const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1];
+           const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1];
+           const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1];
+           oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+         }
+       }
+     }
+     else { // rank 3
+       const sd::LongType dim0 = yTadShape[0];
+       const sd::LongType dim1 = yTadShape[1];
+       const sd::LongType dim2 = yTadShape[2];
+
+       for (sd::LongType j0 = 0; j0 < dim0; j0++) {
+         for (sd::LongType j1 = 0; j1 < dim1; j1++) {
+           PRAGMA_OMP_SIMD
+           for (sd::LongType j2 = 0; j2 < dim2; j2++) {
+             const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1] + j2 * xStrides[2];
+             const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1] + j2 * yTadStrides[2];
+             const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1] + j2 * zTadStrides[2];
+             oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+           }
+         }
+       }
+     }
+   }
+ }
+ else {
+   // Use macros for higher ranks
+   for (auto i = start; i < stop; i++) {
+     auto oZ = z + zTadOffset[i];
+     auto oY = y + tadOffsets[i];
+
+     PRAGMA_OMP_SIMD
+     for (sd::LongType f = 0; f < tadLength; f++) {
+       sd::LongType coords[SD_MAX_RANK];
+       INDEX2COORDS(f, yTadRank, yTadShape, coords);
+
+       sd::LongType xOffset, yOffset, zOffset;
+       COORDS2INDEX(xRank, xStrides, coords, xOffset);
+       COORDS2INDEX(yTadRank, yTadStrides, coords, yOffset);
+       COORDS2INDEX(zTadRank, zTadStrides, coords, zOffset);
+
+       oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
+     }
+   }
+ }
+}
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+static void execRank1(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                      const sd::LongType *zShapeInfo) {
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo,  static_cast<sd::LongType>(0));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo,  static_cast<sd::LongType>(0));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo,  static_cast<sd::LongType>(0));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo,  static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR {
+    if (zStrd0 == 1 && xStrd0 == 1 && yStrd0 == 0) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(x[i0], *y);
+    } else if (zStrd0 == 1 && xStrd0 == 0 && yStrd0 == 1) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(*x, y[i0]);
+    } else if (zStrd0 == 1 && xStrd0 == 1 && yStrd0 == 1) {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0] = OpType::op(x[i0], y[i0]);
+    } else {
+      for (auto i0 = start; i0 < stop; ++i0) z[i0 * zStrd0] = OpType::op(x[i0 * xStrd0], y[i0 * yStrd0]);
+    }
+  };
+  samediff::Threads::parallel_tad(func,  static_cast<sd::LongType>(0), zAxis0);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+static void execRank2(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                      const sd::LongType *zShapeInfo) {
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0)   : static_cast<sd::LongType>(1) );
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0)   : static_cast<sd::LongType>(1) );
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0)   : static_cast<sd::LongType>(1) );
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0)   : static_cast<sd::LongType>(1) );
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1)   : static_cast<sd::LongType>(0) );
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1)   : static_cast<sd::LongType>(0) );
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1)   : static_cast<sd::LongType>(0) );
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1)   : static_cast<sd::LongType>(0) );
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto i0 = start; i0 < stop; ++i0) {
+      auto x0 = x + i0 * xStrd0;
+      auto y0 = y + i0 * yStrd0;
+      auto z0 = z + i0 * zStrd0;
+
+      if (zStrd1 == 1 && xStrd1 == 1 && yStrd1 == 0)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(x0[i1], *y0);
+      else if (zStrd1 == 1 && xStrd1 == 0 && yStrd1 == 1)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(*x0, y0[i1]);
+      else if (zStrd1 == 1 && xStrd1 == 1 && yStrd1 == 1)
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1] = OpType::op(x0[i1], y0[i1]);
+      else
+        for (sd::LongType i1 = 0; i1 < zAxis1; ++i1) z0[i1 * zStrd1] = OpType::op(x0[i1 * xStrd1], y0[i1 * yStrd1]);
+    }
+  };
+
+  samediff::Threads::parallel_tad(func,  static_cast<sd::LongType>(0), zAxis0);
+}
+
+////////////////////////////////////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+static void execRank3(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                      const sd::LongType *zShapeInfo) {
+  // Cache all shape-related values
+  const sd::LongType xRank = shape::rank(xShapeInfo);
+  const sd::LongType yRank = shape::rank(yShapeInfo);
+  const sd::LongType zRank = shape::rank(zShapeInfo);
+
+  auto xStride = shape::stride(xShapeInfo);
+  auto yStride = shape::stride(yShapeInfo);
+  auto zStride = shape::stride(zShapeInfo);
+
+  const sd::LongType zLen0 = shape::sizeAt(zShapeInfo, 0);
+  const sd::LongType zLen1 = shape::sizeAt(zShapeInfo, 1);
+  const sd::LongType zLen2 = shape::sizeAt(zShapeInfo, 2);
+
+  auto func = PRAGMA_THREADS_FOR_2D {
+    sd::LongType coords[3];
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        for (sd::LongType i2 = 0; i2 < zLen2; ++i2) {
+          coords[0] = i0;
+          coords[1] = i1;
+          coords[2] = i2;
+
+          sd::LongType xOffset, yOffset, zOffset;
+          COORDS2INDEX(xRank, xStride, coords, xOffset);
+          COORDS2INDEX(yRank, yStride, coords, yOffset);
+          COORDS2INDEX(zRank, zStride, coords, zOffset);
+
+          z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, zLen0, 1, 0, zLen1, 1);
+}
+
+template <typename X, typename OpType>
+static void execRank4(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                      const sd::LongType *zShapeInfo) {
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(3));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(2));
+
+  sd::LongType zAxis2 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType yStrd2 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+  sd::LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(2) : static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis3 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd3 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd3 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd3 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR_3D {
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        for (auto i2 = start_z; i2 < stop_z; ++i2) {
+          auto x2 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2;
+          auto y2 = y + i0 * yStrd0 + i1 * yStrd1 + i2 * yStrd2;
+          auto z2 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2;
+
+          if (zStrd3 == 1 && xStrd3 == 1 && yStrd3 == 0)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(x2[i3], *y2);
+          else if (zStrd3 == 1 && xStrd3 == 0 && yStrd3 == 1)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(*x2, y2[i3]);
+          else if (zStrd3 == 1 && xStrd3 == 1 && yStrd3 == 1)
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3] = OpType::op(x2[i3], y2[i3]);
+          else
+            for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) z2[i3 * zStrd3] = OpType::op(x2[i3 * xStrd3], y2[i3 * yStrd3]);
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zAxis0, static_cast<sd::LongType>(1), static_cast<sd::LongType>(0), zAxis1, static_cast<sd::LongType>(1), static_cast<sd::LongType>(0), zAxis2, static_cast<sd::LongType>(1));
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+static void execRank5(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                      const sd::LongType *zShapeInfo) {
+  sd::LongType zAxis0 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType yStrd0 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+  sd::LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(0) : static_cast<sd::LongType>(4));
+
+  sd::LongType zAxis1 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType yStrd1 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+  sd::LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(1) : static_cast<sd::LongType>(3));
+
+  sd::LongType zAxis2 = shape::sizeAt(zShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType xStrd2 = shape::strideAt(xShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType yStrd2 = shape::strideAt(yShapeInfo, static_cast<sd::LongType>(2));
+  sd::LongType zStrd2 = shape::strideAt(zShapeInfo, static_cast<sd::LongType>(2));
+
+  sd::LongType zAxis3 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType xStrd3 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType yStrd3 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+  sd::LongType zStrd3 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(3) : static_cast<sd::LongType>(1));
+
+  sd::LongType zAxis4 = shape::sizeAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType xStrd4 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType yStrd4 = shape::strideAt(yShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+  sd::LongType zStrd4 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<sd::LongType>(4) : static_cast<sd::LongType>(0));
+
+  auto func = PRAGMA_THREADS_FOR_3D {
+    for (auto i0 = start_x; i0 < stop_x; ++i0) {
+      for (auto i1 = start_y; i1 < stop_y; ++i1) {
+        for (auto i2 = start_z; i2 < stop_z; ++i2) {
+          for (sd::LongType i3 = 0; i3 < zAxis3; ++i3) {
+            auto x3 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3;
+            auto y3 = y + i0 * yStrd0 + i1 * yStrd1 + i2 * yStrd2 + i3 * yStrd3;
+            auto z3 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2 + i3 * zStrd3;
+
+            if (zStrd4 == 1 && xStrd4 == 1 && yStrd4 == 0)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(x3[i4], *y3);
+            else if (zStrd4 == 1 && xStrd4 == 0 && yStrd4 == 1)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(*x3, y3[i4]);
+            else if (zStrd4 == 1 && xStrd4 == 1 && yStrd4 == 1)
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4) z3[i4] = OpType::op(x3[i4], y3[i4]);
+            else
+              for (sd::LongType i4 = 0; i4 < zAxis4; ++i4)
+                z3[i4 * zStrd4] = OpType::op(x3[i4 * xStrd4], y3[i4 * yStrd4]);
+          }
+        }
+      }
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zAxis0, static_cast<sd::LongType>(1), static_cast<sd::LongType>(0), zAxis1, static_cast<sd::LongType>(1), static_cast<sd::LongType>(0), zAxis2, static_cast<sd::LongType>(1));
+}
+
+
+
+template <typename X, typename OpType>
+static void execDefault(const X *x, const sd::LongType *xShapeInfo, const X *y, const sd::LongType *yShapeInfo, X *z,
+                        const sd::LongType *zShapeInfo) {
+  // Cache all shape-related values
+  const sd::LongType xRank = shape::rank(xShapeInfo);
+  const sd::LongType yRank = shape::rank(yShapeInfo);
+  const sd::LongType zRank = shape::rank(zShapeInfo);
+  const sd::LongType zLen = shape::length(zShapeInfo);
+
+  const bool xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+  const bool yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
+
+  // C-style arrays CANNOT be captured by value in lambdas - they decay to pointers
+  // that point to stack memory. std::array CAN be captured by value, ensuring each
+  // parallel thread gets its own copy of the data with guaranteed lifetime.
+  std::array<sd::LongType, SD_MAX_RANK> xShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> xStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yShapeLocal;
+  std::array<sd::LongType, SD_MAX_RANK> yStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> zStrideLocal;
+  std::array<sd::LongType, SD_MAX_RANK> zShapeLocal;
+
+  std::memcpy(xShapeLocal.data(), shape::shapeOf(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(xStrideLocal.data(), shape::stride(xShapeInfo), xRank * sizeof(sd::LongType));
+  std::memcpy(yShapeLocal.data(), shape::shapeOf(yShapeInfo), yRank * sizeof(sd::LongType));
+  std::memcpy(yStrideLocal.data(), shape::stride(yShapeInfo), yRank * sizeof(sd::LongType));
+  std::memcpy(zStrideLocal.data(), shape::stride(zShapeInfo), zRank * sizeof(sd::LongType));
+  std::memcpy(zShapeLocal.data(), shape::shapeOf(zShapeInfo), zRank * sizeof(sd::LongType));
+
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType xOffset, yOffset, zOffset;
+
+    for (auto i = start; i < stop; ++i) {
+      // Convert linear index to coordinates based on Z (output) shape
+      INDEX2COORDS(i, zRank, zShapeLocal.data(), zCoords);
+      COORDS2INDEX(zRank, zStrideLocal.data(), zCoords, zOffset);
+
+      if (xzSameOffsets) {
+        xOffset = zOffset;
+      } else {
+        // Broadcast Z coordinates to X shape
+        for (sd::LongType d = 0; d < xRank; d++) {
+          xCoords[d] = xShapeLocal[d] == 1 ? 0 : (zCoords[d] % xShapeLocal[d]);
+        }
+        COORDS2INDEX(xRank, xStrideLocal.data(), xCoords, xOffset);
+      }
+
+      if (yzSameOffsets) {
+        yOffset = zOffset;
+      } else {
+        // Broadcast Z coordinates to Y shape
+        for (sd::LongType d = 0; d < yRank; d++) {
+          yCoords[d] = yShapeLocal[d] == 1 ? 0 : (zCoords[d] % yShapeLocal[d]);
+        }
+        COORDS2INDEX(yRank, yStrideLocal.data(), yCoords, yOffset);
+      }
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), zLen);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+void BroadcastInt<X>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                       const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo) {
+const X *x = reinterpret_cast<const X *>(vx);
+const X *y = reinterpret_cast<const X *>(vy);
+X *z = reinterpret_cast<X *>(vz);
+
+const int rank = shape::rank(zShapeInfo);  // xRank = yRank = zRank
+
+switch (rank) {
+case 1:
+  execRank1<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  break;
+case 2:
+  execRank2<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  break;
+case 3:
+  execRank3<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  break;
+case 4:
+  execRank4<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  break;
+case 5:
+  execRank5<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+  break;
+default:
+  execDefault<X, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
+}
+}
+
+}  // namespace broadcast
+}  // namespace functions

libnd4j/include/loops/cpu/comb_compilation_units/broadcast_bool_instantiation_template.cpp.in

@@ -0,0 +1,86 @@
+/* ******************************************************************************
+*
+*
+* 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
+******************************************************************************/
+
+//
+// BroadcastBool template instantiations
+//
+
+#include <loops/pairwise_instantiations.h>
+#include <types/type_promote.h>
+#include <loops/broadcasting_bool.h>
+#include <loops/cpu/broadcasting_bool.hpp>
+/*
+*
+* Function Instantiation:
+* BroadcastBool::exec instantiated for types in @COMB1@, @COMB2@
+*/
+
+ITERATE_COMBINATIONS(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   INSTANT_PROCESS_COMBINATION,
+   functions::broadcast::BroadcastBool,
+   ::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+          void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+          const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+          const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ,
+          sd::LongType start, sd::LongType stop)
+)
+
+/*
+*
+* Function Instantiation:
+* BroadcastBool::exec (simple version) instantiated for types in @COMB1@, @COMB2@
+*/
+
+ITERATE_COMBINATIONS(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   INSTANT_PROCESS_COMBINATION,
+   functions::broadcast::BroadcastBool,
+   ::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+          void *extraParams)
+)
+
+/*
+*
+* Function Instantiation:
+* BroadcastBool::execInverse instantiated for types in @COMB1@, @COMB2@
+*/
+
+ITERATE_COMBINATIONS(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   INSTANT_PROCESS_COMBINATION,
+   functions::broadcast::BroadcastBool,
+   ::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                 const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                 void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+                 const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                 const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ,
+                 sd::LongType start, sd::LongType stop)
+)
+
+/*
+*
+* Explicit instantiation for bool types (since BroadcastBool typically uses bool output)
+*/
+
+template class functions::broadcast::BroadcastBool<bool, bool>;

libnd4j/include/loops/cpu/comb_compilation_units/broadcast_bool_instantiation_template_2.cpp.in

@@ -0,0 +1,31 @@
+/* ******************************************************************************
+*
+*
+* 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
+******************************************************************************/
+
+//
+// BroadcastBool template instantiations - Direct instantiation version
+// This file is designed for 2-type combinations
+//
+
+#include <loops/cpu/broadcasting_bool.hpp>
+#include <types/types.h>
+
+// Direct instantiations will be inserted here by the CMake template system
+// The handler will generate instantiations like:
+// template void functions::broadcast::BroadcastBool<float, bool>::exec(...);
+// template void functions::broadcast::BroadcastBool<double, bool>::exec(...);
+// etc.

libnd4j/include/loops/cpu/comb_compilation_units/broadcast_instantiation_template_3.cpp.in

@@ -0,0 +1,95 @@
+/* ******************************************************************************
+*
+*
+* 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
+******************************************************************************/
+
+//
+// Broadcast template instantiations
+//
+
+#include <loops/pairwise_instantiations.h>
+#include <types/type_promote.h>
+#include <loops/cpu/broadcasting.hpp>
+
+/*
+*
+* Function Instantiation:
+* Broadcast::exec instantiated for types in @COMB1@, @COMB2@, @COMB3@
+*/
+
+ITERATE_COMBINATIONS_3(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   SD_COMMON_TYPES_PART_@COMB3@,
+   INSTANT_PROCESS_COMBINATION_3,
+   functions::broadcast::Broadcast,
+   ::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+          sd::LongType *dimension, sd::LongType dimensionLength,
+          const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+          const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ,
+          sd::LoopKind::Kind loopKind, sd::LongType start, sd::LongType stop)
+)
+
+/*
+*
+* Function Instantiation:
+* Broadcast::exec (simple version) instantiated for types in @COMB1@, @COMB2@, @COMB3@
+*/
+
+ITERATE_COMBINATIONS_3(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   SD_COMMON_TYPES_PART_@COMB3@,
+   INSTANT_PROCESS_COMBINATION_3,
+   functions::broadcast::Broadcast,
+   ::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+          const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo)
+)
+
+/*
+*
+* Function Instantiation:
+* Broadcast::execInverse instantiated for types in @COMB1@, @COMB2@, @COMB3@
+*/
+
+ITERATE_COMBINATIONS_3(
+   SD_COMMON_TYPES_PART_@COMB1@,
+   SD_COMMON_TYPES_PART_@COMB2@,
+   SD_COMMON_TYPES_PART_@COMB3@,
+   INSTANT_PROCESS_COMBINATION_3,
+   functions::broadcast::Broadcast,
+   ::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                 const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                 sd::LongType *dimension, sd::LongType dimensionLength,
+                 const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+                 const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ,
+                 sd::LongType start, sd::LongType stop)
+)
+
+/*
+*
+* Function Instantiation:
+* BroadcastBool::exec instantiated for bool types
+*/
+
+template void functions::broadcast::BroadcastBool<bool, bool>::exec(
+   int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+   const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+   void *extraParams, sd::LongType *dimension, sd::LongType dimensionLength,
+   const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets,
+   const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ,
+   sd::LongType start, sd::LongType stop);

libnd4j/include/loops/cpu/comb_compilation_units/broadcast_int_instantiation_template_1.cpp.in

@@ -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
+******************************************************************************/
+
+//
+// BroadcastInt template instantiations
+//
+
+#include <loops/cpu/broadcasting_int.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+* Macro for instantiating BroadcastInt for a single type
+* This follows the pattern where BroadcastInt is specialized for integer types only
+*/
+
+#define INSTANTIATE_BROADCAST_INT(TYPE_TUPLE) \
+    template void functions::broadcast::BroadcastInt<GET_SECOND(TYPE_TUPLE)>::exec( \
+        int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y, \
+        const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo, \
+        sd::LongType *dimension, sd::LongType dimensionLength, \
+        const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets, \
+        const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ, \
+        sd::LongType start, sd::LongType stop); \
+    \
+    template void functions::broadcast::BroadcastInt<GET_SECOND(TYPE_TUPLE)>::exec( \
+        int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y, \
+        const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo); \
+    \
+    template void functions::broadcast::BroadcastInt<GET_SECOND(TYPE_TUPLE)>::execInverse( \
+        int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y, \
+        const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo, \
+        sd::LongType *dimension, sd::LongType dimensionLength, \
+        const sd::LongType *tadOnlyShapeInfo, const sd::LongType *tadOffsets, \
+        const sd::LongType *tadOnlyShapeInfoZ, const sd::LongType *tadOffsetsZ, \
+        sd::LongType start, sd::LongType stop);
+
+/*
+* Function Instantiation:
+* BroadcastInt instantiated for all integer types in SD_INTEGER_TYPES_PART_@COMB1@
+*/
+ITERATE_LIST((SD_INTEGER_TYPES_PART_@COMB1@), INSTANTIATE_BROADCAST_INT)

libnd4j/include/loops/cpu/comb_compilation_units/indexreduction_instantiation_template_2.cpp.in

@@ -0,0 +1,35 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// Created by agibsonccc on 12/26/24.
+//
+
+#include <helpers/cpu/loops/IndexReductionLoops.hpp>
+#include <system/type_boilerplate.h>
+// Note: Instantiations are generated to prevent compiler memory issues
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    template class sd::IndexReductionLoops,
+    IndexReductionLoops,
+    ;
+)

libnd4j/include/loops/cpu/comb_compilation_units/pairwise_instantiation_template_2.cpp.in

@@ -0,0 +1,40 @@
+// pairwise_instantiation_template_3.cpp.in
+
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 agibsonccc on 11/22/24.
+//
+
+#include <loops/pairwise_instantiations.h>
+#include <types/type_promote.h>
+#include <loops/cpu/pairwise.hpp>
+// Note: Instantiations are generated to prevent compiler memory issues
+
+
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    CALLBACK_INSTANTIATE_PROMOTE,
+    promote,
+    ;
+)

libnd4j/include/loops/cpu/comb_compilation_units/pairwise_instantiation_template_3.cpp.in

@@ -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
+*  *****************************************************************************
+ */
+
+//
+// Created by agibsonccc on 11/22/24.
+//
+
+#include <loops/pairwise_instantiations.h>
+#include <types/type_promote.h>
+#include <loops/cpu/pairwise.hpp>
+
+/*
+*
+* Function Instantiation:
+* PairWiseTransform::exec instantiated for types in @COMB1@, @COMB2@, @COMB3@
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::pairwise_transforms::PairWiseTransform,
+    ::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+     const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+     void *extraParams, sd::LongType start, sd::LongType stop)
+)

libnd4j/include/loops/cpu/comb_compilation_units/random_instantiation_template_1.cpp.in

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/random.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * RandomFunction instantiated for float types in @COMB1@
+ */
+
+#define LIST_CALLBACK_EXEC_TRANSFORM_1(INPUT) \
+    template void functions::random::RandomFunction<GET_SECOND(INPUT)>::execTransform( \
+        int opNum, \
+        sd::Pointer state, \
+        const void *x, \
+        const sd::LongType *xShapeInfo, \
+        void *z, \
+        const sd::LongType *zShapeInfo, \
+        void *extraArguments);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC_TRANSFORM_1)
+
+#define LIST_CALLBACK_EXEC_TRANSFORM_2(INPUT) \
+    template void functions::random::RandomFunction<GET_SECOND(INPUT)>::execTransform( \
+        int opNum, \
+        sd::Pointer state, \
+        const void *x, \
+        const sd::LongType *xShapeInfo, \
+        const void *y, \
+        const sd::LongType *yShapeInfo, \
+        void *z, \
+        const sd::LongType *zShapeInfo, \
+        void *extraArguments);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC_TRANSFORM_2)
+
+#define LIST_CALLBACK_EXEC_TRANSFORM_3(INPUT) \
+    template void functions::random::RandomFunction<GET_SECOND(INPUT)>::execTransform( \
+        int opNum, \
+        sd::Pointer state, \
+        void *z, \
+        const sd::LongType *zShapeInfo, \
+        void *extraArguments);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC_TRANSFORM_3)

libnd4j/include/loops/cpu/comb_compilation_units/reduce3_instantiation_template_2.cpp.in

@@ -0,0 +1,62 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/reduce3.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * Reduce3 instantiated for types @COMB1@, @COMB2@
+ */
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce3::Reduce3,
+    ::execScalar(const int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals, const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce3::Reduce3,
+    ::exec(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals, const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, sd::LongType start, sd::LongType stop)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce3::Reduce3,
+    ::exec(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals, const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffsets, sd::LongType start, sd::LongType stop)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce3::Reduce3,
+    ::execAll(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals, const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo, const sd::LongType *xOffsets, const sd::LongType *yTadShapeInfo, const sd::LongType *yOffsets, sd::LongType start, sd::LongType stop)
+)

libnd4j/include/loops/cpu/comb_compilation_units/reduce_bool_instantiation_template_2.cpp.in

@@ -0,0 +1,54 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/reduce/reduce_bool.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * ReduceBoolFunction instantiated for types @COMB1@, @COMB2@
+ */
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_BOOL_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceBoolFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_BOOL_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceBoolFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams, void *z, const sd::LongType *zShapeInfo)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_BOOL_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceBoolFunction,
+    ::exec(int opNum, sd::memory::Workspace *workspace, const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz, const sd::LongType *zShapeInfo, const sd::LongType *dims)
+)

libnd4j/include/loops/cpu/comb_compilation_units/reduce_float_instantiation_template_2.cpp.in

@@ -0,0 +1,54 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/reduce/reduce_float.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * ReduceFloatFunction instantiated for types @COMB1@, @COMB2@
+ */
+
+ITERATE_COMBINATIONS(
+    SD_NUMERIC_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceFloatFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams)
+)
+
+ITERATE_COMBINATIONS(
+    SD_NUMERIC_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceFloatFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams, void *z, const sd::LongType *zShapeInfo)
+)
+
+ITERATE_COMBINATIONS(
+    SD_NUMERIC_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceFloatFunction,
+    ::exec(int opNum, sd::memory::Workspace *workspace, const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz, const sd::LongType *zShapeInfo, const sd::LongType *dims)
+)

libnd4j/include/loops/cpu/comb_compilation_units/reduce_long_instantiation_template_2.cpp.in

@@ -0,0 +1,54 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/reduce/reduce_long.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * ReduceLongFunction instantiated for types @COMB1@, @COMB2@
+ */
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_LONG_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceLongFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_LONG_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceLongFunction,
+    ::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams, void *z, const sd::LongType *zShapeInfo)
+)
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_LONG_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION_2,
+    functions::reduce::ReduceLongFunction,
+    ::exec(const int opNum, sd::memory::Workspace *workspace, const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz, const sd::LongType *zShapeInfo, sd::LongType *dims)
+)

libnd4j/include/loops/cpu/comb_compilation_units/reduce_same_instantiation_template_1.cpp.in

@@ -0,0 +1,63 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 1/3/25.
+//
+
+#include <loops/cpu/reduce/reduce_same.hpp>
+#include <loops/pairwise_instantiations.h>
+#include <system/type_boilerplate.h>
+
+/*
+ * Function Instantiation:
+ * ReduceSameFunction instantiated for common types in @COMB1@
+ */
+
+#define LIST_CALLBACK_EXEC_SCALAR_1(INPUT) \
+    template void* functions::reduce::ReduceSameFunction<GET_SECOND(INPUT)>::execScalar( \
+        const int opNum, \
+        const void *x, \
+        const sd::LongType *xShapeInfo, \
+        void *extraParams);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC_SCALAR_1)
+
+#define LIST_CALLBACK_EXEC_SCALAR_2(INPUT) \
+    template void functions::reduce::ReduceSameFunction<GET_SECOND(INPUT)>::execScalar( \
+        const int opNum, \
+        const void *x, \
+        const sd::LongType *xShapeInfo, \
+        void *extraParams, \
+        void *z, \
+        const sd::LongType *zShapeInfo);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC_SCALAR_2)
+
+#define LIST_CALLBACK_EXEC(INPUT) \
+    template void functions::reduce::ReduceSameFunction<GET_SECOND(INPUT)>::exec( \
+        int opNum, \
+        sd::memory::Workspace *workspace, \
+        const void *vx, \
+        const sd::LongType *xShapeInfo, \
+        void *vextraParams, \
+        void *vz, \
+        const sd::LongType *zShapeInfo, \
+        const sd::LongType *dims);
+
+ITERATE_LIST((SD_COMMON_TYPES_PART_@COMB1@), LIST_CALLBACK_EXEC)

libnd4j/include/loops/cpu/comb_compilation_units/scalar_instantiation_template_3.cpp.in

@@ -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
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 12/30/24.
+//
+
+#include <loops/pairwise_instantiations.h>
+#include <types/type_promote.h>
+#include <loops/cpu/scalar.hpp>
+
+/*
+ *
+ * Function Instantiation:
+ * ScalarTransform::transform instantiated for types in @COMB1@, @COMB2@, @COMB3@
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::scalar::ScalarTransform,
+    ::transform(int opNum, const void *x, const sd::LongType *xShapeInfo, void *result,
+                const sd::LongType *resultShapeInfo, const void *scalar, void *extraParams, sd::LongType start,
+                sd::LongType stop)
+)

libnd4j/include/loops/cpu/compilation_units/broadcast_bool_p.cpp.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 23/09/18.
+//
+#include <loops/cpu/broadcasting_bool.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace broadcast {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN BroadcastBool, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_BOOL_TYPES);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/broadcast_int_p.cpp.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 23/09/18.
+//
+#include <loops/cpu/broadcasting_int.hpp>
+#cmakedefine SD_INTEGER_TYPES_GEN
+#if defined(SD_INTEGER_TYPES_GEN) && defined(SD_INTEGER_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace broadcast {
+        BUILD_SINGLE_TEMPLATE(template class SD_LIB_HIDDEN BroadcastInt, , SD_INTEGER_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/broadcast_p.cpp.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 23/09/18.
+//
+#include <loops/cpu/broadcasting.hpp>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace broadcast {
+        BUILD_PAIRWISE_TEMPLATE(template class SD_LIB_HIDDEN Broadcast, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/indexreduce_index0.cpp.in

@@ -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 <loops/cpu/indexreduce.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_INDEXING_TYPES_0)
+namespace functions {
+    namespace indexreduce {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN IndexReduce, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_INDEXING_TYPES_0);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/indexreduce_index1.cpp.in

@@ -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 <loops/cpu/indexreduce.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_INDEXING_TYPES_1)
+namespace functions {
+    namespace indexreduce {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN IndexReduce, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_INDEXING_TYPES_1);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/pairwise_p.cpp.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 23/09/18.
+//
+#include <loops/cpu/pairwise.hpp>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace pairwise_transforms {
+        BUILD_PAIRWISE_TEMPLATE(template class SD_LIB_HIDDEN PairWiseTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/random.cpp.in

@@ -0,0 +1,37 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <loops/cpu/random.hpp>
+#cmakedefine SD_FLOAT_TYPES_GEN
+#if defined(SD_FLOAT_TYPES_GEN) && defined(SD_FLOAT_TYPES_@FL_TYPE_INDEX@)
+
+// Helper macros for stringification
+#define STR_HELPER(x) #x
+#define STR(x) STR_HELPER(x)
+
+// Emit the value of SD_PAIRWISE_TYPES_0 during compilation
+#pragma message("SD_PAIRWISE_TYPES_0: " STR((SD_PAIRWISE_TYPES_0)))
+namespace functions {
+    namespace random {
+        BUILD_SINGLE_TEMPLATE(template class SD_LIB_HIDDEN RandomFunction, , SD_FLOAT_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/reduce3_float_types0.cpp.in

@@ -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 <loops/cpu/reduce3.hpp>
+#cmakedefine  SD_COMMON_TYPES_GEN
+#if defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_FLOAT_TYPES_0)
+namespace functions {
+    namespace reduce3 {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_FLOAT_TYPES_0);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/reduce3_float_types1.cpp.in

@@ -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 <loops/cpu/reduce3.hpp>
+#cmakedefine  SD_COMMON_TYPES_GEN
+#if defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_FLOAT_TYPES_1)
+namespace functions {
+    namespace reduce3 {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_FLOAT_TYPES_1);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/reduce3_float_types3.cpp.in

@@ -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 <loops/cpu/reduce3.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_FLOAT_TYPES_3)
+namespace functions {
+    namespace reduce3 {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_FLOAT_TYPES_3);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/reduce3_float_types_2.cpp.in

@@ -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 <loops/cpu/reduce3.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && defined(SD_FLOAT_TYPES_2)
+namespace functions {
+    namespace reduce3 {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES_@FL_TYPE_INDEX@, SD_FLOAT_TYPES_2);
+    }
+}
+#endif

libnd4j/include/loops/cpu/compilation_units/reduce_bool.cpp.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cpu/reduce/reduce_bool.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace reduce {
+        BUILD_DOUBLE_TEMPLATE(template class ReduceBoolFunction, , SD_COMMON_TYPES_@FL_TYPE_INDEX@ , SD_BOOL_TYPES);
+    }
+}
+#endif
+

libnd4j/include/loops/cpu/compilation_units/reduce_float.cpp.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cpu/reduce/reduce_float.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && COUNT_NARG(SD_FLOAT_TYPES) > 0
+namespace functions {
+    namespace reduce {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN ReduceFloatFunction, , SD_COMMON_TYPES_@FL_TYPE_INDEX@ , SD_FLOAT_TYPES);
+    }
+}
+#endif
+

libnd4j/include/loops/cpu/compilation_units/reduce_long.cpp.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cpu/reduce/reduce_long.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@) && COUNT_NARG(SD_LONG_TYPES) > 0
+namespace functions {
+    namespace reduce {
+        BUILD_DOUBLE_TEMPLATE(template class ReduceLongFunction, , SD_COMMON_TYPES_@FL_TYPE_INDEX@ , SD_LONG_TYPES);
+    }
+}
+#endif
+

libnd4j/include/loops/cpu/compilation_units/reduce_same.cpp.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cpu/reduce/reduce_same.hpp>
+#cmakedefine SD_COMMON_TYPES_GEN
+#if  defined(SD_COMMON_TYPES_GEN) && defined(SD_COMMON_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace reduce {
+        BUILD_SINGLE_TEMPLATE(template class ReduceSameFunction, , SD_COMMON_TYPES_@FL_TYPE_INDEX@ );
+    }
+}
+#endif
+

libnd4j/include/loops/cpu/compilation_units/scalar_p.cpp.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver on 9/28/2018.
+//
+#include <loops/cpu/scalar.hpp>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace scalar {
+        BUILD_PAIRWISE_TEMPLATE(template class SD_LIB_HIDDEN ScalarTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cpu/indexreduce.hpp

@@ -0,0 +1,141 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+* License for the specific language governing permissions and limitations
+* under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+//
+// Created by raver on 4/9/2018.
+//
+#include <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <loops/indexreduce.h>
+#include <loops/legacy_ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace indexreduce {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+sd::LongType IndexReduce<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                          void *extraParams) {
+ RETURNING_DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams), INDEX_REDUCE_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+void IndexReduce<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams, void *z,
+                            const sd::LongType *zShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength,
+                            const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {
+ DISPATCH_BY_OPNUM_TT(
+     exec, PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo, dimension, dimensionLength, tadShapeInfo, tadOffset),
+     INDEX_REDUCE_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+sd::LongType IndexReduce<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams) {
+ auto x = reinterpret_cast<const X *>(vx);
+ auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+ auto startingIndex = OpType::startingIndexValue(x);
+ auto len = shape::length(xShapeInfo);
+ sd::OmpLaunchHelper info(len);
+
+ // Cache shape-related values
+ sd::LongType xRank = shape::rank(xShapeInfo);
+ sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+ sd::LongType *xStride = shape::stride(xShapeInfo);
+
+ sd::LongType xShapeInfoCast[SD_MAX_RANK];
+ bool canCastX = sd::DataTypeUtils::castShapeInfo(xShapeInfo, xShapeInfoCast);
+ int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+ IndexValue<X> intermediatery[64];
+ for (int e = 0; e < maxThreads; e++) intermediatery[e].index = -1;
+
+ auto func = PRAGMA_THREADS_FOR {
+   intermediatery[thread_id] = OpType::startingIndexValue(x);
+
+   for (auto i = start; i < stop; i++) {
+     sd::LongType coords[SD_MAX_RANK];
+     INDEX2COORDS(i, xRank, xShape, coords);
+     sd::LongType offset;
+     COORDS2INDEX(xRank, xStride, coords, offset);
+     IndexValue<X> curr(x[offset], i);
+     intermediatery[thread_id] = OpType::update(intermediatery[thread_id], curr, extraParams);
+   }
+ };
+
+ maxThreads = samediff::Threads::parallel_for(func, 0, len, 1, maxThreads);
+
+ for (int e = 0; e < maxThreads; e++) startingIndex = OpType::update(startingIndex, intermediatery[e], extraParams);
+
+ return startingIndex.index;
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void IndexReduce<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                            const sd::LongType *zShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength,
+                            const sd::LongType *tadShapeInfo, const sd::LongType *tadOffset) {
+ auto x = reinterpret_cast<X *>(const_cast<void *>(vx));
+ auto z = reinterpret_cast<Z *>(vz);
+ auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+ const sd::LongType zLen = shape::length(zShapeInfo);
+
+ if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+   if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+   const auto indexValue = OpType::startingIndexValue(x);
+
+   for (sd::LongType i = 0; i < zLen; i++) z[i] = (Z)indexValue.index;
+
+   return;
+ }
+
+ if (shape::isScalar(zShapeInfo)) {
+   z[0] = (Z)execScalar<OpType>(x, xShapeInfo, extraParams);
+   return;
+ }
+
+ auto tadOnlyShapeInfo = tadShapeInfo;
+ auto tadOffsets = tadOffset;
+
+ // When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
+ std::shared_ptr<sd::TadPack> tadPack = nullptr;
+
+ if (tadOnlyShapeInfo == nullptr || tadOffsets == nullptr) {
+   if (dimensionLength < 1) return;
+
+   tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(xShapeInfo), dimension,
+                                                                   dimensionLength);
+   tadOnlyShapeInfo = tadPack->primaryShapeInfo();
+   tadOffsets = tadPack->primaryOffsets();
+ }
+
+ // Let IndexReductionLoops handle the shape caching internally since it's a separate component
+ sd::IndexReductionLoops<X, Z>::template loopIndexReduce<OpType>(x, xShapeInfo, z, zShapeInfo, tadOnlyShapeInfo,
+                                                                 tadOffsets, vextraParams);
+}
+
+}  // namespace indexreduce
+}  // namespace functions

libnd4j/include/loops/cpu/indexreduce_instantiations.cpp

@@ -0,0 +1,599 @@
+#include <helpers/cpu/loops/IndexReductionLoops.hpp>
+#include <types/types.h>
+#include <system/type_boilerplate.h>
+#include "indexreduce.hpp"
+
+// Macro to conditionally instantiate based on HAS_* flags
+#define CONDITIONAL_INSTANTIATE(TYPE_TUPLE, MACRO_CALL) \
+    CONDITIONAL_INSTANTIATE_IMPL(GET_SECOND(TYPE_TUPLE), MACRO_CALL)
+
+#define CONDITIONAL_INSTANTIATE_IMPL(TYPE, MACRO_CALL) \
+    CONDITIONAL_INSTANTIATE_##TYPE(MACRO_CALL)
+
+// Define conditional instantiation for each type
+#ifdef HAS_FLOAT32
+#define CONDITIONAL_INSTANTIATE_float(MACRO_CALL) MACRO_CALL((FLOAT32, float))
+#else
+#define CONDITIONAL_INSTANTIATE_float(MACRO_CALL)
+#endif
+
+#ifdef HAS_DOUBLE
+#define CONDITIONAL_INSTANTIATE_double(MACRO_CALL) MACRO_CALL((DOUBLE, double))
+#else
+#define CONDITIONAL_INSTANTIATE_double(MACRO_CALL)
+#endif
+
+#ifdef HAS_FLOAT16
+#define CONDITIONAL_INSTANTIATE_float16(MACRO_CALL) MACRO_CALL((HALF, float16))
+#else
+#define CONDITIONAL_INSTANTIATE_float16(MACRO_CALL)
+#endif
+
+#ifdef HAS_BFLOAT16
+#define CONDITIONAL_INSTANTIATE_bfloat16(MACRO_CALL) MACRO_CALL((BFLOAT16, bfloat16))
+#else
+#define CONDITIONAL_INSTANTIATE_bfloat16(MACRO_CALL)
+#endif
+
+#ifdef HAS_INT8
+#define CONDITIONAL_INSTANTIATE_int8_t(MACRO_CALL) MACRO_CALL((INT8, int8_t))
+#else
+#define CONDITIONAL_INSTANTIATE_int8_t(MACRO_CALL)
+#endif
+
+#ifdef HAS_INT16
+#define CONDITIONAL_INSTANTIATE_int16_t(MACRO_CALL) MACRO_CALL((INT16, int16_t))
+#else
+#define CONDITIONAL_INSTANTIATE_int16_t(MACRO_CALL)
+#endif
+
+#ifdef HAS_INT32
+#define CONDITIONAL_INSTANTIATE_Int32Type(MACRO_CALL) MACRO_CALL((INT32, Int32Type))
+#else
+#define CONDITIONAL_INSTANTIATE_Int32Type(MACRO_CALL)
+#endif
+
+#ifdef HAS_LONG
+#define CONDITIONAL_INSTANTIATE_LongType(MACRO_CALL) MACRO_CALL((INT64, LongType))
+#else
+#define CONDITIONAL_INSTANTIATE_LongType(MACRO_CALL)
+#endif
+
+#ifdef HAS_UINT8
+#define CONDITIONAL_INSTANTIATE_uint8_t(MACRO_CALL) MACRO_CALL((UINT8, uint8_t))
+#else
+#define CONDITIONAL_INSTANTIATE_uint8_t(MACRO_CALL)
+#endif
+
+#ifdef HAS_UINT16
+#define CONDITIONAL_INSTANTIATE_uint16_t(MACRO_CALL) MACRO_CALL((UINT16, uint16_t))
+#else
+#define CONDITIONAL_INSTANTIATE_uint16_t(MACRO_CALL)
+#endif
+
+#ifdef HAS_UINT32
+#define CONDITIONAL_INSTANTIATE_uint32_t(MACRO_CALL) MACRO_CALL((UINT32, uint32_t))
+#else
+#define CONDITIONAL_INSTANTIATE_uint32_t(MACRO_CALL)
+#endif
+
+#ifdef HAS_UNSIGNEDLONG
+#define CONDITIONAL_INSTANTIATE_unsignedlong(MACRO_CALL) MACRO_CALL((UINT64, unsigned long long))
+#else
+#define CONDITIONAL_INSTANTIATE_unsignedlong(MACRO_CALL)
+#endif
+
+#ifdef HAS_BOOL
+#define CONDITIONAL_INSTANTIATE_bool(MACRO_CALL) MACRO_CALL((BOOL, bool))
+#else
+#define CONDITIONAL_INSTANTIATE_bool(MACRO_CALL)
+#endif
+
+// Now instantiate the operations with int32_t index type (only if HAS_INT32 is defined)
+#ifdef HAS_INT32
+
+#define LIST_CALLBACK_OPS_int32_t(INPUT) \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::IndexMin<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::LastIndex<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::FirstIndex<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), int32_t>::loopIndexReduce<simdOps::IndexMax<GET_SECOND(INPUT), int32_t>>(GET_SECOND(INPUT)*, const sd::LongType*, int32_t*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*);
+
+// Conditionally instantiate for each type
+#ifdef HAS_FLOAT32
+LIST_CALLBACK_OPS_int32_t((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+LIST_CALLBACK_OPS_int32_t((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+LIST_CALLBACK_OPS_int32_t((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+LIST_CALLBACK_OPS_int32_t((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+LIST_CALLBACK_OPS_int32_t((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+LIST_CALLBACK_OPS_int32_t((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+LIST_CALLBACK_OPS_int32_t((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+LIST_CALLBACK_OPS_int32_t((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+LIST_CALLBACK_OPS_int32_t((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+LIST_CALLBACK_OPS_int32_t((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+LIST_CALLBACK_OPS_int32_t((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+LIST_CALLBACK_OPS_int32_t((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+LIST_CALLBACK_OPS_int32_t((BOOL, bool))
+#endif
+
+#endif // HAS_INT32
+
+// And with long long index type (only if HAS_LONG is defined)
+#ifdef HAS_LONG
+
+#define LIST_CALLBACK_OPS_LONG(INPUT) \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::IndexMin<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::LastIndex<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::FirstIndex<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*); \
+    template void sd::IndexReductionLoops<GET_SECOND(INPUT), sd::LongType>::loopIndexReduce<simdOps::IndexMax<GET_SECOND(INPUT), sd::LongType>>(GET_SECOND(INPUT)*, const sd::LongType*, sd::LongType*, const sd::LongType*, const sd::LongType*, const sd::LongType*, void*);
+
+// Conditionally instantiate for each type
+#ifdef HAS_FLOAT32
+LIST_CALLBACK_OPS_LONG((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+LIST_CALLBACK_OPS_LONG((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+LIST_CALLBACK_OPS_LONG((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+LIST_CALLBACK_OPS_LONG((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+LIST_CALLBACK_OPS_LONG((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+LIST_CALLBACK_OPS_LONG((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+LIST_CALLBACK_OPS_LONG((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+LIST_CALLBACK_OPS_LONG((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+LIST_CALLBACK_OPS_LONG((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+LIST_CALLBACK_OPS_LONG((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+LIST_CALLBACK_OPS_LONG((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+LIST_CALLBACK_OPS_LONG((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+LIST_CALLBACK_OPS_LONG((BOOL, bool))
+#endif
+
+#endif // HAS_LONG
+
+// Instantiate execScalar methods with int32_t
+#ifdef HAS_INT32
+
+#define INSTANTIATE_EXEC_SCALAR_INT32(INPUT) \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::IndexMin<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::LastIndex<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::FirstIndex<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar<simdOps::IndexMax<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_EXEC_SCALAR_INT32((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_EXEC_SCALAR_INT32((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_EXEC_SCALAR_INT32((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_EXEC_SCALAR_INT32((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_EXEC_SCALAR_INT32((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_EXEC_SCALAR_INT32((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_EXEC_SCALAR_INT32((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_EXEC_SCALAR_INT32((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_EXEC_SCALAR_INT32((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_EXEC_SCALAR_INT32((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_EXEC_SCALAR_INT32((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_EXEC_SCALAR_INT32((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_EXEC_SCALAR_INT32((BOOL, bool))
+#endif
+
+#endif // HAS_INT32
+
+// Instantiate execScalar methods with sd::LongType
+#ifdef HAS_LONG
+
+#define INSTANTIATE_EXEC_SCALAR_LONG(INPUT) \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::IndexMin<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::LastIndex<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::FirstIndex<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*); \
+    template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar<simdOps::IndexMax<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_EXEC_SCALAR_LONG((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_EXEC_SCALAR_LONG((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_EXEC_SCALAR_LONG((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_EXEC_SCALAR_LONG((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_EXEC_SCALAR_LONG((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_EXEC_SCALAR_LONG((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_EXEC_SCALAR_LONG((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_EXEC_SCALAR_LONG((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_EXEC_SCALAR_LONG((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_EXEC_SCALAR_LONG((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_EXEC_SCALAR_LONG((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_EXEC_SCALAR_LONG((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_EXEC_SCALAR_LONG((BOOL, bool))
+#endif
+
+#endif // HAS_LONG
+
+// Instantiate exec methods with int32_t
+#ifdef HAS_INT32
+
+#define INSTANTIATE_EXEC_INT32(INPUT) \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::IndexMin<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::LastIndex<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::FirstIndex<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec<simdOps::IndexMax<GET_SECOND(INPUT), int32_t>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_EXEC_INT32((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_EXEC_INT32((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_EXEC_INT32((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_EXEC_INT32((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_EXEC_INT32((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_EXEC_INT32((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_EXEC_INT32((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_EXEC_INT32((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_EXEC_INT32((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_EXEC_INT32((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_EXEC_INT32((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_EXEC_INT32((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_EXEC_INT32((BOOL, bool))
+#endif
+
+#endif // HAS_INT32
+
+// Instantiate exec methods with sd::LongType
+#ifdef HAS_LONG
+
+#define INSTANTIATE_EXEC_LONG(INPUT) \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::IndexMin<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::IndexAbsoluteMin<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::LastIndex<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::FirstIndex<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::IndexAbsoluteMax<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*); \
+    template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec<simdOps::IndexMax<GET_SECOND(INPUT), sd::LongType>>(const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_EXEC_LONG((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_EXEC_LONG((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_EXEC_LONG((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_EXEC_LONG((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_EXEC_LONG((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_EXEC_LONG((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_EXEC_LONG((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_EXEC_LONG((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_EXEC_LONG((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_EXEC_LONG((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_EXEC_LONG((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_EXEC_LONG((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_EXEC_LONG((BOOL, bool))
+#endif
+
+#endif // HAS_LONG
+
+// Add these instantiations for the non-templated methods (the ones with int opNum parameter)
+#ifdef HAS_INT32
+
+#define INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32(INPUT) \
+   template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::execScalar(int, const void*, const sd::LongType*, void*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_INT32((BOOL, bool))
+#endif
+
+#endif // HAS_INT32
+
+#ifdef HAS_LONG
+
+#define INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG(INPUT) \
+   template sd::LongType functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::execScalar(int, const void*, const sd::LongType*, void*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_NON_TEMPLATED_EXEC_SCALAR_LONG((BOOL, bool))
+#endif
+
+#endif // HAS_LONG
+
+#ifdef HAS_INT32
+
+#define INSTANTIATE_NON_TEMPLATED_EXEC_INT32(INPUT) \
+   template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), int32_t>::exec(int, const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_NON_TEMPLATED_EXEC_INT32((BOOL, bool))
+#endif
+
+#endif // HAS_INT32
+
+#ifdef HAS_LONG
+
+#define INSTANTIATE_NON_TEMPLATED_EXEC_LONG(INPUT) \
+   template void functions::indexreduce::IndexReduce<GET_SECOND(INPUT), sd::LongType>::exec(int, const void*, const sd::LongType*, void*, void*, const sd::LongType*, sd::LongType*, sd::LongType, const sd::LongType*, const sd::LongType*);
+
+#ifdef HAS_FLOAT32
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((FLOAT32, float))
+#endif
+#ifdef HAS_DOUBLE
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((DOUBLE, double))
+#endif
+#ifdef HAS_FLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((HALF, float16))
+#endif
+#ifdef HAS_BFLOAT16
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((BFLOAT16, bfloat16))
+#endif
+#ifdef HAS_INT8
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((INT8, int8_t))
+#endif
+#ifdef HAS_INT16
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((INT16, int16_t))
+#endif
+#ifdef HAS_INT32
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((INT32, Int32Type))
+#endif
+#ifdef HAS_LONG
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((INT64, LongType))
+#endif
+#ifdef HAS_UINT8
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((UINT8, uint8_t))
+#endif
+#ifdef HAS_UINT16
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((UINT16, uint16_t))
+#endif
+#ifdef HAS_UINT32
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((UINT32, uint32_t))
+#endif
+#ifdef HAS_UNSIGNEDLONG
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((UINT64, uint64_t))
+#endif
+#ifdef HAS_BOOL
+INSTANTIATE_NON_TEMPLATED_EXEC_LONG((BOOL, bool))
+#endif
+
+#endif // HAS_LONG

libnd4j/include/loops/cpu/pairwise.hpp

@@ -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
+ ******************************************************************************/
+
+//
+// Created by remote on 2018-09-20.
+//
+#include <execution/Threads.h>
+#include <helpers/LoopKind.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <helpers/shape.h>
+#include <loops/pairwise_transform.h>
+#include <math/templatemath.h>
+#include <ops/ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <array/ArrayOptions.hXX>
+
+using namespace simdOps;
+
+namespace functions {
+namespace pairwise_transforms {
+
+template <typename X, typename Y, typename Z>
+SD_INLINE void PairWiseTransform<X, Y, Z>::exec(int opNum,
+                                                const void *x,
+                                                sd::LongType xEws,
+                                                const void *y,
+                                                sd::LongType yEws,
+                                                void *z,
+                                                sd::LongType zEws,
+                                                void *extraParams,
+                                                sd::LongType n,
+                                                sd::LongType start,
+                                                sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(exec, PARAMS(x, xEws, y, yEws, z, zEws, extraParams, n, start, stop), PAIRWISE_TRANSFORM_OPS);
+};
+
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+SD_INLINE void PairWiseTransform<X, Y, Z>::exec(const void *vx, sd::LongType xEws, const void *vy, sd::LongType yEws, void *vz,
+                                                sd::LongType zEws, void *vextraParams, sd::LongType n, sd::LongType start,
+                                                sd::LongType stop) {
+
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const Y *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+
+  if (xEws == 1 && yEws == 1 && zEws == 1) {
+    for (sd::LongType i = start; i < stop; i++) {
+      z[i] = OpType::op(x[i], y[i], extraParams);
+    }
+
+  } else {
+    for (sd::LongType i = start; i < stop; i++) z[i * zEws] = OpType::op(x[i * xEws], y[i * yEws], extraParams);
+  }
+
+
+}
+
+template <typename X, typename Y, typename Z>
+SD_INLINE void PairWiseTransform<X, Y, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                                const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                                void *extraParams, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, start, stop),
+                        PAIRWISE_TRANSFORM_OPS);
+};
+
+
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+SD_INLINE void PairWiseTransform<X, Y, Z>::exec(const void *vx,
+                                                const sd::LongType *xShapeInfo,
+                                                const void *vy,
+                                                const sd::LongType *yShapeInfo,
+                                                void *vz,
+                                                const sd::LongType *zShapeInfo,
+                                                void *vextraParams,
+                                                sd::LongType start,
+                                                sd::LongType stop) {
+
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const Y *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+  sd::LongType  xRank = shape::rank(xShapeInfo);
+  sd::LongType  yRank = shape::rank(yShapeInfo);
+  sd::LongType  zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *yShape = shape::shapeOf(yShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *yStride = shape::stride(yShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+  bool allSameOrder = shape::order(xShapeInfo) == shape::order(yShapeInfo) && shape::order(xShapeInfo) == shape::order(zShapeInfo);
+  if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)
+      && shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)
+      && !shape::isViewConst(xShapeInfo)
+      &&  !shape::isViewConst(yShapeInfo) && !shape::isViewConst(zShapeInfo)
+      && allSameOrder) {
+    sd::LongType zCoords[SD_MAX_RANK];
+
+
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, zRank,zShape, zCoords);
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+      z[zOffset] = OpType::op(x[zOffset], y[zOffset], extraParams);
+    }
+  } else if ((shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)
+              || shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo))
+             &&  allSameOrder) {
+    sd::LongType zCoords[SD_MAX_RANK];
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, zRank, zShape, zCoords);
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, zCoords, xOffset);
+      COORDS2INDEX(yRank, yStride, zCoords, yOffset);
+      COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+    }
+  } else if (shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo)
+             && !shape::isViewConst(xShapeInfo)
+             && !shape::isViewConst(yShapeInfo) && !shape::isViewConst(zShapeInfo)
+             && allSameOrder) {
+    sd::LongType zCoords[SD_MAX_RANK];
+
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, zRank, zShape, zCoords);
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, zCoords, xOffset);
+      COORDS2INDEX(yRank, yStride, zCoords, yOffset);
+      COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+    }
+  } else {
+    sd::LongType zCoords[SD_MAX_RANK];
+
+    for (sd::LongType i = start; i < stop; i++) {
+      INDEX2COORDS(i, zRank,xShape, zCoords);
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, zCoords, xOffset);
+      COORDS2INDEX(yRank, yStride, zCoords, yOffset);
+      COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+    }
+  }
+}
+}  // namespace pairwise_transforms
+}  // namespace functions
+

libnd4j/include/loops/cpu/pairwise_bool.cpp

@@ -0,0 +1,151 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+* License for the specific language governing permissions and limitations
+* under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+//
+// Created by remote on 2018-09-20.
+//
+#include <loops/pairwise_bool.h>
+#include <types/types.h>
+#include <helpers/LoopKind.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <execution/Threads.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace pairwise_transforms {
+
+
+template <typename X, typename Z>
+void PairWiseBoolTransform<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                       const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                       void *extraParams, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, start, stop),
+                       PAIRWISE_BOOL_OPS);
+};
+
+template <typename X, typename Z>
+template <typename OpType>
+void PairWiseBoolTransform<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                       const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                       void *vextraParams, sd::LongType start, sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *yShape = shape::shapeOf(yShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *yStride = shape::stride(yShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  sd::LongType n = shape::length(xShapeInfo);
+
+  if (shape::isScalar(yShapeInfo)) {
+    if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, zShape, coords);
+        sd::LongType offset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        z[offset] = OpType::op(x[offset], static_cast<X>(y[0]), static_cast<X*>(extraParams));
+      };
+    } else {
+      PRAGMA_OMP_SIMD
+      for (sd::LongType i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType xOffset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpType::op(x[xOffset], y[0], extraParams);
+      };
+    }
+    return;
+  }
+
+  const sd::LoopKind::Kind kindOfLoop = sd::LoopKind::deduceKindOfLoopXYZ(xShapeInfo, yShapeInfo, zShapeInfo);
+  const bool sameShapesXY = shape::shapeEquals(xShapeInfo, yShapeInfo);
+  const bool isSameLength = shape::length(xShapeInfo) == shape::length(yShapeInfo);
+
+
+  if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo) &&
+      shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i, xRank, xShape, coords);
+      sd::LongType offset;
+      COORDS2INDEX(xRank, xStride, coords, offset);
+      z[offset] = OpType::op(x[offset], y[offset], extraParams);
+    };
+  } else if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i, xRank, xShape, coords);
+      sd::LongType offset, zOffset;
+      COORDS2INDEX(xRank, xStride, coords, offset);
+      COORDS2INDEX(zRank, zStride, coords, zOffset);
+      z[zOffset] = OpType::op(x[offset], y[offset], extraParams);
+    };
+  } else if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i, xRank, xShape, coords);
+      sd::LongType offset, yOffset;
+      COORDS2INDEX(xRank, xStride, coords, offset);
+      COORDS2INDEX(yRank, yStride, coords, yOffset);
+      z[offset] = OpType::op(x[offset], y[yOffset], extraParams);
+    };
+  } else if (shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i, yRank, yShape, coords);
+      sd::LongType xOffset, offset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+      COORDS2INDEX(yRank, yStride, coords, offset);
+      z[offset] = OpType::op(x[xOffset], y[offset], extraParams);
+    };
+  } else {
+    PRAGMA_OMP_SIMD
+    for (sd::LongType i = start; i < stop; i++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i, zRank, zShape, coords);
+      sd::LongType xOffset, yOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+      COORDS2INDEX(yRank, yStride, coords, yOffset);
+      COORDS2INDEX(zRank, zStride, coords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+    };
+  }
+
+}
+
+BUILD_DOUBLE_TEMPLATE( class PairWiseBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+}  // namespace pairwise_transforms
+}  // namespace functions

libnd4j/include/loops/cpu/pairwise_int.cpp

@@ -0,0 +1,153 @@
+/* ******************************************************************************
+*
+*
+* 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 <execution/Threads.h>
+#include <helpers/LoopKind.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <loops/pairwise_int.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace pairwise_transforms {
+
+
+
+template <typename X>
+void PairWiseIntTransform<X>::exec(const int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
+                                   const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
+                                   void *extraParams, const uint64_t start, const uint64_t stop) {
+  DISPATCH_BY_OPNUM_T(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, start, stop),
+                      PAIRWISE_INT_OPS);
+};
+
+template <typename X>
+template <typename OpType>
+void PairWiseIntTransform<X>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                   const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                   void *vextraParams, const uint64_t start, const uint64_t stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *yShape = shape::shapeOf(yShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *yStride = shape::stride(yShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  auto n = shape::length(xShapeInfo);
+  auto xEws = shape::elementWiseStride(xShapeInfo);
+  auto yEws = shape::elementWiseStride(yShapeInfo);
+  auto zEws = shape::elementWiseStride(zShapeInfo);
+
+  if (shape::isScalar(yShapeInfo)) {
+    if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType offset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        z[offset] = OpType::op(x[offset], y[0], extraParams);
+      };
+    } else {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType xOffset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpType::op(x[xOffset], y[0], extraParams);
+      };
+    }
+    return;
+  }
+
+  const sd::LoopKind::Kind kindOfLoop = sd::LoopKind::deduceKindOfLoopXYZ(xShapeInfo, yShapeInfo, zShapeInfo);
+  const bool sameShapesXY = shape::shapeEquals(xShapeInfo, yShapeInfo);
+
+    if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo) &&
+        shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType offset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        z[offset] = OpType::op(x[offset], y[offset], extraParams);
+      }
+    } else if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType offset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpType::op(x[offset], y[offset], extraParams);
+      }
+    } else if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType offset, yOffset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        COORDS2INDEX(yRank, yStride, coords, yOffset);
+        z[offset] = OpType::op(x[offset], y[yOffset], extraParams);
+      }
+    } else if (shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo)) {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, yRank, yShape, coords);
+        sd::LongType xOffset, offset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(yRank, yStride, coords, offset);
+        z[offset] = OpType::op(x[xOffset], y[offset], extraParams);
+      }
+    } else {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, zRank, zShape, coords);
+        sd::LongType xOffset, yOffset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(yRank, yStride, coords, yOffset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+      }
+    }
+
+}
+
+BUILD_SINGLE_TEMPLATE( class PairWiseIntTransform, , SD_COMMON_TYPES);
+}  // namespace pairwise_transforms
+}  // namespace functions

libnd4j/include/loops/cpu/random.hpp

@@ -0,0 +1,186 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+* License for the specific language governing permissions and limitations
+* under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+//
+//  @author raver119@gmail.com, created on 15.12.17.
+//  @author Yurii Shyrma (iuriish@yahoo.com)
+//
+#include <helpers/OmpLaunchHelper.h>
+#include <loops/random.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace randomOps;
+
+namespace functions {
+namespace random {
+
+template <typename X>
+template <typename OpClass>
+void RandomFunction<X>::execTransform(sd::Pointer state, const void *vx, const sd::LongType *xShapeInfo, const void *vy,
+                                      const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                                      void *vextraArguments) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraArguments = reinterpret_cast<X *>(vextraArguments);
+
+  if (OpClass::requiresSpecial) {
+    OpClass::specialOp(state, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraArguments);
+    return;
+  }
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType yRank = shape::rank(yShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *yShape = shape::shapeOf(yShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *yStride = shape::stride(yShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  auto length = shape::length(zShapeInfo);
+  sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
+
+  if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo) &&
+      shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    auto func = PRAGMA_THREADS_FOR {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        z[i] = OpClass::op(x[i], y[i], i, length, rng, extraArguments);
+      }
+    };
+    samediff::Threads::parallel_for(func, 0, length, 1);
+  } else {
+    sd::LongType coords[SD_MAX_RANK];
+    auto func = PRAGMA_THREADS_FOR {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType xOffset, yOffset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(yRank, yStride, coords, yOffset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpClass::op(x[xOffset], y[yOffset], i, length, rng, extraArguments);
+      }
+    };
+    samediff::Threads::parallel_for(func, 0, length, 1);
+  }
+}
+
+template <typename X>
+template <typename OpClass>
+void RandomFunction<X>::execTransform(sd::Pointer state, const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                      const sd::LongType *zShapeInfo, void *vextraArguments) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraArguments = reinterpret_cast<X *>(vextraArguments);
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  auto length = shape::length(zShapeInfo);
+  sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
+
+  if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    sd::LongType coords[SD_MAX_RANK];
+    auto func = PRAGMA_THREADS_FOR {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType offset;
+        COORDS2INDEX(xRank, xStride, coords, offset);
+        z[offset] = OpClass::op(x[offset], i, length, rng, extraArguments);
+      }
+    };
+    samediff::Threads::parallel_for(func, 0, length, 1);
+  } else {
+    sd::LongType coords[SD_MAX_RANK];
+    auto func = PRAGMA_THREADS_FOR {
+      PRAGMA_OMP_SIMD
+      for (auto i = start; i < stop; i++) {
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType xOffset, zOffset;
+        COORDS2INDEX(xRank, xStride, coords, xOffset);
+        COORDS2INDEX(zRank, zStride, coords, zOffset);
+        z[zOffset] = OpClass::op(x[xOffset], i, length, rng, extraArguments);
+      }
+    };
+    samediff::Threads::parallel_for(func, 0, length, 1);
+  }
+}
+
+template <typename X>
+template <typename OpClass>
+void RandomFunction<X>::execTransform(sd::Pointer state, void *vz, const sd::LongType *zShapeInfo,
+                                      void *vextraArguments) {
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraArguments = reinterpret_cast<X *>(vextraArguments);
+
+  // Cache shape-related values
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  auto length = shape::length(zShapeInfo);
+  sd::graph::RandomGenerator *rng = reinterpret_cast<sd::graph::RandomGenerator *>(state);
+
+  sd::LongType coords[SD_MAX_RANK];
+  auto func = PRAGMA_THREADS_FOR {
+    PRAGMA_OMP_SIMD
+    for (auto i = start; i < stop; i++) {
+      INDEX2COORDS(i, zRank, zShape, coords);
+      sd::LongType offset;
+      COORDS2INDEX(zRank, zStride, coords, offset);
+      z[offset] = OpClass::op(i, length, rng, extraArguments);
+    }
+  };
+
+  samediff::Threads::parallel_for(func, 0, length, 1);
+}
+
+// Dispatch functions remain unchanged as they just route to the optimized implementations
+template <typename X>
+void RandomFunction<X>::execTransform(int opNum, sd::Pointer state, const void *x, const sd::LongType *xShapeInfo,
+                                      void *z, const sd::LongType *zShapeInfo, void *extraArguments) {
+  DISPATCH_BY_OPNUM_T(execTransform, PARAMS(state, x, xShapeInfo, z, zShapeInfo, extraArguments), RANDOM_OPS)
+}
+
+template <typename X>
+void RandomFunction<X>::execTransform(int opNum, sd::Pointer state, const void *x, const sd::LongType *xShapeInfo,
+                                      const void *y, const sd::LongType *yShapeInfo, void *z,
+                                      const sd::LongType *zShapeInfo, void *extraArguments) {
+  DISPATCH_BY_OPNUM_T(execTransform, PARAMS(state, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraArguments),
+                      RANDOM_OPS)
+}
+
+template <typename X>
+void RandomFunction<X>::execTransform(int opNum, sd::Pointer state, void *z, const sd::LongType *zShapeInfo,
+                                      void *extraArguments) {
+  DISPATCH_BY_OPNUM_T(execTransform, PARAMS(state, z, zShapeInfo, extraArguments), RANDOM_OPS)
+}
+
+}  // namespace random
+}  // namespace functions

libnd4j/include/loops/cpu/reduce/reduce_bool.hpp

@@ -0,0 +1,198 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_bool.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+    using namespace simdOps;
+
+namespace functions {
+namespace reduce {
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceBoolFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams,
+                                                  void *vz, const sd::LongType *zShapeInfo) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const auto length = shape::length(xShapeInfo);
+
+  if (shape::isEmptyConst(xShapeInfo)) {
+    z[0] = OpType::startingValue(x);
+    return;
+  }
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+    const auto startingVal = OpType::startingValue(x);
+
+    for (sd::LongType i = 0; i < length; i++) {
+      z[i] = startingVal;
+    }
+    return;
+  }
+
+  auto startingValue = OpType::startingValue(x);
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  X intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = startingValue;
+  }
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+  if(shape::isViewConst(xShapeInfo)) {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType indexOffset;
+        COORDS2INDEX(xRank, xStride, coords, indexOffset);
+        intermediate[thread_id] = OpType::update(intermediate[thread_id], OpType::op(x[indexOffset], extraParams), extraParams);
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::update(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  } else {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        intermediate[thread_id] = OpType::update(intermediate[thread_id], OpType::op(x[i], extraParams), extraParams);
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::update(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  }
+}
+template <typename X, typename Z>
+template <typename OpType>
+Z SD_HOST ReduceBoolFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const sd::LongType length = shape::length(xShapeInfo);
+  auto startingValue = OpType::startingValue(x);
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+
+  for (sd::LongType i = 0; i < length; i++) {
+    sd::LongType coords[SD_MAX_RANK];
+    INDEX2COORDS(i, xRank, xShape, coords);
+    sd::LongType offset;
+    COORDS2INDEX(xRank, xStride, coords, offset);
+    startingValue = OpType::update(startingValue, OpType::op(x[offset], extraParams), extraParams);
+  }
+  return OpType::postProcess(startingValue, length, extraParams);
+}
+
+template <typename X, typename Z>
+Z ReduceBoolFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                       void *extraParams) {
+  RETURNING_DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams), REDUCE_BOOL_OPS);
+
+
+}
+
+template <typename X, typename Z>
+void ReduceBoolFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                          void *extraParams, void *z, const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo), REDUCE_BOOL_OPS);
+
+
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceBoolFunction<X, Z>::exec(const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                            void *vresult, const sd::LongType *resultShapeInfo) {
+  auto z = reinterpret_cast<Z *>(vresult);
+  z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceBoolFunction<X, Z>::exec(sd::memory::Workspace *workspace, const void *vx,
+                                            const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                            const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  Z *z = reinterpret_cast<Z *>(vz);
+  X *extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const sd::LongType xRank = shape::rank(xShapeInfo);
+  const sd::LongType zRank = shape::rank(zShapeInfo);
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    const auto startingVal = OpType::startingValue(x);
+    const auto zLen = shape::length(zShapeInfo);
+    if(z != nullptr)
+      for (sd::LongType i = 0; i < zLen; i++) {
+#if defined(PRINT_INDICES)
+        shape::printShapeInfo(xShapeInfo);
+        shape::printShapeInfo(zShapeInfo);
+        printf("i: %lld\n  ReduceBoolFunction<X, Z>::exec", i);
+#endif
+        z[i] = startingVal;
+      }
+    return;
+  }
+
+  if (shape::length(zShapeInfo) == 1) {
+    z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+    return;
+  }
+
+#ifdef SD_LOOPS_INLINED
+  sd::ReductionLoops<X, Z, X>::template loopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, extraParams);
+#else
+  sd::ReductionBoolLoops<X, Z>::template innerloopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims,
+                                                                 extraParams);
+#endif
+}
+
+template <typename X, typename Z>
+void ReduceBoolFunction<X, Z>::exec(int opNum, sd::memory::Workspace *workspace, const void *vx,
+                                    const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                    const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(workspace, vx, xShapeInfo, vextraParams, vz, zShapeInfo, dims), REDUCE_BOOL_OPS);
+}
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cpu/reduce/reduce_float.hpp

@@ -0,0 +1,380 @@
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <helpers/ShapeBuilders.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_float.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <algorithm>
+
+using namespace simdOps;
+
+namespace functions {
+namespace reduce {
+
+// =============================================================================
+// TYPE-SAFE UTILITIES FOR ALL NUMERIC TYPES INCLUDING FLOAT16
+// =============================================================================
+
+namespace SafeTypeUtils {
+
+/**
+ * @brief Type-safe array initialization that works with float16 and all other types
+ */
+template<typename T>
+SD_INLINE void initializeArray(T* array, size_t count) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    // For arithmetic types including float16, use loop initialization
+    for (size_t i = 0; i < count; i++) {
+      array[i] = static_cast<T>(0);
+    }
+  } else {
+    // For non-arithmetic types, use default initialization
+    std::fill_n(array, count, T{});
+  }
+}
+
+/**
+ * @brief Safe type conversion for mixed-type operations
+ */
+template<typename From, typename To>
+SD_INLINE constexpr To safeCast(const From& value) {
+  return static_cast<To>(value);
+}
+
+/**
+ * @brief Convert parameter arrays between types safely
+ */
+template<typename SourceType, typename TargetType>
+SD_INLINE void convertParams(const SourceType* source, TargetType* target, size_t count = 8) {
+  if (source && target) {
+    for (size_t i = 0; i < count; i++) {
+      target[i] = safeCast<SourceType, TargetType>(source[i]);
+    }
+  }
+}
+
+/**
+ * @brief Determine appropriate parameter type for mixed operations
+ * For float16, use float for better compatibility; otherwise use Z type
+ */
+template<typename X, typename Z>
+struct CompatibleParamType {
+  using type = typename std::conditional_t<
+      std::is_same_v<Z, float16> || std::is_same_v<Z, bfloat16>,
+      float,  // Use float for half-precision types
+      Z       // Use Z for all other types
+      >;
+};
+
+} // namespace SafeTypeUtils
+
+// =============================================================================
+// REDUCE FLOAT FUNCTION IMPLEMENTATION WITH FLOAT16 SUPPORT
+// =============================================================================
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceFloatFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams,
+                                                   void *vz, const sd::LongType *zShapeInfo) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+
+  // Convert to Z* for consistency with macro expectations
+  Z *extraParams = nullptr;
+  Z convertedParams[8];
+
+  if (vextraParams != nullptr) {
+    if constexpr (std::is_same_v<Z, X>) {
+      extraParams = reinterpret_cast<Z*>(vextraParams);
+    } else {
+      // Convert parameters to Z type
+      auto originalParams = reinterpret_cast<X*>(vextraParams);
+      for (int i = 0; i < 8; ++i) {
+        convertedParams[i] = static_cast<Z>(originalParams[i]);
+      }
+      extraParams = convertedParams;
+    }
+  }
+
+  const auto length = shape::length(xShapeInfo);
+
+  if (shape::isEmptyConst(xShapeInfo)) {
+    z[0] = static_cast<Z>(OpType::startingValue(x));
+    return;
+  }
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+    const auto startingVal = static_cast<Z>(OpType::startingValue(x));
+
+    for (sd::LongType i = 0; i < length; i++) {
+      z[i] = startingVal;
+    }
+    return;
+  }
+
+  auto startingValue = static_cast<typename OpType::InterType>(OpType::startingValue(x));
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  typename OpType::InterType intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = startingValue;
+  }
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+
+  if(shape::isViewConst(xShapeInfo)) {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType indexOffset;
+        COORDS2INDEX(xRank, xStride, coords, indexOffset);
+
+        auto opResult = OpType::op(x[indexOffset], extraParams);
+        intermediate[thread_id] = OpType::update(
+            intermediate[thread_id],
+            opResult,
+            extraParams
+        );
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::merge(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  } else {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        auto opResult = OpType::op(x[i], extraParams);
+        intermediate[thread_id] = OpType::update(
+            intermediate[thread_id],
+            opResult,
+            extraParams
+        );
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::merge(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  }
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+Z SD_HOST ReduceFloatFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+
+  // Convert to Z* for compatibility with OpType::op
+  Z *extraParams = nullptr;
+  Z convertedParams[8];
+
+  if (vextraParams != nullptr) {
+    if constexpr (std::is_same_v<Z, X>) {
+      extraParams = reinterpret_cast<Z*>(vextraParams);
+    } else {
+      // Convert the parameters to Z type
+      auto originalParams = reinterpret_cast<X*>(vextraParams);
+      for (int i = 0; i < 8; ++i) {
+        convertedParams[i] = static_cast<Z>(originalParams[i]);
+      }
+      extraParams = convertedParams;
+    }
+  }
+
+  const sd::LongType length = shape::length(xShapeInfo);
+  auto startingValue = static_cast<typename OpType::InterType>(OpType::startingValue(x));
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+
+  for (sd::LongType i = 0; i < length; i++) {
+    sd::LongType coords[SD_MAX_RANK];
+    INDEX2COORDS(i, xRank, xShape, coords);
+    sd::LongType offset;
+    COORDS2INDEX(xRank, xStride, coords, offset);
+
+    auto opResult = OpType::op(x[offset], extraParams);
+    startingValue = OpType::update(startingValue, opResult, extraParams);
+  }
+
+  return OpType::postProcess(startingValue, length, extraParams);
+}
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceFloatFunction<X, Z>::exec(sd::memory::Workspace *workspace, const void *vx,
+                                             const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                             const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  Z *z = reinterpret_cast<Z *>(vz);
+
+  using CompatibleParamType = typename SafeTypeUtils::CompatibleParamType<X, Z>::type;
+  CompatibleParamType *compatibleExtraParams = nullptr;
+  CompatibleParamType convertedParams[8];
+  SafeTypeUtils::initializeArray(convertedParams, 8);
+
+  if (vextraParams != nullptr) {
+    if constexpr (std::is_same_v<X, CompatibleParamType>) {
+      compatibleExtraParams = reinterpret_cast<CompatibleParamType*>(vextraParams);
+    } else {
+      SafeTypeUtils::convertParams(reinterpret_cast<X*>(vextraParams), convertedParams, 8);
+      compatibleExtraParams = convertedParams;
+    }
+  }
+
+  const int xRank = shape::rank(xShapeInfo);
+  const int zRank = shape::rank(zShapeInfo);
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    const auto startingVal = std::is_same<OpType, simdOps::Mean<X, Z>>::value
+                                 ? sd::DataTypeUtils::nanOrZero<Z>()
+                                 : SafeTypeUtils::safeCast<X, Z>(OpType::startingValue(x));
+    const auto zLen = shape::length(zShapeInfo);
+    if (z != nullptr)
+      for (sd::LongType i = 0; i < zLen; i++) z[i] = startingVal;
+    return;
+  }
+
+  if (shape::length(zShapeInfo) == 1) {
+    z[0] = execScalar<OpType>(x, xShapeInfo, compatibleExtraParams);
+    return;
+  }
+
+  if (OpType::requiresSpecialAccumulation) {
+    // FIXED: Handle execSpecial with flexible parameter types
+    // The enhanced macro provides template overloads that accept any arithmetic type
+        
+    if constexpr (std::is_same_v<CompatibleParamType, sd::LongType>) {
+      // Direct call for sd::LongType parameters
+      OpType::execSpecial(x, xShapeInfo, compatibleExtraParams, z, zShapeInfo,
+                          const_cast<sd::LongType *>(dims) + zRank, xRank - zRank,
+                          nullptr, nullptr);
+    } else {
+      // Convert to sd::LongType for operations that specifically need it
+      sd::LongType longExtraParams[8];
+      SafeTypeUtils::initializeArray(longExtraParams, 8);
+            
+      if (compatibleExtraParams != nullptr) {
+        SafeTypeUtils::convertParams(compatibleExtraParams, longExtraParams, 8);
+      }
+            
+      // Use template overload that accepts sd::LongType*
+      OpType::execSpecial(x, xShapeInfo, longExtraParams, z, zShapeInfo,
+                          const_cast<sd::LongType *>(dims) + zRank, xRank - zRank,
+                          nullptr, nullptr);
+    }
+    return;
+  }
+
+#ifdef SD_LOOPS_INLINED
+  sd::ReductionLoops<X, Z, CompatibleParamType>::template loopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, compatibleExtraParams);
+#else
+  sd::ReductionFloatLoops<X, Z>::template innerloopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, compatibleExtraParams);
+#endif
+}
+
+
+template <typename X, typename Z>
+Z ReduceFloatFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                        void *extraParams) {
+  RETURNING_DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams), REDUCE_FLOAT_OPS);
+}
+
+template <typename X, typename Z>
+void ReduceFloatFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                           void *extraParams, void *z, const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo), REDUCE_FLOAT_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceFloatFunction<X, Z>::exec(const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                             void *vresult, const sd::LongType *resultShapeInfo) {
+  auto z = reinterpret_cast<Z *>(vresult);
+  z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+Z SD_HOST ReduceFloatFunction<X, Z>::execScalar(const void *vx, sd::LongType xEws, sd::LongType length,
+                                                void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+
+  using CompatibleParamType = typename SafeTypeUtils::CompatibleParamType<X, Z>::type;
+  CompatibleParamType *compatibleExtraParams = nullptr;
+  CompatibleParamType convertedParams[8];
+  SafeTypeUtils::initializeArray(convertedParams, 8);
+
+  if (vextraParams != nullptr) {
+    if constexpr (std::is_same_v<X, CompatibleParamType>) {
+      compatibleExtraParams = reinterpret_cast<CompatibleParamType*>(vextraParams);
+    } else {
+      SafeTypeUtils::convertParams(reinterpret_cast<X*>(vextraParams), convertedParams, 8);
+      compatibleExtraParams = convertedParams;
+    }
+  }
+
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  using InterType = typename OpType::InterType;
+  InterType intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = SafeTypeUtils::safeCast<X, InterType>(OpType::startingValue(x));
+  }
+
+  auto func = PRAGMA_THREADS_FOR {
+    if (xEws == 1) {
+      for (auto i = start; i < stop; i++) {
+        auto opResult = OpType::op(x[i], compatibleExtraParams);
+        intermediate[thread_id] = OpType::update(
+            intermediate[thread_id],
+            SafeTypeUtils::safeCast<decltype(opResult), InterType>(opResult),
+            compatibleExtraParams
+        );
+      }
+    } else {
+      for (auto i = start; i < stop; i++) {
+        auto opResult = OpType::op(x[i * xEws], compatibleExtraParams);
+        intermediate[thread_id] = OpType::update(
+            intermediate[thread_id],
+            SafeTypeUtils::safeCast<decltype(opResult), InterType>(opResult),
+            compatibleExtraParams
+        );
+      }
+    }
+  };
+
+  maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+  for (int e = 1; e < maxThreads; e++)
+    intermediate[0] = OpType::update(intermediate[0], intermediate[e], compatibleExtraParams);
+
+  return SafeTypeUtils::safeCast<InterType, Z>(OpType::postProcess(intermediate[0], length, compatibleExtraParams));
+}
+
+template <typename X, typename Z>
+void ReduceFloatFunction<X, Z>::exec(int opNum, sd::memory::Workspace *workspace, const void *vx,
+                                     const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                     const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(workspace, vx, xShapeInfo, vextraParams, vz, zShapeInfo, dims), REDUCE_FLOAT_OPS);
+}
+
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cpu/reduce/reduce_long.hpp

@@ -0,0 +1,246 @@
+/* ******************************************************************************
+*
+*
+* 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 <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_long.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace reduce {
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceLongFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams,
+                                                  void *vz, const sd::LongType *zShapeInfo) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+
+  // For reduce_long operations like MatchCondition, extraParams contain comparison values
+  // (e.g., compare value, epsilon) that must remain in input type X, not be converted to Z.
+  // Converting double→LongType would truncate 0.5→0, breaking comparisons.
+  auto compatibleExtraParams = reinterpret_cast<X*>(vextraParams);
+
+  const sd::LongType length = shape::length(xShapeInfo);
+
+  if (shape::isEmptyConst(xShapeInfo)) {
+    z[0] = static_cast<Z>(OpType::startingValue(x));
+    return;
+  }
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+    const auto startingVal = static_cast<Z>(OpType::startingValue(x));
+
+    for (sd::LongType i = 0; i < length; i++) z[i] = startingVal;
+    return;
+  }
+
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  typename OpType::InterType intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = static_cast<typename OpType::InterType>(OpType::startingValue(x));
+  }
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+
+
+  auto func = PRAGMA_THREADS_FOR {
+    sd::LongType coords[SD_MAX_RANK];
+    for (auto i = start; i < stop; i++) {
+      INDEX2COORDS(i, xRank, xShape, coords);
+      sd::LongType indexOffset;
+      COORDS2INDEX(xRank, xStride, coords, indexOffset);
+      
+      intermediate[thread_id] = OpType::update(
+          intermediate[thread_id],
+          OpType::op(x[indexOffset], compatibleExtraParams),
+          compatibleExtraParams);
+    }
+  };
+  maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+  for (int e = 1; e < maxThreads; e++)
+    intermediate[0] = OpType::update(intermediate[0], intermediate[e], compatibleExtraParams);
+
+  z[0] = OpType::postProcess(intermediate[0], length, compatibleExtraParams);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+Z SD_HOST ReduceLongFunction<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+
+
+  // For reduce_long operations like MatchCondition, extraParams contain comparison values
+  // (e.g., compare value, epsilon) that must remain in input type X, not be converted to Z.
+  // Converting double→LongType would truncate 0.5→0, breaking comparisons.
+  auto compatibleExtraParams = reinterpret_cast<X*>(vextraParams);
+
+  const sd::LongType length = shape::length(xShapeInfo);
+
+  auto startingValue = static_cast<typename OpType::InterType>(OpType::startingValue(x));
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+
+
+  sd::LongType coords[SD_MAX_RANK];
+  for (sd::LongType i = 0; i < length; i++) {
+    INDEX2COORDS(i, xRank, xShape, coords);
+    sd::LongType indexOffset;
+    COORDS2INDEX(xRank, xStride, coords, indexOffset);
+
+    
+    startingValue = OpType::update(startingValue, OpType::op(x[indexOffset], compatibleExtraParams), compatibleExtraParams);
+  }
+  
+  Z result = OpType::postProcess(startingValue, length, compatibleExtraParams);
+  return result;
+}
+
+template <typename X, typename Z>
+Z ReduceLongFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                       void *extraParams) {
+  RETURNING_DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams), REDUCE_LONG_OPS);
+}
+
+template <typename X, typename Z>
+void ReduceLongFunction<X, Z>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                          void *extraParams, void *z, const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo), REDUCE_LONG_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceLongFunction<X, Z>::exec(const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                            void *vresult, const sd::LongType *resultShapeInfo) {
+
+  auto z = reinterpret_cast<Z *>(vresult);
+  z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+Z SD_HOST ReduceLongFunction<X, Z>::execScalar(const void *vx, sd::LongType xEws, sd::LongType length,
+                                               void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+
+
+  // For reduce_long operations like MatchCondition, extraParams contain comparison values
+  // (e.g., compare value, epsilon) that must remain in input type X, not be converted to Z.
+  // Converting double→LongType would truncate 0.5→0, breaking comparisons.
+  auto compatibleExtraParams = reinterpret_cast<X*>(vextraParams);
+
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  typename OpType::InterType intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = static_cast<typename OpType::InterType>(OpType::startingValue(x));
+  }
+
+
+  auto func = PRAGMA_THREADS_FOR {
+    if (xEws == 1) {
+      for (auto i = start; i < stop; i++) {
+        intermediate[thread_id] = OpType::update(intermediate[thread_id], OpType::op(x[i], compatibleExtraParams), compatibleExtraParams);
+      }
+    } else {
+      for (auto i = start; i < stop; i++) {
+        intermediate[thread_id] =
+            OpType::update(intermediate[thread_id], OpType::op(x[i * xEws], compatibleExtraParams), compatibleExtraParams);
+      }
+    }
+  };
+
+  maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+  for (int e = 1; e < maxThreads; e++)
+    intermediate[0] = OpType::update(intermediate[0], intermediate[e], compatibleExtraParams);
+
+  Z result = OpType::postProcess(intermediate[0], length, compatibleExtraParams);
+  return result;
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST ReduceLongFunction<X, Z>::exec(sd::memory::Workspace *workspace, const void *vx,
+                                            const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                            const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  Z *z = reinterpret_cast<Z *>(vz);
+
+  // For reduce_long operations like MatchCondition, extraParams contain comparison values
+  // (e.g., compare value, epsilon) that must remain in input type X, not be converted to Z.
+  // Converting double→LongType would corrupt the comparison values.
+  auto compatibleExtraParams = reinterpret_cast<X*>(vextraParams);
+
+  const int xRank = shape::rank(xShapeInfo);
+  const int zRank = shape::rank(zShapeInfo);
+
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    const auto startingVal = static_cast<Z>(OpType::startingValue(x));
+    const auto zLen = shape::length(zShapeInfo);
+
+    for (sd::LongType i = 0; i < zLen; i++) z[i] = startingVal;
+    return;
+  }
+
+  if (shape::length(zShapeInfo) == 1) {
+    z[0] = execScalar<OpType>(x, xShapeInfo, vextraParams);
+    return;
+  }
+
+  if (OpType::requiresSpecialAccumulation) {
+    // For execSpecial, use original LongType* extraParams as expected by that API
+    auto originalParams = reinterpret_cast<sd::LongType *>(vextraParams);
+    OpType::execSpecial(x, xShapeInfo, originalParams, z, zShapeInfo, const_cast<sd::LongType *>(dims) + zRank,
+                        xRank - zRank, nullptr, nullptr);
+    return;
+  }
+  // Call ReductionLongLoops with properly typed X* extraParams
+  sd::ReductionLongLoops<X, Z>::template innerloopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, compatibleExtraParams);
+
+}
+
+template <typename X, typename Z>
+void ReduceLongFunction<X, Z>::exec(const int opNum, sd::memory::Workspace *workspace, const void *vx,
+                                    const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                    const sd::LongType *zShapeInfo,  sd::LongType *dims) {
+
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(workspace, vx, xShapeInfo, vextraParams, vz, zShapeInfo, dims), REDUCE_LONG_OPS);
+}
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cpu/reduce/reduce_same.hpp

@@ -0,0 +1,204 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <helpers/OmpLaunchHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_same.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include <chrono>
+
+using namespace simdOps;
+
+namespace functions {
+namespace reduce {
+template <typename X>
+template <typename OpType>
+void SD_HOST ReduceSameFunction<X>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams,
+                                               void *vz, const sd::LongType *zShapeInfo) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const auto length = shape::length(xShapeInfo);
+
+  if (shape::isEmptyConst(xShapeInfo)) {
+    z[0] = OpType::startingValue(x);
+    return;
+  }
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+    const auto startingVal = OpType::startingValue(x);
+
+    for (sd::LongType i = 0; i < length; i++) {
+      z[i] = startingVal;
+    }
+    return;
+  }
+
+  auto startingValue = OpType::startingValue(x);
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  X intermediate[64];
+
+  PRAGMA_OMP_SIMD
+  for (auto e = 0; e < maxThreads; e++) {
+    intermediate[e] = startingValue;
+  }
+
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+  if(shape::isViewConst(xShapeInfo)) {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        sd::LongType coords[SD_MAX_RANK];
+        INDEX2COORDS(i, xRank, xShape, coords);
+        sd::LongType indexOffset;
+        COORDS2INDEX(xRank, xStride, coords, indexOffset);
+        intermediate[thread_id] = OpType::update(intermediate[thread_id], OpType::op(x[indexOffset], extraParams), extraParams);
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::update(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  } else {
+    auto func = PRAGMA_THREADS_FOR {
+      for (auto i = start; i < stop; i++) {
+        intermediate[thread_id] = OpType::update(intermediate[thread_id], OpType::op(x[i], extraParams), extraParams);
+      }
+    };
+    maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+    PRAGMA_OMP_SIMD
+    for (int e = 1; e < maxThreads; e++) {
+      intermediate[0] = OpType::update(intermediate[0], intermediate[e], extraParams);
+    }
+
+    z[0] = OpType::postProcess(intermediate[0], length, extraParams);
+  }
+
+}
+
+template <typename X>
+template <typename OpType>
+X SD_HOST ReduceSameFunction<X>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const sd::LongType length = shape::length(xShapeInfo);
+
+  if (shape::isEmptyConst(xShapeInfo)) {
+    return OpType::startingValue(x);
+  }
+
+  auto startingValue = OpType::startingValue(x);
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType* xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType* xStride = shape::stride(xShapeInfo);
+
+  for (sd::LongType i = 0; i < length; i++) {
+    sd::LongType coords[SD_MAX_RANK];
+    INDEX2COORDS(i, xRank, xShape, coords);
+    sd::LongType indexOffset;
+    COORDS2INDEX(xRank, xStride, coords, indexOffset);
+    startingValue = OpType::update(startingValue, OpType::op(x[indexOffset], extraParams), extraParams);
+  }
+  return OpType::postProcess(startingValue, length, extraParams);
+}
+
+template <typename X>
+X ReduceSameFunction<X>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams) {
+  RETURNING_DISPATCH_BY_OPNUM_T(execScalar, PARAMS(x, xShapeInfo, extraParams), REDUCE_SAME_OPS);
+}
+
+template <typename X>
+void ReduceSameFunction<X>::execScalar(const int opNum, const void *x, const sd::LongType *xShapeInfo,
+                                       void *extraParams, void *z, const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_T(execScalar, PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo), REDUCE_SAME_OPS);
+}
+
+template <typename X>
+template <typename OpType>
+void SD_HOST ReduceSameFunction<X>::exec(const void *x, const sd::LongType *xShapeInfo, void *extraParams, void *vz,
+                                         const sd::LongType *zShapeInfo) {
+  auto z = reinterpret_cast<X *>(vz);
+  z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+}
+
+template <typename X>
+template <typename OpType>
+void SD_HOST ReduceSameFunction<X>::exec(sd::memory::Workspace *workspace, const void *vx,
+                                         const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                         const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  const X *x = reinterpret_cast<const X *>(vx);
+  X *z = reinterpret_cast<X *>(vz);
+  X *extraParams = reinterpret_cast<X *>(vextraParams);
+
+  const sd::LongType xRank = shape::rank(xShapeInfo);
+  const sd::LongType zRank = shape::rank(zShapeInfo);
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+    const auto startingVal = OpType::startingValue(x);
+    const auto zLen = shape::length(zShapeInfo);
+    if (z != nullptr) {
+      for (sd::LongType i = 0; i < zLen; i++) {
+        z[i] = startingVal;
+      }
+    }
+    return;
+  }
+
+  if (shape::length(zShapeInfo) == 1) {
+    z[0] = execScalar<OpType>(x, xShapeInfo, extraParams);
+    return;
+  }
+
+  if (OpType::requiresSpecialAccumulation) {
+    OpType::execSpecial(x, xShapeInfo, extraParams, z, zShapeInfo, const_cast<sd::LongType *>(dims) + zRank,
+                        xRank - zRank, nullptr, nullptr);
+    return;
+  }
+
+#ifdef SD_LOOPS_INLINED
+  sd::ReductionLoops<X, X, X>::template loopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, extraParams);
+#else
+  sd::ReductionSameLoops<X>::template innerloopReduce<OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, extraParams);
+#endif
+}
+
+template <typename X>
+void ReduceSameFunction<X>::exec(int opNum, sd::memory::Workspace *workspace, const void *vx,
+                                 const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                 const sd::LongType *zShapeInfo, const sd::LongType *dims) {
+  DISPATCH_BY_OPNUM_T(exec, PARAMS(workspace, vx, xShapeInfo, vextraParams, vz, zShapeInfo, dims), REDUCE_SAME_OPS);
+}
+
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cpu/reduce3.hpp

@@ -0,0 +1,235 @@
+/* ******************************************************************************
+*
+*
+* 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), created on 19.11.2018
+
+#include <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce3.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace reduce3 {
+
+
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void Reduce3<X, Z>::execScalar(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, const void *vy,
+                               const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+  auto length = shape::length(xShapeInfo);
+
+  if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY ||
+      sd::ArrayOptions::arrayType(yShapeInfo) == sd::ArrayType::EMPTY) {
+    if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+    const auto startingVal = OpType::startingValue(x);
+
+    for (sd::LongType i = 0; i < length; i++) z[i] = startingVal;
+    return;
+  }
+
+  // Cache shape-related values
+  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);
+
+  Z extraParamsVals[3] = {(Z)0.0f, (Z)0.0f, (Z)0.0f};
+  Z startingVal = OpType::startingValue(x);
+  int maxThreads = sd::math::sd_min<int>(64, sd::Environment::getInstance().maxThreads());
+  Z intermediate[64];
+  Z extraParamsLocal[3 * 64];
+
+  PRAGMA_OMP_SIMD
+  for (int e = 0; e < maxThreads; e++) intermediate[e] = startingVal;
+
+  sd::ops::safe_zero(extraParamsLocal, 3);
+  if (extraParams != nullptr) {
+    PRAGMA_OMP_SIMD
+    for (int e = 0; e < maxThreads; e++) {
+      extraParamsLocal[3 * e] = extraParams[0];
+      extraParamsLocal[3 * e + 1] = extraParams[1];
+      extraParamsLocal[3 * e + 2] = extraParams[2];
+    }
+  }
+
+  auto func = PRAGMA_THREADS_FOR {
+    for (auto i = start; i < stop; i++) {
+      sd::LongType xCoords[SD_MAX_RANK];
+      sd::LongType yCoords[SD_MAX_RANK];
+
+      INDEX2COORDS(i, xRank, xShape, xCoords);
+      INDEX2COORDS(i, yRank, yShape, yCoords);
+      sd::LongType xOffset = 0, yOffset = 0;
+      COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+      COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+
+      intermediate[thread_id] = OpType::update(intermediate[thread_id],
+                                               OpType::op(x[xOffset], y[yOffset], extraParamsLocal + 3 * thread_id),
+                                               extraParamsLocal + 3 * thread_id);
+    }
+  };
+
+  maxThreads = samediff::Threads::parallel_for(func, 0, length, 1, maxThreads);
+
+
+  // merge step
+  for (int e = 0; e < maxThreads; e++) OpType::aggregateExtraParams(extraParamsVals, extraParamsLocal + 3 * e);
+
+  for (int e = 0; e < maxThreads; e++) startingVal = OpType::update(startingVal, intermediate[e], extraParamsVals);
+
+  // writing out result
+  z[0] = OpType::postProcess(startingVal, length, extraParamsVals);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+void Reduce3<X, Z>::execScalar(const int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals,
+                               const void *vy, const sd::LongType *yShapeInfo, void *vz,
+                               const sd::LongType *zShapeInfo) {
+  DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(vx, xShapeInfo, extraParamsVals, vy, yShapeInfo, vz, zShapeInfo),
+                       REDUCE3_OPS);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void Reduce3<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, const void *vy,
+                         const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                         sd::LongType *dimension,
+                         sd::LongType dimensionLength, sd::LongType start, sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+  if (shape::isScalar(zShapeInfo)) {
+    execScalar<OpType>(vx, xShapeInfo, vextraParams, vy, yShapeInfo, vz, zShapeInfo);
+    return;
+  }
+
+#ifdef SD_LOOPS_INLINED
+  sd::Reduction3Loops<X, Z>::template loopReduce3<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension,
+                                                          dimensionLength, extraParams, start, stop);
+#else
+  sd::Reduction3Loops<X, Z>::template innerloopReduce3<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension,
+                                                               dimensionLength, extraParams, start, stop);
+#endif
+}
+
+//////////////////////////////////////////////////////////////////////////
+// Rest of the functions remain the same as they primarily dispatch to other functions
+// or use the Reduction3Loops class which handles its own shape caching internally
+
+template <typename X, typename Z>
+template <typename OpType>
+void Reduce3<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, const void *vy,
+                         const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                         sd::LongType *dimension,
+                         sd::LongType dimensionLength, const sd::LongType *tadShapeInfo, const sd::LongType *tadOffsets,
+                         sd::LongType start, sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+#ifdef SD_LOOPS_INLINED
+  sd::Reduction3Loops<X, Z>::template loopReduce3<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension,
+                                                          dimensionLength, extraParams, start, stop);
+#else
+  sd::Reduction3Loops<X, Z>::template innerloopReduce3<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension,
+                                                               dimensionLength, extraParams, start, stop);
+#endif
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void Reduce3<X, Z>::execAll(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, const void *vy,
+                            const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                            sd::LongType *dimension,
+                            sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo, const sd::LongType *xOffsets,
+                            const sd::LongType *yTadShapeInfo, const sd::LongType *yOffsets, sd::LongType start,
+                            sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto y = reinterpret_cast<const X *>(vy);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+#ifdef SD_LOOPS_INLINED
+  sd::Reduction3Loops<X, Z>::template loopReduce3All<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, xTadShapeInfo,
+                                                             xOffsets, yTadShapeInfo, yOffsets, extraParams, start,
+                                                             stop);
+#else
+  sd::Reduction3Loops<X, Z>::template innerloopReduce3All<OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo,
+                                                                  xTadShapeInfo, xOffsets, yTadShapeInfo, yOffsets,
+                                                                  extraParams, start, stop);
+#endif
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+void Reduce3<X, Z>::exec(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals,
+                         const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                         sd::LongType *dimension, sd::LongType dimensionLength, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(
+      exec,
+      PARAMS(vx, xShapeInfo, extraParamsVals, vy, yShapeInfo, vz, zShapeInfo, dimension, dimensionLength, start, stop),
+      REDUCE3_OPS);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+void Reduce3<X, Z>::exec(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals,
+                         const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                         sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *tadShapeInfo,
+                         const sd::LongType *tadOffsets, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(exec,
+                       PARAMS(vx, xShapeInfo, extraParamsVals, vy, yShapeInfo, vz, zShapeInfo, dimension,
+                              dimensionLength, tadShapeInfo, tadOffsets, start, stop),
+                       REDUCE3_OPS);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+void Reduce3<X, Z>::execAll(int opNum, const void *vx, const sd::LongType *xShapeInfo, void *extraParamsVals,
+                            const void *vy, const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
+                            sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                            const sd::LongType *xOffsets, const sd::LongType *yTadShapeInfo,
+                            const sd::LongType *yOffsets, sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(execAll,
+                       PARAMS(vx, xShapeInfo, extraParamsVals, vy, yShapeInfo, vz, zShapeInfo, dimension,
+                              dimensionLength, xTadShapeInfo, xOffsets, yTadShapeInfo, yOffsets, start, stop),
+                       REDUCE3_OPS);
+}
+
+}  // namespace reduce3
+}  // namespace functions

libnd4j/include/loops/cpu/scalar.hpp

@@ -0,0 +1,143 @@
+/* ******************************************************************************
+*
+*
+* 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 08.10.2017.
+//
+#include <execution/Threads.h>
+#include <helpers/LoopKind.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include <cstdint>
+
+#include "../legacy_ops.h"
+#include "../scalar.h"
+
+using namespace simdOps;
+
+namespace functions {
+namespace scalar {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void ScalarTransform<X, Y, Z>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                         const sd::LongType *zShapeInfo, const void *vscalars, sd::LongType *dimension,
+                                         sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                         const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                         const sd::LongType *zTadOffsets, sd::LongType start, sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto scalars = reinterpret_cast<const Y *>(vscalars);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+  if (zTadShapeInfo == nullptr) {
+    zTadShapeInfo = xTadShapeInfo;
+    zTadOffsets = xTadOffsets;
+  }
+
+  // Cache shape-related values for TAD operations
+  sd::LongType xTadRank = shape::rank(xTadShapeInfo);
+  sd::LongType zTadRank = shape::rank(zTadShapeInfo);
+  sd::LongType *xTadShape = shape::shapeOf(xTadShapeInfo);
+  sd::LongType *zTadShape = shape::shapeOf(zTadShapeInfo);
+  sd::LongType *xTadStride = shape::stride(xTadShapeInfo);
+  sd::LongType *zTadStride = shape::stride(zTadShapeInfo);
+
+  const int tadLength = shape::tadLength(xShapeInfo, dimension, dimensionLength);
+  for (auto r = start; r < stop; r++) {
+    auto oZ = z + zTadOffsets[r];
+    auto oX = x + xTadOffsets[r];
+    PRAGMA_OMP_SIMD
+    for (int f = 0; f < tadLength; f++) {
+      sd::LongType coords[SD_MAX_RANK];
+      sd::LongType xOffset, zOffset;
+      INDEX2COORDS(f, xTadRank, xTadShape, coords);
+      INDEX2COORDS(f, zTadRank, zTadShape, coords);
+      COORDS2INDEX(xTadRank, xTadStride, coords, xOffset);
+      COORDS2INDEX(zTadRank, zTadStride, coords, zOffset);
+      oZ[zOffset] = OpType::op(oX[xOffset], scalars[0], extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void ScalarTransform<X, Y, Z>::transform(int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                         void *z, const sd::LongType *zShapeInfo, const void *scalars,
+                                         sd::LongType *dimension,
+                                         sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                         const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                         const sd::LongType *zTadOffsets,
+                                         sd::LongType start, sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(transform,
+                        PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo, scalars, dimension, dimensionLength,
+                               xTadShapeInfo, xTadOffsets, zTadShapeInfo, zTadOffsets, start, stop),
+                        SCALAR_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void ScalarTransform<X, Y, Z>::transform(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                                         const sd::LongType *zShapeInfo, const void *scalar, void *extraParams,
+                                         const sd::LongType start, const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TTT(transform, PARAMS(x, xShapeInfo, z, zShapeInfo, scalar, extraParams, start, stop), SCALAR_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void ScalarTransform<X, Y, Z>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                         const sd::LongType *zShapeInfo, const void *vscalar,
+                                         void *vextraParams,
+                                         const sd::LongType start, const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto scalar = reinterpret_cast<const Y *>(vscalar);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+  //need special handling for scalars as strides may not be set for scalars
+  if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
+    z[0] = OpType::op(x[0], scalar[0], extraParams);
+    return;
+
+  }
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  PRAGMA_OMP_SIMD
+  for (auto i = start; i < stop; i++) {
+    sd::LongType coords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset, zOffset;
+    INDEX2COORDS(i, xRank, xShape, coords);
+    INDEX2COORDS(i, zRank, zShape, zCoords);
+    COORDS2INDEX(xRank, xStride, coords, xOffset);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+    z[zOffset] = OpType::op(x[xOffset], scalar[0], extraParams);
+  };
+}
+////////////////////////////////////////////////////////////////////////
+
+}  // namespace scalar
+}  // namespace functions

libnd4j/include/loops/cpu/scalar_bool.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
+******************************************************************************/
+
+//
+// Created by raver119 on 08.10.2017.
+//
+#include "../scalar_bool.h"
+
+#include <execution/Threads.h>
+#include <helpers/LoopKind.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include "../legacy_ops.h"
+
+using namespace simdOps;
+
+namespace functions {
+namespace scalar {
+
+template <typename X, typename Z>
+template <typename OpType>
+void ScalarBoolTransform<X, Z>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                          const sd::LongType *zShapeInfo, const void *vscalars,
+                                          sd::LongType *dimension,
+                                          sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                          const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                          const sd::LongType *zTadOffsets, const sd::LongType start, const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto scalars = reinterpret_cast<const X *>(vscalars);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  if (zTadShapeInfo == nullptr) {
+    zTadShapeInfo = xTadShapeInfo;
+    zTadOffsets = xTadOffsets;
+  }
+
+  // Cache shape-related values for TAD operations
+  sd::LongType xTadRank = shape::rank(xTadShapeInfo);
+  sd::LongType zTadRank = shape::rank(zTadShapeInfo);
+  sd::LongType *xTadShape = shape::shapeOf(xTadShapeInfo);
+  sd::LongType *zTadShape = shape::shapeOf(zTadShapeInfo);
+  sd::LongType *xTadStride = shape::stride(xTadShapeInfo);
+  sd::LongType *zTadStride = shape::stride(zTadShapeInfo);
+
+  const int tadLength = shape::tadLength(xShapeInfo, dimension, dimensionLength);
+  const int numTads = shape::length(xShapeInfo) / tadLength;
+  int num_threads = sd::math::sd_min<int>(numTads, sd::Environment::getInstance().maxThreads());
+
+  for (auto r = start; r < stop; r++) {
+    auto oZ = z + zTadOffsets[r];
+    auto oX = x + xTadOffsets[r];
+
+    PRAGMA_OMP_SIMD
+    for (int f = 0; f < tadLength; f++) {
+      sd::LongType coords[SD_MAX_RANK];
+      sd::LongType xOffset, zOffset;
+      INDEX2COORDS(f, xTadRank, xTadShape, coords);
+      COORDS2INDEX(xTadRank, xTadStride, coords, xOffset);
+      COORDS2INDEX(zTadRank, zTadStride, coords, zOffset);
+      oZ[zOffset] = OpType::op(oX[xOffset], scalars[0], extraParams);
+    }
+  }
+}
+
+template <typename X, typename Z>
+void ScalarBoolTransform<X, Z>::transform(int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                          void *z, const sd::LongType *zShapeInfo, const void *scalars,
+                                          sd::LongType *dimension,
+                                          sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                          const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                          const sd::LongType *zTadOffsets, const sd::LongType start, const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(transform,
+                       PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo, scalars, dimension, dimensionLength,
+                              xTadShapeInfo, xTadOffsets, zTadShapeInfo, zTadOffsets, start, stop),
+                       SCALAR_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+void ScalarBoolTransform<X, Z>::transform(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                                          const sd::LongType *zShapeInfo, const void *scalar, void *extraParams,
+                                          const sd::LongType start, const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_TT(transform, PARAMS(x, xShapeInfo, z, zShapeInfo, scalar, extraParams, start, stop),
+                       SCALAR_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void ScalarBoolTransform<X, Z>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                          const sd::LongType *zShapeInfo, const void *vscalar, void *vextraParams,
+                                          const sd::LongType start, const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto scalar = reinterpret_cast<const X *>(vscalar)[0];
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (auto i2 = start; i2 < stop; i2++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i2, xRank, xShape, coords);
+      sd::LongType offset;
+      COORDS2INDEX(xRank, xStride, coords, offset);
+      z[offset] = OpType::op(x[offset], scalar, extraParams);
+    };
+  } else {
+    PRAGMA_OMP_SIMD
+    for (auto i2 = start; i2 < stop; i2++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i2, xRank, xShape, coords);
+      sd::LongType xOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+      COORDS2INDEX(zRank, zStride, coords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], scalar, extraParams);
+    };
+  }
+}
+
+BUILD_DOUBLE_TEMPLATE( class ScalarBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+
+}  // namespace scalar
+}  // namespace functions

libnd4j/include/loops/cpu/scalar_int.cpp

@@ -0,0 +1,174 @@
+/* ******************************************************************************
+*
+*
+* 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 "../scalar_int.h"
+
+#include <execution/Threads.h>
+#include <helpers/LoopKind.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include "../legacy_ops.h"
+
+using namespace simdOps;
+
+namespace functions {
+namespace scalar {
+
+template <typename X>
+template <typename OpType>
+void ScalarIntTransform<X>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vextraParams, void *vz,
+                                      const sd::LongType *zShapeInfo, const void *vscalars, sd::LongType *dimension,
+                                      sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                      const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                      const sd::LongType *zTadOffsets, const sd::LongType start, const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto scalars = reinterpret_cast<const X *>(vscalars);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  if (zTadShapeInfo == nullptr) {
+    zTadShapeInfo = xTadShapeInfo;
+    zTadOffsets = xTadOffsets;
+  }
+
+  // Cache shape-related values
+  sd::LongType xTadRank = shape::rank(xTadShapeInfo);
+  sd::LongType zTadRank = shape::rank(zTadShapeInfo);
+  sd::LongType *xTadShape = shape::shapeOf(xTadShapeInfo);
+  sd::LongType *zTadShape = shape::shapeOf(zTadShapeInfo);
+  sd::LongType *xTadStride = shape::stride(xTadShapeInfo);
+  sd::LongType *zTadStride = shape::stride(zTadShapeInfo);
+
+  const int tadLength = shape::tadLength(xShapeInfo, dimension, dimensionLength);
+
+  for (auto r = start; r < stop; r++) {
+    auto oZ = z + zTadOffsets[r];
+    auto oX = x + xTadOffsets[r];
+
+    PRAGMA_OMP_SIMD
+    for (int f = 0; f < tadLength; f++) {
+      sd::LongType coords[SD_MAX_RANK];
+      sd::LongType xOffset, zOffset;
+      INDEX2COORDS(f, xTadRank, xTadShape, coords);
+      COORDS2INDEX(xTadRank, xTadStride, coords, xOffset);
+      COORDS2INDEX(zTadRank, zTadStride, coords, zOffset);
+      oZ[zOffset] = OpType::op(oX[xOffset], scalars[0], extraParams);
+    }
+  }
+}
+
+template <typename X>
+void ScalarIntTransform<X>::transform(int opNum, const void *x, const sd::LongType *xShapeInfo, void *extraParams,
+                                      void *z, const sd::LongType *zShapeInfo, const void *scalars,
+                                      sd::LongType *dimension,
+                                      sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
+                                      const sd::LongType *xTadOffsets, const sd::LongType *zTadShapeInfo,
+                                      const sd::LongType *zTadOffsets, const sd::LongType start, const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_T(transform,
+                      PARAMS(x, xShapeInfo, extraParams, z, zShapeInfo, scalars, dimension, dimensionLength,
+                             xTadShapeInfo, xTadOffsets, zTadShapeInfo, zTadOffsets, start, stop),
+                      SCALAR_INT_OPS);
+}
+
+template <typename X>
+void ScalarIntTransform<X>::transform(const int opNum, const void *x, sd::LongType xEws, void *z, sd::LongType zEws,
+                                      const void *scalar, void *extraParams, const sd::LongType n, const sd::LongType start,
+                                      const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_T(transform, PARAMS(x, xEws, z, zEws, scalar, extraParams, n, start, stop), SCALAR_INT_OPS);
+}
+
+template <typename X>
+void ScalarIntTransform<X>::transform(const int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                                      const sd::LongType *zShapeInfo, const void *scalar, void *extraParams,
+                                      const sd::LongType start, const sd::LongType stop) {
+  DISPATCH_BY_OPNUM_T(transform, PARAMS(x, xShapeInfo, z, zShapeInfo, scalar, extraParams, start, stop),
+                      SCALAR_INT_OPS);
+}
+
+template <typename X>
+template <typename OpType>
+void ScalarIntTransform<X>::transform(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                      const sd::LongType *zShapeInfo, const void *vscalar, void *vextraParams,
+                                      const sd::LongType start, const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto scalar = reinterpret_cast<const X *>(vscalar)[0];
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  // Cache shape-related values
+  sd::LongType xRank = shape::rank(xShapeInfo);
+  sd::LongType zRank = shape::rank(zShapeInfo);
+  sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+  sd::LongType *zShape = shape::shapeOf(zShapeInfo);
+  sd::LongType *xStride = shape::stride(xShapeInfo);
+  sd::LongType *zStride = shape::stride(zShapeInfo);
+
+  if (shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo)) {
+    PRAGMA_OMP_SIMD
+    for (auto i3 = start; i3 < stop; i3++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i3, xRank, xShape, coords);
+      sd::LongType offset;
+      COORDS2INDEX(xRank, xStride, coords, offset);
+      z[offset] = OpType::op(x[offset], scalar, extraParams);
+    };
+  } else {
+    PRAGMA_OMP_SIMD
+    for (auto i2 = start; i2 < stop; i2++) {
+      sd::LongType coords[SD_MAX_RANK];
+      INDEX2COORDS(i2, xRank, xShape, coords);
+      sd::LongType xOffset, zOffset;
+      COORDS2INDEX(xRank, xStride, coords, xOffset);
+      COORDS2INDEX(zRank, zStride, coords, zOffset);
+      z[zOffset] = OpType::op(x[xOffset], scalar, extraParams);
+    };
+  }
+}
+
+template <typename X>
+template <typename OpType>
+void ScalarIntTransform<X>::transform(const void *vx, sd::LongType xEws, void *vz, sd::LongType zEws,
+                                      const void *vscalar, void *vextraParams, const sd::LongType len, const sd::LongType start,
+                                      const sd::LongType stop) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto scalar = reinterpret_cast<const X *>(vscalar)[0];
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  if (scalar < static_cast<X>((sizeof(X) * 8))) {
+    PRAGMA_OMP_SIMD
+    for (auto i = start; i < stop; i++) {
+      auto xi = i * xEws;
+      auto zi = i * zEws;
+      z[zi] = OpType::op(x[xi], scalar, extraParams);
+    }
+  }
+}
+
+BUILD_SINGLE_TEMPLATE(class ScalarIntTransform, , SD_INTEGER_TYPES);
+
+// Note: Platform-specific instantiations for "unsigned long" and "unsigned int" on Linux
+// are now generated by cmake/TemplateProcessing.cmake which preserves these raw types
+// instead of canonicalizing them to typedefs. This fixes JavaCPP symbol resolution issues.
+
+}  // namespace scalar
+}  // namespace functions

libnd4j/include/loops/cpu/summarystatsreduce.cpp

@@ -0,0 +1,165 @@
+/* ******************************************************************************
+*
+*
+* 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 18.12.17.
+//
+#include <execution/Threads.h>
+#include <helpers/ConstantTadHelper.h>
+
+#include <helpers/shape.h>
+#include <loops/summarystatsreduce.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace summarystats {
+
+template <typename X, typename Z>
+Z SummaryStatsReduce<X, Z>::execScalar(const int opNum, const bool biasCorrected,  void *x,
+                                       sd::LongType *xShapeInfo, void *extraParams) {
+ RETURNING_DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(biasCorrected, x, xShapeInfo, extraParams), SUMMARY_STATS_OPS);
+}
+
+template <typename X, typename Z>
+void SummaryStatsReduce<X, Z>::execScalar(const int opNum, const bool biasCorrected,  void *x,
+                                          sd::LongType *xShapeInfo, void *extraParams, void *z,
+                                          sd::LongType *zShapeInfo) {
+ DISPATCH_BY_OPNUM_TT(execScalar, PARAMS(biasCorrected, x, xShapeInfo, extraParams, z, zShapeInfo), SUMMARY_STATS_OPS);
+}
+
+template <typename X, typename Z>
+void SummaryStatsReduce<X, Z>::exec(int opNum, bool biasCorrected,  void *x,
+                                    sd::LongType *xShapeInfo, void *extraParams, void *z,
+                                    sd::LongType *zShapeInfo,
+                                   sd::LongType *dimension, sd::LongType dimensionLength) {
+ DISPATCH_BY_OPNUM_TT(exec,
+                      PARAMS(biasCorrected, x, xShapeInfo, extraParams, z, zShapeInfo, dimension, dimensionLength),
+                      SUMMARY_STATS_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SummaryStatsReduce<X, Z>::execScalar(const bool biasCorrected,  void *vx,  sd::LongType *xShapeInfo,
+                                         void *vextraParams, void *vz,  sd::LongType *zShapeInfo) {
+ auto z = reinterpret_cast<Z *>(vz);
+ z[0] = execScalar<OpType>(biasCorrected, vx, xShapeInfo, vextraParams);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+Z SummaryStatsReduce<X, Z>::execScalar(const bool biasCorrected,  void *vx,  sd::LongType *xShapeInfo,
+                                      void *vextraParams) {
+ auto x = reinterpret_cast<const X *>(vx);
+ auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+ // Cache shape-related values
+ sd::LongType xRank = shape::rank(xShapeInfo);
+ sd::LongType *xShape = shape::shapeOf(xShapeInfo);
+ sd::LongType *xStride = shape::stride(xShapeInfo);
+
+ SummaryStatsData<X> startingIndex;
+ startingIndex.initialize();
+ auto length = shape::length(xShapeInfo);
+
+ for (sd::LongType i = 0; i < length; i++) {
+   sd::LongType coords[SD_MAX_RANK];
+   INDEX2COORDS(i, xRank, xShape, coords);
+   sd::LongType xOffset;
+   COORDS2INDEX(xRank, xStride, coords, xOffset);
+
+   SummaryStatsData<X> curr;
+   curr.initWithValue(x[xOffset]);
+   startingIndex = update(startingIndex, curr, extraParams);
+ }
+
+ return OpType::getValue(biasCorrected, startingIndex);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SummaryStatsReduce<X, Z>::exec(bool biasCorrected,  void *vx,  sd::LongType *xShapeInfo,
+                                   void *vextraParams, void *vz,  sd::LongType *zShapeInfo,
+                                   sd::LongType *dimension, sd::LongType dimensionLength) {
+ auto x = reinterpret_cast< X *>(vx);
+ auto z = reinterpret_cast<Z *>(vz);
+ auto extraParams = reinterpret_cast<Z *>(vextraParams);
+ auto resultLength = shape::length(zShapeInfo);
+
+ if (sd::ArrayOptions::arrayType(xShapeInfo) == sd::ArrayType::EMPTY) {
+   if (sd::ArrayOptions::arrayType(zShapeInfo) == sd::ArrayType::EMPTY) return;
+   SummaryStatsData<X> comp;
+   comp.initWithValue(x[0]);
+
+   for (sd::LongType i = 0; i < resultLength; i++) z[i] = OpType::getValue(biasCorrected, comp);
+   return;
+ }
+
+ if (shape::isScalar(zShapeInfo)) {
+   z[0] = execScalar<OpType>(biasCorrected, (void *)x, xShapeInfo, extraParams);
+   return;
+ }
+
+ if (dimensionLength < 1) return;
+
+ // When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
+ auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(xShapeInfo), dimension, dimensionLength);
+ if (resultLength == 1 || dimensionLength == shape::rank(xShapeInfo) || tadPack->numberOfTads() == 1) {
+   z[0] = execScalar<OpType>(biasCorrected, x, xShapeInfo, extraParams);
+   return;
+ }
+
+ auto tadShapeShapeInfo = tadPack->primaryShapeInfo();
+ auto tadLength = shape::length(tadPack->primaryShapeInfo());
+
+ // Cache TAD shape-related values
+ sd::LongType tadRank = shape::rank(tadShapeShapeInfo);
+ sd::LongType *tadShape = shape::shapeOf(tadShapeShapeInfo);
+ sd::LongType *tadStride = shape::stride(tadShapeShapeInfo);
+
+ auto func = PRAGMA_THREADS_FOR {
+   for (auto r = start; r < stop; r++) {
+     auto tadOffsetForBlock = tadPack->primaryOffsets()[r];
+     auto tx = x + tadOffsetForBlock;
+     SummaryStatsData<X> comp;
+     comp.initWithValue(tx[0]);
+
+     for (sd::LongType i = 1; i < tadLength; i++) {
+       sd::LongType coords[SD_MAX_RANK];
+       INDEX2COORDS(i, tadRank, tadShape, coords);
+       sd::LongType xOffset;
+       COORDS2INDEX(tadRank, tadStride, coords, xOffset);
+
+       SummaryStatsData<X> indexVal2;
+       indexVal2.initWithValue(tx[xOffset]);
+
+       comp = update(comp, OpType::op(indexVal2, extraParams), extraParams);
+     }
+
+     z[r] = OpType::getValue(biasCorrected, comp);
+   }
+ };
+
+ samediff::Threads::parallel_tad(func, 0, resultLength, 1);
+}
+
+BUILD_DOUBLE_TEMPLATE( class SummaryStatsReduce, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
+}  // namespace summarystats
+}  // namespace functions

libnd4j/include/loops/cpu/transform/transform_any.cpp

@@ -0,0 +1,73 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+// Selective rendering - MUST be included before types.h to define HAS_* flags
+#include <system/selective_rendering/core.h>
+#include <system/selective_rendering/bool_types.h>
+#include <system/selective_rendering/float_types.h>
+#include <system/selective_rendering/bfloat_types.h>
+#include <system/selective_rendering/int_types.h>
+#include <system/selective_rendering/uint_types.h>
+
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_any.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace transform {
+
+template <typename X, typename Z>
+void TransformAny<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                              const sd::LongType *zShapeInfo, void *extraParams, sd::LongType threadId,
+                              sd::LongType numThreads) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(x, xShapeInfo, z, zShapeInfo, extraParams, threadId, numThreads),
+                       TRANSFORM_ANY_OPS);
+}
+
+/////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST TransformAny<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                      const sd::LongType *zShapeInfo, void *vextraParams, sd::LongType threadId,
+                                      sd::LongType numThreads) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  sd::TransformLoops<X, Z, X>::template loopTransform<OpType>(x, xShapeInfo, z, zShapeInfo, extraParams, threadId,
+                                                              numThreads);
+}
+
+BUILD_DOUBLE_TEMPLATE(class TransformAny, , SD_NUMERIC_TYPES, SD_NUMERIC_TYPES);
+// Also instantiate numeric types -> bool for comparison operations
+BUILD_DOUBLE_TEMPLATE(class TransformAny, , SD_NUMERIC_TYPES, SD_BOOL_TYPES);
+// And bool -> numeric types for casting operations
+BUILD_DOUBLE_TEMPLATE(class TransformAny, , SD_BOOL_TYPES, SD_NUMERIC_TYPES);
+// And bool -> bool for boolean operations
+BUILD_DOUBLE_TEMPLATE(class TransformAny, , SD_BOOL_TYPES, SD_BOOL_TYPES);
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cpu/transform/transform_bool.cpp

@@ -0,0 +1,56 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author  raver119@gmail.com
+//
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_bool.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace transform {
+
+template <typename X, typename Z>
+void TransformBool<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                               const sd::LongType *zShapeInfo, void *extraParams, sd::LongType threadId,
+                               sd::LongType numThreads) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(x, xShapeInfo, z, zShapeInfo, extraParams, threadId, numThreads),
+                       TRANSFORM_BOOL_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST TransformBool<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                       const sd::LongType *zShapeInfo, void *vextraParams, sd::LongType threadId,
+                                       sd::LongType numThreads) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  sd::TransformLoops<X, Z, X>::template loopTransform<OpType>(x, xShapeInfo, z, zShapeInfo, extraParams, threadId,
+                                                              numThreads);
+}
+
+BUILD_DOUBLE_TEMPLATE( class TransformBool, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cpu/transform/transform_float.cpp

@@ -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
+//
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_float.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace transform {
+template <typename X, typename Z>
+void TransformFloat<X, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                                const sd::LongType *zShapeInfo, void *extraParams, sd::LongType threadId,
+                                sd::LongType numThreads) {
+  DISPATCH_BY_OPNUM_TT(exec, PARAMS(x, xShapeInfo, z, zShapeInfo, extraParams, threadId, numThreads),
+                       TRANSFORM_FLOAT_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST TransformFloat<X, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                        const sd::LongType *zShapeInfo, void *vextraParams, sd::LongType threadId,
+                                        sd::LongType numThreads) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<Z *>(vz);
+  auto extraParams = reinterpret_cast<Z *>(vextraParams);
+
+  sd::TransformLoops<X, Z, Z>::template loopTransform<OpType>(x, xShapeInfo, z, zShapeInfo, extraParams, threadId,
+                                                              numThreads);
+}
+
+BUILD_DOUBLE_TEMPLATE( class TransformFloat, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
+
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cpu/transform/transform_same.cpp

@@ -0,0 +1,58 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author  raver119@gmail.com
+//
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_same.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace transform {
+
+template <typename X>
+void TransformSame<X>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                            const sd::LongType *zShapeInfo, void *extraParams, sd::LongType threadId,
+                            sd::LongType numThreads) {
+  DISPATCH_BY_OPNUM_T(exec, PARAMS(x, xShapeInfo, z, zShapeInfo, extraParams, threadId, numThreads),
+                      TRANSFORM_SAME_OPS);
+}
+
+template <typename X>
+template <typename OpType>
+void SD_HOST TransformSame<X>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                    const sd::LongType *zShapeInfo, void *vextraParams, sd::LongType threadId,
+                                    sd::LongType numThreads) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+
+  sd::TransformLoops<X, X, X>::template loopTransform<OpType>(x, xShapeInfo, z, zShapeInfo, extraParams, threadId,
+                                                              numThreads);
+}
+
+BUILD_SINGLE_TEMPLATE( class TransformSame, , SD_NUMERIC_TYPES);
+BUILD_SINGLE_TEMPLATE( class TransformSame, , SD_BOOL_TYPES);
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cpu/transform/transform_strict.cpp

@@ -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
+//
+#include <helpers/Loops.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_strict.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+namespace functions {
+namespace transform {
+
+template <typename X>
+void TransformStrict<X>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, void *z,
+                              const sd::LongType *zShapeInfo, void *extraParams, sd::LongType threadId,
+                              sd::LongType numThreads) {
+  DISPATCH_BY_OPNUM_T(exec, PARAMS(x, xShapeInfo, z, zShapeInfo, extraParams, threadId, numThreads),
+                      TRANSFORM_STRICT_OPS);
+}
+
+template <typename X>
+template <typename OpType>
+void SD_HOST TransformStrict<X>::exec(const void *vx, const sd::LongType *xShapeInfo, void *vz,
+                                      const sd::LongType *zShapeInfo, void *vextraParams, sd::LongType threadId,
+                                      sd::LongType numThreads) {
+  auto x = reinterpret_cast<const X *>(vx);
+  auto z = reinterpret_cast<X *>(vz);
+  auto extraParams = reinterpret_cast<X *>(vextraParams);
+  sd::TransformLoops<X, X, X>::template loopTransform<OpType>(x, xShapeInfo, z, zShapeInfo, extraParams, threadId,
+                                                              numThreads);
+}
+
+BUILD_SINGLE_TEMPLATE( class TransformStrict, , SD_FLOAT_TYPES);
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cuda/broadcasting.chpp

@@ -0,0 +1,441 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+ * the License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#include <loops/broadcasting.h>
+#include <loops/legacy_ops.h>
+#include <types/types.h>
+#include <system/Environment.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <string>
+#include <stdexcept>
+#include <helpers/StringUtils.h>
+#include <ops/specials_cuda.h>
+#include <loops/pairwise_instantiations_single.h>
+
+using namespace simdOps;
+
+template<typename X, typename Y, typename Z, typename OpClass>
+static SD_KERNEL void broadcastSimple(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  functions::broadcast::Broadcast<X,Y,Z>::template transformCuda<OpClass>(
+      x, xShapeInfo,
+      y, yShapeInfo,
+      z, zShapeInfo,
+      dimension, dimensionLength,
+      tadOnlyShapeInfo, tadOffsets,
+      tadOnlyShapeInfoZ, tadOffsetsZ);
+}
+
+template<typename X, typename Y, typename Z, typename OpClass>
+static SD_KERNEL void broadcastSimple(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const void* y,
+    const sd::LongType* yShapeInfo,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::broadcast::Broadcast<X,Y,Z>::template transformCuda<OpClass>(
+      x, xShapeInfo,
+      y, yShapeInfo,
+      z, zShapeInfo);
+}
+
+template<typename X, typename Y, typename Z, typename OpClass>
+static SD_KERNEL void broadcastInverseSimple(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, long long int dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  functions::broadcast::Broadcast<X,Y,Z>::template transformInverseCuda<OpClass>(
+      x, xShapeInfo,
+      y, yShapeInfo,
+      z, zShapeInfo,
+      dimension, dimensionLength,
+      tadOnlyShapeInfo, tadOffsets,
+      tadOnlyShapeInfoZ, tadOffsetsZ);
+}
+
+namespace functions {
+namespace broadcast {
+
+template<typename X, typename Y, typename Z>
+template <typename OpClass>
+SD_HOST void Broadcast<X,Y,Z>::intermediateBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    void const* x, sd::LongType const* xShapeInfo,
+    void const* y, sd::LongType const* yShapeInfo,
+    void* z, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  broadcastSimple<X, Y, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x, xShapeInfo,
+          y, yShapeInfo,
+          z, zShapeInfo,
+          dimension, dimensionLength,
+          tadOnlyShapeInfo, tadOffsets,
+          tadOnlyShapeInfoZ, tadOffsetsZ);
+}
+
+template<typename X, typename Y, typename Z>
+template <typename OpClass>
+SD_HOST void Broadcast<X,Y,Z>::intermediateBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    const void* x, const sd::LongType* xShapeInfo,
+    const void* y, const sd::LongType* yShapeInfo,
+    void* z, const sd::LongType* zShapeInfo) {
+
+  broadcastSimple<X, Y, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x, xShapeInfo,
+          y, yShapeInfo,
+          z, zShapeInfo);
+}
+
+template<typename X, typename Y, typename Z>
+SD_HOST void Broadcast<X,Y,Z>::execBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    int opNum,
+    void const* x, sd::LongType const* xShapeInfo,
+    void const* y, sd::LongType const* yShapeInfo,
+    void* z, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateBroadcast,
+      PARAMS(launchDims, stream,
+             x, xShapeInfo,
+             y, yShapeInfo,
+             z, zShapeInfo,
+             dimension, dimensionLength,
+             tadOnlyShapeInfo, tadOffsets,
+             tadOnlyShapeInfoZ, tadOffsetsZ),
+      OPS_A(BROADCAST_OPS))
+
+  sd::DebugHelper::checkErrorCode(stream, "execBroadcast(...) failed");
+}
+
+template<typename X, typename Y, typename Z>
+SD_HOST void Broadcast<X,Y,Z>::execBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    const int opNum,
+    const void* x, const sd::LongType* xShapeInfo,
+    const void* y, const sd::LongType* yShapeInfo,
+    void* z, const sd::LongType* zShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateBroadcast,
+      PARAMS(launchDims, stream,
+             x, xShapeInfo,
+             y, yShapeInfo,
+             z, zShapeInfo),
+      OPS_A(BROADCAST_OPS))
+
+  sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcast(...) failed");
+}
+
+template<typename X, typename Y, typename Z>
+template <typename OpClass>
+SD_HOST void Broadcast<X,Y,Z>::intermediateInverseBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    void const* x, sd::LongType const* xShapeInfo,
+    void const* y, sd::LongType const* yShapeInfo,
+    void* z, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  broadcastInverseSimple<X, Y, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x, xShapeInfo,
+          y, yShapeInfo,
+          z, zShapeInfo,
+          dimension, dimensionLength,
+          tadOnlyShapeInfo, tadOffsets,
+          tadOnlyShapeInfoZ, tadOffsetsZ);
+}
+
+template<typename X, typename Y, typename Z>
+SD_HOST void Broadcast<X,Y,Z>::execInverseBroadcast(
+    dim3 launchDims, cudaStream_t* stream,
+    int opNum,
+    void const* x, sd::LongType const* xShapeInfo,
+    void const* y, sd::LongType const* yShapeInfo,
+    void* z, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateInverseBroadcast,
+      PARAMS(launchDims, stream,
+             x, xShapeInfo,
+             y, yShapeInfo,
+             z, zShapeInfo,
+             dimension, dimensionLength,
+             tadOnlyShapeInfo, tadOffsets,
+             tadOnlyShapeInfoZ, tadOffsetsZ),
+      OPS_A(BROADCAST_OPS))
+
+  sd::DebugHelper::checkErrorCode(stream, "execInverseBroadcast(...) failed");
+}
+
+template<typename X, typename Y, typename Z>
+template <typename OpType>
+SD_DEVICE void Broadcast<X,Y,Z>::transformInverseCuda(
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  if (tadOnlyShapeInfoZ == nullptr) {
+    tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+    tadOffsetsZ       = tadOffsets;
+  }
+
+  auto x = reinterpret_cast<X const*>(vx);
+  auto y = reinterpret_cast<Y const*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+
+  __shared__ sd::LongType tadLength;
+  __shared__ int numTads;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShape;
+  __shared__ const sd::LongType* xStride;
+
+  __shared__ sd::LongType yRank;
+  __shared__ const sd::LongType* yShape;
+  __shared__ const sd::LongType* yStride;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShape;
+  __shared__ const sd::LongType* zStride;
+
+  if (threadIdx.x == 0) {
+    tadLength = shape::length(tadOnlyShapeInfo);
+    numTads   = shape::length(yShapeInfo) / tadLength;
+
+    xRank   = shape::rank(xShapeInfo);
+    xShape  = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+
+    yRank   = shape::rank(tadOnlyShapeInfo);
+    yShape  = shape::shapeOf(tadOnlyShapeInfo);
+    yStride = shape::stride(tadOnlyShapeInfo);
+
+    zRank   = shape::rank(tadOnlyShapeInfoZ);
+    zShape  = shape::shapeOf(tadOnlyShapeInfoZ);
+    zStride = shape::stride(tadOnlyShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
+    auto rY = y + tadOffsets[r];
+    auto rZ = z + tadOffsetsZ[r];
+
+    for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.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(i, xRank, xShape, xCoords);
+      COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+
+      INDEX2COORDS(i, yRank, yShape, yCoords);
+      COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+
+      INDEX2COORDS(i, zRank, zShape, zCoords);
+      COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+      rZ[zOffset] = OpType::op(x[xOffset], rY[yOffset]);
+    }
+  }
+}
+
+template<typename X, typename Y, typename Z>
+template <typename OpType>
+SD_DEVICE void Broadcast<X,Y,Z>::transformCuda(
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
+
+  if (tadOnlyShapeInfoZ == nullptr) {
+    tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+    tadOffsetsZ       = tadOffsets;
+  }
+
+  auto x = reinterpret_cast<X const*>(vx);
+  auto y = reinterpret_cast<Y const*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+
+  __shared__ sd::LongType tadLength;
+  __shared__ sd::LongType numTads;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShape;
+  __shared__ const sd::LongType* xStride;
+
+  __shared__ sd::LongType yRank;
+  __shared__ const sd::LongType* yShape;
+  __shared__ const sd::LongType* yStride;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShape;
+  __shared__ const sd::LongType* zStride;
+
+  if (threadIdx.x == 0) {
+    tadLength = shape::length(tadOnlyShapeInfo);
+    numTads   = shape::length(xShapeInfo) / tadLength;
+
+    xRank   = shape::rank(tadOnlyShapeInfo);
+    xShape  = shape::shapeOf(tadOnlyShapeInfo);
+    xStride = shape::stride(tadOnlyShapeInfo);
+
+    yRank   = shape::rank(yShapeInfo);
+    yShape  = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+
+    zRank   = shape::rank(tadOnlyShapeInfoZ);
+    zShape  = shape::shapeOf(tadOnlyShapeInfoZ);
+    zStride = shape::stride(tadOnlyShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    auto rX = x + tadOffsets[r];
+    auto rZ = z + tadOffsetsZ[r];
+
+    for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
+      sd::LongType coordsX[SD_MAX_RANK];
+      sd::LongType coordsY[SD_MAX_RANK];
+      sd::LongType coordsZ[SD_MAX_RANK];
+      sd::LongType xOffset;
+      sd::LongType yOffset;
+      sd::LongType zOffset;
+
+      INDEX2COORDS(i, xRank, xShape, coordsX);
+      COORDS2INDEX(xRank, xStride, coordsX, xOffset);
+
+      INDEX2COORDS(i, yRank, yShape, coordsY);
+      COORDS2INDEX(yRank, yStride, coordsY, yOffset);
+
+      INDEX2COORDS(i, zRank, zShape, coordsZ);
+      COORDS2INDEX(zRank, zStride, coordsZ, zOffset);
+
+      rZ[zOffset] = OpType::op(rX[xOffset], y[yOffset]);
+    }
+  }
+}
+
+template<typename X, typename Y, typename Z>
+template <typename OpType>
+SD_DEVICE void Broadcast<X,Y,Z>::transformCuda(
+    const void* vx, const sd::LongType* xShapeInfo,
+    const void* vy, const sd::LongType* yShapeInfo,
+    void* vz, const sd::LongType* zShapeInfo) {
+
+  const X* x = reinterpret_cast<const X*>(vx);
+  const Y* y = reinterpret_cast<const Y*>(vy);
+  Z* z       = reinterpret_cast<Z*>(vz);
+
+  __shared__ sd::LongType zLen;
+  __shared__ int rank;
+  __shared__ bool xzSameOffsets, yzSameOffsets;
+
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ const sd::LongType* yStridePtr;
+  __shared__ const sd::LongType* zStridePtr;
+  __shared__ const sd::LongType* zShapePtr;
+
+  if (threadIdx.x == 0) {
+    zLen          = shape::length(zShapeInfo);
+    rank          = shape::rank(zShapeInfo);
+    xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+    yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
+
+    xStridePtr    = shape::stride(xShapeInfo);
+    yStridePtr    = shape::stride(yShapeInfo);
+    zStridePtr    = shape::stride(zShapeInfo);
+    zShapePtr     = shape::shapeOf(zShapeInfo);
+  }
+  __syncthreads();
+
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  sd::LongType coords[SD_MAX_RANK];
+
+  for (sd::LongType i = tid; i < zLen; i += blockDim.x * gridDim.x) {
+    INDEX2COORDS(i, rank, zShapePtr, coords);
+
+    sd::LongType zOffset;
+    COORDS2INDEX(rank, zStridePtr, coords, zOffset);
+
+    sd::LongType xOffset, yOffset;
+    if (xzSameOffsets) {
+      xOffset = zOffset;
+    } else {
+      COORDS2INDEX(rank, xStridePtr, coords, xOffset);
+    }
+
+    if (yzSameOffsets) {
+      yOffset = zOffset;
+    } else {
+      COORDS2INDEX(rank, yStridePtr, coords, yOffset);
+    }
+
+    z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
+  }
+}
+
+
+} // namespace broadcast
+} // namespace functions

libnd4j/include/loops/cuda/broadcasting_bool.cu

@@ -0,0 +1,570 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <cuda.h>
+#include <cuda_runtime.h>
+#include <helpers/StringUtils.h>
+#include <loops/broadcasting_bool.h>
+#include <loops/legacy_ops.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+    using namespace simdOps;
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpClass>
+static SD_KERNEL void broadcastBoolSimple(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  functions::broadcast::BroadcastBool<X, Z>::template transformCuda<OpClass>(
+      x,
+      xShapeInfo,
+      y,
+      yShapeInfo,
+      z,
+      zShapeInfo,
+      extraParams,
+      dimension,
+      dimensionLength,
+      tadOnlyShapeInfo,
+      tadOffsets,
+      tadOnlyShapeInfoZ,
+      tadOffsetsZ);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpClass>
+static SD_KERNEL void broadcastBoolSimple(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const void* y,
+    const sd::LongType* yShapeInfo,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    void* extraParams) {
+
+  functions::broadcast::BroadcastBool<X, Z>::template transformCuda<OpClass>(
+      x,
+      xShapeInfo,
+      y,
+      yShapeInfo,
+      z,
+      zShapeInfo,
+      extraParams);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpClass>
+static SD_KERNEL void broadcastBoolInverseSimple(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  functions::broadcast::BroadcastBool<X, Z>::template transformInverseCuda<OpClass>(
+      x,
+      xShapeInfo,
+      y,
+      yShapeInfo,
+      z,
+      zShapeInfo,
+      extraParams,
+      dimension,
+      dimensionLength,
+      tadOnlyShapeInfo,
+      tadOffsets,
+      tadOnlyShapeInfoZ,
+      tadOffsetsZ);
+}
+
+namespace functions {
+namespace broadcast {
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpClass>
+SD_HOST void BroadcastBool<X, Z>::intermediateBroadcast(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  broadcastBoolSimple<X, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x,
+          xShapeInfo,
+          y,
+          yShapeInfo,
+          z,
+          zShapeInfo,
+          extraParams,
+          dimension,
+          dimensionLength,
+          tadOnlyShapeInfo,
+          tadOffsets,
+          tadOnlyShapeInfoZ,
+          tadOffsetsZ);
+
+  sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpClass>
+SD_HOST void BroadcastBool<X, Z>::intermediateBroadcast(
+    dim3 launchDims,
+    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) {
+
+  broadcastBoolSimple<X, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x,
+          xShapeInfo,
+          y,
+          yShapeInfo,
+          z,
+          zShapeInfo,
+          extraParams);
+
+  sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void BroadcastBool<X, Y>::execBroadcast(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateBroadcast,
+      PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams,
+             dimension, dimensionLength, tadOnlyShapeInfo, tadOffsets,
+             tadOnlyShapeInfoZ, tadOffsetsZ),
+      OPS_A(BROADCAST_BOOL_OPS));
+
+  sd::DebugHelper::checkErrorCode(stream, "execBroadcast(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void BroadcastBool<X, Y>::execBroadcast(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const void* y,
+    const sd::LongType* yShapeInfo,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    void* extraParams) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateBroadcast,
+      PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams),
+      OPS_A(BROADCAST_BOOL_OPS));
+
+  DEBUG_KERNEL(stream, opNum);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpClass>
+SD_HOST void BroadcastBool<X, Z>::intermediateInverseBroadcast(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  broadcastBoolInverseSimple<X, Z, OpClass>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x,
+          xShapeInfo,
+          y,
+          yShapeInfo,
+          z,
+          zShapeInfo,
+          extraParams,
+          dimension,
+          dimensionLength,
+          tadOnlyShapeInfo,
+          tadOffsets,
+          tadOnlyShapeInfoZ,
+          tadOffsetsZ);
+
+  sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void BroadcastBool<X, Y>::execInverseBroadcast(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void const* y,
+    sd::LongType const* yShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateInverseBroadcast,
+      PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams,
+             dimension, dimensionLength, tadOnlyShapeInfo, tadOffsets,
+             tadOnlyShapeInfoZ, tadOffsetsZ),
+      OPS_A(BROADCAST_BOOL_OPS));
+
+  sd::DebugHelper::checkErrorCode(stream, "execInverseBroadcast(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void BroadcastBool<X, Z>::transformInverseCuda(
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  if (tadOnlyShapeInfoZ == nullptr) {
+    tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+    tadOffsetsZ       = tadOffsets;
+  }
+
+  auto x          = reinterpret_cast<const X*>(vx);
+  auto y          = reinterpret_cast<const X*>(vy);
+  auto z          = reinterpret_cast<Z*>(vz);
+  auto extraParams= reinterpret_cast<X*>(vextraParams);
+
+  __shared__ sd::LongType tadLength;
+  __shared__ int numTads;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType tadRank;
+  __shared__ const sd::LongType* tadShapePtr;
+  __shared__ const sd::LongType* tadStridePtr;
+
+  __shared__ sd::LongType tadRankZ;
+  __shared__ const sd::LongType* tadShapePtrZ;
+  __shared__ const sd::LongType* tadStridePtrZ;
+
+  if (threadIdx.x == 0) {
+    tadLength = shape::length(tadOnlyShapeInfo);
+    numTads   = shape::length(yShapeInfo) / tadLength;
+
+    xRank     = shape::rank(xShapeInfo);
+    xShapePtr = shape::shapeOf(xShapeInfo);
+    xStridePtr= shape::stride(xShapeInfo);
+
+    tadRank     = shape::rank(tadOnlyShapeInfo);
+    tadShapePtr = shape::shapeOf(tadOnlyShapeInfo);
+    tadStridePtr= shape::stride(tadOnlyShapeInfo);
+
+    tadRankZ     = shape::rank(tadOnlyShapeInfoZ);
+    tadShapePtrZ = shape::shapeOf(tadOnlyShapeInfoZ);
+    tadStridePtrZ= shape::stride(tadOnlyShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
+    auto rZ = z + tadOffsetsZ[r];
+    auto rY = y + tadOffsets[r];
+
+    for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
+      sd::LongType coordsX[SD_MAX_RANK];
+      sd::LongType coordsY[SD_MAX_RANK];
+      sd::LongType coordsZ[SD_MAX_RANK];
+      sd::LongType xOffset;
+      sd::LongType yOffset;
+      sd::LongType zOffset;
+
+      // for x
+      INDEX2COORDS(i, xRank, xShapePtr, coordsX);
+      COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
+
+      // for y (tad)
+      INDEX2COORDS(i, tadRank, tadShapePtr, coordsY);
+      COORDS2INDEX(tadRank, tadStridePtr, coordsY, yOffset);
+
+      // for z (tadZ)
+      INDEX2COORDS(i, tadRankZ, tadShapePtrZ, coordsZ);
+      COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, zOffset);
+
+      rZ[zOffset] = OpType::op(x[xOffset], rY[yOffset], extraParams);
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void BroadcastBool<X, Z>::transformCuda(
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadOnlyShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  if (tadOnlyShapeInfoZ == nullptr) {
+    tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+    tadOffsetsZ       = tadOffsets;
+  }
+
+  auto x          = reinterpret_cast<const X*>(vx);
+  auto y          = reinterpret_cast<const X*>(vy);
+  auto z          = reinterpret_cast<Z*>(vz);
+  auto extraParams= reinterpret_cast<X*>(vextraParams);
+
+  __shared__ sd::LongType tadLength;
+  __shared__ sd::LongType numTads;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType yRank;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    tadLength = shape::length(tadOnlyShapeInfo);
+    numTads   = shape::length(xShapeInfo) / tadLength;
+
+    xRank     = shape::rank(tadOnlyShapeInfo);
+    xShapePtr = shape::shapeOf(tadOnlyShapeInfo);
+    xStridePtr= shape::stride(tadOnlyShapeInfo);
+
+    yRank     = shape::rank(yShapeInfo);
+    yShapePtr = shape::shapeOf(yShapeInfo);
+    yStridePtr= shape::stride(yShapeInfo);
+
+    zRank     = shape::rank(tadOnlyShapeInfoZ);
+    zShapePtr = shape::shapeOf(tadOnlyShapeInfoZ);
+    zStridePtr= shape::stride(tadOnlyShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    auto rX = x + tadOffsets[r];
+    auto rZ = z + tadOffsetsZ[r];
+
+    for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
+      sd::LongType coordsX[SD_MAX_RANK];
+      sd::LongType coordsY[SD_MAX_RANK];
+      sd::LongType coordsZ[SD_MAX_RANK];
+      sd::LongType xOffset;
+      sd::LongType yOffset;
+      sd::LongType zOffset;
+
+      INDEX2COORDS(i, xRank, xShapePtr, coordsX);
+      COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
+
+      INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+      COORDS2INDEX(yRank, yStridePtr, coordsY, yOffset);
+
+      INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+      COORDS2INDEX(zRank, zStridePtr, coordsZ, zOffset);
+
+      rZ[zOffset] = OpType::op(rX[xOffset], y[yOffset], extraParams);
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void BroadcastBool<X, Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    const void* vy,
+    const sd::LongType* yShapeInfo,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vextraParams) {
+
+  const auto x = reinterpret_cast<const X*>(vx);
+  const auto y = reinterpret_cast<const X*>(vy);
+  auto z       = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ sd::LongType zLen;
+  __shared__ int xRank, yRank, zRank;
+  __shared__ bool xzSameOffsets, yzSameOffsets;
+
+  // We'll store stride data for each shape in shared mem
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ const sd::LongType* yStridePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* zShapePtr;
+
+  if (threadIdx.x == 0) {
+    zLen           = shape::length(zShapeInfo);
+
+    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);
+
+    xzSameOffsets  = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+    yzSameOffsets  = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
+  }
+  __syncthreads();
+
+  const auto tid          = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto totalThreads = blockDim.x * gridDim.x;
+
+  for (sd::LongType i2 = tid; i2 < zLen; i2 += totalThreads) {
+    sd::LongType coords[SD_MAX_RANK];
+    INDEX2COORDS(i2, zRank, zShapePtr, coords);
+
+    sd::LongType zOffset;
+    COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
+
+    sd::LongType xOffset;
+    if (xzSameOffsets) {
+      xOffset = zOffset;
+    } else {
+      COORDS2INDEX(xRank, xStridePtr, coords, xOffset);
+    }
+
+    sd::LongType yOffset;
+    if (yzSameOffsets) {
+      yOffset = zOffset;
+    } else {
+      COORDS2INDEX(yRank, yStridePtr, coords, yOffset);
+    }
+
+    z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+  }
+}
+
+// build the class
+BUILD_DOUBLE_TEMPLATE( class BroadcastBool, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+
+}  // namespace broadcast
+}  // namespace functions

libnd4j/include/loops/cuda/broadcasting_int.cu

@@ -0,0 +1,476 @@
+/******************************************************************************
+*
+*
+* 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 <cuda.h>
+#include <cuda_runtime.h>
+#include <helpers/StringUtils.h>
+#include <loops/broadcasting_int.h>
+#include <loops/legacy_ops.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <loops/pairwise_instantiations.h>
+#include <stdexcept>
+#include <string>
+
+using namespace simdOps;
+
+//////////////////////////////////////////////////////////////////////////////
+// Cached kernel that caches shape info in shared memory and performs the broadcast
+template <typename X, typename OpClass>
+__global__ void broadcastIntSimpleCached(
+   void const* x,
+   sd::LongType const* xShapeInfo,
+   void const* y,
+   sd::LongType const* yShapeInfo,
+   void* z,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // Delegate the broadcast operation to the transformCuda method with cached shape info
+ functions::broadcast::BroadcastInt<X>::template transformCuda<OpClass>(
+     x,
+     xShapeInfo,
+     y,
+     yShapeInfo,
+     z,
+     zShapeInfo,
+     dimension,
+     dimensionLength,
+     tadOnlyShapeInfo,
+     tadOffsets,
+     tadOnlyShapeInfoZ,
+     tadOffsetsZ);
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// Cached kernel that caches shape info in shared memory and performs the inverse broadcast
+template <typename X, typename OpClass>
+__global__ void broadcastIntInverseSimpleCached(
+   void const* x,
+   sd::LongType const* xShapeInfo,
+   void const* y,
+   sd::LongType const* yShapeInfo,
+   void* z,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // Delegate the inverse broadcast operation to the transformInverseCuda method with cached shape info
+ functions::broadcast::BroadcastInt<X>::template transformInverseCuda<OpClass>(
+     x,
+     xShapeInfo,
+     y,
+     yShapeInfo,
+     z,
+     zShapeInfo,
+     dimension,
+     dimensionLength,
+     tadOnlyShapeInfo,
+     tadOffsets,
+     tadOnlyShapeInfoZ,
+     tadOffsetsZ);
+}
+
+namespace functions {
+namespace broadcast {
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the intermediateBroadcast function that launches the cached kernel with dimensions
+template <typename X>
+template <typename OpClass>
+SD_HOST void BroadcastInt<X>::intermediateBroadcast(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   void const* x,
+   sd::LongType const* xShapeInfo,
+   void const* y,
+   sd::LongType const* yShapeInfo,
+   void* z,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // Launch the cached broadcastIntSimpleCached kernel with all parameters
+ broadcastIntSimpleCached<X, OpClass>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShapeInfo,
+         y,
+         yShapeInfo,
+         z,
+         zShapeInfo,
+         dimension,
+         dimensionLength,
+         tadOnlyShapeInfo,
+         tadOffsets,
+         tadOnlyShapeInfoZ,
+         tadOffsetsZ);
+
+ // Check for any errors during kernel execution
+ sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcast(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the intermediateBroadcast function that launches the cached kernel without dimensions
+template <typename X>
+template <typename OpClass>
+SD_HOST void BroadcastInt<X>::intermediateBroadcast(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   const void* x,
+   const sd::LongType* xShapeInfo,
+   const void* y,
+   const sd::LongType* yShapeInfo,
+   void* z,
+   const sd::LongType* zShapeInfo)
+{
+ // Launch the cached broadcastIntSimpleCached kernel without dimensions
+ broadcastIntSimpleCached<X, OpClass>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShapeInfo,
+         y,
+         yShapeInfo,
+         z,
+         zShapeInfo,
+         nullptr, // dimension
+         0,       // dimensionLength
+         nullptr, // tadOnlyShapeInfo
+         nullptr, // tadOffsets
+         nullptr, // tadOnlyShapeInfoZ
+         nullptr  // tadOffsetsZ
+     );
+
+ // Check for any errors during kernel execution
+ sd::DebugHelper::checkGlobalErrorCode("broadcastIntSimpleCached(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the intermediateInverseBroadcast function that launches the cached inverse kernel with dimensions
+template <typename X>
+template <typename OpClass>
+SD_HOST void BroadcastInt<X>::intermediateInverseBroadcast(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   void const* x,
+   sd::LongType const* xShapeInfo,
+   void const* y,
+   sd::LongType const* yShapeInfo,
+   void* z,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // Launch the cached broadcastIntInverseSimpleCached kernel with all parameters
+ broadcastIntInverseSimpleCached<X, OpClass>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShapeInfo,
+         y,
+         yShapeInfo,
+         z,
+         zShapeInfo,
+         dimension,
+         dimensionLength,
+         tadOnlyShapeInfo,
+         tadOffsets,
+         tadOnlyShapeInfoZ,
+         tadOffsetsZ);
+
+ // Check for any errors during kernel execution
+ sd::DebugHelper::checkGlobalErrorCode("broadcastIntInverseSimpleCached(...) failed");
+}
+
+
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the transformCuda device function for BroadcastInt with cached shape info
+template <typename X>
+template <typename OpClass>
+SD_DEVICE void BroadcastInt<X>::transformCuda(
+   void const* vx,
+   sd::LongType const* xShapeInfo,
+   void const* vy,
+   sd::LongType const* yShapeInfo,
+   void* vz,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // If tadOnlyShapeInfoZ is null, set it to tadOnlyShapeInfo and tadOffsetsZ to tadOffsets
+ if (tadOnlyShapeInfoZ == nullptr) {
+   tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+   tadOffsetsZ       = tadOffsets;
+ }
+
+ // Cast pointers to appropriate types
+ auto x = reinterpret_cast<const X*>(vx);
+ auto y = reinterpret_cast<const X*>(vy);
+ auto z = reinterpret_cast<X*>(vz);
+
+ // Shared memory for caching shape information
+ __shared__ sd::LongType tadLength;
+ __shared__ int numTads;
+ __shared__ int xRank;
+ __shared__ int yRank;
+ __shared__ int zRank;
+
+ __shared__ const sd::LongType* tadShape;
+ __shared__ const sd::LongType* tadStride;
+ __shared__ const sd::LongType* tadShapeZ;
+ __shared__ const sd::LongType* tadStrideZ;
+
+ if (threadIdx.x == 0) {
+   // Cache essential shape information
+   tadLength = shape::length(tadOnlyShapeInfo);
+   numTads   = shape::length(xShapeInfo) / tadLength;
+
+   xRank = shape::rank(xShapeInfo);
+   yRank = shape::rank(yShapeInfo);
+   zRank = shape::rank(zShapeInfo);
+
+   tadShape  = shape::shapeOf(tadOnlyShapeInfo);
+   tadStride = shape::stride(tadOnlyShapeInfo);
+
+   tadShapeZ  = shape::shapeOf(tadOnlyShapeInfoZ);
+   tadStrideZ = shape::stride(tadOnlyShapeInfoZ);
+ }
+ __syncthreads();
+
+ // Each block handles a subset of TADs
+ for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+   auto xTad = x + tadOffsets[r];
+   auto zTad = z + tadOffsetsZ[r];
+
+   // Loop over TAD elements
+   for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
+     sd::LongType coords[SD_MAX_RANK];
+     sd::LongType xOffset, yOffset, zOffset;
+
+     // Convert index to coordinates using cached shape info
+     INDEX2COORDS(i, xRank, tadShape, coords);
+     COORDS2INDEX(xRank, tadStride, coords, xOffset);
+
+     COORDS2INDEX(yRank, shape::stride(yShapeInfo), coords, yOffset);
+
+     INDEX2COORDS(i, zRank, tadShapeZ, coords);
+     COORDS2INDEX(zRank, tadStrideZ, coords, zOffset);
+
+     // Apply the operation
+     zTad[zOffset] = OpClass::op(x[xOffset], y[yOffset]);
+   }
+ }
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the transformInverseCuda device function for BroadcastInt with cached shape info
+template <typename X>
+template <typename OpClass>
+SD_DEVICE void BroadcastInt<X>::transformInverseCuda(
+   void const* vx,
+   sd::LongType const* xShapeInfo,
+   void const* vy,
+   sd::LongType const* yShapeInfo,
+   void* vz,
+   sd::LongType const* zShapeInfo,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   sd::LongType const* tadOnlyShapeInfo,
+   sd::LongType const* tadOffsets,
+   sd::LongType const* tadOnlyShapeInfoZ,
+   sd::LongType const* tadOffsetsZ)
+{
+ // If tadOnlyShapeInfoZ is null, set it to tadOnlyShapeInfo and tadOffsetsZ to tadOffsets
+ if (tadOnlyShapeInfoZ == nullptr) {
+   tadOnlyShapeInfoZ = tadOnlyShapeInfo;
+   tadOffsetsZ       = tadOffsets;
+ }
+
+ // Cast pointers to appropriate types
+ auto x = reinterpret_cast<const X*>(vx);
+ auto y = reinterpret_cast<const X*>(vy);
+ auto z = reinterpret_cast<X*>(vz);
+
+ // Shared memory for caching shape information
+ __shared__ sd::LongType tadLength;
+ __shared__ int numTads;
+ __shared__ int xRank;
+ __shared__ int yRank;
+ __shared__ int zRank;
+
+ __shared__ const sd::LongType* tadShape;
+ __shared__ const sd::LongType* tadStride;
+ __shared__ const sd::LongType* tadShapeZ;
+ __shared__ const sd::LongType* tadStrideZ;
+
+ if (threadIdx.x == 0) {
+   // Cache essential shape information
+   tadLength = shape::length(tadOnlyShapeInfo);
+   numTads   = shape::length(yShapeInfo) / tadLength;
+
+   xRank = shape::rank(xShapeInfo);
+   yRank = shape::rank(yShapeInfo);
+   zRank = shape::rank(zShapeInfo);
+
+   tadShape  = shape::shapeOf(tadOnlyShapeInfo);
+   tadStride = shape::stride(tadOnlyShapeInfo);
+
+   tadShapeZ  = shape::shapeOf(tadOnlyShapeInfoZ);
+   tadStrideZ = shape::stride(tadOnlyShapeInfoZ);
+ }
+ __syncthreads();
+
+ // Each block handles a subset of TADs
+ for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
+   auto zTad = z + tadOffsetsZ[r];
+   auto yTad = y + tadOffsets[r];
+
+   // Loop over TAD elements
+   for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
+     // Derive coordinates and offsets
+     sd::LongType coords[SD_MAX_RANK];
+     sd::LongType xOffset, yOffset, zOffset;
+
+     // Convert index to coordinates using cached shape info
+     INDEX2COORDS(i, xRank, tadShape, coords);
+     COORDS2INDEX(xRank, tadStride, coords, xOffset);
+
+     COORDS2INDEX(yRank, shape::stride(yShapeInfo), coords, yOffset);
+
+     INDEX2COORDS(i, zRank, tadShapeZ, coords);
+     COORDS2INDEX(zRank, tadStrideZ, coords, zOffset);
+
+     // Apply the inverse operation
+     zTad[zOffset] = OpClass::op(x[xOffset], yTad[yOffset]);
+   }
+ }
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// Implementation of the transformCuda device function for BroadcastInt without dimensions
+template <typename X>
+template <typename OpClass>
+SD_DEVICE void BroadcastInt<X>::transformCuda(
+   const void* vx,
+   const sd::LongType* xShapeInfo,
+   const void* vy,
+   const sd::LongType* yShapeInfo,
+   void* vz,
+   const sd::LongType* zShapeInfo)
+{
+ const X* x = reinterpret_cast<const X*>(vx);
+ const X* y = reinterpret_cast<const X*>(vy);
+ X* z       = reinterpret_cast<X*>(vz);
+
+ // Shared memory for caching shape information
+ __shared__ sd::LongType zLen;
+ __shared__ int rank;
+ __shared__ bool xzSameOffsets, yzSameOffsets;
+
+ __shared__ const sd::LongType* xShapeCached;
+ __shared__ const sd::LongType* yShapeCached;
+ __shared__ const sd::LongType* zShapeCached;
+
+ __shared__ const sd::LongType* xStrideCached;
+ __shared__ const sd::LongType* yStrideCached;
+ __shared__ const sd::LongType* zStrideCached;
+
+ if (threadIdx.x == 0) {
+   // Cache essential shape information
+   zLen           = shape::length(zShapeInfo);
+   rank           = shape::rank(zShapeInfo);
+
+   xzSameOffsets  = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
+   yzSameOffsets  = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
+
+   xShapeCached   = shape::shapeOf(xShapeInfo);
+   yShapeCached   = shape::shapeOf(yShapeInfo);
+   zShapeCached   = shape::shapeOf(zShapeInfo);
+
+   xStrideCached  = shape::stride(xShapeInfo);
+   yStrideCached  = shape::stride(yShapeInfo);
+   zStrideCached  = shape::stride(zShapeInfo);
+ }
+ __syncthreads();
+
+ const auto tid          = blockIdx.x * blockDim.x + threadIdx.x;
+ const auto totalThreads = blockDim.x * gridDim.x;
+
+ sd::LongType coords[SD_MAX_RANK];
+
+ for (sd::LongType i = tid; i < zLen; i += totalThreads) {
+   // Quick coordinate transform
+   INDEX2COORDS(i, rank, zShapeCached, coords);
+
+   sd::LongType zOffset, xOffset, yOffset;
+   COORDS2INDEX(rank, zStrideCached, coords, zOffset);
+
+   if (xzSameOffsets) {
+     xOffset = zOffset;
+   } else {
+     COORDS2INDEX(rank, xStrideCached, coords, xOffset);
+   }
+
+   if (yzSameOffsets) {
+     yOffset = zOffset;
+   } else {
+     COORDS2INDEX(rank, yStrideCached, coords, yOffset);
+   }
+
+   z[zOffset] = OpClass::op(x[xOffset], y[yOffset]);
+ }
+}
+
+
+//////////////////////////////////////////////////////////////////////////////
+// Instantiate templates for common integer types
+BUILD_SINGLE_TEMPLATE(
+   template class BroadcastInt, ,
+   SD_INTEGER_TYPES);
+
+}  // namespace broadcast
+}  // namespace functions
+
+
+
+

libnd4j/include/loops/cuda/comb_compilation_units/broadcast_exec_broadcast_with_dimension.cu.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// Created by agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/broadcasting.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function instantiation for Broadcast::execBroadcast with dimension
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* tadOnlyShapeInfoZ,
+        sd::LongType const* tadOffsetsZ);
+);
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo);
+);
+

libnd4j/include/loops/cuda/comb_compilation_units/broadcast_exec_inverse_broadcast.cu.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// Created by agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/broadcasting.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function instantiation for Broadcast::execInverseBroadcast
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execInverseBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* tadOnlyShapeInfoZ,
+        sd::LongType const* tadOffsetsZ);
+);

libnd4j/include/loops/cuda/comb_compilation_units/broadcast_instantiation_template_3.cu.in

@@ -0,0 +1,96 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/broadcasting.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * Broadcast::execBroadcast instantiated for types in @COMB1@, @COMB2@, @COMB3@
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* tadOnlyShapeInfoZ,
+        sd::LongType const* tadOffsetsZ);
+);
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo);
+);
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::broadcast::Broadcast,
+    ::execInverseBroadcast(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* x,
+        sd::LongType const* xShapeInfo,
+        void const* y,
+        sd::LongType const* yShapeInfo,
+        void* z,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* tadOnlyShapeInfoZ,
+        sd::LongType const* tadOffsetsZ);
+);

libnd4j/include/loops/cuda/comb_compilation_units/indexreduce_instantiation_template_2.cu.in

@@ -0,0 +1,81 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/indexreduce.cu>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * IndexReduce::executeIndexReduce instantiated for types in @COMB1@, @COMB2@
+ * Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
+ */
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_INDEXING_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::indexreduce::IndexReduce,
+    ::executeIndexReduce(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        void const* dx,
+        sd::LongType const* xShapeInfo,
+        sd::LongType xRank,
+        void* extraParams,
+        void* result,
+        sd::LongType const* zShapeInfo,
+        sd::LongType zRank,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        int postProcessOrNot,
+        sd::LongType* allocationBuffer,
+        void* reductionBuffer,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets);
+);
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_INDEXING_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::indexreduce::IndexReduce,
+    ::executeIndexReduceScalar(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        void const* dx,
+        sd::LongType const* xShapeInfo,
+        sd::LongType xRank,
+        void* extraParams,
+        void* result,
+        sd::LongType const* zShapeInfo,
+        sd::LongType zRank,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        int postProcessOrNot,
+        sd::LongType* allocationBuffer,
+        void* reductionBuffer,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets);
+);

libnd4j/include/loops/cuda/comb_compilation_units/pairwise_exec_cuda_shaped.cu.in

@@ -0,0 +1,50 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/pairwise.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function instantiation for PairWiseTransform::executeCudaShaped
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::pairwise_transforms::PairWiseTransform,
+    ::executeCudaShaped(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        const void *vy,
+        const sd::LongType *yShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        void *vextraParams);
+);

libnd4j/include/loops/cuda/comb_compilation_units/pairwise_instantiation_template_3.cu.in

@@ -0,0 +1,49 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/pairwise.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * PairWiseTransform::executeCudaShaped instantiated for types in @COMB1@, @COMB2@, @COMB3@
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::pairwise_transforms::PairWiseTransform,
+    ::executeCudaShaped(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        const void *vy,
+        const sd::LongType *yShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        void *vextraParams);
+);

libnd4j/include/loops/cuda/comb_compilation_units/reduce3_instantiation_template_2.cu.in

@@ -0,0 +1,104 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/reduce3.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * Reduce3::exec instantiated for types in @COMB1@, @COMB2@
+ * Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
+ */
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce3::Reduce3,
+    ::exec(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* vx,
+        sd::LongType const* xShapeInfo,
+        void const* vy,
+        sd::LongType const* yShapeInfo,
+        void* extraParams,
+        void* vz,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        int postProcessOrNot,
+        sd::LongType* allocationPointer,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* yTadOnlyShapeInfo,
+        sd::LongType const* yTadOffsets);
+);
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce3::Reduce3,
+    ::execAll(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* vx,
+        sd::LongType const* xShapeInfo,
+        void const* vy,
+        sd::LongType const* yShapeInfo,
+        void* extraParams,
+        void* vz,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* dimension,
+        long long int dimensionLength,
+        int postProcessOrNot,
+        sd::LongType* allocationPointer,
+        sd::LongType const* tadOnlyShapeInfo,
+        sd::LongType const* tadOffsets,
+        sd::LongType const* yTadOnlyShapeInfo,
+        sd::LongType const* yTadOffsets);
+);
+
+ITERATE_COMBINATIONS(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce3::Reduce3,
+    ::execScalar(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        void const* vx,
+        sd::LongType const* xShapeInfo,
+        void const* vy,
+        sd::LongType const* yShapeInfo,
+        void* extraParams,
+        void* vz,
+        sd::LongType const* zShapeInfo,
+        sd::LongType* allocationPointer,
+        void* reductionBuffer,
+        sd::LongType const* tadOnlyShapeInfo);
+);

libnd4j/include/loops/cuda/comb_compilation_units/reduce_float_instantiation_template_2.cu.in

@@ -0,0 +1,73 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/reduce/reduce_float.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * ReduceFloatFunction::execReduce instantiated for types in @COMB1@, @COMB2@
+ * Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
+ */
+
+ITERATE_COMBINATIONS(
+    SD_NUMERIC_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceFloatFunction,
+    ::execReduce(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        const void* x,
+        const sd::LongType* dXShapeInfo,
+        const sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* vreductionBuffer,
+        void* z,
+        const sd::LongType* dZShapeInfo,
+        const sd::LongType* hZShapeInfo,
+        const sd::LongType* dims);
+);
+
+ITERATE_COMBINATIONS(
+    SD_NUMERIC_TYPES_PART_@COMB1@,
+    SD_FLOAT_TYPES_PART_@COMB2@,
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceFloatFunction,
+    ::execReduceScalar(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        const void* x,
+        const sd::LongType* xShapeInfo,
+        const sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* z,
+        const sd::LongType* dZShapeInfo,
+        const sd::LongType* hZShapeInfo,
+        sd::LongType* dimension,
+        long long int dimensionLength,
+        void* reductionBuffer,
+        const sd::LongType* tadOnlyShapeInfo);
+);

libnd4j/include/loops/cuda/comb_compilation_units/scalar_exec_cuda_along_dimension.cu.in

@@ -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
+ *  *****************************************************************************
+ */
+
+//
+// Created by agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/scalar.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function instantiation for ScalarTransform::executeCudaAlongDimension
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::scalar::ScalarTransform,
+    ::executeCudaAlongDimension(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        const void *vscalars,
+        void *vextraParams,
+        sd::LongType *dimension,
+        sd::LongType dimensionLength,
+        const sd::LongType *tadShapeInfo,
+        const sd::LongType *tadOffsets,
+        const sd::LongType *tadShapeInfoZ,
+        const sd::LongType *tadOffsetsZ);
+);

libnd4j/include/loops/cuda/comb_compilation_units/scalar_exec_cuda_shaped.cu.in

@@ -0,0 +1,51 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/scalar.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function instantiation for ScalarTransform::executeCudaShaped
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_SINGLE_TYPE_@COMB1@,
+    SD_SINGLE_TYPE_@COMB2@,
+    SD_SINGLE_TYPE_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::scalar::ScalarTransform,
+    ::executeCudaShaped(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        const sd::LongType *hxShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        const sd::LongType *hzShapeInfo,
+        const void *vscalar,
+        void *vextraParams);
+);

libnd4j/include/loops/cuda/comb_compilation_units/scalar_instantiation_template_3.cu.in

@@ -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
+ ******************************************************************************/
+
+//
+// Created by agibsonccc on 3/10/25.
+//
+
+#include <loops/cuda/scalar.chpp>
+#include <system/type_boilerplate.h>
+#include <loops/pairwise_instantiations_single.h>
+
+/*
+ * Function Instantiation:
+ * ScalarTransform::executeCudaShaped instantiated for types in @COMB1@, @COMB2@, @COMB3@
+ */
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::scalar::ScalarTransform,
+    ::executeCudaShaped(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        const sd::LongType *hxShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        const sd::LongType *hzShapeInfo,
+        const void *vscalar,
+        void *vextraParams);
+);
+
+ITERATE_COMBINATIONS_3(
+    SD_COMMON_TYPES_PART_@COMB1@,
+    SD_COMMON_TYPES_PART_@COMB2@,
+    SD_COMMON_TYPES_PART_@COMB3@,
+    INSTANT_PROCESS_COMBINATION_3,
+    functions::scalar::ScalarTransform,
+    ::executeCudaAlongDimension(
+        dim3& launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void *vx,
+        const sd::LongType *xShapeInfo,
+        void *vz,
+        const sd::LongType *zShapeInfo,
+        const void *vscalars,
+        void *vextraParams,
+        sd::LongType *dimension,
+        sd::LongType dimensionLength,
+        const sd::LongType *tadShapeInfo,
+        const sd::LongType *tadOffsets,
+        const sd::LongType *tadShapeInfoZ,
+        const sd::LongType *tadOffsetsZ);
+);

libnd4j/include/loops/cuda/compilation_units/broadcasting.cu.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cuda/broadcasting.chpp>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace broadcast {
+        BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN Broadcast, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cuda/compilation_units/pairwise.cu.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cuda/pairwise.chpp>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace pairwise_transforms {
+        BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN PairWiseTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cuda/compilation_units/reduce3.cu.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cuda/reduce3.chpp>
+#cmakedefine SD_FLOAT_TYPES_GEN
+#if defined(SD_FLOAT_TYPES_GEN) && defined(SD_FLOAT_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace reduce3 {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES, SD_FLOAT_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cuda/compilation_units/reduce_float.cu.in

@@ -0,0 +1,30 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include <loops/cuda/reduce/reduce_float.chpp>
+#cmakedefine SD_FLOAT_TYPES_GEN
+#if defined(SD_FLOAT_TYPES_GEN) && defined(SD_FLOAT_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace reduce {
+        BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN ReduceFloatFunction, , SD_COMMON_TYPES, SD_FLOAT_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cuda/compilation_units/scalar.cu.in

@@ -0,0 +1,31 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <loops/cuda/scalar.chpp>
+#include <system/type_boilerplate.h>
+#cmakedefine SD_PAIRWISE_TYPES_GEN
+#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
+namespace functions {
+    namespace scalar {
+    BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN ScalarTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
+    }
+}
+#endif

libnd4j/include/loops/cuda/indexreduce.cu

@@ -0,0 +1,331 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing,
+* software distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+* the License for the specific language governing permissions and
+* limitations under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+//
+// Created by raver on 4/9/2018.
+//
+#include <helpers/DebugHelper.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+//note: keep this. It's required for proper linker work
+
+#include "../indexreduce.h"
+#include "../legacy_ops.h"
+
+
+using namespace simdOps;
+
+template <typename X, typename Z>
+static SD_KERNEL void simpleIndexReduceGeneric(const int op, void const *dx, sd::LongType const *xShapeInfo,
+                                              sd::LongType xRank,
+                                              void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType zRank,
+                                              sd::LongType *dimension, sd::LongType dimensionLength, int postProcessOrNot, sd::LongType *allocationBuffer, void *reductionBuffer,
+                                              sd::LongType const *tadOnlyShapeInfo, sd::LongType const *tadOffsets) {
+ functions::indexreduce::IndexReduce<X, Z>::transform(op, dx, xShapeInfo, extraParams, result, zShapeInfo, dimension,
+                                                      dimensionLength, postProcessOrNot, allocationBuffer,
+                                                      reductionBuffer, tadOnlyShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace indexreduce {
+
+template <typename X, typename Z>
+SD_HOST void IndexReduce<X, Z>::executeIndexReduceScalar(
+   dim3 launchDims, cudaStream_t *stream, const int opNum, void const *dx, sd::LongType const *xShapeInfo,
+   sd::LongType xRank,
+   void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType zRank,
+   sd::LongType *dimension, sd::LongType dimensionLength,
+   int postProcessOrNot,sd::LongType *allocationBuffer, void *reductionBuffer, sd::LongType const *tadOnlyShapeInfo,
+   sd::LongType const *tadOffsets) {
+ simpleIndexReduceGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+     opNum, dx, xShapeInfo, xRank, extraParams, result, zShapeInfo, 0, nullptr, 0, 1, allocationBuffer,
+     reductionBuffer, tadOnlyShapeInfo, tadOffsets);
+ sd::DebugHelper::checkErrorCode(stream, "executeIndexReduceScalar(...) failed");
+
+}
+
+template <typename X, typename Z>
+SD_HOST void IndexReduce<X, Z>::executeIndexReduce(dim3 launchDims,
+                                                  cudaStream_t *stream,
+                                                  const int opNum,
+                                                  void const *dx,
+                                                  sd::LongType const *xShapeInfo,
+                                                  sd::LongType xRank,
+                                                  void *extraParams,
+                                                  void *result,
+                                                  sd::LongType const *zShapeInfo,
+                                                  sd::LongType zRank,
+                                                  sd::LongType *dimension,
+                                                  sd::LongType dimensionLength,
+                                                  int postProcessOrNot,
+                                                  sd::LongType *allocationBuffer,
+                                                  void *reductionBuffer,
+                                                  sd::LongType const *tadOnlyShapeInfo,
+                                                  sd::LongType const *tadOffsets) {
+ simpleIndexReduceGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+     opNum, dx, xShapeInfo, xRank, extraParams, result, zShapeInfo, zRank, dimension, dimensionLength, postProcessOrNot,
+     allocationBuffer, reductionBuffer, tadOnlyShapeInfo, tadOffsets);
+ sd::DebugHelper::checkErrorCode(stream, "executeIndexReduce(...) failed");
+
+}
+
+// This is the un-specialized struct.  Note that we prevent instantiation of this
+// struct by putting an undefined symbol in the function body so it won't compile.
+template <typename T>
+struct SharedIndexValue {
+ // Ensure that we won't compile any un-specialized types
+ SD_DEVICE T *getPointer() {
+   extern SD_DEVICE void error(void);
+   error();
+   return 0;
+ }
+};
+
+// Following are the specializations for the following types.
+// int, sd::Unsigned, char, uchar, short, ushort, long long, ulong long, bool, float, and double
+// One could also specialize it for user-defined types.
+
+template <>
+struct SharedIndexValue<float> {
+ SD_DEVICE IndexValue<float> *getPointer() {
+   extern __shared__ IndexValue<float> s_int2[];
+   return s_int2;
+ }
+};
+// Following are the specializations for the following types.
+// int, sd::Unsigned, char, uchar, short, ushort, long long, ulong long, bool, float, and double
+// One could also specialize it for user-defined types.
+
+template <>
+struct SharedIndexValue<double> {
+ SD_DEVICE IndexValue<double> *getPointer() {
+   extern __shared__ IndexValue<double> s_int6[];
+   return s_int6;
+ }
+};
+
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void IndexReduce<X, Z>::aggregatePartials(IndexValue<X> *sPartials, sd::LongType tid,
+                                                   sd::LongType numElements, void *vextraParams) {
+ // start the shared memory loop on the next power of 2 less
+ // than the block size.  If block size is not a power of 2,
+ // accumulate the intermediate sums in the remainder range.
+ auto extraParams = static_cast<X *>(vextraParams);
+ sd::LongType floorPow2 = static_cast<sd::LongType>(blockDim.x);
+
+ if (floorPow2 & (floorPow2 - 1)) {
+   while (floorPow2 & (floorPow2 - 1)) {
+     floorPow2 &= floorPow2 - 1;
+   }
+
+   if (tid >= floorPow2) {
+     IndexValue<X> prev = sPartials[tid - floorPow2];
+     IndexValue<X> curr = sPartials[tid];
+     sPartials[tid - floorPow2] = OpType::update(prev, curr, extraParams);
+   }
+   __syncthreads();
+ }
+
+ for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
+   if (tid < activeThreads && tid + activeThreads < numElements) {
+     IndexValue<X> curr = sPartials[tid];
+     IndexValue<X> next = sPartials[tid + activeThreads];
+     sPartials[tid] = OpType::update(curr, next, extraParams);
+   }
+   __syncthreads();
+ }
+}
+
+template <typename X, typename Y>
+SD_DEVICE void IndexReduce<X, Y>::transform(int opNum, void const *x, sd::LongType const *xShapeInfo,
+                                           void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType *dimension,
+                                           sd::LongType dimensionLength, int postProcessOrNot,
+                                           sd::LongType *allocationBuffer, void *reductionBuffer,
+                                           sd::LongType const *tadShapeInfo, sd::LongType const *tadOffset) {
+ DISPATCH_BY_OPNUM_TT(transform,
+                      PARAMS(x, xShapeInfo, extraParams, result, zShapeInfo, dimension, dimensionLength,
+                             postProcessOrNot, allocationBuffer, reductionBuffer, tadShapeInfo, tadOffset),
+                      INDEX_REDUCE_OPS);
+}
+
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void IndexReduce<X, Z>::transform(void const *vdx, sd::LongType const *xShapeInfo, void *vextraParams,
+                                           void *vz, sd::LongType const *zShapeInfo, sd::LongType *dimension,
+                                           sd::LongType dimensionLength, int postProcessOrNot,
+                                           sd::LongType *allocationBuffer,
+                                           void *vreductionBuffer, sd::LongType const *tadOnlyShapeInfo,
+                                           sd::LongType const *tadOffsets) {
+ auto dx = reinterpret_cast<X const *>(vdx);
+ auto z = reinterpret_cast<Z *>(vz);
+ auto extraParams = static_cast<X *>(vextraParams);
+ auto reductionBuffer = static_cast<unsigned int *>(vreductionBuffer);
+ auto order = shape::order(xShapeInfo);
+ sd::LongType tid = static_cast<sd::LongType>(blockIdx.x * blockDim.x + threadIdx.x);
+ __shared__ volatile bool resultScalar;
+
+ __shared__ IndexValue<X> sPartials[SD_CUDA_BLOCK_SIZE];
+
+ sPartials[threadIdx.x] = OpType::startingIndexValue(dx);
+
+ __shared__ volatile sd::LongType xLength;
+
+ __shared__ volatile sd::LongType zLen;
+
+ IndexValue<X> reduction = OpType::startingIndexValue(dx);
+ sd::LongType threadIdxX = static_cast<sd::LongType>(threadIdx.x);
+ sd::LongType blockDimX = static_cast<sd::LongType>(blockDim.x);
+ sd::LongType blockIdxX = static_cast<sd::LongType>(blockIdx.x);
+ sd::LongType gridDimX = static_cast<sd::LongType>(gridDim.x);
+
+ if (threadIdxX == 0) {
+   if (zShapeInfo != nullptr)
+     zLen = shape::length(zShapeInfo);
+   else
+     zLen = 1;
+
+   if (zLen == 1)
+     resultScalar = true;
+   else
+     resultScalar = false;
+
+   xLength = shape::length(xShapeInfo);
+ }
+ __syncthreads();
+
+ if (!resultScalar) {
+   __shared__ sd::LongType tadLength;
+   __shared__ sd::LongType tadEWS;
+   __shared__ sd::LongType numTads;
+
+   if (threadIdx.x == 0) {
+     tadLength = shape::length(tadOnlyShapeInfo);
+     tadEWS = shape::elementWiseStride(tadOnlyShapeInfo);
+     numTads = shape::length(xShapeInfo) / tadLength;
+   }
+   __syncthreads();
+
+   if (dimensionLength > 1 || tadEWS < 1) {
+     for (sd::LongType r = blockIdxX; r < numTads; r += gridDimX) {
+       auto tadOffsetForBlock = tadOffsets[r];
+       sPartials[threadIdxX] = OpType::startingIndexValue(dx);
+
+       for (sd::LongType i = threadIdxX; i < tadLength; i += blockDimX) {
+         sd::LongType coords[SD_MAX_RANK];
+         sd::LongType xOffset;
+         INDEX2COORDS(i, shape::rank(tadOnlyShapeInfo), shape::shapeOf(tadOnlyShapeInfo), coords);
+         COORDS2INDEX(shape::rank(tadOnlyShapeInfo),  shape::stride(tadOnlyShapeInfo), coords, xOffset);
+         xOffset += tadOffsetForBlock;
+         IndexValue<X> comp{dx[xOffset], i};
+         sPartials[threadIdxX] = OpType::update(sPartials[threadIdxX], comp, extraParams);
+       }
+
+       __syncthreads();
+       aggregatePartials<OpType>(sPartials,threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDimX, tadLength), extraParams);
+
+       __syncthreads();
+       if (threadIdxX == 0) {
+         z[r] = static_cast<Z>(sPartials[threadIdxX].index);
+       }
+       __syncthreads();
+     }
+   } else {
+     for (sd::LongType i = blockIdxX; i < numTads; i += gridDimX) {
+       sd::LongType tadOffsetForBlock = tadOffsets[i];
+
+       sPartials[threadIdxX] = OpType::startingIndexValue(dx);
+
+       for (sd::LongType x = threadIdxX; x < tadLength; x += blockDimX) {
+         sd::LongType coords[SD_MAX_RANK];
+         sd::LongType xOffset;
+         INDEX2COORDS(x, shape::rank(tadOnlyShapeInfo), shape::shapeOf(tadOnlyShapeInfo), coords);
+         COORDS2INDEX(shape::rank(tadOnlyShapeInfo), shape::stride(tadOnlyShapeInfo), coords, xOffset);
+         IndexValue<X> comp{dx[tadOffsetForBlock + xOffset], x};
+         sPartials[threadIdxX] = OpType::update(sPartials[threadIdxX], comp, extraParams);
+       }
+
+       __syncthreads();
+       aggregatePartials<OpType>(sPartials, threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDim.x, tadLength), extraParams);
+
+       __syncthreads();
+       if (threadIdxX == 0) {
+         z[i] = static_cast<Z>(sPartials[threadIdxX].index);
+       }
+       __syncthreads();
+     }
+   }
+ } else {
+   auto n = shape::length(xShapeInfo);
+
+   for (sd::LongType i = tid; i < n; i += (gridDimX * blockDimX)) {
+     sd::LongType coords[SD_MAX_RANK];
+     sd::LongType xOffset;
+     INDEX2COORDS(i, shape::rank(xShapeInfo), shape::shapeOf(xShapeInfo), coords);
+     COORDS2INDEX(shape::rank(xShapeInfo), shape::stride(xShapeInfo), coords, xOffset);
+     IndexValue<X> comp{dx[xOffset], i};
+     reduction = OpType::update(reduction, comp, extraParams);
+   }
+
+   sPartials[threadIdxX] = reduction;
+   __syncthreads();
+   aggregatePartials<OpType>(sPartials, threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDim.x, n), extraParams);
+   if (gridDimX > 1) {
+     __shared__ bool amLast;
+     unsigned int *unsignedSharedMemory = (unsigned int *)reductionBuffer;
+     tid = threadIdx.x;
+     if (threadIdx.x == 0)
+       reductionBuffer[blockIdx.x] = sPartials[threadIdx.x].index;
+
+     __threadfence();
+     __syncthreads();
+
+     if (threadIdx.x == 0) {
+       unsigned int ticket = atomicInc(&unsignedSharedMemory[16384], gridDim.x);
+       amLast = (ticket == gridDim.x - 1);
+     }
+
+     __syncthreads();
+
+     if (amLast) {
+       sPartials[threadIdx.x] = OpType::startingIndexValue(dx);
+       for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+         IndexValue<X> comp{static_cast<X>(0), reductionBuffer[i]};
+         sPartials[threadIdx.x] = OpType::update(sPartials[threadIdx.x], comp, extraParams);
+       }
+       __syncthreads();
+       aggregatePartials<OpType>(sPartials, threadIdxX, gridDim.x, extraParams);
+
+       if (threadIdx.x == 0) {
+         z[0] = static_cast<Z>(sPartials[threadIdx.x].index);
+         unsignedSharedMemory[16384] = 0;
+       }
+     }
+   } else {
+     if (threadIdx.x == 0) {
+       z[0] = static_cast<Z>(sPartials[threadIdx.x].index);
+     }
+
+   }
+ }
+}
+
+}  // namespace indexreduce
+}  // namespace functions

libnd4j/include/loops/cuda/inplace_loops/reduce_same_inplace.h

@@ -0,0 +1,242 @@
+/* ******************************************************************************
+ *
+ *
+ * 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_REDUCE_SAME_LOOPS_H
+#define DEV_TESTS_REDUCE_SAME_LOOPS_H
+
+#include <cuda_runtime.h>
+#include <helpers/shape.h>
+#include <ops/ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include <execution/LaunchContext.h>
+#include "execution/cuda/LaunchDims.h"
+using namespace simdOps;
+
+namespace functions {
+namespace reduce {
+
+template <typename X>
+class ReduceSameInplace {
+ public:
+  // Dispatch method for old-style calls
+  static SD_INLINE void SD_DEVICE execScalarCudaLegacy(
+      int opNum,
+      void* vx,
+      sd::LongType* xShapeInfo,
+      void* vextraParams,
+      void* vz,
+      sd::LongType* zShapeInfo,
+      void* vreductionBuffer,
+      sd::LongType* tadOnlyShapeInfo);
+
+  // Template for real call
+  template <typename OpClass>
+  static SD_INLINE void SD_DEVICE execScalarCuda(
+      void* vx,
+      sd::LongType* xShapeInfo,
+      void* vextraParams,
+      void* vz,
+      sd::LongType* zShapeInfo,
+      void* vreductionBuffer,
+      sd::LongType* tadOnlyShapeInfo);
+
+  template <typename OpClass>
+  static SD_INLINE void SD_DEVICE aggregatePartials(
+      void* vsPartials,
+      sd::LongType tid,
+      sd::LongType numItems,
+      void* vextraParams);
+};
+
+template <typename X>
+template <typename OpClass>
+SD_INLINE void SD_DEVICE ReduceSameInplace<X>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  auto sPartials   = static_cast<X*>(vsPartials);
+  auto extraParams = static_cast<X*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= (floorPow2 - 1);
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpClass::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] =
+          OpClass::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+template <typename X>
+SD_INLINE void SD_DEVICE ReduceSameInplace<X>::execScalarCudaLegacy(
+    int opNum,
+    void* vx,
+    sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_T(
+      execScalarCuda,
+      PARAMS(vx, xShapeInfo, vextraParams, vz, zShapeInfo, vreductionBuffer, tadOnlyShapeInfo),
+      REDUCE_SAME_OPS);
+}
+
+template <typename X>
+template <typename OpClass>
+SD_INLINE void SD_DEVICE ReduceSameInplace<X>::execScalarCuda(
+    void* vx,
+    sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    sd::LongType* tadOnlyShapeInfo) {
+
+  auto x             = reinterpret_cast<X*>(vx);
+  auto z             = reinterpret_cast<X*>(vz);
+  auto extraParams   = reinterpret_cast<X*>(vextraParams);
+  auto reductionBuff = reinterpret_cast<X*>(vreductionBuffer);
+
+  // We'll cache relevant shape info in shared memory so we don't call them repeatedly
+  __shared__ sd::LongType length;
+  __shared__ int rank;
+  __shared__ const sd::LongType* shapePtr;
+  __shared__ const sd::LongType* stridePtr;
+
+  if (threadIdx.x == 0) {
+    length    = shape::length(xShapeInfo);
+    rank      = shape::rank(xShapeInfo);
+    shapePtr  = shape::shapeOf(xShapeInfo);
+    stridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  const auto tid      = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto gridSize = gridDim.x * blockDim.x;
+
+  // We'll use some shared memory for partial sums
+  __shared__ X* sPartials;
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char shmem[];
+    sPartials = reinterpret_cast<X*>(shmem);
+  }
+  __syncthreads();
+
+  // Each thread gets a starting value
+  sPartials[threadIdx.x] = OpClass::startingValue(x);
+
+  // We'll stride over the entire array
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType coords[SD_MAX_RANK];
+    sd::LongType offset;
+
+    INDEX2COORDS(i, rank, shapePtr, coords);
+    COORDS2INDEX(rank, stridePtr, coords, offset);
+
+    sPartials[threadIdx.x] =
+        OpClass::update(sPartials[threadIdx.x], OpClass::op(x[offset], extraParams), extraParams);
+  }
+  __syncthreads();
+
+  // Next: reduce partial sums in the block
+  aggregatePartials<OpClass>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  // If gridDim.x > 1, we do a multi-block reduce using the global buffer
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(reductionBuff);
+    __shared__ bool amLast;
+
+    // each block's sum is stored in the 'reductionBuff'
+    if (threadIdx.x == 0) {
+      reductionBuff[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    // The 16384 is a special "counter" location in the reductionBuff
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == gridDim.x - 1);
+    }
+    __syncthreads();
+
+    if (amLast) {
+      // We do a final block-level reduce
+      tc[16384] = 0; // reset
+      sPartials[threadIdx.x] = OpClass::startingValue(x);
+
+      // accumulate partial sums from each block
+      for (int i = threadIdx.x; i < static_cast<int>(gridDim.x); i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpClass::update(sPartials[threadIdx.x], reductionBuff[i], extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpClass>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpClass::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  }
+  else {
+    // single-block case
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(reductionBuff);
+      tc[16384] = 0;
+      z[0]      = OpClass::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+}  // namespace reduce
+}  // namespace functions
+
+#endif  // DEV_TESTS_REDUCE_SAME_LOOPS_H

libnd4j/include/loops/cuda/inplace_loops/scalar_inplace.h

@@ -0,0 +1,138 @@
+/* ******************************************************************************
+ *
+ *
+ * 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_SCALAR_INPLACE_H
+#define DEV_TESTS_SCALAR_INPLACE_H
+
+#include <helpers/shape.h>
+#include <ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+    using namespace simdOps;
+
+namespace functions {
+namespace scalar {
+
+template <typename X, typename Y, typename Z>
+class ScalarInplace {
+ public:
+  static SD_INLINE SD_DEVICE void transformCudaLegacy(
+      int opNum,
+      void* vscalar,
+      void* vy,
+      sd::LongType* yShapeInfo,
+      void* vparams,
+      void* vz,
+      sd::LongType* zShapeInfo,
+      int* allocationBuffer);
+
+  template <typename OpClass>
+  static SD_INLINE SD_DEVICE void transformCuda(
+      void* vscalar,
+      void* vy,
+      sd::LongType* yShapeInfo,
+      void* vparams,
+      void* vz,
+      sd::LongType* zShapeInfo,
+      int* allocationBuffer);
+};
+
+template <typename X, typename Y, typename Z>
+SD_INLINE SD_DEVICE void ScalarInplace<X, Y, Z>::transformCudaLegacy(
+    int opNum,
+    void* vscalar,
+    void* vy,
+    sd::LongType* yShapeInfo,
+    void* vparams,
+    void* vz,
+    sd::LongType* zShapeInfo,
+    int* allocationBuffer) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      transformCuda,
+      PARAMS(vscalar, vy, yShapeInfo, vparams, vz, zShapeInfo, allocationBuffer),
+      SCALAR_OPS);
+}
+
+template <typename X, typename Y, typename Z>
+template <typename OpClass>
+SD_INLINE SD_DEVICE void ScalarInplace<X, Y, Z>::transformCuda(
+    void* vscalar,
+    void* vy,
+    sd::LongType* yShapeInfo,
+    void* vparams,
+    void* vz,
+    sd::LongType* zShapeInfo,
+    int* allocationBuffer) {
+
+  auto scalar  = reinterpret_cast<X*>(vscalar)[0];
+  auto y       = reinterpret_cast<Y*>(vy);
+  auto params  = reinterpret_cast<Z*>(vparams);
+  auto z       = reinterpret_cast<Z*>(vz);
+
+  // We'll store relevant shape info for y & z in shared memory
+  __shared__ sd::LongType yLen;
+  __shared__ int yRank;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ int zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    yLen       = shape::length(yShapeInfo);
+    yRank      = shape::rank(yShapeInfo);
+    yShapePtr  = shape::shapeOf(yShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+
+    zRank      = shape::rank(zShapeInfo);
+    zShapePtr  = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  const int totalThreads = gridDim.x * blockDim.x;
+  const int tid          = blockIdx.x * blockDim.x + threadIdx.x;
+
+  // Each thread does part of the length
+  for (sd::LongType i = tid; i < yLen; i += totalThreads) {
+    sd::LongType coordsY[SD_MAX_RANK];
+    sd::LongType coordsZ[SD_MAX_RANK];
+    sd::LongType offsetY;
+    sd::LongType offsetZ;
+
+    INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+    COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
+
+    INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+    COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+    // apply scalar op in place to produce z
+    z[offsetZ] = OpClass::op(y[offsetY], scalar, params);
+  }
+}
+
+}  // namespace scalar
+}  // namespace functions
+
+#endif  // DEV_TESTS_SCALAR_INPLACE_H

libnd4j/include/loops/cuda/inplace_loops/transform_strict_inplace.h

@@ -0,0 +1,148 @@
+/* ******************************************************************************
+ *
+ *
+ * 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_TRANSFORM_FLOAT_INPLACE_H
+#define DEV_TESTS_TRANSFORM_FLOAT_INPLACE_H
+
+#include <helpers/shape.h>
+#include <ops.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+    using namespace simdOps;
+
+#define LOCAL_TRANSFORM_STRICT_OPS (23, Exp), (24, Log)
+
+namespace functions {
+namespace transform {
+
+template <typename X>
+class TransformStrictInplace {
+ public:
+  static SD_INLINE SD_DEVICE void transformCudaLegacy(
+      int opNum,
+      void* dy,
+      sd::LongType* shapeInfo,
+      void* params,
+      void* result,
+      sd::LongType* zShapeInfo,
+      int* allocationPointer,
+      void* reductionPointer,
+      sd::LongType* tadShapeInfo,
+      sd::LongType* tadOffsets);
+
+  template <typename OpClass>
+  static SD_INLINE SD_DEVICE void transformCuda(
+      void* vdy,
+      sd::LongType* shapeInfo,
+      void* vparams,
+      void* vresult,
+      sd::LongType* zShapeInfo,
+      int* allocationPointer,
+      void* vreductionPointer,
+      sd::LongType* tadShapeInfo,
+      sd::LongType* tadOffsets);
+};
+
+template <typename X>
+template <typename OpClass>
+SD_INLINE SD_DEVICE void TransformStrictInplace<X>::transformCuda(
+    void* vdy,
+    sd::LongType* shapeInfo,
+    void* vparams,
+    void* vresult,
+    sd::LongType* zShapeInfo,
+    int* allocationPointer,
+    void* vreductionPointer,
+    sd::LongType* tadShapeInfo,
+    sd::LongType* tadOffsets) {
+
+  auto dy    = static_cast<X*>(vdy);
+  auto result= static_cast<X*>(vresult);
+  auto params= static_cast<X*>(vparams);
+  // reductionPointer is unused in this basic transform?
+
+  const int tid          = blockIdx.x * blockDim.x + threadIdx.x;
+  const int totalThreads = gridDim.x * blockDim.x;
+
+  // Cache shape info in shared memory to avoid repeated calls
+  __shared__ sd::LongType length;
+  __shared__ int rank;
+  __shared__ const sd::LongType* shapePtr;
+  __shared__ const sd::LongType* stridePtr;
+
+  __shared__ int zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length    = shape::length(shapeInfo);
+    rank      = shape::rank(shapeInfo);
+    shapePtr  = shape::shapeOf(shapeInfo);
+    stridePtr = shape::stride(shapeInfo);
+
+    zRank      = shape::rank(zShapeInfo);
+    zShapePtr  = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType inCoords[SD_MAX_RANK];
+    sd::LongType outCoords[SD_MAX_RANK];
+    sd::LongType inOffset;
+    sd::LongType outOffset;
+
+    INDEX2COORDS(i, rank, shapePtr, inCoords);
+    COORDS2INDEX(rank, stridePtr, inCoords, inOffset);
+
+    INDEX2COORDS(i, zRank, zShapePtr, outCoords);
+    COORDS2INDEX(zRank, zStridePtr, outCoords, outOffset);
+
+    // Apply transform op in place
+    result[outOffset] = OpClass::op(dy[inOffset], params);
+  }
+}
+
+template <typename X>
+SD_INLINE SD_DEVICE void TransformStrictInplace<X>::transformCudaLegacy(
+    int opNum,
+    void* dy,
+    sd::LongType* shapeInfo,
+    void* params,
+    void* result,
+    sd::LongType* zShapeInfo,
+    int* allocationPointer,
+    void* reductionPointer,
+    sd::LongType* tadShapeInfo,
+    sd::LongType* tadOffsets) {
+
+  DISPATCH_BY_OPNUM_T(
+      transformCuda,
+      PARAMS(dy, shapeInfo, params, result, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets),
+      LOCAL_TRANSFORM_STRICT_OPS);
+}
+
+}  // namespace transform
+}  // namespace functions
+
+#undef LOCAL_TRANSFORM_STRICT_OPS
+#endif  // DEV_TESTS_TRANSFORM_FLOAT_INPLACE_H

libnd4j/include/loops/cuda/pairwise.chpp

@@ -0,0 +1,137 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+ * the License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//  @author raver119@gmail.com
+//  @author Yurii Shyrma (iuriish@yahoo.com)
+
+#ifndef PAIRWISE_CU
+#define PAIRWISE_CU
+#include <loops/pairwise_instantiations.h>
+
+#include "../pairwise_transform.h"
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z, typename OpType>
+SD_KERNEL static void pairwiseSimpleShaped(void const* vx, sd::LongType const* xShapeInfo,
+                                           void const* vy, sd::LongType const* yShapeInfo,
+                                           void* vz, sd::LongType const* zShapeInfo,
+                                           void* vextraParams) {
+  auto x          = static_cast<X const*>(vx);
+  auto y          = static_cast<Y const*>(vy);
+  auto z          = static_cast<Z*>(vz);
+  auto extraParams= static_cast<Z*>(vextraParams);
+
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  __shared__ sd::LongType len;
+  __shared__ sd::LongType xRank;
+  __shared__ sd::LongType yRank;
+  __shared__ sd::LongType zRank;
+
+  __shared__ const sd::LongType* xShape;
+  __shared__ const sd::LongType* xStride;
+  __shared__ const sd::LongType* yShape;
+  __shared__ const sd::LongType* yStride;
+  __shared__ const sd::LongType* zShape;
+  __shared__ const sd::LongType* zStride;
+
+  if (threadIdx.x == 0) {
+    len     = shape::length(xShapeInfo);
+
+    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 i = tid; i < len; i += (gridDim.x * blockDim.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(i, xRank, xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+
+    INDEX2COORDS(i, yRank, yShape, yCoords);
+    COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+
+    INDEX2COORDS(i, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+  }
+}
+
+namespace functions {
+namespace pairwise_transforms {
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void SD_HOST PairWiseTransform<X, Y, Z>::intermediateShaped(dim3& launchDims,
+                                                            cudaStream_t* stream,
+                                                            void const* vx,
+                                                            sd::LongType const* xShapeInfo,
+                                                            void const* vy,
+                                                            sd::LongType const* yShapeInfo,
+                                                            void* vz,
+                                                            sd::LongType const* zShapeInfo,
+                                                            void* vextraParams) {
+  pairwiseSimpleShaped<X, Y, Z, OpType>
+  <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      vx, xShapeInfo,
+      vy, yShapeInfo,
+      vz, zShapeInfo,
+      vextraParams);
+  sd::DebugHelper::checkErrorCode(stream, "PairWiseTransform intermediateShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void SD_HOST PairWiseTransform<X, Y, Z>::executeCudaShaped(dim3& launchDims,
+                                                           cudaStream_t* stream,
+                                                           int opNum,
+                                                           void const* vx,
+                                                           sd::LongType const* xShapeInfo,
+                                                           void const* vy,
+                                                           sd::LongType const* yShapeInfo,
+                                                           void* vz,
+                                                           sd::LongType const* zShapeInfo,
+                                                           void* vextraParams) {
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
+      PAIRWISE_TRANSFORM_OPS);
+}
+
+}  // namespace pairwise_transforms
+}  // namespace functions
+
+#endif  // PAIRWISE_CU

libnd4j/include/loops/cuda/pairwise_bool.cu

@@ -0,0 +1,158 @@
+/********************************************************************************
+ *
+ * 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), created on 08.11.2018
+//
+#ifndef PAIRWISE_BOOL_CU
+#define PAIRWISE_BOOL_CU
+
+#include "../pairwise_bool.h"
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL static void pairwiseSimpleShaped(
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  auto x            = reinterpret_cast<const X*>(vx);
+  auto y            = reinterpret_cast<const X*>(vy);
+  auto z            = reinterpret_cast<Z*>(vz);
+  auto extraParams  = reinterpret_cast<X*>(vextraParams);
+
+  const int tid     = blockIdx.x * blockDim.x + threadIdx.x;
+  __shared__ sd::LongType length;
+
+  // Cache shape info
+  __shared__ int xRank;
+  __shared__  sd::LongType* xShapePtr;
+  __shared__  sd::LongType* xStridePtr;
+
+  __shared__ int yRank;
+  __shared__  sd::LongType* yShapePtr;
+  __shared__  sd::LongType* yStridePtr;
+
+  __shared__ int zRank;
+  __shared__  sd::LongType* zShapePtr;
+  __shared__  sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length      = shape::length(xShapeInfo);
+
+    xRank       = shape::rank(xShapeInfo);
+    xShapePtr   = shape::shapeOf(xShapeInfo);
+    xStridePtr  = shape::stride(xShapeInfo);
+
+    yRank       = shape::rank(yShapeInfo);
+    yShapePtr   = shape::shapeOf(yShapeInfo);
+    yStridePtr  = shape::stride(yShapeInfo);
+
+    zRank       = shape::rank(zShapeInfo);
+    zShapePtr   = shape::shapeOf(zShapeInfo);
+    zStridePtr  = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  const auto totalThreads = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType coordsX[SD_MAX_RANK];
+    sd::LongType coordsY[SD_MAX_RANK];
+    sd::LongType coordsZ[SD_MAX_RANK];
+
+    sd::LongType offsetX;
+    sd::LongType offsetY;
+    sd::LongType offsetZ;
+
+    INDEX2COORDS(i, xRank, xShapePtr, coordsX);
+    COORDS2INDEX(xRank, xStridePtr, coordsX, offsetX);
+
+    INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+    COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
+
+    INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+    COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+    z[offsetZ] = OpType::op(x[offsetX], y[offsetY], extraParams);
+  }
+}
+
+namespace functions {
+namespace pairwise_transforms {
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+void SD_HOST PairWiseBoolTransform<X,Z>::intermediateShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  pairwiseSimpleShaped<X, Z, OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          vy,
+          yShapeInfo,
+          vz,
+          zShapeInfo,
+          vextraParams);
+
+  sd::DebugHelper::checkErrorCode(
+      stream,
+      "PairWiseBoolTransform intermediateShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+void PairWiseBoolTransform<X,Y>::executeCudaShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
+      PAIRWISE_BOOL_OPS);
+}
+
+BUILD_DOUBLE_TEMPLATE( class PairWiseBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+
+}  // namespace pairwise_transforms
+}  // namespace functions
+
+#endif  // PAIRWISE_BOOL_CU

libnd4j/include/loops/cuda/pairwise_int.cu

@@ -0,0 +1,163 @@
+/********************************************************************************
+ *
+ *
+ * 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), created on 08.11.2018
+//
+#ifndef PAIRWISE_INT_CU
+#define PAIRWISE_INT_CU
+
+#include "../pairwise_int.h"
+
+    using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+SD_KERNEL static void pairwiseSimpleShaped(
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto y           = reinterpret_cast<const X*>(vy);
+  auto z           = reinterpret_cast<X*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  const int tid     = blockIdx.x * blockDim.x + threadIdx.x;
+
+  __shared__ sd::LongType length;
+  __shared__ int xRank;
+  __shared__ int yRank;
+  __shared__ int zRank;
+
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length      = shape::length(xShapeInfo);
+
+    xRank       = shape::rank(xShapeInfo);
+    xShapePtr   = shape::shapeOf(xShapeInfo);
+    xStridePtr  = shape::stride(xShapeInfo);
+
+    yRank       = shape::rank(yShapeInfo);
+    yShapePtr   = shape::shapeOf(yShapeInfo);
+    yStridePtr  = shape::stride(yShapeInfo);
+
+    zRank       = shape::rank(zShapeInfo);
+    zShapePtr   = shape::shapeOf(zShapeInfo);
+    zStridePtr  = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  const auto totalThreads = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType coordsX[SD_MAX_RANK];
+    sd::LongType coordsY[SD_MAX_RANK];
+    sd::LongType coordsZ[SD_MAX_RANK];
+
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, coordsX);
+    COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
+
+    INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+    COORDS2INDEX(yRank, yStridePtr, coordsY, yOffset);
+
+    INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+    COORDS2INDEX(zRank, zStridePtr, coordsZ, zOffset);
+
+    z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+  }
+}
+
+namespace functions {
+namespace pairwise_transforms {
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+void SD_HOST PairWiseIntTransform<X>::intermediateShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  pairwiseSimpleShaped<X, OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          vy,
+          yShapeInfo,
+          vz,
+          zShapeInfo,
+          vextraParams);
+
+  sd::DebugHelper::checkErrorCode(
+      stream,
+      "PairWiseIntTransform intermediateShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X>
+void PairWiseIntTransform<X>::executeCudaShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void* vextraParams) {
+
+  DISPATCH_BY_OPNUM_T(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
+      PAIRWISE_INT_OPS);
+
+  sd::DebugHelper::checkErrorCode(
+      stream,
+      "PairWiseIntTransform executeCudaShaped(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( class PairWiseIntTransform, , SD_INTEGER_TYPES);
+
+}  // namespace pairwise_transforms
+}  // namespace functions
+
+#endif  // PAIRWISE_INT_CU

libnd4j/include/loops/cuda/random.cu

@@ -0,0 +1,491 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+* License for the specific language governing permissions and limitations
+* under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+#include <helpers/DebugHelper.h>
+#include <loops/random.h>
+#include <ops/specials_cuda.h>
+#include <system/common.h>
+#include <system/op_boilerplate.h>
+
+
+using namespace randomOps;
+
+template <typename T, typename OpClass>
+static SD_INLINE SD_DEVICE void randomSingleGeneric(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer,
+                                                   void* extraArguments) {
+ functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, z, zShapeBuffer, extraArguments);
+}
+
+template <typename T, typename OpClass>
+static SD_INLINE SD_DEVICE void randomDoubleGeneric(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer,
+                                                   void* z, sd::LongType const* zShapeBuffer, void* extraArguments) {
+ functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, x, xShapeBuffer, z, zShapeBuffer,
+                                                                           extraArguments);
+}
+
+template <typename T, typename OpClass>
+static SD_INLINE SD_DEVICE void randomTripleGeneric(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer,
+                                                   void const* y, sd::LongType const* yShapeBuffer, void* z,
+                                                   sd::LongType const* zShapeBuffer, void* extraArguments) {
+ functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, x, xShapeBuffer, y, yShapeBuffer, z,
+                                                                           zShapeBuffer, extraArguments);
+}
+
+// here we generate kernels for target operations
+DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, float,
+                      INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, double,
+                      INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, float16,
+                      INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, bfloat16,
+                      INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, float,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
+                            sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, double,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
+                            sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, float16,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
+                            sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, bfloat16,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
+                            sd::LongType const* zShapeBuffer, void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, float,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
+                            sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
+                            void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                      OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, double,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
+                            sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
+                            void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                      OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, float16,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
+                            sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
+                            void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                      OPS_A(RANDOM_OPS))
+DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, bfloat16,
+                      INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
+                            sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
+                            void* extraArguments),
+                      PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                      OPS_A(RANDOM_OPS))
+
+
+namespace functions {
+namespace random {
+template <typename T>
+template <typename OpClass>
+void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void const* vx, sd::LongType const* xShapeBuffer,
+                                                    void const* vy, sd::LongType const* yShapeBuffer, void* vz,
+                                                    sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<T const*>(vx);
+ auto y = reinterpret_cast<T const*>(vy);
+ auto z = reinterpret_cast<T*>(vz);
+ auto extraArguments = reinterpret_cast<T*>(vextraArguments);
+
+ if (OpClass::requiresSpecial) {
+   OpClass::specialOpCuda(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments);
+   return;
+ } else {
+   __shared__ sd::LongType length;
+   __shared__ char xOrder;
+   __shared__ char yOrder;
+   __shared__ char zOrder;
+
+   // Cache shape information for x buffer
+   __shared__ sd::LongType xRank;
+   __shared__ const sd::LongType* xShapePtr;
+   __shared__ const sd::LongType* xStridePtr;
+
+   // Cache shape information for y buffer
+   __shared__ sd::LongType yRank;
+   __shared__ const sd::LongType* yShapePtr;
+   __shared__ const sd::LongType* yStridePtr;
+
+   // Cache shape information for z buffer
+   __shared__ sd::LongType zRank;
+   __shared__ const sd::LongType* zShapePtr;
+   __shared__ const sd::LongType* zStridePtr;
+
+   __shared__ sd::graph::RandomGenerator* buffer;
+   __shared__ unsigned char* cB;
+   __shared__ unsigned char* dB;
+   sd::graph::RandomGenerator* devBuffer;
+
+   if (threadIdx.x == 0) {
+     length = shape::length(zShapeBuffer);
+     xOrder = shape::order(xShapeBuffer);
+     yOrder = shape::order(yShapeBuffer);
+     zOrder = shape::order(zShapeBuffer);
+
+     // Cache all shape information in thread 0
+     xRank = shape::rank(xShapeBuffer);
+     xShapePtr = shape::shapeOf(xShapeBuffer);
+     xStridePtr = shape::stride(xShapeBuffer);
+
+     yRank = shape::rank(yShapeBuffer);
+     yShapePtr = shape::shapeOf(yShapeBuffer);
+     yStridePtr = shape::stride(yShapeBuffer);
+
+     zRank = shape::rank(zShapeBuffer);
+     zShapePtr = shape::shapeOf(zShapeBuffer);
+     zStridePtr = shape::stride(zShapeBuffer);
+
+     extern __shared__ unsigned char shmem[];
+     buffer = (sd::graph::RandomGenerator*)shmem;
+     cB = shmem;
+     devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
+     dB = reinterpret_cast<unsigned char*>(state);
+   }
+   __syncthreads();
+
+   // using this loop instead of memcpy
+   for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x)
+     cB[e] = dB[e];
+
+   __syncthreads();
+
+   int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+   for (sd::LongType i = tid; i < length; i += 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(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);
+
+     z[zOffset] = OpClass::op(x[xOffset], y[yOffset], i, length, buffer, extraArguments);
+   }
+ }
+}
+
+template <typename T>
+template <typename OpClass>
+void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void const* vx, sd::LongType const* xShapeBuffer,
+                                                   void* vz, sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<T const*>(vx);
+ auto z = reinterpret_cast<T*>(vz);
+ auto extraArguments = reinterpret_cast<T*>(vextraArguments);
+
+ __shared__ sd::LongType length;
+ __shared__ char xOrder;
+ __shared__ char zOrder;
+
+ __shared__ sd::graph::RandomGenerator* buffer;
+ __shared__ unsigned char* cB;
+ __shared__ unsigned char* dB;
+ __shared__ sd::graph::RandomGenerator* devBuffer;
+
+ if (threadIdx.x == 0) {
+   extern __shared__ unsigned char shmem[];
+   buffer = (sd::graph::RandomGenerator*)shmem;
+   cB = shmem;
+   devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
+   dB = reinterpret_cast<unsigned char*>(state);
+
+   length = shape::length(zShapeBuffer);
+   xOrder = shape::order(xShapeBuffer);
+   zOrder = shape::order(zShapeBuffer);
+ }
+ __syncthreads();
+
+ // using this loop instead of memcpy
+ for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x) cB[e] = dB[e];
+
+ __syncthreads();
+
+ for (sd::LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < length; i += blockDim.x * gridDim.x) {
+   sd::LongType xCoords[SD_MAX_RANK];
+   sd::LongType zCoords[SD_MAX_RANK];
+   sd::LongType xOffset;
+   sd::LongType zOffset;
+
+   INDEX2COORDS(i, shape::rank(xShapeBuffer), shape::shapeOf(xShapeBuffer), xCoords);
+   COORDS2INDEX(shape::rank(xShapeBuffer), shape::stride(xShapeBuffer), xCoords, xOffset);
+   INDEX2COORDS(i, shape::rank(zShapeBuffer), shape::shapeOf(zShapeBuffer), zCoords);
+   COORDS2INDEX(shape::rank(zShapeBuffer), shape::stride(zShapeBuffer), zCoords, zOffset);
+
+   z[zOffset] = OpClass::op(x[xOffset], i, length, buffer, extraArguments);
+ }
+}
+template <typename T>
+template <typename OpClass>
+void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void* vz, sd::LongType const* zShapeBuffer,
+                                                   void* vextraArguments) {
+ auto z = reinterpret_cast<T*>(vz);
+ auto extraArguments = reinterpret_cast<T*>(vextraArguments);
+
+ __shared__ sd::LongType length;
+ __shared__ sd::graph::RandomGenerator* buffer;
+ __shared__ unsigned char* cB;
+ __shared__ unsigned char* dB;
+ __shared__ sd::graph::RandomGenerator* devBuffer;
+
+ if (threadIdx.x == 0) {
+   extern __shared__ unsigned char shmem[];
+   buffer = (sd::graph::RandomGenerator*)shmem;
+   cB = shmem;
+   devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
+   dB = reinterpret_cast<unsigned char*>(state);
+   length = shape::length(zShapeBuffer);
+ }
+ __syncthreads();
+
+ // using this loop instead of memcpy
+ for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x) cB[e] = dB[e];
+
+ __syncthreads();
+
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+ for (sd::LongType i = tid; i < length; i += blockDim.x * gridDim.x) {
+   sd::LongType zCoords[SD_MAX_RANK];
+   sd::LongType zOffset;
+
+   INDEX2COORDS(i, shape::rank(zShapeBuffer), shape::shapeOf(zShapeBuffer), zCoords);
+   COORDS2INDEX(shape::rank(zShapeBuffer), shape::stride(zShapeBuffer), zCoords, zOffset);
+
+   z[zOffset] = OpClass::op(i, length, buffer, extraArguments);
+ }
+}
+
+template <>
+SD_HOST void RandomFunction<float>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                     sd::Pointer stateHost, void* vz, sd::LongType const* zShapeBuffer,
+                                                     void* vextraArguments) {
+ auto z = reinterpret_cast<float*>(vz);
+ auto extraArguments = reinterpret_cast<float*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomSingle, float, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<float16>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                       sd::Pointer stateHost, void* vz,
+                                                       sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto z = reinterpret_cast<float16*>(vz);
+ auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomSingle, float16, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<bfloat16>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                        sd::Pointer stateHost, void* vz,
+                                                        sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto z = reinterpret_cast<bfloat16*>(vz);
+ auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomSingle, bfloat16, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<double>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                      sd::Pointer stateHost, void* vz,
+                                                      sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto z = reinterpret_cast<double*>(vz);
+ auto extraArguments = reinterpret_cast<double*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomSingle, double, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<float>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                     sd::Pointer stateHost, void const* vx,
+                                                     sd::LongType const* xShapeBuffer, void* vz,
+                                                     sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<float const*>(vx);
+ auto z = reinterpret_cast<float*>(vz);
+ auto extraArguments = reinterpret_cast<float*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomDouble, float, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<float16>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                       sd::Pointer stateHost, void const* vx,
+                                                       sd::LongType const* xShapeBuffer, void* vz,
+                                                       sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<float16 const*>(vx);
+ auto z = reinterpret_cast<float16*>(vz);
+ auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomDouble, float16, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<bfloat16>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                        sd::Pointer stateHost, void const* vx,
+                                                        sd::LongType const* xShapeBuffer, void* vz,
+                                                        sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<bfloat16 const*>(vx);
+ auto z = reinterpret_cast<bfloat16*>(vz);
+ auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomDouble, bfloat16, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<double>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                      sd::Pointer stateHost, void const* vx,
+                                                      sd::LongType const* xShapeBuffer, void* vz,
+                                                      sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<double const*>(vx);
+ auto z = reinterpret_cast<double*>(vz);
+ auto extraArguments = reinterpret_cast<double*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomDouble, double, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<float>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                     sd::Pointer stateHost, void const* vx,
+                                                     sd::LongType const* xShapeBuffer, void const* vy,
+                                                     sd::LongType const* yShapeBuffer, void* vz,
+                                                     sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<float const*>(vx);
+ auto y = reinterpret_cast<float const*>(vy);
+ auto z = reinterpret_cast<float*>(vz);
+ auto extraArguments = reinterpret_cast<float*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomTriple, float,
+                 PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<float16>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                       sd::Pointer stateHost, void const* vx,
+                                                       sd::LongType const* xShapeBuffer, void const* vy,
+                                                       sd::LongType const* yShapeBuffer, void* vz,
+                                                       sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<float16 const*>(vx);
+ auto y = reinterpret_cast<float16 const*>(vy);
+ auto z = reinterpret_cast<float16*>(vz);
+ auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomTriple, float16,
+                 PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<bfloat16>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                        sd::Pointer stateHost, void const* vx,
+                                                        sd::LongType const* xShapeBuffer, void const* vy,
+                                                        sd::LongType const* yShapeBuffer, void* vz,
+                                                        sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<bfloat16 const*>(vx);
+ auto y = reinterpret_cast<bfloat16 const*>(vy);
+ auto z = reinterpret_cast<bfloat16*>(vz);
+ auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomTriple, bfloat16,
+                 PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+template <>
+SD_HOST void RandomFunction<double>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
+                                                      sd::Pointer stateHost, void const* vx,
+                                                      sd::LongType const* xShapeBuffer, void const* vy,
+                                                      sd::LongType const* yShapeBuffer, void* vz,
+                                                      sd::LongType const* zShapeBuffer, void* vextraArguments) {
+ auto x = reinterpret_cast<double const*>(vx);
+ auto y = reinterpret_cast<double const*>(vy);
+ auto z = reinterpret_cast<double*>(vz);
+ auto extraArguments = reinterpret_cast<double*>(vextraArguments);
+
+ // this macro builds bunch of IF/ELSE selectors for kernel launch
+ DISPATCH_SIMPLE(randomTriple, double,
+                 PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
+                 OPS_A(RANDOM_OPS))
+
+ sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( class RandomFunction, , SD_COMMON_TYPES);
+
+
+}  // namespace random
+}  // namespace functions

libnd4j/include/loops/cuda/reduce/reduce_bool.cu

@@ -0,0 +1,558 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <exceptions/cuda_exception.h>
+#include <execution/LaunchContext.h>
+#include <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_bool.h>
+#include <loops/scalar.h>
+#include <system/op_boilerplate.h>
+#include <system/common.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL SD_INLINE void simpleReduce(
+    const void* x,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::reduce::ReduceBoolFunction<X, Z>::template transformCuda<OpType>(
+      x,
+      outerXTadShapeInfo,
+      innerXTadShapeInfo,
+      vreductionBuffer,
+      extraParams,
+      z,
+      zShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL SD_INLINE void simpleScalar(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  functions::reduce::ReduceBoolFunction<X, Z>::template execScalarCuda<OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+}
+
+namespace functions {
+namespace reduce {
+
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE  void ReduceBoolFunction<X, Z>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  auto sPartials   = reinterpret_cast<Z*>(vsPartials);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= (floorPow2 - 1);
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] = OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE void ReduceBoolFunction<X, Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* vreductionBuffer,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ int tadLen;
+  __shared__ int numTads;
+  __shared__ bool sameOffsets;
+
+  // Cache shape info
+  __shared__ sd::LongType outerRank;
+  __shared__ sd::LongType innerRank;
+  __shared__ sd::LongType zRank;
+
+  __shared__ const sd::LongType* outerShapePtr;
+  __shared__ const sd::LongType* outerStridePtr;
+
+  __shared__ const sd::LongType* innerShapePtr;
+  __shared__ const sd::LongType* innerStridePtr;
+
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    outerRank      = shape::rank(outerXTadShapeInfo);
+    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
+    outerStridePtr = shape::stride(outerXTadShapeInfo);
+
+    innerRank      = shape::rank(innerXTadShapeInfo);
+    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
+    innerStridePtr = shape::stride(innerXTadShapeInfo);
+
+    zRank          = shape::rank(zShapeInfo);
+    zShapePtr      = shape::shapeOf(zShapeInfo);
+    zStridePtr     = shape::stride(zShapeInfo);
+
+    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
+    tadLen         = shape::length(innerXTadShapeInfo);
+    numTads        = shape::length(outerXTadShapeInfo);
+  }
+  __syncthreads();
+
+  sd::LongType coords[SD_MAX_RANK];
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    INDEX2COORDS(r, outerRank, outerShapePtr, coords);
+
+    sd::LongType outerOffset;
+    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
+
+    sd::LongType zOffset;
+    if (sameOffsets) {
+      zOffset = outerOffset;
+    } else {
+      INDEX2COORDS(r, zRank, zShapePtr, coords);
+      COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
+    }
+
+    const X* xTad = x + outerOffset;
+    sPartials[threadIdx.x] = OpType::startingValue(xTad);
+
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType iCoords[SD_MAX_RANK];
+      sd::LongType innerOffset;
+
+      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
+      COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);
+
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(xTad[innerOffset], extraParams),
+          extraParams);
+    }
+    __syncthreads();
+
+    aggregatePartials<OpType>(
+        sPartials,
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraParams);
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE void ReduceBoolFunction<X, Z>::execScalarCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType length;
+
+  // Cache shape info
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+  auto reductionBuffer = reinterpret_cast<Z *>(vreductionBuffer);
+
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    sPartials[threadIdx.x] = OpType::update(
+        sPartials[threadIdx.x],
+        OpType::op(x[xOffset], extraParams),
+        extraParams);
+  }
+  __syncthreads();
+
+  aggregatePartials<OpType>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
+    __shared__ bool amLast;
+
+    if (threadIdx.x == 0) {
+      reductionBuffer[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == (gridDim.x - 1));
+    }
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384] = 0;
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] = OpType::update(
+            sPartials[threadIdx.x],
+            reinterpret_cast<Z*>(vreductionBuffer)[i],
+            extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(vreductionBuffer);
+      tc[16384] = 0;
+      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST SD_INLINE void ReduceBoolFunction<X, Z>::intermediate(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    sd::LongType* dXShapeInfo,
+    sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    sd::LongType* dZShapeInfo,
+    sd::LongType* hZShapeInfo,
+    sd::LongType* dims) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal =
+        static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+
+    auto res = cudaMemcpyAsync(
+        sd::LaunchContext::defaultContext()->getScalarPointer(),
+        &startingVal,
+        sizeof(Z),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceBoolFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
+    }
+
+    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
+    // scalar assign
+    scalar::ScalarTransform<Z,Z,Z>::executeCudaShaped(
+        launchDims,
+        stream,
+        14,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        ptr,
+        nullptr);
+    sd::DebugHelper::checkErrorCode(stream, "reduceBoolDim empty(...) failed");
+  } else {
+    const sd::LongType zRank   = shape::rank(hZShapeInfo);
+    const sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
+
+    auto outerPack = sd::ConstantShapeHelper::getInstance()
+                         .createSubArrShapeInfo(hXShapeInfo, dims, zRank);
+    auto innerPack = sd::ConstantShapeHelper::getInstance()
+                         .createSubArrShapeInfo(hXShapeInfo, dims + zRank, tadRank);
+
+    simpleReduce<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            outerPack->special(),
+            innerPack->special(),
+            extraParams,
+            vreductionBuffer,
+            z,
+            dZShapeInfo);
+    sd::DebugHelper::checkErrorCode(stream, "reduceBoolDim(...) failed");
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST SD_INLINE void ReduceBoolFunction<X, Z>::intermediateScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal =
+        static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+    auto res = cudaMemcpyAsync(
+        z,
+        &startingVal,
+        sizeof(Z),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceBoolFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
+    }
+    sd::DebugHelper::checkErrorCode(stream, "reduceBoolScalar empty(...) failed");
+  } else {
+    simpleScalar<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            xShapeInfo,
+            extraParams,
+            z,
+            zShapeInfo,
+            dimension,
+            dimensionLength,
+            reductionBuffer,
+            tadOnlyShapeInfo);
+    sd::DebugHelper::checkErrorCode(stream, "reduceBoolScalar(...) failed");
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST  SD_INLINE void ReduceBoolFunction<X, Y>::execReduceScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateScalar,
+      PARAMS(
+          launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z,
+          zShapeInfo, hZShapeInfo, dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
+      OPS_A(REDUCE_BOOL_OPS));
+  sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST SD_INLINE void ReduceBoolFunction<X, Y>::execReduce(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    sd::LongType* dXShapeInfo,
+    sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    sd::LongType* dZShapeInfo,
+    sd::LongType* hZShapeInfo,
+    sd::LongType* dims) {
+
+  if (shape::length(hZShapeInfo) == 1) {
+    execReduceScalar(
+        launchDims, stream, opNum,
+        x, dXShapeInfo, hXShapeInfo,
+        extraParams, z,
+        dZShapeInfo, hZShapeInfo,
+        nullptr, 0,
+        vreductionBuffer,
+        nullptr);
+  } else {
+    DISPATCH_BY_OPNUM_TT(
+        intermediate,
+        PARAMS(
+            launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams,
+            vreductionBuffer, z, dZShapeInfo, hZShapeInfo, dims),
+        OPS_A(REDUCE_BOOL_OPS));
+  }
+  DEBUG_KERNEL(stream, opNum);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
+  int sPartialsLength = sMemSize / sizeof(X);
+  X* sPartialsDeref   = reinterpret_cast<X*>(*sPartials);
+  for (int i = 0; i < sPartialsLength; i++) {
+    sPartialsDeref[i] = extraParams[0];
+  }
+}
+
+
+ITERATE_COMBINATIONS(
+    (SD_COMMON_TYPES),
+    (SD_BOOL_TYPES),
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceBoolFunction,
+    ::execReduce(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        const void* x,
+        sd::LongType* dXShapeInfo,
+        sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* vreductionBuffer,
+        void* z,
+        sd::LongType* dZShapeInfo,
+        sd::LongType* hZShapeInfo,
+        sd::LongType* dims);
+);
+
+ITERATE_COMBINATIONS(
+    (SD_COMMON_TYPES),
+    (SD_BOOL_TYPES),
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceBoolFunction,
+    ::execReduceScalar(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        const void* x,
+        const sd::LongType* xShapeInfo,
+        const sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* z,
+        const sd::LongType* zShapeInfo,
+        const sd::LongType* hZShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        void* reductionBuffer,
+        const sd::LongType* tadOnlyShapeInfo);
+);
+
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cuda/reduce/reduce_float.chpp

@@ -0,0 +1,502 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <execution/LaunchContext.h>
+#include <exceptions/cuda_exception.h>
+#include <system/op_boilerplate.h>
+#include <loops/reduce_float.h>
+#include <loops/scalar.h>
+#include <loops/legacy_ops.h>
+#include <helpers/DebugHelper.h>
+#include <types/types.h>
+#include <ops/specials_cuda.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+    using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL void simpleReduce(
+    const void* x,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::reduce::ReduceFloatFunction<X,Z>::template transformCuda<OpType>(
+      x,
+      outerXTadShapeInfo,
+      innerXTadShapeInfo,
+      extraParams,
+      vreductionBuffer,
+      z,
+      zShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL void simpleScalar(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  functions::reduce::ReduceFloatFunction<X, Z>::template execScalarCuda<OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+}
+
+namespace functions {
+namespace reduce {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X,Z>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  using Y             = typename OpType::InterType;
+  auto sPartials      = reinterpret_cast<Y*>(vsPartials);
+  auto extraParams    = reinterpret_cast<Z*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= (floorPow2 - 1);
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] =
+          OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X,Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* vextraParams,
+    void* vreductionBuffer,
+    void* vz,
+    const sd::LongType* zShapeInfo) {
+
+  auto x          = reinterpret_cast<const X*>(vx);
+  auto z          = reinterpret_cast<Z*>(vz);
+  auto extraParams= reinterpret_cast<Z*>(vextraParams);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType tadLen;
+  __shared__ sd::LongType numTads;
+  __shared__ bool sameOffsets;
+
+  // Cache ranks/shape/stride for outer and z
+  __shared__ sd::LongType outerRank;
+  __shared__ const sd::LongType* outerStridePtr;
+  __shared__ const sd::LongType* outerShapePtr;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zStridePtr;
+  __shared__ const sd::LongType* zShapePtr;
+
+  // Cache ranks/shape/stride for inner as well if needed
+  __shared__ sd::LongType innerRank;
+  __shared__ const sd::LongType* innerStridePtr;
+  __shared__ const sd::LongType* innerShapePtr;
+
+  if (threadIdx.x == 0) {
+    outerRank      = shape::rank(outerXTadShapeInfo);
+    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
+    outerStridePtr = shape::stride(outerXTadShapeInfo);
+
+    zRank          = shape::rank(zShapeInfo);
+    zShapePtr      = shape::shapeOf(zShapeInfo);
+    zStridePtr     = shape::stride(zShapeInfo);
+
+    innerRank      = shape::rank(innerXTadShapeInfo);
+    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
+    innerStridePtr = shape::stride(innerXTadShapeInfo);
+
+    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
+    tadLen         = shape::length(innerXTadShapeInfo);
+    numTads        = shape::length(outerXTadShapeInfo);
+  }
+  __syncthreads();
+
+  sd::LongType coords[SD_MAX_RANK];
+  sd::LongType zCoords[SD_MAX_RANK];
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    INDEX2COORDS(r, outerRank, outerShapePtr, coords);
+
+    sd::LongType outerOffset;
+    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
+
+    sd::LongType zOffset;
+    if (sameOffsets) {
+      zOffset = outerOffset;
+    } else {
+      INDEX2COORDS(r, zRank, zShapePtr, zCoords);
+      COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+    }
+
+    const X* xTad = x + outerOffset;
+    sPartials[threadIdx.x] = OpType::startingValue(xTad);
+
+    // For the inner dimension
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType iCoords[SD_MAX_RANK];
+      sd::LongType innerOffset;
+
+      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
+      COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);
+
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(xTad[innerOffset], extraParams),
+          extraParams);
+    }
+    __syncthreads();
+
+    aggregatePartials<OpType>(
+        sPartials,
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraParams);
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      z[zOffset] =
+          OpType::postProcess(sPartials[0], tadLen, extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X, Z>::execScalarCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<Z*>(vextraParams);
+  auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);
+
+  using Y = typename OpType::InterType;
+
+  __shared__ Y sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType length;
+
+  // Cache rank/shape/stride
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    if(xOffset < length) {
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(x[xOffset], extraParams),
+          extraParams);
+    }
+
+  }
+  __syncthreads();
+
+  aggregatePartials<OpType>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
+    __shared__ bool amLast;
+
+    if (threadIdx.x == 0) {
+      reductionBuffer[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == (gridDim.x - 1));
+    }
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384]           = 0;
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::update(sPartials[threadIdx.x],
+                           reductionBuffer[i],
+                           extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(reductionBuffer);
+      tc[16384] = 0;
+      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template<typename OpType>
+SD_HOST void ReduceFloatFunction<X,Z>::intermediate(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
+                                 ? sd::DataTypeUtils::nanOrZero<Z>()
+                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+
+    auto res = cudaMemcpyAsync(
+        sd::LaunchContext::defaultContext()->getScalarPointer(),
+        &startingVal,
+        sizeof(Z),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceFloatFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
+    }
+
+    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
+    // scalar assign
+    functions::scalar::ScalarTransform<Z, Z, Z>::executeCudaShaped(
+        launchDims,
+        stream,
+        14,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        ptr,
+        nullptr);
+  } else {
+    const int zRank   = shape::rank(hZShapeInfo);
+    const int tadRank = shape::rank(hXShapeInfo) - zRank;
+
+    auto outerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
+        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
+    auto innerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
+        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);
+
+    simpleReduce<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            reinterpret_cast<const sd::LongType*>(outerPack->special()),
+            reinterpret_cast<const sd::LongType*>(innerPack->special()),
+            extraParams,
+            vreductionBuffer,
+            z,
+            dZShapeInfo);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template<typename OpType>
+SD_HOST void ReduceFloatFunction<X,Z>::intermediateScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
+                                 ? sd::DataTypeUtils::nanOrZero<Z>()
+                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+
+    auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z), cudaMemcpyHostToDevice, *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceFloatFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
+    }
+  } else {
+    simpleScalar<X, Z, OpType><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x,
+        xShapeInfo,
+        extraParams,
+        z,
+        dZShapeInfo,
+        dimension,
+        dimensionLength,
+        reductionBuffer,
+        tadOnlyShapeInfo);
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction intermediateScalar(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void ReduceFloatFunction<X,Y>::execReduceScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateScalar,
+      PARAMS(
+          launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z, dZShapeInfo, hZShapeInfo, dimension,
+          dimensionLength, reductionBuffer, tadOnlyShapeInfo),
+      OPS_A(REDUCE_FLOAT_OPS));
+  sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void ReduceFloatFunction<X,Y>::execReduce(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::length(hZShapeInfo) == 1) {
+    ReduceFloatFunction<X,Y>::execReduceScalar(
+        launchDims, stream, opNum, x, dXShapeInfo, hXShapeInfo,
+        extraParams, z, dZShapeInfo, hZShapeInfo,
+        nullptr, 0, vreductionBuffer, nullptr);
+  } else {
+    DISPATCH_BY_OPNUM_TT(
+        intermediate,
+        PARAMS(
+            launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
+            hZShapeInfo, dims),
+        OPS_A(REDUCE_FLOAT_OPS));
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction execReduce(...) failed");
+}
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cuda/reduce/reduce_long.cu

@@ -0,0 +1,552 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <exceptions/cuda_exception.h>
+#include <execution/LaunchContext.h>
+#include <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_long.h>
+#include <loops/scalar.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+    using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL SD_INLINE void simpleReduce(
+    const void* x,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::reduce::ReduceLongFunction<X, Z>::template transformCuda<OpType>(
+      x,
+      outerXTadShapeInfo,
+      innerXTadShapeInfo,
+      extraParams,
+      vreductionBuffer,
+      z,
+      zShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_DEVICE SD_INLINE void reduceScalarGeneric(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  functions::reduce::ReduceLongFunction<X, Z>::template execScalarCuda<OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL SD_INLINE  void simpleScalar(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  reduceScalarGeneric<X, Z, OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      dimension,
+      dimensionLength,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+}
+
+namespace functions {
+namespace reduce {
+
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  auto sPartials   = reinterpret_cast<Z*>(vsPartials);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= floorPow2 - 1;
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] = OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* vextraParams,
+    void* vreductionBuffer,
+    void* vz,
+    const sd::LongType* zShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ int tadLen;
+  __shared__ int numTads;
+  __shared__ bool sameOffsets;
+
+  // Cache shape info for outer/inner and z if needed
+  __shared__ sd::LongType outerRank;
+  __shared__ sd::LongType innerRank;
+  __shared__ sd::LongType zRank;
+
+  __shared__ const sd::LongType* outerShapePtr;
+  __shared__ const sd::LongType* outerStridePtr;
+
+  __shared__ const sd::LongType* innerShapePtr;
+  __shared__ const sd::LongType* innerStridePtr;
+
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    outerRank      = shape::rank(outerXTadShapeInfo);
+    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
+    outerStridePtr = shape::stride(outerXTadShapeInfo);
+
+    innerRank      = shape::rank(innerXTadShapeInfo);
+    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
+    innerStridePtr = shape::stride(innerXTadShapeInfo);
+
+    zRank          = shape::rank(zShapeInfo);
+    zShapePtr      = shape::shapeOf(zShapeInfo);
+    zStridePtr     = shape::stride(zShapeInfo);
+
+    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
+    tadLen         = shape::length(innerXTadShapeInfo);
+    numTads        = shape::length(outerXTadShapeInfo);
+  }
+  __syncthreads();
+
+  sd::LongType coords[SD_MAX_RANK];
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    INDEX2COORDS(r, outerRank, outerShapePtr, coords);
+
+    sd::LongType outerOffset;
+    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
+
+    sd::LongType zOffset;
+    if (sameOffsets) {
+      zOffset = outerOffset;
+    } else {
+      INDEX2COORDS(r, zRank, zShapePtr, coords);
+      COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
+    }
+
+    const X* xTad = x + outerOffset;
+    sPartials[threadIdx.x] = OpType::startingValue(xTad);
+
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType iCoords[SD_MAX_RANK];
+      sd::LongType xOffset;
+      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
+      COORDS2INDEX(innerRank, innerStridePtr, iCoords, xOffset);
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(xTad[xOffset], extraParams),
+          extraParams);
+    }
+    __syncthreads();
+
+    aggregatePartials<OpType>(
+        sPartials,
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraParams);
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::execScalarCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+  auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType length;
+
+  // Cache shape info
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    sPartials[threadIdx.x] = OpType::update(
+        sPartials[threadIdx.x],
+        OpType::op(x[xOffset], extraParams),
+        extraParams);
+  }
+  __syncthreads();
+
+  aggregatePartials<OpType>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
+    __shared__ bool amLast;
+
+    if (threadIdx.x == 0) {
+      reductionBuffer[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == gridDim.x - 1);
+    }
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384] = 0;
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::update(sPartials[threadIdx.x], reductionBuffer[i], extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(reductionBuffer);
+      tc[16384] = 0;
+      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST  SD_INLINE void ReduceLongFunction<X, Z>::intermediate(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+     sd::LongType* dXShapeInfo,
+     sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    sd::LongType* dZShapeInfo,
+    sd::LongType* hZShapeInfo,
+    sd::LongType* dims) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    // If input is empty, we skip unless z is also empty
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal = static_cast<Z>(
+        OpType::startingValue(reinterpret_cast<const X*>(x)));
+    auto res = cudaMemcpyAsync(
+        sd::LaunchContext::defaultContext()->getScalarPointer(),
+        &startingVal,
+        sizeof(Z),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceLongFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
+    }
+
+    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
+    // scalar assign
+    scalar::ScalarTransform<Z,Z,Z>::executeCudaShaped(
+        launchDims, stream, 14, z, dZShapeInfo, hXShapeInfo,
+        z, dZShapeInfo, hZShapeInfo,
+        ptr, nullptr);
+  } else {
+     sd::LongType zRank   = shape::rank(hZShapeInfo);
+     sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
+
+    auto outerPack = sd::ConstantShapeHelper::getInstance()
+                         .createSubArrShapeInfo(hXShapeInfo, dims, zRank);
+    auto innerPack = sd::ConstantShapeHelper::getInstance()
+                         .createSubArrShapeInfo(hXShapeInfo, dims + zRank, tadRank);
+
+    simpleReduce<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            outerPack->special(),
+            innerPack->special(),
+            extraParams,
+            vreductionBuffer,
+            z,
+            dZShapeInfo);
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction intermediate(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST  SD_INLINE void ReduceLongFunction<X, Z>::intermediateScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal =
+        static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+    auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z),
+                               cudaMemcpyHostToDevice, *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceLongFunction<X,Z>::intermediateScalar: failed to copy resulting scalar",
+          res);
+    }
+  } else {
+    simpleScalar<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            xShapeInfo,
+            extraParams,
+            z,
+            zShapeInfo,
+            dimension,
+            dimensionLength,
+            reductionBuffer,
+            tadOnlyShapeInfo);
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction intermediateScalar(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST SD_INLINE void ReduceLongFunction<X, Y>::execReduceScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+     sd::LongType* xShapeInfo,
+     sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+     sd::LongType* zShapeInfo,
+     sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+     sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateScalar,
+      PARAMS(launchDims, stream, x, xShapeInfo, hXShapeInfo,
+             extraParams, z, zShapeInfo, hZShapeInfo,
+             dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
+      OPS_A(REDUCE_LONG_OPS));
+  sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction execReduceScalar(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST SD_INLINE void ReduceLongFunction<X, Y>::execReduce(dim3 launchDims, cudaStream_t *stream, int opNum, const void *vx,
+                                                    sd::LongType *dXShapeInfo,  sd::LongType *hXShapeInfo, void *extraParams,
+                                                    void *vreductionBuffer, void *vz,  sd::LongType *dZShapeInfo,
+                                                    sd::LongType *hZShapeInfo,  sd::LongType *dims) {
+
+  if (shape::length(hZShapeInfo) == 1) {
+    execReduceScalar(
+        launchDims, stream, opNum,
+        vx, dXShapeInfo, hXShapeInfo,
+        extraParams, vz,
+        dZShapeInfo, hZShapeInfo,
+        nullptr, 0, vreductionBuffer, nullptr);
+  } else {
+    DISPATCH_BY_OPNUM_TT(
+        intermediate,
+        PARAMS(launchDims, stream, vx, dXShapeInfo, hXShapeInfo,
+               extraParams, vreductionBuffer, vz, dZShapeInfo, hZShapeInfo, dims),
+        OPS_A(REDUCE_LONG_OPS));
+  }
+  DEBUG_KERNEL(stream, opNum);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
+  int sPartialsLength = sMemSize / sizeof(X);
+  X* sPartialsDeref   = reinterpret_cast<X*>(*sPartials);
+  for (int i = 0; i < sPartialsLength; i++) {
+    sPartialsDeref[i] = extraParams[0];
+  }
+}
+
+
+ITERATE_COMBINATIONS(
+    (SD_COMMON_TYPES),
+    (SD_LONG_TYPES),
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceLongFunction,
+    ::execReduce(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        int opNum,
+        const void* vx,
+        sd::LongType* dXShapeInfo,
+        sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* vreductionBuffer,
+        void* vz,
+        sd::LongType* dZShapeInfo,
+        sd::LongType* hZShapeInfo,
+        sd::LongType* dims);
+);
+
+ITERATE_COMBINATIONS(
+    (SD_COMMON_TYPES),
+    (SD_LONG_TYPES),
+    INSTANT_PROCESS_COMBINATION,
+    functions::reduce::ReduceLongFunction,
+    ::execReduceScalar(
+        dim3 launchDims,
+        cudaStream_t* stream,
+        const int opNum,
+        const void* x,
+        sd::LongType* xShapeInfo,
+        sd::LongType* hXShapeInfo,
+        void* extraParams,
+        void* z,
+        sd::LongType* zShapeInfo,
+        sd::LongType* hZShapeInfo,
+        sd::LongType* dimension,
+        sd::LongType dimensionLength,
+        void* reductionBuffer,
+        sd::LongType* tadOnlyShapeInfo);
+);
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cuda/reduce/reduce_same.cu

@@ -0,0 +1,553 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <exceptions/cuda_exception.h>
+#include <execution/LaunchContext.h>
+#include <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/reduce_same.h>
+#include <loops/scalar.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+SD_KERNEL SD_INLINE void simpleReduce(
+    const void* x,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::reduce::ReduceSameFunction<X>::template transformCuda<OpType>(
+      x,
+      outerXTadShapeInfo,
+      innerXTadShapeInfo,
+      extraParams,
+      vreductionBuffer,
+      z,
+      zShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename OpType>
+SD_KERNEL SD_INLINE void simpleScalar(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  functions::reduce::ReduceSameFunction<X>::template execScalarCuda<OpType>(
+      x, xShapeInfo, extraParams, z, zShapeInfo, reductionBuffer, tadOnlyShapeInfo);
+}
+
+namespace functions {
+namespace reduce {
+
+template <typename X>
+template <typename OpType>
+SD_DEVICE SD_INLINE  void ReduceSameFunction<X>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  auto sPartials   = reinterpret_cast<X*>(vsPartials);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= floorPow2 - 1;
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] =
+          OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+SD_DEVICE SD_INLINE  void ReduceSameFunction<X>::transformCuda(
+    const void* vx,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* vextraParams,
+    void* vreductionBuffer,
+    void* vz,
+    const sd::LongType* zShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<X*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ X sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ int tadLen;
+  __shared__ int numTads;
+  __shared__ bool sameOffsets;
+
+  // Cache shape info for outer/inner and z if needed
+  __shared__ sd::LongType outerRank;
+  __shared__ sd::LongType innerRank;
+  __shared__ sd::LongType zRank;
+
+  __shared__ const sd::LongType* outerShapePtr;
+  __shared__ const sd::LongType* outerStridePtr;
+
+  __shared__ const sd::LongType* innerShapePtr;
+  __shared__ const sd::LongType* innerStridePtr;
+
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    outerRank      = shape::rank(outerXTadShapeInfo);
+    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
+    outerStridePtr = shape::stride(outerXTadShapeInfo);
+
+    innerRank      = shape::rank(innerXTadShapeInfo);
+    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
+    innerStridePtr = shape::stride(innerXTadShapeInfo);
+
+    zRank          = shape::rank(zShapeInfo);
+    zShapePtr      = shape::shapeOf(zShapeInfo);
+    zStridePtr     = shape::stride(zShapeInfo);
+
+    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
+    tadLen         = shape::length(innerXTadShapeInfo);
+    numTads        = shape::length(outerXTadShapeInfo);
+  }
+  __syncthreads();
+
+  sd::LongType coords[SD_MAX_RANK];
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    INDEX2COORDS(r, outerRank, outerShapePtr, coords);
+
+    sd::LongType outerOffset;
+    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
+
+    sd::LongType zOffset;
+    if (sameOffsets) {
+      zOffset = outerOffset;
+    } else {
+      INDEX2COORDS(r, zRank, zShapePtr, coords);
+      COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
+    }
+
+    const X* xTad = x + outerOffset;
+    sPartials[threadIdx.x] = OpType::startingValue(xTad);
+
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType iCoords[SD_MAX_RANK];
+      sd::LongType xOffset;
+
+      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
+      COORDS2INDEX(innerRank, innerStridePtr, iCoords, xOffset);
+
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(xTad[xOffset], extraParams),
+          extraParams);
+    }
+    __syncthreads();
+
+    aggregatePartials<OpType>(
+        sPartials,
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraParams);
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_DEVICE SD_INLINE  void ReduceSameFunction<X>::execScalarCudaLegacy(
+    int opNum,
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_T(
+      execScalarCuda,
+      PARAMS(vx, xShapeInfo, vextraParams, vz, zShapeInfo, vreductionBuffer, tadOnlyShapeInfo),
+      REDUCE_SAME_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+SD_DEVICE SD_INLINE  void ReduceSameFunction<X>::execScalarCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<X*>(vz);
+  auto extraParams = reinterpret_cast<X*>(vextraParams);
+
+  __shared__ X sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType length;
+  auto reductionBuffer = reinterpret_cast<X *>(vreductionBuffer);
+
+  // Cache shape info
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    sPartials[threadIdx.x] = OpType::update(
+        sPartials[threadIdx.x],
+        OpType::op(x[xOffset], extraParams),
+        extraParams);
+  }
+  __syncthreads();
+
+  aggregatePartials<OpType>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
+    __shared__ bool amLast;
+
+    if (threadIdx.x == 0) {
+      reductionBuffer[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == (gridDim.x - 1));
+    }
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384]           = 0;
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::update(sPartials[threadIdx.x], reinterpret_cast<X*>(vreductionBuffer)[i], extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(vreductionBuffer);
+      tc[16384] = 0;
+      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+SD_HOST SD_INLINE void ReduceSameFunction<X>::intermediate(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    // If input is empty, skip unless z is also empty
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal =
+        static_cast<X>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+    auto res = cudaMemcpyAsync(
+        sd::LaunchContext::defaultContext()->getScalarPointer(),
+        &startingVal,
+        sizeof(X),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceSameFunction<X>::intermediate: failed to copy temporary scalar", res);
+    }
+
+    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
+    // scalar assign
+    scalar::ScalarTransform<X,X,X>::executeCudaShaped(
+        launchDims,
+        stream,
+        14,
+        z,
+        dZShapeInfo,
+        hXShapeInfo,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        ptr,
+        nullptr);
+  } else {
+    const sd::LongType zRank   = shape::rank(hZShapeInfo);
+    const sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
+
+    auto outerPack = sd::ConstantShapeHelper::getInstance()
+        .createSubArrShapeInfo(
+            const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
+    auto innerPack = sd::ConstantShapeHelper::getInstance()
+        .createSubArrShapeInfo(
+            const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);
+
+    simpleReduce<X, OpType>
+    <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x,
+        outerPack->special(),
+        innerPack->special(),
+        extraParams,
+        vreductionBuffer,
+        z,
+        dZShapeInfo);
+    sd::DebugHelper::checkErrorCode(stream, "ReduceSameFunction intermediate(...) failed");
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+SD_HOST SD_INLINE void ReduceSameFunction<X>::intermediateScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    if (shape::isEmptyConst(hZShapeInfo)) return;
+
+    const auto startingVal =
+        static_cast<X>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+    auto res = cudaMemcpyAsync(z, &startingVal, sizeof(X),
+                               cudaMemcpyHostToDevice, *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceSameFunction<X>::intermediateScalar: failed to copy resulting scalar",
+          res);
+    }
+  } else {
+    simpleScalar<X, OpType>
+    <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x,
+        xShapeInfo,
+        extraParams,
+        z,
+        zShapeInfo,
+        dimension,
+        dimensionLength,
+        reductionBuffer,
+        tadOnlyShapeInfo);
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceSameFunction intermediateScalar(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_HOST SD_INLINE  void ReduceSameFunction<X>::execReduceScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_T(
+      intermediateScalar,
+      PARAMS(
+          launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z,
+          zShapeInfo, hZShapeInfo, dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
+      REDUCE_SAME_OPS);
+  sd::DebugHelper::checkErrorCode(stream, "execReduceScalarSame(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_HOST SD_INLINE  void ReduceSameFunction<X>::execReduce(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::length(hZShapeInfo) == 1) {
+    execReduceScalar(
+        launchDims, stream, opNum,
+        x, dXShapeInfo, hXShapeInfo,
+        extraParams, z,
+        dZShapeInfo, hZShapeInfo,
+        nullptr, 0, vreductionBuffer, nullptr);
+  } else {
+    DISPATCH_BY_OPNUM_T(
+        intermediate,
+        PARAMS(
+            launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
+            hZShapeInfo, dims),
+        REDUCE_SAME_OPS);
+  }
+  DEBUG_KERNEL(stream, opNum);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
+  int sPartialsLength = sMemSize / sizeof(X);
+  X* sPartialsDeref   = reinterpret_cast<X*>(*sPartials);
+  for (int i = 0; i < sPartialsLength; i++) {
+    sPartialsDeref[i] = extraParams[0];
+  }
+}
+
+
+#define INSTANT_PROCESS_SINGLE(a1) \
+    template void functions::reduce::ReduceSameFunction<GET_SECOND(a1)>::execReduce( \
+        dim3 launchDims, \
+        cudaStream_t* stream, \
+        const int opNum, \
+        const void* x, \
+        const sd::LongType* dXShapeInfo, \
+        const sd::LongType* hXShapeInfo, \
+        void* extraParams, \
+        void* vreductionBuffer, \
+        void* z, \
+        const sd::LongType* dZShapeInfo, \
+        const sd::LongType* hZShapeInfo, \
+        const sd::LongType* dims); \
+    \
+    template void functions::reduce::ReduceSameFunction<GET_SECOND(a1)>::execReduceScalar( \
+        dim3 launchDims, \
+        cudaStream_t* stream, \
+        int opNum, \
+        const void* x, \
+        const sd::LongType* xShapeInfo, \
+        const sd::LongType* hXShapeInfo, \
+        void* extraParams, \
+        void* z, \
+        const sd::LongType* zShapeInfo, \
+        const sd::LongType* hZShapeInfo, \
+        sd::LongType* dimension, \
+        sd::LongType dimensionLength, \
+        void* reductionBuffer, \
+        const sd::LongType* tadOnlyShapeInfo);
+
+ITERATE_LIST((SD_COMMON_TYPES), INSTANT_PROCESS_SINGLE)
+
+}  // namespace reduce
+}  // namespace functions

libnd4j/include/loops/cuda/reduce3.chpp

@@ -0,0 +1,703 @@
+/* ******************************************************************************
+ *
+ *
+ * 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), created on 19.11.2018
+
+#include <system/op_boilerplate.h>
+#include <loops/reduce3.h>
+#include <loops/legacy_ops.h>
+#include <types/types.h>
+#include <ops/specials_cuda.h>
+
+    using namespace simdOps;
+
+namespace functions {
+namespace reduce3 {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_KERNEL void execScalarGeneric(const int opNum,
+                                 void const* vx, sd::LongType const* xShapeInfo,
+                                 void const* vy, sd::LongType const* yShapeInfo,
+                                 void *extraParams,
+                                 void *vz, sd::LongType const* zShapeInfo,
+                                 sd::LongType* allocationPointer,
+                                 void *reductionBuffer,
+                                 sd::LongType const* tadOnlyShapeInfo) {
+
+  Reduce3<X, Z>::execScalarCuda(opNum,
+                                vx, xShapeInfo,
+                                vy, yShapeInfo,
+                                extraParams,
+                                vz, zShapeInfo,
+                                allocationPointer,
+                                reductionBuffer,
+                                tadOnlyShapeInfo);
+}
+
+template <typename X, typename Z>
+SD_KERNEL void execAllGeneric(const int opNum,
+                              void const* vx, sd::LongType const* xShapeInfo,
+                              void const* vy, sd::LongType const* yShapeInfo,
+                              void *extraParams,
+                              void *vz, sd::LongType const* zShapeInfo,
+                              sd::LongType* dimension,
+                              long long int dimensionLength,
+                              int postProcessOrNot,
+                              sd::LongType* allocationPointer,
+                              sd::LongType const* tadOnlyShapeInfo,
+                              sd::LongType const* tadOffsets,
+                              sd::LongType const* yTadOnlyShapeInfo,
+                              sd::LongType const* yTadOffsets) {
+
+  Reduce3<X, Z>::execAllCuda(opNum,
+                             vx, xShapeInfo,
+                             vy, yShapeInfo,
+                             extraParams,
+                             vz, zShapeInfo,
+                             dimension, dimensionLength,
+                             postProcessOrNot,
+                             allocationPointer,
+                             tadOnlyShapeInfo, tadOffsets,
+                             yTadOnlyShapeInfo, yTadOffsets);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_KERNEL void execGeneric(const int opNum,
+                           void const* vx, sd::LongType const* xShapeInfo,
+                           void const* vy, sd::LongType const* yShapeInfo,
+                           void *extraParams,
+                           void *vz, sd::LongType const* zShapeInfo,
+                           sd::LongType *dimension,
+                           sd::LongType dimensionLength,
+                           int postProcessOrNot,
+                           sd::LongType* allocationPointer,
+                           sd::LongType const* tadOnlyShapeInfo,
+                           sd::LongType const* tadOffsets,
+                           sd::LongType const* yTadOnlyShapeInfo,
+                           sd::LongType const* yTadOffsets) {
+
+  Reduce3<X, Z>::execCuda(opNum,
+                          vx, xShapeInfo,
+                          vy, yShapeInfo,
+                          extraParams,
+                          vz, zShapeInfo,
+                          dimension, dimensionLength,
+                          postProcessOrNot,
+                          allocationPointer,
+                          tadOnlyShapeInfo, tadOffsets,
+                          yTadOnlyShapeInfo, yTadOffsets);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void Reduce3<X, Z>::aggregatePartials(void* vsPartials,
+                                                sd::LongType tid,
+                                                sd::LongType numItems,
+                                                void* vextraParams) {
+
+  auto sPartials   = reinterpret_cast<Z*>(vsPartials);
+  auto extraParams = reinterpret_cast<Z*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1))
+      floorPow2 &= floorPow2 - 1;
+
+    if (tid >= floorPow2)
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads) {
+      sPartials[tid] =
+          OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void Reduce3<X, Z>::execScalarCuda(
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* allocationPointer,
+    void* reductionBuffer,
+    sd::LongType const* tadOnlyShapeInfo) {
+
+  auto x = reinterpret_cast<const X*>(vx);
+  auto y = reinterpret_cast<const X*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+
+  __shared__ sd::LongType length;
+  __shared__ Z extraZ[3];  // just 3 values used in logic below
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+
+  // Cache rank/shape/stride for x and y in __shared__
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType yRank;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length        = shape::length(xShapeInfo);
+
+    xRank         = shape::rank(xShapeInfo);
+    xShapePtr     = shape::shapeOf(xShapeInfo);
+    xStridePtr    = shape::stride(xShapeInfo);
+
+    yRank         = shape::rank(yShapeInfo);
+    yShapePtr     = shape::shapeOf(yShapeInfo);
+    yStridePtr    = shape::stride(yShapeInfo);
+
+    extraZ[0] = (Z) 0.0f;
+    extraZ[1] = (Z) 0.0f;
+    extraZ[2] = extraParams != nullptr
+                    ? (reinterpret_cast<Z*>(extraParams))[2]
+                    : (Z) 0.0f;
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  char xOrder = shape::order(xShapeInfo);
+  char yOrder = shape::order(yShapeInfo);
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+
+  // fill partial sums
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    INDEX2COORDS(i, yRank, yShapePtr, yCoords);
+    COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
+
+    sPartials[threadIdx.x] =
+        OpType::update(sPartials[threadIdx.x],
+                       OpType::opAtomic(x[xOffset], y[yOffset], extraZ),
+                       extraZ);
+  }
+
+  __syncthreads();
+  // reduce partial sums
+  aggregatePartials<OpType>(
+      reinterpret_cast<void*>(sPartials),
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraZ);
+  __syncthreads();
+
+  // multi-block reduce
+  if (gridDim.x > 1) {
+    auto tc        = reinterpret_cast<unsigned int*>(reductionBuffer);
+    __shared__ bool amLast;
+    tid = threadIdx.x;
+    Z* extraBuffer = reinterpret_cast<Z*>(allocationPointer);
+
+    if (threadIdx.x == 0) {
+      reinterpret_cast<Z*>(reductionBuffer)[blockIdx.x] = sPartials[0];
+      extraBuffer[blockIdx.x]                           = extraZ[0];
+      extraBuffer[gridDim.x + blockIdx.x]               = extraZ[1];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == gridDim.x - 1);
+    }
+    sPartials[tid] = OpType::startingValue(x);
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384] = 0; // reset
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      if (tid == 0 && extraZ[0] != static_cast<Z>(0) && extraZ[1] != static_cast<Z>(0)) {
+        extraZ[0] = 0.0;
+        extraZ[1] = 0.0;
+        for (int i = 0; i < gridDim.x; i++) {
+          extraZ[0] += extraBuffer[i];
+          extraZ[1] += extraBuffer[gridDim.x + i];
+        }
+      }
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::update(sPartials[threadIdx.x],
+                           (reinterpret_cast<Z*>(reductionBuffer))[i],
+                           extraZ);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          reinterpret_cast<void*>(sPartials),
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraZ);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraZ);
+      }
+    }
+  } else {
+    if (tid == 0) {
+      auto tc     = reinterpret_cast<unsigned int*>(reductionBuffer);
+      tc[16384]   = 0;
+      z[0]        = OpType::postProcess(sPartials[0], length, extraZ);
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void Reduce3<X, Z>::transformAll(
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    int postProcessOrNot, sd::LongType* allocationPointer,
+    sd::LongType const* xTadShapeInfo, sd::LongType const* xOffsets,
+    sd::LongType const* yTadShapeInfo, sd::LongType const* yOffsets) {
+
+  auto dx = reinterpret_cast<const X*>(vx);
+  auto dy = reinterpret_cast<const X*>(vy);
+  auto z  = reinterpret_cast<Z*>(vz);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ Z extraZ[OpType::extraParamsLen > 0 ? OpType::extraParamsLen : 1];
+
+  __shared__ int xTadLength;
+  __shared__ int yTadLength;
+  __shared__ int xTads;
+  __shared__ int yTads;
+
+  // Cache shape info for xTadShapeInfo and yTadShapeInfo
+  __shared__ sd::LongType xTadRank;
+  __shared__ const sd::LongType* xTadShapePtr;
+  __shared__ const sd::LongType* xTadStridePtr;
+
+  __shared__ sd::LongType yTadRank;
+  __shared__ const sd::LongType* yTadShapePtr;
+  __shared__ const sd::LongType* yTadStridePtr;
+
+  if (threadIdx.x == 0) {
+    sPartials[threadIdx.x] = OpType::startingValue(dx);
+
+    xTadLength = shape::length(xTadShapeInfo);
+    yTadLength = shape::length(yTadShapeInfo);
+
+    xTads = shape::length(xShapeInfo) / xTadLength;
+    yTads = shape::length(yShapeInfo) / yTadLength;
+
+    xTadRank      = shape::rank(xTadShapeInfo);
+    xTadShapePtr  = shape::shapeOf(xTadShapeInfo);
+    xTadStridePtr = shape::stride(xTadShapeInfo);
+
+    yTadRank      = shape::rank(yTadShapeInfo);
+    yTadShapePtr  = shape::shapeOf(yTadShapeInfo);
+    yTadStridePtr = shape::stride(yTadShapeInfo);
+  }
+  __syncthreads();
+
+  Z startingVal = OpType::startingValue(dx);
+
+  const int maxBlock = blockDim.x;
+  const int limit    = (xTadLength + maxBlock - 1) / maxBlock; // ceiling division
+
+  for (int r = blockIdx.x; r < xTads; r += (blockDim.x * gridDim.x)) {
+    // load partial x
+    auto xLocal  = dx + xOffsets[r];
+
+    // Fill tile for x once if thread < xTadLength
+    __shared__ X sXCache[SD_CUDA_BLOCK_SIZE];
+    if (threadIdx.x < xTadLength && threadIdx.x < maxBlock) {
+      sd::LongType xCoords[SD_MAX_RANK];
+      sd::LongType xOff;
+      INDEX2COORDS(threadIdx.x, xTadRank, xTadShapePtr, xCoords);
+      COORDS2INDEX(xTadRank, xTadStridePtr, xCoords, xOff);
+      sXCache[threadIdx.x] = xLocal[xOff];
+    }
+    __syncthreads();
+
+    for (int g = 0; g < yTads; g++) {
+      auto yLocal = dy + yOffsets[g];
+      int ri       = (r * yTads) + g;
+
+      sPartials[threadIdx.x] = startingVal;
+      if (OpType::extraParamsLen > 0 && threadIdx.x < OpType::extraParamsLen) {
+        extraZ[threadIdx.x] = startingVal;
+      }
+      __syncthreads();
+
+      // Possibly multiple tiles per x
+      for (int t = 0; t < limit; t++) {
+        // re-fetch x tile if needed
+        if (t >= 1 && (threadIdx.x + t * maxBlock < xTadLength)) {
+          sd::LongType xCoords[SD_MAX_RANK];
+          sd::LongType xOff;
+          INDEX2COORDS(threadIdx.x + t * maxBlock,
+                       xTadRank, xTadShapePtr, xCoords);
+          COORDS2INDEX(xTadRank, xTadStridePtr, xCoords, xOff);
+          sXCache[threadIdx.x] = xLocal[xOff];
+        }
+        __syncthreads();
+
+        // compute partials
+        for (int f = threadIdx.x + t * maxBlock; (f < xTadLength) && (f < (t + 1) * maxBlock); f += (blockDim.x * gridDim.x)) {
+          sd::LongType yCoords[SD_MAX_RANK];
+          sd::LongType yOff;
+          INDEX2COORDS(f, yTadRank, yTadShapePtr, yCoords);
+          COORDS2INDEX(yTadRank, yTadStridePtr, yCoords, yOff);
+
+          sPartials[threadIdx.x] = OpType::update(
+              sPartials[threadIdx.x],
+              OpType::opAtomic(sXCache[threadIdx.x], yLocal[yOff], extraZ),
+              extraZ);
+        }
+        __syncthreads();
+      }
+
+      // reduce partials
+      aggregatePartials<OpType>(
+          reinterpret_cast<void*>(sPartials),
+          threadIdx.x,
+          sd::math::sd_min<int>(blockDim.x, xTadLength),
+          extraZ);
+      __syncthreads();
+
+      // store final
+      if (threadIdx.x == 0) {
+        z[ri] = OpType::postProcess(sPartials[0], xTadLength, extraZ);
+      }
+      __syncthreads();
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void Reduce3<X, Z>::transform(
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    int postProcessOrNot,
+    sd::LongType* allocationPointer,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* yTadOnlyShapeInfo,
+    sd::LongType const* yTadOffsets) {
+
+  if (shape::isScalar(zShapeInfo))
+    return; // no-op if scalar
+
+  if (yTadOnlyShapeInfo == nullptr)
+    yTadOnlyShapeInfo = yShapeInfo; // execReduce3TAD case
+
+  auto x = reinterpret_cast<const X*>(vx);
+  auto y = reinterpret_cast<const X*>(vy);
+  auto z = reinterpret_cast<Z*>(vz);
+
+  // For partial sums
+  __shared__ Z extraZ[OpType::extraParamsLen > 0 ? OpType::extraParamsLen : 1];
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+
+  __shared__ sd::LongType tadLen;
+  __shared__ sd::LongType zLen;
+  __shared__ sd::LongType yTadNum;
+  __shared__ sd::LongType xTadEws;
+  __shared__ sd::LongType yTadEws;
+  __shared__ char xTadOrder;
+  __shared__ char yTadOrder;
+
+  // Cache shape info
+  __shared__ sd::LongType xTadRank;
+  __shared__ const sd::LongType* xTadShape;
+  __shared__ const sd::LongType* xTadStride;
+
+  __shared__ sd::LongType yTadRank;
+  __shared__ const sd::LongType* yTadShape;
+  __shared__ const sd::LongType* yTadStride;
+
+  __shared__ sd::LongType zRank;  // Might not be used, but let's keep consistent
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    tadLen     = shape::length(tadOnlyShapeInfo);
+    zLen       = shape::length(zShapeInfo);
+    yTadNum    = shape::length(yShapeInfo) / tadLen;
+    xTadEws    = shape::elementWiseStride(tadOnlyShapeInfo);
+    yTadEws    = shape::elementWiseStride(yTadOnlyShapeInfo);
+    xTadOrder  = shape::order(tadOnlyShapeInfo);
+    yTadOrder  = shape::order(yTadOnlyShapeInfo);
+
+    xTadRank   = shape::rank(tadOnlyShapeInfo);
+    xTadShape  = shape::shapeOf(tadOnlyShapeInfo);
+    xTadStride = shape::stride(tadOnlyShapeInfo);
+
+    yTadRank   = shape::rank(yTadOnlyShapeInfo);
+    yTadShape  = shape::shapeOf(yTadOnlyShapeInfo);
+    yTadStride = shape::stride(yTadOnlyShapeInfo);
+
+    zRank      = shape::rank(zShapeInfo);
+    zShapePtr  = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+
+    sPartials[threadIdx.x] = OpType::startingValue(x);
+  }
+  __syncthreads();
+
+  Z startingVal = OpType::startingValue(x);
+
+  // either a direct ews approach or fallback
+  for (int i = blockIdx.x; i < zLen; i += gridDim.x) {
+    sd::LongType xBaseOffset = tadOffsets[i];
+    sd::LongType yBaseOffset = (yTadNum == 1 ? 0 : yTadOffsets[i]);
+
+    if (OpType::extraParamsLen > 0 && threadIdx.x < OpType::extraParamsLen) {
+      extraZ[threadIdx.x] = startingVal;
+    }
+    __syncthreads();
+
+    // partial sums
+    for (int j = threadIdx.x; j < tadLen; j += blockDim.x) {
+      sd::LongType xCoords[SD_MAX_RANK];
+      sd::LongType yCoords[SD_MAX_RANK];
+      sd::LongType xOff;
+      sd::LongType yOff;
+
+      INDEX2COORDS(j, xTadRank, xTadShape, xCoords);
+      COORDS2INDEX(xTadRank, xTadStride, xCoords, xOff);
+
+      INDEX2COORDS(j, yTadRank, yTadShape, yCoords);
+      COORDS2INDEX(yTadRank, yTadStride, yCoords, yOff);
+
+      // update partial
+      if (j < blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::opAtomic(x[xBaseOffset + xOff],
+                             y[yBaseOffset + yOff],
+                             extraZ);
+      } else {
+        sPartials[threadIdx.x] = OpType::update(
+            sPartials[threadIdx.x],
+            OpType::opAtomic(x[xBaseOffset + xOff],
+                             y[yBaseOffset + yOff],
+                             extraZ),
+            extraZ);
+      }
+    }
+    __syncthreads();
+
+    // reduce partials
+    aggregatePartials<OpType>(
+        reinterpret_cast<void*>(sPartials),
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraZ);
+    __syncthreads();
+
+    // write final
+    if (threadIdx.x == 0) {
+      z[i] = OpType::postProcess(sPartials[0], tadLen, extraZ);
+    }
+    __syncthreads();
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_DEVICE void Reduce3<X, Y>::execCuda(
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    int postProcessOrNot, sd::LongType* allocationPointer,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* yTadOnlyShapeInfo, sd::LongType const* yTadOffsets) {
+
+  DISPATCH_BY_OPNUM_TT(
+      transform,
+      PARAMS(vx, xShapeInfo, vy, yShapeInfo,
+             extraParams, vz, zShapeInfo,
+             dimension, dimensionLength, postProcessOrNot,
+             allocationPointer, tadOnlyShapeInfo, tadOffsets,
+             yTadOnlyShapeInfo, yTadOffsets),
+      REDUCE3_OPS);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_DEVICE void Reduce3<X, Y>::execAllCuda(
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, sd::LongType dimensionLength,
+    int postProcessOrNot, sd::LongType* allocationPointer,
+    sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
+    sd::LongType const* yTadOnlyShapeInfo, sd::LongType const* yTadOffsets) {
+
+  DISPATCH_BY_OPNUM_TT(
+      transformAll,
+      PARAMS(vx, xShapeInfo, vy, yShapeInfo,
+             extraParams, vz, zShapeInfo,
+             dimension, dimensionLength,
+             postProcessOrNot, allocationPointer,
+             tadOnlyShapeInfo, tadOffsets,
+             yTadOnlyShapeInfo, yTadOffsets),
+      REDUCE3_OPS);
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_DEVICE void Reduce3<X, Y>::execScalarCuda(
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* allocationPointer,
+    void* reductionBuffer,
+    sd::LongType const* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TT(
+      execScalarCuda,
+      PARAMS(vx, xShapeInfo, vy, yShapeInfo,
+             extraParams, vz, zShapeInfo,
+             allocationPointer, reductionBuffer, tadOnlyShapeInfo),
+      REDUCE3_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_HOST void Reduce3<X, Z>::exec(
+    dim3 launchDims, cudaStream_t* stream,
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    int postProcessOrNot,
+    sd::LongType* allocationPointer,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* yTadOnlyShapeInfo,
+    sd::LongType const* yTadOffsets) {
+
+  execGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      opNum,
+      vx, xShapeInfo,
+      vy, yShapeInfo,
+      extraParams,
+      vz, zShapeInfo,
+      dimension, dimensionLength,
+      postProcessOrNot,
+      allocationPointer,
+      tadOnlyShapeInfo, tadOffsets,
+      yTadOnlyShapeInfo, yTadOffsets);
+  sd::DebugHelper::checkErrorCode(stream, "reduce3exec(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_HOST void Reduce3<X, Z>::execAll(
+    dim3 launchDims, cudaStream_t* stream,
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* dimension, long long int dimensionLength,
+    int postProcessOrNot,
+    sd::LongType* allocationPointer,
+    sd::LongType const* tadOnlyShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* yTadOnlyShapeInfo,
+    sd::LongType const* yTadOffsets) {
+
+  execAllGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      opNum, vx, xShapeInfo,
+      vy, yShapeInfo,
+      extraParams,
+      vz, zShapeInfo,
+      dimension, dimensionLength,
+      postProcessOrNot,
+      allocationPointer,
+      tadOnlyShapeInfo, tadOffsets,
+      yTadOnlyShapeInfo, yTadOffsets);
+  sd::DebugHelper::checkErrorCode(stream, "execAllGeneric(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+SD_HOST void Reduce3<X, Z>::execScalar(
+    dim3 launchDims, cudaStream_t* stream,
+    int opNum,
+    void const* vx, sd::LongType const* xShapeInfo,
+    void const* vy, sd::LongType const* yShapeInfo,
+    void* extraParams,
+    void* vz, sd::LongType const* zShapeInfo,
+    sd::LongType* allocationPointer,
+    void* reductionBuffer,
+    sd::LongType const* tadOnlyShapeInfo) {
+
+  execScalarGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      opNum,
+      vx, xShapeInfo,
+      vy, yShapeInfo,
+      extraParams,
+      vz, zShapeInfo,
+      allocationPointer,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+  sd::DebugHelper::checkErrorCode(stream, "execScalarGeneric(...) failed");
+}
+
+}  // namespace reduce3
+}  // namespace functions

libnd4j/include/loops/cuda/scalar.chpp

@@ -0,0 +1,304 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 SCALAR_CU
+#define SCALAR_CU
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <loops/pairwise_instantiations.h>
+
+#include "../../types/types.h"
+#include "loops/scalar.h"
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z, typename OpType>
+SD_KERNEL static void scalarSimpleShaped(
+    void const* x,
+    void const* y,
+    sd::LongType const* xShapeInfo,
+    void* params,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    sd::LongType* allocationBuffer) {
+
+  auto xVals    = reinterpret_cast<X const*>(x);
+  auto scalar   = reinterpret_cast<Y const*>(y)[0];
+  auto pVals    = reinterpret_cast<Z*>(params);
+  auto zVals    = reinterpret_cast<Z*>(z);
+
+  const int totalThreads = gridDim.x * blockDim.x;
+  const int tid          = blockIdx.x * blockDim.x + threadIdx.x;
+
+  __shared__ sd::LongType length;
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+    zRank      = shape::rank(zShapeInfo);
+    zShapePtr  = shape::shapeOf(zShapeInfo);
+    zStridePtr = shape::stride(zShapeInfo);
+
+    //need special handling for scalars as strides may not be set for scalars
+    if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
+      z[0] = OpType::op(x[0], scalar[0], extraParams);
+      return;
+
+    }
+  }
+  __syncthreads();
+
+  //need special handling for scalars as strides may not be set for scalars
+  if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
+    return;
+
+  }
+  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);
+
+    zVals[zOffset] = OpType::op(xVals[xOffset], scalar, pVals);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z, typename OpType>
+SD_KERNEL static void scalarAlongDimension(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void* extraParams,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void const* scalars,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  auto xVals     = reinterpret_cast<X const*>(x);
+  auto zVals     = reinterpret_cast<Z*>(z);
+  auto pVals     = reinterpret_cast<Z*>(extraParams);
+  auto sVals     = reinterpret_cast<Y const*>(scalars);
+
+  if (tadShapeInfoZ == nullptr) {
+    tadShapeInfoZ = tadShapeInfo;
+    tadOffsetsZ   = tadOffsets;
+  }
+
+  __shared__ sd::LongType tadLength;
+  __shared__ sd::LongType numTads;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType tadRank;
+  __shared__ const sd::LongType* tadShapePtr;
+  __shared__ const sd::LongType* tadStridePtr;
+
+  __shared__ sd::LongType tadRankZ;
+  __shared__ const sd::LongType* tadShapePtrZ;
+  __shared__ const sd::LongType* tadStridePtrZ;
+
+  if (threadIdx.x == 0) {
+    tadLength      = shape::length(tadShapeInfo);
+    numTads        = shape::length(xShapeInfo) / tadLength;
+
+    xRank          = shape::rank(xShapeInfo);
+    xShapePtr      = shape::shapeOf(xShapeInfo);
+    xStridePtr     = shape::stride(xShapeInfo);
+
+    tadRank        = shape::rank(tadShapeInfo);
+    tadShapePtr    = shape::shapeOf(tadShapeInfo);
+    tadStridePtr   = shape::stride(tadShapeInfo);
+
+    tadRankZ       = shape::rank(tadShapeInfoZ);
+    tadShapePtrZ   = shape::shapeOf(tadShapeInfoZ);
+    tadStridePtrZ  = shape::stride(tadShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
+    auto zPart   = zVals + tadOffsetsZ[r];
+    auto xPart   = xVals + tadOffsets[r];
+    auto scalarV = sVals[r];
+
+    for (sd::LongType f = threadIdx.x; f < tadLength; f += blockDim.x) {
+      sd::LongType xCoords[SD_MAX_RANK];
+      sd::LongType zCoords[SD_MAX_RANK];
+      sd::LongType xOffset;
+      sd::LongType zOffset;
+
+      INDEX2COORDS(f, tadRank, tadShapePtr, xCoords);
+      COORDS2INDEX(tadRank, tadStridePtr, xCoords, xOffset);
+
+      INDEX2COORDS(f, tadRankZ, tadShapePtrZ, zCoords);
+      COORDS2INDEX(tadRankZ, tadStridePtrZ, zCoords, zOffset);
+
+      zPart[zOffset] = OpType::op(xPart[xOffset], scalarV, pVals);
+    }
+  }
+}
+
+namespace functions {
+namespace scalar {
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void SD_HOST ScalarTransform<X, Y, Z>::intermediateShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    sd::LongType const* hxShapeInfo,  // Unused in the kernel
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    sd::LongType const* hzShapeInfo,  // Unused in the kernel
+    void const* vscalar,
+    void* vextraParams,
+    sd::LongType* allocPointer) {
+
+  scalarSimpleShaped<X, Y, Z, OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          vscalar,
+          xShapeInfo,
+          vextraParams,
+          vz,
+          zShapeInfo,
+          allocPointer);
+  sd::DebugHelper::checkErrorCode(stream, "scalarSimpleShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void SD_HOST ScalarTransform<X, Y, Z>::intermediateAlongDimension(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void const* scalars,
+    void* extraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  scalarAlongDimension<X, Y, Z, OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x,
+          xShapeInfo,
+          extraParams,
+          z,
+          zShapeInfo,
+          scalars,
+          dimension,
+          dimensionLength,
+          tadShapeInfo,
+          tadOffsets,
+          tadShapeInfoZ,
+          tadOffsetsZ);
+  sd::DebugHelper::checkErrorCode(stream, "scalarAlongDimension(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void ScalarTransform<X, Y, Z>::executeCudaShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    sd::LongType const* hxShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    sd::LongType const* hzShapeInfo,
+    void const* vscalar,
+    void* vextraParams) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, hxShapeInfo, vz, zShapeInfo, hzShapeInfo, vscalar, vextraParams, nullptr),
+      SCALAR_OPS);
+}
+
+
+        ////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void ScalarTransform<X, Y, Z>::executeCudaAlongDimension(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void const* vscalars,
+    void* vextraParams,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ) {
+
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateAlongDimension,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo, vscalars, vextraParams, dimension, dimensionLength,
+             tadShapeInfo, tadOffsets, tadShapeInfoZ, tadOffsetsZ),
+      SCALAR_OPS);
+}
+
+}  // namespace scalar
+}  // namespace functions
+
+#endif // SCALAR_CU

libnd4j/include/loops/cuda/scalar_bool.cu

@@ -0,0 +1,364 @@
+/******************************************************************************
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+* See the NOTICE file distributed with this work for additional
+* information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing,
+* software distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+// @author Yurii Shyrma (iuriish@yahoo.com), created on 08.11.2018
+// @author raver119@gmail.com
+
+#ifndef SCALAR_BOOL_CU
+#define SCALAR_BOOL_CU
+
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include "../legacy_ops.h"
+#include "../scalar_bool.h"
+#include <helpers/DebugHelper.h>
+#include <system/Environment.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+// A kernel that applies a scalar bool transform along a specific dimension (TAD).
+// It uses shared memory caching for relevant shape information to reduce overhead.
+template <typename X, typename Z, typename OpType>
+__global__ void scalarAlongDimensionCachedKernel(
+   void const* x,
+   const sd::LongType* xShapeInfo,
+   void* extraParams,
+   void* z,
+   const sd::LongType* zShapeInfo,
+   void const* scalars,                     // per-TAD scalars
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets,
+   const sd::LongType* tadShapeInfoZ,       // if null, same as x TAD
+   const sd::LongType* tadOffsetsZ) {
+
+ auto xTyped         = reinterpret_cast<const X*>(x);
+ auto zTyped         = reinterpret_cast<Z*>(z);
+ auto extra          = reinterpret_cast<X*>(extraParams);
+ auto scalarsTyped   = reinterpret_cast<const X*>(scalars);
+
+ // If not provided, fallback
+ const auto* actualTadShapeInfoZ = (tadShapeInfoZ == nullptr ? tadShapeInfo : tadShapeInfoZ);
+ const auto* actualTadOffsetsZ   = (tadShapeInfoZ == nullptr ? tadOffsets   : tadOffsetsZ);
+
+ // Cache shape info in shared memory
+ __shared__ sd::LongType tadLen;
+ __shared__ sd::LongType numTads;
+
+ __shared__ int   tadRank;
+ __shared__ const sd::LongType* tadShapePtr;
+ __shared__ const sd::LongType* tadStridePtr;
+
+ __shared__ int   tadRankZ;
+ __shared__ const sd::LongType* tadShapePtrZ;
+ __shared__ const sd::LongType* tadStridePtrZ;
+
+ if (threadIdx.x == 0) {
+   tadLen      = shape::length(tadShapeInfo);
+   numTads     = shape::length(xShapeInfo) / tadLen;
+
+   tadRank     = shape::rank(tadShapeInfo);
+   tadShapePtr = shape::shapeOf(tadShapeInfo);
+   tadStridePtr= shape::stride(tadShapeInfo);
+
+   tadRankZ      = shape::rank(actualTadShapeInfoZ);
+   tadShapePtrZ  = shape::shapeOf(actualTadShapeInfoZ);
+   tadStridePtrZ = shape::stride(actualTadShapeInfoZ);
+ }
+ __syncthreads();
+
+ // Each block handles multiple TADs
+ for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+   Z* zTad    = zTyped + actualTadOffsetsZ[r];
+   const X* xTad = xTyped + tadOffsets[r];
+   X  scalar   = scalarsTyped[r];
+
+   // Each thread processes part of a single TAD
+   for (sd::LongType f = threadIdx.x; f < tadLen; f += blockDim.x) {
+     sd::LongType coordsX[SD_MAX_RANK];
+     sd::LongType coordsZ[SD_MAX_RANK];
+
+     sd::LongType offsetX;
+     sd::LongType offsetZ;
+
+     // Compute offset for X TAD
+     INDEX2COORDS(f, tadRank, tadShapePtr, coordsX);
+     COORDS2INDEX(tadRank, tadStridePtr, coordsX, offsetX);
+
+     // Compute offset for Z TAD
+     INDEX2COORDS(f, tadRankZ, tadShapePtrZ, coordsZ);
+     COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, offsetZ);
+
+     zTad[offsetZ] = OpType::op(xTad[offsetX], scalar, extra);
+   }
+ }
+}
+
+////////////////////////////////////////////////////////////////////////
+// A kernel to apply a scalar transform to a shaped buffer, with caching logic
+// for shape info in shared memory to reduce overhead.
+template <typename X, typename Z, typename OpType>
+__global__ void scalarSimpleShapedCachedKernel(
+   void const* x,                 // the "scalar" input
+   void const* y,                 // the "array" input
+   const sd::LongType* xShapeInfo,// we just read rank from here if needed
+   void* params,
+   void* z,
+   const sd::LongType* zShapeInfo,
+   sd::LongType* allocationBuffer) {
+
+ auto scalar = reinterpret_cast<const X*>(x)[0];
+ auto yTyped = reinterpret_cast<const X*>(y);
+ auto zTyped = reinterpret_cast<Z*>(z);
+ auto extra  = reinterpret_cast<X*>(params);
+
+ __shared__ sd::LongType length;
+ __shared__ int yRank;
+ __shared__ const sd::LongType* yShapePtr;
+ __shared__ const sd::LongType* yStridePtr;
+
+ __shared__ int zRank;
+ __shared__ const sd::LongType* zShapePtr;
+ __shared__ const sd::LongType* zStridePtr;
+
+ if (threadIdx.x == 0) {
+   length      = shape::length(xShapeInfo);  // or maybe shape::length(zShapeInfo)
+   // Actually we only need length from either input array's shape
+   // but let's assume x is scalar, so let's do it from z shape if that's a shaped array
+   // For now we keep as is.
+
+   yRank       = shape::rank(xShapeInfo);    // or we do: shape::rank(some-other-shape)
+   yShapePtr   = shape::shapeOf(xShapeInfo);
+   yStridePtr  = shape::stride(xShapeInfo);
+
+   zRank       = shape::rank(zShapeInfo);
+   zShapePtr   = shape::shapeOf(zShapeInfo);
+   zStridePtr  = shape::stride(zShapeInfo);
+ }
+ __syncthreads();
+
+ const auto tid          = blockDim.x * blockIdx.x + threadIdx.x;
+ const auto totalThreads = gridDim.x * blockDim.x;
+
+ for (sd::LongType i = tid; i < length; i += totalThreads) {
+   sd::LongType coordsY[SD_MAX_RANK];
+   sd::LongType coordsZ[SD_MAX_RANK];
+   sd::LongType offsetY;
+   sd::LongType offsetZ;
+
+   // get offset for Y
+   INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+   COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
+
+   // get offset for Z
+   INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+   COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+   zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, extra);
+ }
+}
+
+// *********************************************************************//
+// *********************************************************************//
+namespace functions {
+namespace scalar {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+__device__ void ScalarBoolTransform<X,Z>::transformCuda(
+   void const* vscalar,
+   void const* vy,
+   const sd::LongType* yShapeInfo,
+   void* vparams,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   sd::LongType* allocationBuffer) {
+
+ auto scalar = reinterpret_cast<const X*>(vscalar)[0];
+ auto yTyped = reinterpret_cast<const X*>(vy);
+ auto zTyped = reinterpret_cast<Z*>(vz);
+ auto extra  = reinterpret_cast<const X*>(vparams);
+
+ // store shape info in shared memory
+ __shared__ sd::LongType length;
+ __shared__ int yRank;
+ __shared__ const sd::LongType* yShapePtr;
+ __shared__ const sd::LongType* yStridePtr;
+
+ __shared__ int zRank;
+ __shared__ const sd::LongType* zShapePtr;
+ __shared__ const sd::LongType* zStridePtr;
+
+ if (threadIdx.x == 0) {
+   length      = shape::length(yShapeInfo);
+
+   yRank       = shape::rank(yShapeInfo);
+   yShapePtr   = shape::shapeOf(yShapeInfo);
+   yStridePtr  = shape::stride(yShapeInfo);
+
+   zRank       = shape::rank(zShapeInfo);
+   zShapePtr   = shape::shapeOf(zShapeInfo);
+   zStridePtr  = shape::stride(zShapeInfo);
+ }
+ __syncthreads();
+
+ const auto tid          = blockIdx.x * blockDim.x + threadIdx.x;
+ const auto totalThreads = blockDim.x * gridDim.x;
+
+ for (sd::LongType i = tid; i < length; i += totalThreads) {
+   sd::LongType coordsY[SD_MAX_RANK];
+   sd::LongType coordsZ[SD_MAX_RANK];
+   sd::LongType offsetY;
+   sd::LongType offsetZ;
+
+   INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+   COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
+
+   INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+   COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+   zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, const_cast<X*>(extra));
+ }
+}
+
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+__host__ void ScalarBoolTransform<X,Z>::intermediateAlongDimension(
+   dim3& launchDims,
+   cudaStream_t* stream,
+   void const* x,
+   const sd::LongType* xShapeInfo,
+   void* z,
+   const sd::LongType* zShapeInfo,
+   void const* scalars,
+   void* extraParams,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets,
+   const sd::LongType* tadShapeInfoZ,
+   const sd::LongType* tadOffsetsZ)
+{
+ scalarAlongDimensionCachedKernel<X,Z,OpType>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShapeInfo,
+         extraParams,
+         z,
+         zShapeInfo,
+         scalars,
+         dimension,
+         dimensionLength,
+         tadShapeInfo,
+         tadOffsets,
+         tadShapeInfoZ,
+         tadOffsetsZ);
+
+ sd::DebugHelper::checkErrorCode(stream, "scalarAlongDimensionCachedKernel(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+__host__ void ScalarBoolTransform<X,Z>::intermediateShaped(
+   dim3& launchDims,
+   cudaStream_t* stream,
+   void const* vx,
+   const sd::LongType* xShapeInfo,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   void const* vscalar,
+   void* vextraParams,
+   sd::LongType* allocPointer)
+{
+ scalarSimpleShapedCachedKernel<X,Z,OpType>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         vx,
+         vscalar,
+         xShapeInfo,
+         vextraParams,
+         vz,
+         zShapeInfo,
+         allocPointer);
+
+ sd::DebugHelper::checkErrorCode(stream, "scalarSimpleShapedCachedKernel(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+__host__ void ScalarBoolTransform<X, Y>::executeCudaShaped(
+   dim3& launchDims,
+   cudaStream_t* stream,
+   int opNum,
+   void const* vx,
+   const sd::LongType* xShapeInfo,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   void const* vscalar,
+   void const* vextraParams)
+{
+ if (sd::Environment::getInstance().isDebugAndVerbose()) {
+   printf("H14 opNum:[%i]\n", opNum);
+ }
+
+ DISPATCH_BY_OPNUM_TT(
+     intermediateShaped,
+     PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
+            vscalar, const_cast<void*>(vextraParams), nullptr),
+     SCALAR_BOOL_OPS);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+__host__ void ScalarBoolTransform<X, Y>::executeCudaAlongDimension(
+   dim3& launchDims,
+   cudaStream_t* stream,
+   int opNum,
+   void const* vx,
+   const sd::LongType* xShapeInfo,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   void const* vscalars,
+   void* vextraParams,
+   sd::LongType* dimension,
+   sd::LongType dimensionLength,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets,
+   const sd::LongType* tadShapeInfoZ,
+   const sd::LongType* tadOffsetsZ)
+{
+ DISPATCH_BY_OPNUM_TT(
+     intermediateAlongDimension,
+     PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
+            vscalars, vextraParams, dimension, dimensionLength,
+            tadShapeInfo, tadOffsets,
+            tadShapeInfoZ, tadOffsetsZ),
+     SCALAR_BOOL_OPS);
+}
+
+BUILD_DOUBLE_TEMPLATE( class ScalarBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+
+}  // namespace scalar
+}  // namespace functions
+
+#endif  // SCALAR_BOOL_CU

libnd4j/include/loops/cuda/scalar_int.cu

@@ -0,0 +1,388 @@
+/******************************************************************************
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com), created on 08.11.2018
+// @author raver119@gmail.com
+//
+#ifndef SCALAR_INT_CU
+#define SCALAR_INT_CU
+
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+#include "../legacy_ops.h"
+#include "../scalar_int.h"
+#include <helpers/DebugHelper.h>
+#include <system/Environment.h>
+
+    using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+// A kernel that applies an integer-based scalar transform along specified dimension (via TAD).
+// We'll cache shape info in shared memory to reduce repeated calls to shapeOf, strideOf, etc.
+template <typename X, typename OpType>
+__global__ void scalarAlongDimensionCachedKernel(
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void* extraParams,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void const* scalars,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ)
+{
+  // delegate the actual transform to the transformCuda method,
+  // caching shape info inside that method.
+  functions::scalar::ScalarIntTransform<X>::template transformCuda<OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      scalars,
+      dimension,
+      dimensionLength,
+      tadShapeInfo,
+      tadOffsets,
+      tadShapeInfoZ,
+      tadOffsetsZ
+  );
+}
+
+////////////////////////////////////////////////////////////////////////
+// A kernel to handle shaped transforms: x is "scalar," y is "array," but
+// we store shape data in shared memory.
+template <typename X, typename OpType>
+__global__ void scalarSimpleShapedCachedKernel(
+    void const* x,
+    void const* y,
+    sd::LongType const* xShapeInfo,
+    void* params,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    sd::LongType* allocationBuffer)
+{
+  // We'll call the transformCuda method (which will do caching).
+  functions::scalar::ScalarIntTransform<X>::template transformCuda<OpType>(
+      y,
+      x,
+      xShapeInfo,
+      params,
+      z,
+      zShapeInfo,
+      allocationBuffer
+  );
+}
+
+namespace functions {
+namespace scalar {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+__device__ void ScalarIntTransform<X>::transformCuda(
+    void const* vscalar,
+    void const* vy,
+    sd::LongType const* yShapeInfo,
+    void* vparams,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    sd::LongType* allocationBuffer)
+{
+  auto scalar = reinterpret_cast<const X*>(vscalar)[0];
+  auto yTyped = reinterpret_cast<const X*>(vy);
+  auto zTyped = reinterpret_cast<X*>(vz);
+  auto extra  = reinterpret_cast<X*>(vparams);
+
+  // cache shape info in shared memory
+  __shared__ sd::LongType length;
+
+  __shared__ int yRank;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ int zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length    = shape::length(yShapeInfo);
+
+    yRank     = shape::rank(yShapeInfo);
+    yShapePtr = shape::shapeOf(yShapeInfo);
+    yStridePtr= shape::stride(yShapeInfo);
+
+    zRank     = shape::rank(zShapeInfo);
+    zShapePtr = shape::shapeOf(zShapeInfo);
+    zStridePtr= shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  const auto tid          = blockDim.x * blockIdx.x + threadIdx.x;
+  const auto totalThreads = gridDim.x * blockDim.x;
+
+  // now we do the transform
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType coordsY[SD_MAX_RANK];
+    sd::LongType coordsZ[SD_MAX_RANK];
+
+    sd::LongType offsetY;
+    sd::LongType offsetZ;
+
+    INDEX2COORDS(i, yRank, yShapePtr, coordsY);
+    COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
+
+    INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+    COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+    zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, extra);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+__device__ void ScalarIntTransform<X>::transformCuda(
+    sd::LongType len,
+    void const* vx,
+    void const* vy,
+    sd::LongType yEWS,
+    void* vparams,
+    void* vz,
+    sd::LongType zEWS,
+    sd::LongType* allocationBuffer)
+{
+  auto x     = reinterpret_cast<const X*>(vx)[0];  // scalar
+  auto yTyped= reinterpret_cast<const X*>(vy);
+  auto zTyped= reinterpret_cast<X*>(vz);
+  auto extra = reinterpret_cast<X*>(vparams);
+
+  const int tid          = blockDim.x * blockIdx.x + threadIdx.x;
+  const int totalThreads = blockDim.x * gridDim.x;
+
+  for (sd::LongType i = tid; i < len; i += totalThreads) {
+    zTyped[i * zEWS] = OpType::op(yTyped[i * yEWS], x, extra);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+__device__ void ScalarIntTransform<X>::transformCuda(
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void const* vscalars,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ)
+{
+  auto x        = reinterpret_cast<const X*>(vx);
+  auto scalars  = reinterpret_cast<const X*>(vscalars);
+  auto zTyped   = reinterpret_cast<X*>(vz);
+  auto extra    = reinterpret_cast<X*>(vextraParams);
+
+  // if z TAD not provided, fallback
+  const auto* actualTadShapeInfoZ = (tadShapeInfoZ == nullptr ? tadShapeInfo : tadShapeInfoZ);
+  const auto* actualTadOffsetsZ   = (tadShapeInfoZ == nullptr ? tadOffsets : tadOffsetsZ);
+
+  // cache shape info in shared memory
+  __shared__ sd::LongType tadLen;
+  __shared__ sd::LongType numTads;
+
+  __shared__ int tadRank;
+  __shared__ const sd::LongType* tadShapePtr;
+  __shared__ const sd::LongType* tadStridePtr;
+
+  __shared__ int tadRankZ;
+  __shared__ const sd::LongType* tadShapePtrZ;
+  __shared__ const sd::LongType* tadStridePtrZ;
+
+  if (threadIdx.x == 0) {
+    tadLen       = shape::length(tadShapeInfo);
+    numTads      = shape::length(xShapeInfo) / tadLen;
+
+    tadRank      = shape::rank(tadShapeInfo);
+    tadShapePtr  = shape::shapeOf(tadShapeInfo);
+    tadStridePtr = shape::stride(tadShapeInfo);
+
+    tadRankZ      = shape::rank(actualTadShapeInfoZ);
+    tadShapePtrZ  = shape::shapeOf(actualTadShapeInfoZ);
+    tadStridePtrZ = shape::stride(actualTadShapeInfoZ);
+  }
+  __syncthreads();
+
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    X* zTad       = zTyped + actualTadOffsetsZ[r];
+    const X* xTad = x       + tadOffsets[r];
+    X    scalar   = scalars[r];
+
+    // each thread processes part of TAD
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType coordsX[SD_MAX_RANK];
+      sd::LongType coordsZ[SD_MAX_RANK];
+
+      sd::LongType offsetX;
+      sd::LongType offsetZ;
+
+      INDEX2COORDS(i, tadRank, tadShapePtr, coordsX);
+      COORDS2INDEX(tadRank, tadStridePtr, coordsX, offsetX);
+
+      INDEX2COORDS(i, tadRankZ, tadShapePtrZ, coordsZ);
+      COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, offsetZ);
+
+      zTad[offsetZ] = OpType::op(xTad[offsetX], scalar, extra);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+__host__ void ScalarIntTransform<X>::intermediateAlongDimension(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* x,
+    sd::LongType const* xShapeInfo,
+    void* z,
+    sd::LongType const* zShapeInfo,
+    void const* scalars,
+    void* extraParams,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ)
+{
+  // we use the new, cached version
+  scalarAlongDimensionCachedKernel<X,OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          x,
+          xShapeInfo,
+          extraParams,
+          z,
+          zShapeInfo,
+          scalars,
+          dimension,
+          dimensionLength,
+          tadShapeInfo,
+          tadOffsets,
+          tadShapeInfoZ,
+          tadOffsetsZ);
+
+  sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform intermediateAlongDimension(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+template <typename OpType>
+__host__ void ScalarIntTransform<X>::intermediateShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void const* vscalar,
+    void* vextraParams,
+    sd::LongType* allocPointer)
+{
+  // call the new cached kernel
+  scalarSimpleShapedCachedKernel<X,OpType>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          vscalar,
+          xShapeInfo,
+          vextraParams,
+          vz,
+          zShapeInfo,
+          allocPointer);
+
+  sd::DebugHelper::checkGlobalErrorCode("scalarSimpleShapedCachedKernel(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+__host__ void ScalarIntTransform<X>::executeCudaShaped(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void const* vscalar,
+    void* vextraParams)
+{
+  if (sd::Environment::getInstance().isDebugAndVerbose()) {
+    printf("H14 scalar int transform opNum:[%i]\n", opNum);
+  }
+
+  DISPATCH_BY_OPNUM_T(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo, vscalar, vextraParams, nullptr),
+      SCALAR_INT_OPS);
+
+  sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform executeCudaShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X>
+__host__ void ScalarIntTransform<X>::executeCudaAlongDimension(
+    dim3& launchDims,
+    cudaStream_t* stream,
+    int opNum,
+    void const* vx,
+    sd::LongType const* xShapeInfo,
+    void* vz,
+    sd::LongType const* zShapeInfo,
+    void const* vscalars,
+    void* vextraParams,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    sd::LongType const* tadShapeInfo,
+    sd::LongType const* tadOffsets,
+    sd::LongType const* tadShapeInfoZ,
+    sd::LongType const* tadOffsetsZ)
+{
+  DISPATCH_BY_OPNUM_T(
+      intermediateAlongDimension,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
+             vscalars, vextraParams, dimension, dimensionLength,
+             tadShapeInfo, tadOffsets, tadShapeInfoZ, tadOffsetsZ),
+      SCALAR_INT_OPS);
+
+  sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform executeCudaAlongDimension(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( class ScalarIntTransform, , SD_INTEGER_TYPES);
+
+}  // namespace scalar
+}  // namespace functions
+
+#endif  // SCALAR_INT_CU

libnd4j/include/loops/cuda/specials/accumulateKernel.cu

@@ -0,0 +1,88 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 15.11.2018
+//
+#include <loops/special_kernels.h>
+
+
+namespace sd {
+
+///////////////////////////////////////////////////////////////////////
+/**
+ * This kernel accumulates X arrays, and stores z into Z
+ *
+ * @tparam T
+ * @param x
+ * @param z
+ * @param n
+ * @param length
+ */
+template <typename T>
+SD_DEVICE void accumulateKernel(void **vx, void *vz, int n, const LongType length) {
+  auto x = reinterpret_cast<T **>(vx);
+  auto z = reinterpret_cast<T *>(vz);
+
+  __shared__ T *shmem;
+
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char sharedmem[];
+    shmem = (T *)sharedmem;
+  }
+  __syncthreads();
+
+  for (int r = blockDim.x * blockIdx.x; r < length; r += blockDim.x * gridDim.x) {
+    shmem[threadIdx.x] = 0.0f;
+
+    LongType baseIdx = r;
+
+    // aggregation step, we roll over all arrays
+    for (int ar = 0; ar < n; ar++) {
+      T *cdata = (T *)x[ar];
+      cdata += baseIdx;
+
+      if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] += cdata[threadIdx.x];
+    }
+
+    T *wdata = z + baseIdx;
+
+    // saving accumulated values
+    if (baseIdx + threadIdx.x < length) wdata[threadIdx.x] = shmem[threadIdx.x];
+  }
+}
+
+///////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL void execAccumulateKernel(void **vx, void *vz, int n, const LongType length) {
+  accumulateKernel<T>(vx, vz, n, length);
+}
+
+///////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_HOST void accumulateKernelGeneric(dim3 &launchDims, cudaStream_t *stream, void **vx, void *vz, int n,
+                                     const LongType length) {
+  execAccumulateKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vx, vz, n, length);
+  DebugHelper::checkErrorCode(stream, "accumulate(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( void accumulateKernelGeneric,
+                      (dim3 & launchDims, cudaStream_t *stream, void **vx, void *vz, int n, const sd::LongType length),
+                      SD_COMMON_TYPES);
+}  // namespace sd

libnd4j/include/loops/cuda/specials/averagingKernel.cu

@@ -0,0 +1,97 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 15.11.2018
+//
+#include <loops/special_kernels.h>
+
+
+namespace sd {
+
+///////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_DEVICE void averagingKernel(void **vdx, void *vdz, int n, LongType length, bool propagate) {
+  auto dx = reinterpret_cast<T **>(vdx);
+  auto dz = reinterpret_cast<T *>(vdz);
+
+  __shared__ T *shmem;
+
+  if (threadIdx.x == 0) {
+    extern __shared__ unsigned char sharedmem[];
+    shmem = (T *)sharedmem;
+  }
+  __syncthreads();
+
+  // each block cycles over it's own part of arrays
+  for (int r = blockDim.x * blockIdx.x; r < length; r += blockDim.x * gridDim.x) {
+    shmem[threadIdx.x] = (T)0.0f;
+
+    LongType baseIdx = r;
+
+    // aggregation step, we roll over all arrays
+    for (int ar = 0; ar < n; ar++) {
+      T *cdata = (T *)dx[ar];
+      cdata += baseIdx;
+
+      if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] += cdata[threadIdx.x];
+    }
+
+    // average data in shared memory
+    if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] /= n;
+
+    // div step & write out step
+    if (dz != nullptr) {
+      T *wdata = dz + baseIdx;
+
+      if (baseIdx + threadIdx.x < length) {
+        wdata[threadIdx.x] = shmem[threadIdx.x];
+      }
+    }
+
+    // propagate averaged data to all arrays
+    if (propagate)
+      for (int ar = 0; ar < n; ar++) {
+        T *cdata = (T *)dx[ar];
+        cdata += baseIdx;
+
+        if (baseIdx + threadIdx.x < length) cdata[threadIdx.x] = shmem[threadIdx.x];
+      }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL void execAveragingKernel(void **vdx, void *vdz, int n, LongType length, bool propagate) {
+  averagingKernel<T>(vdx, vdz, n, length, propagate);
+}
+
+///////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_HOST void averagingKernelGeneric(dim3 &launchDims, cudaStream_t *stream, void **vdx, void *vdz, int n,
+                                    LongType length, bool propagate) {
+  execAveragingKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vdx, vdz, n, length, propagate);
+  DebugHelper::checkErrorCode(stream, "averaging(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( void averagingKernelGeneric,
+                      (dim3 & launchDims, cudaStream_t *stream, void **vdx, void *vdz, int n, sd::LongType length,
+                       bool propagate),
+                      SD_COMMON_TYPES);
+}  // namespace sd

libnd4j/include/loops/cuda/specials/bitonicArbitraryStep.cu

@@ -0,0 +1,294 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 28.11.2018
+//
+#include <ops/specials_cuda.h>
+
+    //////////////////////////////////////////////////////////////////////////
+    template <typename X, typename Y>
+    SD_KERNEL void bitonicArbitraryStepKernelKey(
+        void* vx,
+        const sd::LongType* xShapeInfo,
+        void* vy,
+        const sd::LongType* yShapeInfo,
+        int window,
+        int length,
+        int reverse,
+        bool descending) {
+
+  auto x         = static_cast<X*>(vx);
+  auto y         = static_cast<Y*>(vy);
+  const int tid  = threadIdx.x + blockDim.x * blockIdx.x;
+  const int half = window >> 1;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType yRank;    // Potentially unused for direct indexing, but let's keep the pattern consistent
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ sd::LongType xLength;
+
+  if (threadIdx.x == 0) {
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+    yRank      = shape::rank(yShapeInfo);
+    yShapePtr  = shape::shapeOf(yShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+
+    xLength    = shape::length(xShapeInfo);
+  }
+  __syncthreads();
+
+  const int WARP_SIZE = 32;
+  const int numWarps  = (gridDim.x * blockDim.x) / WARP_SIZE;
+  const int warpId    = tid / WARP_SIZE;
+  const int warpIdx   = tid % WARP_SIZE;
+
+  int firstPosition;
+  int firstStep;
+  int secondPosition;
+  int secondStep;
+
+  if (half >= 128) {
+    firstPosition = blockIdx.x * window;
+    firstStep     = gridDim.x * window;
+
+    secondPosition = threadIdx.x;
+    secondStep     = blockDim.x;
+  }
+  else if (half >= 32) {
+    firstPosition = warpId * window;
+    firstStep     = numWarps * window;
+
+    secondPosition = warpIdx;
+    secondStep     = WARP_SIZE;
+  }
+  else {
+    firstPosition = tid * window;
+    firstStep     = blockDim.x * gridDim.x * window;
+
+    secondPosition = 0;
+    secondStep     = 1;
+  }
+
+  for (int i = firstPosition; i < length; i += firstStep) {
+    for (int j = secondPosition; j < half; j += secondStep) {
+      const int it = (reverse) ? i + j + half : i + window - j - 1;
+      const int ij = i + j;
+      if (it < length && ij < length) {
+        sd::LongType itCoords[SD_MAX_RANK];
+        sd::LongType ijCoords[SD_MAX_RANK];
+        sd::LongType itOffset;
+        sd::LongType ijOffset;
+
+        INDEX2COORDS(it, xRank, xShapePtr, itCoords);
+        COORDS2INDEX(xRank, xStridePtr, itCoords, itOffset);
+
+        INDEX2COORDS(ij, xRank, xShapePtr, ijCoords);
+        COORDS2INDEX(xRank, xStridePtr, ijCoords, ijOffset);
+
+        X v0 = x[ijOffset];
+        X v1 = x[itOffset];
+
+        const bool condition = (!descending == (v0 > v1));
+        if (condition) {
+          x[ijOffset] = v1;
+          x[itOffset] = v0;
+
+          sd::LongType itCoordsY[SD_MAX_RANK];
+          sd::LongType ijCoordsY[SD_MAX_RANK];
+          sd::LongType itOffsetY;
+          sd::LongType ijOffsetY;
+
+          INDEX2COORDS(it, yRank, yShapePtr, itCoordsY);
+          COORDS2INDEX(yRank, yStridePtr, itCoordsY, itOffsetY);
+
+          INDEX2COORDS(ij, yRank, yShapePtr, ijCoordsY);
+          COORDS2INDEX(yRank, yStridePtr, ijCoordsY, ijOffsetY);
+
+          Y ytemp        = y[ijOffsetY];
+          y[ijOffsetY]   = y[itOffsetY];
+          y[itOffsetY]   = ytemp;
+        }
+      }
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL void execBitonicArbitraryStepKernel(
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    int window,
+    int length,
+    int reverse,
+    bool descending) {
+
+  auto x         = static_cast<T*>(vx);
+  const int tid  = threadIdx.x + blockDim.x * blockIdx.x;
+  const int half = window >> 1;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType xLength;
+
+  // We'll omit using shared memory for x data except for small merges,
+  // but keep the pattern of caching shape info
+  if (threadIdx.x == 0) {
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+    xLength    = shape::length(xShapeInfo);
+  }
+  __syncthreads();
+
+  const int WARP_SIZE = 32;
+  const int numWarps  = (gridDim.x * blockDim.x) / WARP_SIZE;
+  const int warpId    = tid / WARP_SIZE;
+  const int warpIdx   = tid % WARP_SIZE;
+
+  int firstPosition;
+  int firstStep;
+  int secondPosition;
+  int secondStep;
+
+  if (half >= 128) {
+    firstPosition  = blockIdx.x * window;
+    firstStep      = gridDim.x * window;
+
+    secondPosition = threadIdx.x;
+    secondStep     = blockDim.x;
+  }
+  else if (half >= 32) {
+    firstPosition  = warpId * window;
+    firstStep      = numWarps * window;
+
+    secondPosition = warpIdx;
+    secondStep     = WARP_SIZE;
+  }
+  else {
+    firstPosition  = tid * window;
+    firstStep      = blockDim.x * gridDim.x * window;
+
+    secondPosition = 0;
+    secondStep     = 1;
+  }
+
+  for (int i = firstPosition; i < length; i += firstStep) {
+    for (int j = secondPosition; j < half; j += secondStep) {
+      const int it = (reverse) ? i + j + half : i + window - j - 1;
+      const int ij = i + j;
+      if (it < length && ij < length) {
+        sd::LongType itCoords[SD_MAX_RANK];
+        sd::LongType ijCoords[SD_MAX_RANK];
+        sd::LongType itOffset;
+        sd::LongType ijOffset;
+
+        INDEX2COORDS(it, xRank, xShapePtr, itCoords);
+        COORDS2INDEX(xRank, xStridePtr, itCoords, itOffset);
+
+        INDEX2COORDS(ij, xRank, xShapePtr, ijCoords);
+        COORDS2INDEX(xRank, xStridePtr, ijCoords, ijOffset);
+
+        T v0 = x[ijOffset];
+        T v1 = x[itOffset];
+
+        const bool condition = (!descending == (v0 > v1));
+        if (condition) {
+          x[ijOffset] = v1;
+          x[itOffset] = v0;
+        }
+      }
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_HOST void bitonicArbitraryStepGeneric(
+    dim3 &launchDims,
+    cudaStream_t *stream,
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    int window,
+    int length,
+    int reverse,
+    bool descending) {
+
+  execBitonicArbitraryStepKernel<T>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          window,
+          length,
+          reverse,
+          descending);
+
+  sd::DebugHelper::checkErrorCode(stream, "execBitonicArbitraryStepKernel  failed");
+}
+
+template <typename X, typename Y>
+SD_HOST void bitonicArbitraryStepGenericKey(
+    dim3 &launchDims,
+    cudaStream_t *stream,
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vy,
+    const sd::LongType* yShapeInfo,
+    int window,
+    int length,
+    int reverse,
+    bool descending) {
+
+  bitonicArbitraryStepKernelKey<X, Y>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          vy,
+          yShapeInfo,
+          window,
+          length,
+          reverse,
+          descending);
+
+  sd::DebugHelper::checkErrorCode(stream, "bitonicArbitraryStepKernelKey failed");
+}
+
+BUILD_SINGLE_TEMPLATE(
+     void bitonicArbitraryStepGeneric,
+    (dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo, int window,
+     int length, int reverse, bool descending),
+    SD_COMMON_TYPES);
+
+BUILD_DOUBLE_TEMPLATE(
+     void bitonicArbitraryStepGenericKey,
+    (dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo, void *vy,
+     sd::LongType const *yShapeInfo, int window, int length, int reverse, bool descending),
+    SD_COMMON_TYPES, SD_COMMON_TYPES);

libnd4j/include/loops/cuda/specials/bitonicSortStep.cu

@@ -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 Yurii Shyrma, created on 28.11.2018
+//
+#include <ops/specials_cuda.h>
+
+    //////////////////////////////////////////////////////////////////////////
+    template <typename X, typename Y>
+    SD_KERNEL void bitonicSortStepKernelKey(
+        void* vx,
+        const sd::LongType* xShapeInfo,
+        void* vy,
+        const sd::LongType* yShapeInfo,
+        int j,
+        int k,
+        int length,
+        bool descending) {
+
+  auto x           = static_cast<X*>(vx);
+  auto y           = static_cast<Y*>(vy);
+  const unsigned int i = threadIdx.x + blockDim.x * blockIdx.x;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ sd::LongType yRank;
+  __shared__ const sd::LongType* yShapePtr;
+  __shared__ const sd::LongType* yStridePtr;
+
+  __shared__ sd::LongType xLength;
+
+  if (threadIdx.x == 0) {
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+    yRank      = shape::rank(yShapeInfo);
+    yShapePtr  = shape::shapeOf(yShapeInfo);
+    yStridePtr = shape::stride(yShapeInfo);
+
+    xLength    = shape::length(xShapeInfo);
+  }
+  __syncthreads();
+
+  if (i >= static_cast<unsigned int>(length)) return;
+
+  const unsigned int ixj = i ^ j;
+  if (ixj <= i) return;
+
+  sd::LongType iCoords[SD_MAX_RANK];
+  sd::LongType ixjCoords[SD_MAX_RANK];
+  sd::LongType iOffset;
+  sd::LongType ixjOffset;
+
+  INDEX2COORDS(i, xRank, xShapePtr, iCoords);
+  COORDS2INDEX(xRank, xStridePtr, iCoords, iOffset);
+
+  INDEX2COORDS(ixj, xRank, xShapePtr, ixjCoords);
+  COORDS2INDEX(xRank, xStridePtr, ixjCoords, ixjOffset);
+
+  const bool ascending = ((i & k) == 0);
+  X xi = x[iOffset];
+  X xixj = x[ixjOffset];
+
+  if (ascending) {
+    // Sort ascending
+    if (!descending == (xi > xixj)) {
+      x[iOffset]      = xixj;
+      x[ixjOffset]    = xi;
+
+      sd::LongType iCoordsY[SD_MAX_RANK];
+      sd::LongType ixjCoordsY[SD_MAX_RANK];
+      sd::LongType iOffsetY;
+      sd::LongType ixjOffsetY;
+
+      INDEX2COORDS(i, yRank, yShapePtr, iCoordsY);
+      COORDS2INDEX(yRank, yStridePtr, iCoordsY, iOffsetY);
+
+      INDEX2COORDS(ixj, yRank, yShapePtr, ixjCoordsY);
+      COORDS2INDEX(yRank, yStridePtr, ixjCoordsY, ixjOffsetY);
+
+      Y yi   = y[iOffsetY];
+      Y yixj = y[ixjOffsetY];
+      y[iOffsetY]   = yixj;
+      y[ixjOffsetY] = yi;
+    }
+  }
+  else {
+    // Sort descending
+    if (!descending == (xi < xixj)) {
+      x[iOffset]      = xixj;
+      x[ixjOffset]    = xi;
+
+      sd::LongType iCoordsY[SD_MAX_RANK];
+      sd::LongType ixjCoordsY[SD_MAX_RANK];
+      sd::LongType iOffsetY;
+      sd::LongType ixjOffsetY;
+
+      INDEX2COORDS(i, yRank, yShapePtr, iCoordsY);
+      COORDS2INDEX(yRank, yStridePtr, iCoordsY, iOffsetY);
+
+      INDEX2COORDS(ixj, yRank, yShapePtr, ixjCoordsY);
+      COORDS2INDEX(yRank, yStridePtr, ixjCoordsY, ixjOffsetY);
+
+      Y yi   = y[iOffsetY];
+      Y yixj = y[ixjOffsetY];
+      y[iOffsetY]   = yixj;
+      y[ixjOffsetY] = yi;
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_KERNEL void bitonicSortStepKernel(
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    int j,
+    int k,
+    int length,
+    bool descending) {
+
+  auto x           = static_cast<T*>(vx);
+  const unsigned int i = threadIdx.x + blockDim.x * blockIdx.x;
+
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+  __shared__ sd::LongType xLength;
+
+  if (threadIdx.x == 0) {
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+
+    xLength    = shape::length(xShapeInfo);
+  }
+  __syncthreads();
+
+  if (i >= static_cast<unsigned int>(length)) return;
+
+  const unsigned int ixj = i ^ j;
+  if (ixj <= i) return;
+
+  sd::LongType iCoords[SD_MAX_RANK];
+  sd::LongType ixjCoords[SD_MAX_RANK];
+  sd::LongType iOffset;
+  sd::LongType ixjOffset;
+
+  INDEX2COORDS(i, xRank, xShapePtr, iCoords);
+  COORDS2INDEX(xRank, xStridePtr, iCoords, iOffset);
+
+  INDEX2COORDS(ixj, xRank, xShapePtr, ixjCoords);
+  COORDS2INDEX(xRank, xStridePtr, ixjCoords, ixjOffset);
+
+  const bool ascending = ((i & k) == 0);
+  T xi   = x[iOffset];
+  T xixj = x[ixjOffset];
+
+  if (ascending) {
+    // Sort ascending
+    if (!descending == (xi > xixj)) {
+      x[iOffset]    = xixj;
+      x[ixjOffset]  = xi;
+    }
+  }
+  else {
+    // Sort descending
+    if (!descending == (xi < xixj)) {
+      x[iOffset]    = xixj;
+      x[ixjOffset]  = xi;
+    }
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename T>
+SD_HOST void bitonicSortStepGeneric(
+    dim3 &launchDims,
+    cudaStream_t *stream,
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    int j,
+    int k,
+    int length,
+    bool descending) {
+
+  bitonicSortStepKernel<T>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          j,
+          k,
+          length,
+          descending);
+
+  sd::DebugHelper::checkErrorCode(stream, "bitonicSortStepGeneric failed");
+}
+
+//////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void bitonicSortStepGenericKey(
+    dim3 &launchDims,
+    cudaStream_t *stream,
+    void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vy,
+    const sd::LongType* yShapeInfo,
+    int j,
+    int k,
+    int length,
+    bool descending) {
+
+  bitonicSortStepKernelKey<X, Y>
+      <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+          vx,
+          xShapeInfo,
+          vy,
+          yShapeInfo,
+          j,
+          k,
+          length,
+          descending);
+
+  sd::DebugHelper::checkErrorCode(stream, "bitonicSortStepGenericKey failed");
+}
+
+BUILD_SINGLE_TEMPLATE(
+    template void bitonicSortStepGeneric,
+    (dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
+     int j, int k, int length, bool descending),
+    SD_COMMON_TYPES);
+
+BUILD_DOUBLE_TEMPLATE(
+    template void bitonicSortStepGenericKey,
+    (dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
+     void *vy, sd::LongType const *yShapeInfo, int j, int k, int length, bool descending),
+    SD_COMMON_TYPES, SD_COMMON_TYPES);

libnd4j/include/loops/cuda/specials/concatKernel.cu

@@ -0,0 +1,300 @@
+/* ******************************************************************************
+*
+*
+* 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, created on 15.11.2018
+//
+#include <loops/special_kernels.h>
+
+namespace sd {
+
+template <typename T>
+SD_DEVICE void concatKernel(
+   int numArrays,
+   Pointer* data,
+   Pointer* inputShapeInfos,
+   void* vz,
+   LongType* resultShapeInfo,
+   Pointer* tadPointers,
+   Pointer* offsetPointers,
+   LongType* zTadShape,
+   LongType* zOffsets) {
+
+ const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+ // Shared variables for shape data
+ __shared__ int zRank;
+ __shared__ const LongType* zShape;
+ __shared__ const LongType* zStride;
+ __shared__ int zTadRank;
+ __shared__ const LongType* zTadShapeOf;
+ __shared__ const LongType* zTadStride;
+
+ if (threadIdx.x == 0) {
+   zRank = shape::rank(resultShapeInfo);
+   zShape = shape::shapeOf(resultShapeInfo);
+   zStride = shape::stride(resultShapeInfo);
+   zTadRank = shape::rank(zTadShape);
+   zTadShapeOf = shape::shapeOf(zTadShape);
+   zTadStride = shape::stride(zTadShape);
+ }
+ __syncthreads();
+
+ auto result = reinterpret_cast<T*>(vz);
+ auto dataT = reinterpret_cast<T**>(data);
+ auto shapeInfoPtrs = reinterpret_cast<LongType**>(inputShapeInfos);
+ auto tadShapes = reinterpret_cast<LongType**>(tadPointers);
+ auto tadOffsets = reinterpret_cast<LongType**>(offsetPointers);
+
+ __shared__ bool _vec;
+ __shared__ int baseIdx;
+ __shared__ int arrOffset;
+ __shared__ int yLength;
+ __shared__ char yOrder;
+ __shared__ int numTads;
+ __shared__ char zOrder;
+
+ if (threadIdx.x == 0) {
+   zOrder = shape::order(resultShapeInfo);
+   _vec = shape::isVector(resultShapeInfo);
+ }
+ __syncthreads();
+
+ const int zLength = shape::length(resultShapeInfo);
+
+ // Special case for when result is a vector
+ if (_vec) {
+   for (int r = blockIdx.x; r < numArrays; r += gridDim.x) {
+     auto currShapeInfo = shapeInfoPtrs[r];
+     if (shape::isVector(currShapeInfo) || shape::order(currShapeInfo) == zOrder) {
+       if (threadIdx.x == 0) {
+         yLength = shape::length(currShapeInfo);
+         baseIdx = 0;
+         for (int f = 0; f < r; f++) {
+           baseIdx += shape::length(shapeInfoPtrs[f]);
+         }
+       }
+       __syncthreads();
+
+       for (int i = threadIdx.x; i < yLength && baseIdx + i < zLength; i += blockDim.x) {
+         result[baseIdx + i] = dataT[r][i];
+       }
+       __syncthreads();
+     }
+   }
+   return;
+ }
+
+ // For all non-vector results
+ for (int r = 0; r < numArrays; r++) {
+   auto currShape = shapeInfoPtrs[r];
+   auto currData = dataT[r];
+   auto currTad = tadShapes[r];
+   auto currOffsets = tadOffsets[r];
+
+   // Shared variables for current TAD shape data
+   __shared__ int currTadRank;
+   __shared__ const LongType* currTadShape;
+   __shared__ const LongType* currTadStride;
+
+   if (threadIdx.x == 0) {
+     currTadRank = shape::rank(currTad);
+     currTadShape = shape::shapeOf(currTad);
+     currTadStride = shape::stride(currTad);
+   }
+   __syncthreads();
+
+   if (threadIdx.x == 0) {
+     yLength = shape::length(currTad);
+     yOrder = shape::order(currTad);
+     numTads = shape::length(currShape) / yLength;
+
+     arrOffset = 0;
+     for (int f = 0; f < r; f++) {
+       arrOffset += shape::length(tadShapes[f]);
+     }
+   }
+   __syncthreads();
+
+   // Case for single element
+   if (yLength == 1 && _vec) {
+     for (LongType j = tid; j < numTads; j += blockDim.x * gridDim.x) {
+       LongType inputOffset = currOffsets[j];
+       LongType resultOffset = zOffsets[j];
+
+       T* dataTAD = currData + inputOffset;
+       T* resultTAD = result + resultOffset;
+
+       LongType sub[SD_MAX_RANK];
+       INDEX2COORDS(arrOffset, zTadRank, zTadShapeOf, sub);
+
+       LongType baseOffset;
+       COORDS2INDEX(zTadRank, zTadStride, sub, baseOffset);
+       resultTAD += baseOffset;
+
+       INDEX2COORDS(0, currTadRank, currTadShape, sub);
+
+       LongType yOffset;
+       COORDS2INDEX(currTadRank, currTadStride, sub, yOffset);
+
+       COORDS2INDEX(zTadRank, zTadStride, sub, resultOffset);
+       resultTAD[resultOffset] = dataTAD[yOffset];
+     }
+   }
+   else {
+     for (LongType j = blockIdx.x; j < numTads; j += gridDim.x) {
+       const LongType inputOffset = currOffsets[j];
+       const LongType resultOffset = zOffsets[j];
+
+       T* dataTAD = currData + inputOffset;
+       T* resultTAD = result + resultOffset;
+
+       LongType sub[SD_MAX_RANK];
+       INDEX2COORDS(arrOffset, zTadRank, zTadShapeOf, sub);
+
+       LongType baseOffset;
+       COORDS2INDEX(zTadRank, zTadStride, sub, baseOffset);
+       resultTAD += baseOffset;
+
+       if (zOrder == yOrder) {
+         for (int i = threadIdx.x; i < yLength; i += blockDim.x) {
+           resultTAD[i] = dataTAD[i];
+         }
+       }
+       else {
+         if (shape::order(resultShapeInfo) == shape::order(currTad)) {
+           if (threadIdx.x == 0) {
+             baseIdx = 0;
+             for (int f = 0; f < r; f++) {
+               baseIdx += shape::length(shapeInfoPtrs[f]);
+             }
+           }
+           __syncthreads();
+
+           if (numTads == 1) {
+             for (int k = threadIdx.x; k < yLength; k += blockDim.x) {
+               resultTAD[baseIdx + k] = dataTAD[k];
+             }
+           }
+           else {
+             LongType yIdx[SD_MAX_RANK];
+             for (LongType i = threadIdx.x; i < yLength; i += blockDim.x) {
+               INDEX2COORDS(i, currTadRank, currTadShape, yIdx);
+
+               LongType yOffset;
+               COORDS2INDEX(currTadRank, currTadStride, yIdx, yOffset);
+
+               resultTAD[baseIdx + i] = dataTAD[yOffset];
+             }
+           }
+           __syncthreads();
+         }
+         else {
+           LongType zIdx[SD_MAX_RANK];
+           LongType yIdx[SD_MAX_RANK];
+
+           for (LongType i = threadIdx.x; i < yLength; i += blockDim.x) {
+             INDEX2COORDS(i, currTadRank, currTadShape, yIdx);
+             INDEX2COORDS(i, zTadRank, zTadShapeOf, zIdx);
+
+             LongType yOffset;
+             COORDS2INDEX(currTadRank, currTadStride, yIdx, yOffset);
+
+             LongType rOffset;
+             COORDS2INDEX(zTadRank, zTadStride, zIdx, rOffset);
+             resultTAD[rOffset] = dataTAD[yOffset];
+           }
+         }
+       }
+       __syncthreads();
+     }
+   }
+   __syncthreads();
+ }
+}
+
+template <typename T>
+SD_KERNEL void execConcatKernel(
+   int numArrays,
+   Pointer* data,
+   Pointer* inputShapeInfos,
+   void* vz,
+   LongType* zShapeInfo,
+   Pointer* tadPointers,
+   Pointer* offsetPointers,
+   LongType* zTadShape,
+   LongType* zOffsets) {
+
+ concatKernel<T>(
+     numArrays,
+     data,
+     inputShapeInfos,
+     vz,
+     zShapeInfo,
+     tadPointers,
+     offsetPointers,
+     zTadShape,
+     zOffsets);
+}
+
+template <typename T>
+SD_HOST void concatKernelGeneric(
+   dim3 &launchDims,
+   cudaStream_t *stream,
+   int numArrays,
+   Pointer* data,
+   Pointer* inputShapeInfos,
+   void* vz,
+   LongType* zShapeInfo,
+   Pointer* tadPointers,
+   Pointer* offsetPointers,
+   LongType* zTadShape,
+   LongType* zOffsets) {
+
+ execConcatKernel<T>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         numArrays,
+         data,
+         inputShapeInfos,
+         vz,
+         zShapeInfo,
+         tadPointers,
+         offsetPointers,
+         zTadShape,
+         zOffsets);
+
+ DebugHelper::checkErrorCode(stream, "concatGenericLegacy(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE(
+    void concatKernelGeneric,
+   (dim3 &launchDims,
+    cudaStream_t *stream,
+    int numArrays,
+    sd::Pointer* data,
+    sd::Pointer* inputShapeInfos,
+    void* vz,
+    sd::LongType* zShapeInfo,
+    sd::Pointer* tadPointers,
+    sd::Pointer* offsetPointers,
+    sd::LongType* zTadShape,
+    sd::LongType* zOffsets),
+   SD_COMMON_TYPES);
+
+}  // namespace sd

libnd4j/include/loops/cuda/specials/concatKernelHStack.cu

@@ -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
+// @author Yurii Shyrma, created on 15.11.2018
+//
+#include <loops/special_kernels.h>
+
+namespace sd {
+template <typename T>
+SD_DEVICE void concatKernelHStack(int numArrays, Pointer *data, Pointer *inputShapeInfos, void *vz,
+                                  LongType *zShapeInfo) {
+  // We expect each input array to be a vector, and the result (z) to be a 2D matrix for horizontal stacking.
+  // The row dimension is presumably 1 for each input, with the column dimension the length of each vector,
+  // stacked horizontally.
+
+  auto z = reinterpret_cast<T*>(vz);
+  auto inputData = reinterpret_cast<T**>(data);
+  auto shapes = reinterpret_cast<LongType**>(inputShapeInfos);
+
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+  // Output shape data in shared memory
+  __shared__ int zRank;
+  __shared__ const LongType* zShape;
+  __shared__ const LongType* zStride;
+
+  // Current input array shape data in shared memory
+  __shared__ int inRank;
+  __shared__ const LongType* inShape;
+  __shared__ const LongType* inStride;
+
+  // Working variables in shared memory
+  __shared__ int inputLength;
+  __shared__ int baseIdx;
+
+  // Initialize output shape data once
+  if (threadIdx.x == 0) {
+    zRank = shape::rank(zShapeInfo);
+    zShape = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  // Loop over all input arrays
+  for (int r = blockIdx.x; r < numArrays; r += gridDim.x) {
+    // Cache the current input array's shape data and compute offsets
+    if (threadIdx.x == 0) {
+      // Cache current input shape data
+      inRank = shape::rank(shapes[r]);
+      inShape = shape::shapeOf(shapes[r]);
+      inStride = shape::stride(shapes[r]);
+
+      // Compute base offset
+      baseIdx = 0;
+      for (int f = 0; f < r; f++) {
+        baseIdx += shape::length(shapes[f]);
+      }
+      inputLength = shape::length(shapes[r]);
+    }
+    __syncthreads();
+
+    // Each thread will copy a subset of data
+    for (int i = tid; i < inputLength; i += blockDim.x * gridDim.x) {
+      // Coordinates in the input vector
+      LongType inCoords[SD_MAX_RANK];
+      // Coordinates in the output 2D shape
+      LongType outCoords[SD_MAX_RANK];
+
+      // 1) Get input coordinates using cached shape data
+      INDEX2COORDS(i, inRank, inShape, inCoords);
+      LongType inOffset;
+      COORDS2INDEX(inRank, inStride, inCoords, inOffset);
+
+      // 2) The output coordinate index is baseIdx + i in the horizontal dimension
+      const LongType outIndex = baseIdx + i;
+
+      // Get output coordinates using cached shape data
+      INDEX2COORDS(outIndex, zRank, zShape, outCoords);
+      LongType outOffset;
+      COORDS2INDEX(zRank, zStride, outCoords, outOffset);
+
+      z[outOffset] = inputData[r][inOffset];
+    }
+    __syncthreads();
+  }
+}
+
+template <typename T>
+SD_KERNEL void execConcatKernelHStack(int numArrays, Pointer *data, Pointer *inputShapeInfos, void *vz,
+                                      LongType *zShapeInfo) {
+  concatKernelHStack<T>(numArrays, data, inputShapeInfos, vz, zShapeInfo);
+}
+
+template <typename T>
+SD_HOST void concatKernelHStackGeneric(dim3 &launchDims, cudaStream_t *stream,
+                                       int numArrays,
+                                       Pointer *data,
+                                       Pointer *inputShapeInfos,
+                                       void *vz,
+                                       LongType *zShapeInfo) {
+  execConcatKernelHStack<T>
+  <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      numArrays, data, inputShapeInfos, vz, zShapeInfo);
+  DebugHelper::checkErrorCode(stream, "concatHStack(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE(
+ void concatKernelHStackGeneric,
+(dim3 &launchDims, cudaStream_t *stream, int numArrays, sd::Pointer *data,
+sd::Pointer *inputShapeInfos, void *vz, sd::LongType *zShapeInfo),
+SD_COMMON_TYPES);
+
+}  // namespace sd