chore: import upstream snapshot with attribution

libnd4j/include/ops/declarable/generic/nn/activations/crelu.cpp

@@ -0,0 +1,110 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_crelu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+#include <ops/declarable/helpers/transforms.h>
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(crelu, 1, 1, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);
+
+  REQUIRE_TRUE(x->isR(), 0, "CRELU: input must be real type");
+
+  auto tmp = x->dup(x->ordering());
+  tmp->applyTransform(transform::Neg, tmp);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::concat(block.launchContext(), {x, tmp}, *z, x->rankOf() - 1);
+
+  // TODO: make this configurable?
+  double threshold = 0.0;
+  z->applyScalar(scalar::RELU, threshold, z);
+
+  STORE_RESULT(z);
+
+  delete tmp;
+  return Status::OK;
+}
+
+DECLARE_TYPES(crelu) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+DECLARE_SHAPE_FN(crelu) {
+  auto inShape = inputShape->at(0);
+  std::vector<LongType> shape;
+  for (int e = 0; e < shape::rank(inShape); e++) shape.emplace_back(shape::shapeOf(inShape)[e]);
+
+  shape[shape.size() - 1] *= 2;
+  auto newShape =
+      ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(inShape), shape::order(inShape), shape);
+
+  return SHAPELIST(newShape);
+}
+
+CUSTOM_OP_IMPL(crelu_bp, 2, 1, false, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilonNext = INPUT_VARIABLE(1);
+  auto epsilon = OUTPUT_VARIABLE(0);
+
+  // at first step we build fwd activation
+  crelu op;
+  auto tmpResult = op.evaluate({input});
+  if (tmpResult.status() != Status::OK) return tmpResult.status();
+
+  auto actv = tmpResult.at(0);
+
+  // now we do RELU backward pass
+  helpers::reluDerivative(block.launchContext(), actv, epsilonNext);
+  // now we split updated array into 2 chunks along last dimension
+  concat_bp opc;
+  auto dec = opc.evaluate({input, input, actv}, {-1});
+  if (dec.status() != Status::OK) return dec.status();
+
+  // and now we subtract two parts of epsilons and pass result out
+  auto pos = dec.at(0);
+  auto neg = dec.at(1);
+
+  pos->applyPairwiseTransform(pairwise::Subtract, neg, epsilon);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(crelu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+
+DECLARE_SHAPE_FN(crelu_bp) {
+  auto inShape = inputShape->at(0);
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(inShape)->primary());
+  return ret;
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/cube.cpp

@@ -0,0 +1,61 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_cube)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(cube, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyTransform(transform::Cube, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(cube) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+CONFIGURABLE_OP_IMPL(cube_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+  helpers::cubeDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(cube_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/elu.cpp

@@ -0,0 +1,67 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_elu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(elu, 1, 1, true, -2, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  const auto alpha = block.numT() > 0 ? T_ARG(0) : 1.f;
+
+  input->applyScalar(scalar::ELU, alpha, output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(elu) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(elu_bp, 2, 1, true, -2, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto output = OUTPUT_VARIABLE(0);
+
+  const auto alpha = block.numT() > 0 ? T_ARG(0) : 1.f;
+
+  helpers::eluDerivative(block.launchContext(), input, epsilon, output, alpha);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(elu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/hardsigmoid.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_hardsigmoid)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(hardsigmoid, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyTransform(transform::HardSigmoid, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(hardsigmoid) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(hardsigmoid_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::hardSigmoidDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(hardsigmoid_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/hardtanh.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_hardtanh)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(hardtanh, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyTransform(transform::HardTanh, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(hardtanh) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(hardtanh_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::hardTanhDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(hardtanh_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/identity.cpp

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_identity)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+OP_IMPL(identity, 1, 1, true) {
+  auto z = OUTPUT_VARIABLE(0);
+
+  if (!block.isInplace()) {
+    auto first = INPUT_VARIABLE(0);
+
+    // we hope for memcpy here
+    z->assign(first);
+  }
+
+  return Status::OK;
+}
+DECLARE_SYN(linear, identity);
+
+DECLARE_TYPES(identity) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+OP_IMPL(identity_bp, 2, 1, true) {
+  auto first = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+  auto z = OUTPUT_VARIABLE(0);
+
+  z->assign(epsilon);
+
+  return Status::OK;
+}
+DECLARE_SYN(LinearGrad, identity_bp);
+
+DECLARE_TYPES(identity_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/identity_n.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 sgazeos@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_identity_n)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(identity_n, 1, 1, true, 0, 0) {
+  if (!block.isInplace()) {
+    for (size_t i = 0; i < block.width(); ++i) {
+      auto x = INPUT_VARIABLE(i);
+      auto z = OUTPUT_VARIABLE(i);
+
+      x->applyTransform(transform::Identity, z);
+    }
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(identity_n) {
+  auto shapes = SHAPELIST();
+  for (int i = 0; i < inputShape->size(); ++i) {
+    shapes->push_back(CONSTANT(inputShape->at(i)));
+  }
+  return shapes;
+}
+
+DECLARE_TYPES(identity_n) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes(ANY);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/lrelu.cpp

@@ -0,0 +1,67 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_lrelu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(lrelu, 1, 1, true, -2, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  float alpha = block.numT() > 0 ? T_ARG(0) : 0.01f;
+
+  input->applyScalar(scalar::LeakyRELU, alpha, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lrelu) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(lrelu_bp, 2, 1, true, -2, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  float alpha = block.numT() > 0 ? T_ARG(0) : 0.01f;
+
+  helpers::leakyReluDerivative(block.launchContext(), input, epsilon, z, alpha);
+  return Status::OK;
+}
+
+DECLARE_TYPES(lrelu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/prelu.cpp

@@ -0,0 +1,171 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 24.07.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_prelu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/activations.h>
+
+#include <numeric>
+
+namespace sd {
+namespace ops {
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(prelu, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto alpha = INPUT_VARIABLE(1);
+  auto output = OUTPUT_VARIABLE(0);
+
+  std::vector<LongType> sharedAxes = *block.getIArguments();
+
+  const int inputRank = input->rankOf();
+  const int numSharedAxes = sharedAxes.size();  // can be zero as well
+  const LongType inputLen = input->lengthOf();
+  const LongType alphaLen = alpha->lengthOf();
+  auto* inputShapeVec = input->getShapeAsVector();
+  auto* alphaShapeVec = alpha->getShapeAsVector();
+  const std::vector<LongType> inputShape = *inputShapeVec;
+  const std::vector<LongType> alphaShape = *alphaShapeVec;
+  delete inputShapeVec;
+  delete alphaShapeVec;
+
+  //***** input validation *****//
+  std::vector<LongType> expectedAlphaShape(&inputShape[1], &inputShape[inputRank]);
+
+  REQUIRE_TRUE(inputRank > 1, 0,
+               "PRELU OP: wrong rank of input array, expected rank should be > 1, but got %i instead !", inputRank);
+
+  for (int i = 0; i < numSharedAxes; ++i) {
+    if (sharedAxes[i] <= 0) sharedAxes[i] += inputRank - 1;
+    REQUIRE_TRUE(1 <= sharedAxes[i] && sharedAxes[i] <= inputRank - 1, 0,
+                 "PRELU OP: wrong axis value %i in sharedAxes at position %i, axis value must be within range [1, "
+                 "input_rank-1] !",
+                 sharedAxes[i], i);
+    expectedAlphaShape[sharedAxes[i] - 1] = 1;
+  }
+
+
+  NDArray *alpha2 =  alphaShape != expectedAlphaShape ? alpha->reshape(alpha->ordering(), expectedAlphaShape) : alpha;
+  helpers::prelu(block.launchContext(), input,
+                 alpha2,
+                 output);
+
+  if(alpha2 != alpha) delete alpha2;
+  return Status::OK;
+}
+
+DECLARE_TYPES(prelu) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(prelu_bp, 3, 2, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto alpha = INPUT_VARIABLE(1);
+  auto dLdO = INPUT_VARIABLE(2);
+
+  auto dLdI = OUTPUT_VARIABLE(0);
+  auto dLdA = OUTPUT_VARIABLE(1);
+
+  std::vector<LongType> sharedAxes = *block.getIArguments();
+
+  const int inputRank = input->rankOf();
+  const int numSharedAxes = sharedAxes.size();  // can be zero as well
+  const LongType inputLen = input->lengthOf();
+  const LongType alphaLen = alpha->lengthOf();
+  auto* inputShapeVec = input->getShapeAsVector();
+  auto* alphaShapeVec = alpha->getShapeAsVector();
+  const std::vector<LongType> inputShape = *inputShapeVec;
+  const std::vector<LongType> alphaShape = *alphaShapeVec;
+  delete inputShapeVec;
+  delete alphaShapeVec;
+
+  //***** input validation *****//
+
+  // temporary limitation imposed by Yurii
+  REQUIRE_TRUE(inputRank <= SD_MAX_RANK / 2, 0, "rank of input array should be <= SD_MAX_RANK/2, but got %i instead!",
+               inputRank);
+  REQUIRE_TRUE(input->lengthOf() / alpha->lengthOf() <= SD_MAX_RANK * 2, 0,
+               "the length of input array should be no more than SD_MAX_RANK*2 times the alpha array length, but got "
+               "%lld and %lld correspondingly!",
+               input->lengthOf(), alpha->lengthOf());
+
+  std::vector<LongType> expectedAlphaShape(&inputShape[1], &inputShape[inputRank]);
+
+  REQUIRE_TRUE(inputRank > 1, 0,
+               "PRELU_BP OP: wrong rank of input array, expected rank should be > 1, but got %i instead !", inputRank);
+
+  for (int i = 0; i < numSharedAxes; ++i) {
+    if (sharedAxes[i] <= 0) sharedAxes[i] += inputRank - 1;
+    REQUIRE_TRUE(1 <= sharedAxes[i] && sharedAxes[i] <= inputRank - 1, 0,
+                 "PRELU_BP OP: wrong axis value %i in sharedAxes at position %i, axis value must be within range [1, "
+                 "input_rank-1] !",
+                 sharedAxes[i], i);
+    expectedAlphaShape[sharedAxes[i] - 1] = 1;
+  }
+
+  LongType product = 1;
+  for (const auto& item : expectedAlphaShape) product *= item;
+
+  REQUIRE_TRUE(product == alphaLen, 0, "PRELU_BP OP: wrong shape of alpha array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedAlphaShape).c_str(), ShapeUtils::shapeAsString(alphaShape).c_str());
+  // ***** end of validation ***** //
+
+  NDArray* alphaReshaped = nullptr;
+  NDArray* dLdAReshaped = nullptr;
+  
+  if (alphaShape != expectedAlphaShape) {
+    alphaReshaped = alpha->reshape(alpha->ordering(), expectedAlphaShape);
+    dLdAReshaped = dLdA->reshape(dLdA->ordering(), expectedAlphaShape);
+  }
+
+  helpers::preluBP(block.launchContext(), input, 
+                   alphaReshaped != nullptr ? alphaReshaped : alpha, 
+                   dLdO, dLdI, 
+                   dLdAReshaped != nullptr ? dLdAReshaped : dLdA);
+
+  if (alphaReshaped != nullptr) {
+    delete alphaReshaped;
+    delete dLdAReshaped;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(prelu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedInputTypes(2, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(1, {FLOAT32, DOUBLE, HALF});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/rationaltanh.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_rationaltanh)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(rationaltanh, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyTransform(transform::RationalTanh, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(rationaltanh) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(rationaltanh_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::rationalTanhDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(rationaltanh_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/rectifiedtanh.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_rectifiedtanh)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(rectifiedtanh, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyTransform(transform::RectifiedTanh, output);
+  STORE_RESULT(output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(rectifiedtanh) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(rectifiedtanh_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::rectifiedTanhDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(rectifiedtanh_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/relu.cpp

@@ -0,0 +1,65 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_relu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(relu, 1, 1, true, 1, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  auto scalar = block.numT() > 0 ? block.getTArguments()->at(0) : 0.0;
+
+  first->applyScalar(scalar::RELU, scalar, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(relu) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+CONFIGURABLE_OP_IMPL(relu_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+  helpers::reluDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+DECLARE_SYN(ReluGrad, relu_bp);
+
+DECLARE_TYPES(relu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/relu6.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma, created on 16.02.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_relu6)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(relu6, 1, 1, true, 1, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  input->applyScalar(scalar::RELU6, T_ARG(0), output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(relu6) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(relu6_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto gradO = INPUT_VARIABLE(1);
+  auto gradI = OUTPUT_VARIABLE(0);
+
+  helpers::relu6Derivative(block.launchContext(), input, gradO, gradI);
+  return Status::OK;
+}
+
+DECLARE_TYPES(relu6_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/selu.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_selu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(selu, 1, 1, true, 0, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  first->applyTransform(transform::SELU, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(selu) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(selu_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+  helpers::seluDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(selu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/sigmoid.cpp

@@ -0,0 +1,64 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_sigmoid)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(sigmoid, 1, 1, true, 0, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  first->applyTransform(transform::Sigmoid, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(sigmoid) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(sigmoid_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::sigmoidDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+
+DECLARE_TYPES(sigmoid_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/softplus.cpp

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_softplus)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(softplus, 1, 1, true, 0, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  first->applyTransform(transform::SoftPlus, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(softplus) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(softplus_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+
+  helpers::softPlusDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+DECLARE_SYN(SoftplusGrad, softplus_bp);
+
+DECLARE_TYPES(softplus_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/softsign.cpp

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_softsign)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(softsign, 1, 1, true, 0, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  first->applyTransform(transform::SoftSign, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(softsign) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(softsign_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+  auto z = OUTPUT_VARIABLE(0);
+  helpers::softSignDerivative(block.launchContext(), input, epsilon, z);
+
+  return Status::OK;
+}
+DECLARE_SYN(SoftsignGrad, softsign_bp);
+
+DECLARE_TYPES(softsign_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/tanh.cpp

@@ -0,0 +1,66 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_tanh)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(tanh, 1, 1, true, 0, 0) {
+  auto first = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  first->applyTransform(transform::Tanh, z);
+
+  STORE_RESULT(*z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(tanh) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(tanh_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto epsilon = INPUT_VARIABLE(1);
+
+
+  auto z = OUTPUT_VARIABLE(0);
+  helpers::tanhDerivative(block.launchContext(), input, epsilon, z);
+  return Status::OK;
+}
+DECLARE_SYN(TanhGrad, tanh_bp);
+
+DECLARE_TYPES(tanh_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/activations/thresholdedrelu.cpp

@@ -0,0 +1,69 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 24.07.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_thresholdedrelu)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/activations.h>
+#include <ops/declarable/helpers/legacy_helpers.h>
+namespace sd {
+namespace ops {
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(thresholdedrelu, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  auto scalar = block.numT() > 0 ? block.getTArguments()->at(0) : 0.0;
+
+  helpers::thresholdRelu(block.launchContext(), input, scalar, output);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(thresholdedrelu) { getOpDescriptor()->setAllowedInputTypes(0, ANY)->setSameMode(true); }
+
+////////////////////////////////////////////////////////////////////////
+CONFIGURABLE_OP_IMPL(thresholdedrelu_bp, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto dLdO = INPUT_VARIABLE(1);
+
+  auto dLdI = OUTPUT_VARIABLE(0);
+  auto threshold = block.numT() > 0 ? block.getTArguments()->at(0) : 0.0;
+
+  helpers::thresholdReluDerivative(block.launchContext(), input, threshold, dLdO, dLdI);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(thresholdedrelu_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
+      ->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/apply_sgd.cpp

@@ -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
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_apply_sgd)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/gradient.h>
+
+namespace sd {
+namespace ops {
+CONFIGURABLE_OP_IMPL(apply_sgd, 2, 1, true, -2, 0) {
+  auto parameters = INPUT_VARIABLE(0);
+  auto gradients = INPUT_VARIABLE(1);
+
+  double lr = 0.0;
+
+  REQUIRE_TRUE(
+      parameters->isSameShape(gradients), 0,
+      "ApplySGD: parameters and gradients should have the same shape, but got parameters = %s and gradients = %s !",
+      ShapeUtils::shapeAsString(parameters).c_str(), ShapeUtils::shapeAsString(gradients).c_str());
+
+  if (block.width() == 3) {
+    auto tarr = INPUT_VARIABLE(2);
+    lr = tarr->e<double>(0);
+  } else if (block.getTArguments()->size() == 1) {
+    lr = T_ARG(0);
+  } else {
+    REQUIRE_TRUE(false, 0, "ApplyGradients op should have LR announced either es T argument or additional NDArray!");
+  }
+
+  auto Z = OUTPUT_VARIABLE(0);
+
+  helpers::applyGradientDescent(block.launchContext(), parameters, gradients, lr, Z);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(ApplyGradientDescent, apply_sgd);
+}  // namespace ops
+
+DECLARE_TYPES(apply_sgd) { getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS})->setAllowedOutputTypes({ALL_FLOATS}); }
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/batchnorm.cpp

@@ -0,0 +1,335 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// @author raver119@gmail.com, created on 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_batchnorm)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/batchnorm.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(batchnorm, 3, 1, false, 1, 2) {
+  auto input = INPUT_VARIABLE(0);
+  auto mean = INPUT_VARIABLE(1);
+  auto variance = INPUT_VARIABLE(2);
+  NDArray* gamma = nullptr;
+  NDArray* beta = nullptr;
+
+  auto output = OUTPUT_VARIABLE(0);
+
+  const bool applyScale = (bool)INT_ARG(0);
+  const bool applyOffset = (bool)INT_ARG(1);
+  const double epsilon = T_ARG(0);
+
+  if (applyScale) gamma = INPUT_VARIABLE(3);
+  if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
+
+  const int numOfIntArgs = block.getIArguments()->size();
+  const int inRank = input->rankOf();
+
+  // get axes args to normalize input array over
+  std::vector<sd::LongType> axes;
+  if (numOfIntArgs > 2)
+    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
+  else
+    axes.push_back(inRank - 1);  // default dimension to reduce along is last dimension
+
+  const sd::LongType numOfAxes = axes.size();
+  REQUIRE_TRUE(numOfAxes <= inRank, 0,
+               "BATCHNORM op: too big number of input axes to normalize over, expected number should be less or equal "
+               "to rank of input array, but got %i and %i correspondingly !",
+               numOfAxes, inRank);
+
+  // evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
+  // for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
+  // {3}, then expected shape would be {5}
+  std::vector<sd::LongType> expShape;
+  if (numOfAxes == 1)
+    expShape.push_back(input->sizeAt(axes[0]));
+  else {  // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
+    expShape = std::vector<sd::LongType>(inRank, 1);
+    for (sd::LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
+  }
+
+  REQUIRE_TRUE(mean->isSameShape(expShape), 0,
+               "BATCHNORM op: wrong shape of mean array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
+  REQUIRE_TRUE(variance->isSameShape(expShape), 0,
+               "BATCHNORM op: wrong shape of variance array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
+  if (gamma)
+    REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
+                 "BATCHNORM op: wrong shape of gamma array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
+  if (beta)
+    REQUIRE_TRUE(beta->isSameShape(expShape), 0,
+                 "BATCHNORM op: wrong shape of beta array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
+
+  // types of all input arrays should be the same
+  for (unsigned long i = 1; i < block.width(); ++i)
+    REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
+                 "BATCHNORM op: types of all input arrays should be the same !");
+
+  sd_debug("MKL-DNN is not used for batchnorm!\n", 0);
+
+  // formula: output = gamma * ((input - mean) / sqrt(variance + epsilon)) + beta
+  helpers::batchnorm(input, mean, variance, gamma, beta, output, axes, epsilon);
+
+
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(batchnorm) { getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS})->setSameMode(true); }
+
+DECLARE_SHAPE_FN(batchnorm) {
+  auto inShapeInfo = inputShape->at(0);
+  DataType outType = DataTypeUtils::pickFloatingType(ArrayOptions::dataType(inShapeInfo));
+
+  auto outShapeInfo = ShapeBuilders::copyShapeInfoAndType(
+      inShapeInfo, outType, false, block.getWorkspace());  // output shape is identical to input shape
+
+  return SHAPELIST(CONSTANT(outShapeInfo));
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(batchnorm_bp, 4, 3, false, 1, 2) {
+  NDArray* input = INPUT_VARIABLE(0);
+  NDArray* mean = INPUT_VARIABLE(1);
+  NDArray* variance = INPUT_VARIABLE(2);
+  NDArray* gamma = nullptr;
+  NDArray* beta = nullptr;
+  NDArray* dLdO = INPUT_VARIABLE(block.width() - 1);  // next epsilon
+
+  NDArray* dLdI = OUTPUT_VARIABLE(0);
+  NDArray* dLdM = OUTPUT_VARIABLE(1);
+  NDArray* dLdV = OUTPUT_VARIABLE(2);
+  NDArray* dLdG = nullptr;
+  NDArray* dLdB = nullptr;
+
+  const bool applyScale = (bool)INT_ARG(0);
+  const bool applyOffset = (bool)INT_ARG(1);
+  const float epsilon = T_ARG(0);
+
+  if (applyScale) {
+    gamma = INPUT_VARIABLE(3);
+    dLdG = OUTPUT_VARIABLE(3);
+  }
+  if (applyOffset) {
+    beta = INPUT_VARIABLE(3 + (int)applyScale);
+    dLdB = OUTPUT_VARIABLE(3 + (int)applyScale);
+  }
+
+  const int numOfIntArgs = block.getIArguments()->size();
+  const int inRank = input->rankOf();
+
+  // get axes args to normalize input array over
+  std::vector<LongType> axes;
+  if (numOfIntArgs > 2)
+    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
+  else
+    axes.push_back(inRank - 1);  // default dimension to reduce along is last dimension
+
+  const sd::LongType numOfAxes = axes.size();
+  REQUIRE_TRUE(numOfAxes <= inRank, 0,
+               "BATCHNORM_BP op: too big number of input axes to normalize over, expected number should be less or "
+               "equal to rank of input array, but got %i and %i correspondingly !",
+               numOfAxes, inRank);
+
+  // evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
+  // for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
+  // {3}, then expected shape would be {5}
+  std::vector<sd::LongType> expShape;
+  if (numOfAxes == 1)
+    expShape.push_back(input->sizeAt(axes[0]));
+  else {  // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
+    expShape = std::vector<sd::LongType>(inRank, 1);
+    for (sd::LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
+  }
+
+  REQUIRE_TRUE(mean->isSameShape(expShape), 0,
+               "BATCHNORM_BP op: wrong shape of mean array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
+  REQUIRE_TRUE(variance->isSameShape(expShape), 0,
+               "BATCHNORM_BP op: wrong shape of variance array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
+  if (gamma)
+    REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
+                 "BATCHNORM_BP op: wrong shape of gamma array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
+  if (beta)
+    REQUIRE_TRUE(beta->isSameShape(expShape), 0,
+                 "BATCHNORM_BP op: wrong shape of beta array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
+
+  REQUIRE_TRUE(input->isSameShape(dLdO), 0,
+               "BATCHNORM_BP op: wrong shape of output gradients array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(dLdO).c_str());
+
+  // types of all input arrays should be the same (except dLdO)
+  for (unsigned long i = 1; i < block.width() - 2; ++i)
+    REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
+                 "BATCHNORM_BP op: types of arrays (input, mean, variance, gamma, beta) should be the same !");
+
+  // ***** calculations ***** //
+
+  // notations:
+  // f = g * (gamma * ((x - m) / (v + eps)^0.5) + beta) -> means dLdO * ff_output, g = dLdO
+  // stdInv = 1 / (v + eps)^0.5
+  // N - batch size (product of spatial dimensions)
+
+  // derivatives:
+  // dLdI = dfdx + dfdm*dmdx + dfdv*(dvdm*dmdx + dvdx)
+
+  // dfdx =  gamma*stdInv*g;
+  // dfdm = -gamma*stdInv*g_sum;
+  // dmdx  = 1/N;
+  // dvdx  = 2 *  (x - m) / N
+  // dvdm  = -2 * [(x - m)]_sum / N
+  // dfdv  = -0.5 * [g*(x - m)]_sum * stdInv^3, drop gamma here for calc convenience
+
+  // finally:
+  // dLdI = gamma * (  stdInv * (g - g_sum/N) + (2/N) * dfdv * (dvdm/2  + (x - m))  )
+
+  // dLdG = (g * (x - m))_sum * stdInv
+  // dLdB = g_sum
+
+  // variance = input->varianceAlongDimension(variance::SummaryStatsVariance, false,
+
+
+  const auto excludedAxes = ShapeUtils::evalDimsToExclude(inRank, axes.size(),axes.data());
+  const bool keepUnitiesInShape = inRank == mean->rankOf();
+
+  // inverse batch size 1/N
+  const float Ninv = 1.f * shape::tadLength(input->shapeInfo(), (axes.data()), axes.size()) / input->lengthOf();
+
+  // input - mean
+  NDArray xMinusMean(input);  // empty array with same shape as input
+  input->applyBroadcast(sd::broadcast::Subtract, &axes, mean, &xMinusMean);
+
+  // stdInv = 1 / (variance + epsilon)^0.5
+  NDArray* stdInv = (*variance) + epsilon;
+  stdInv->applyTransform(transform::Reciprocal, stdInv);  // 1 / (variance + epsilon)
+  stdInv->applyTransform(transform::Sqrt, stdInv);        // 1 / (variance + epsilon)^0.5
+
+  // dvdm (use dLdM as storage for dvdm)
+  xMinusMean.reduceAlongDimension(sd::reduce::Sum, dLdM, excludedAxes, keepUnitiesInShape);
+  *dLdM *= -Ninv;
+
+  // g_sum
+  auto* gSum = dLdO->reduceAlongDimension(sd::reduce::Sum, excludedAxes, keepUnitiesInShape);
+
+  // dLdB
+  if (applyOffset) dLdB->assign(gSum);
+
+  // stdInv * (g - g_sum/N) (use dLdI as storage for this expression)
+  *gSum *= Ninv;
+  dLdO->applyBroadcast(sd::broadcast::Subtract, &axes, gSum, dLdI);
+  delete gSum;
+
+  dLdI->applyBroadcast(sd::broadcast::Multiply, &axes, stdInv, dLdI);
+
+  // dLdV <- [g*(x - m)]_sum
+  auto* xMinusMeanTimesDLdO = xMinusMean * (*dLdO);
+  xMinusMeanTimesDLdO->reduceAlongDimension(sd::reduce::Sum, dLdV, excludedAxes, keepUnitiesInShape);
+  delete xMinusMeanTimesDLdO;
+
+  // dLdG
+  *dLdV *= (*stdInv);
+  if (applyScale) dLdG->assign(dLdV);
+
+  // (2 / N) * dfdv (use dLdV as storage for dfdv)
+  // dLdV *= stdInv * stdInv becomes dLdV *= stdInv^2
+  auto* stdInvSquared = (*stdInv) * (*stdInv);
+  *dLdV *= (*stdInvSquared);
+  delete stdInvSquared;
+  *dLdV *= -Ninv;            // -0.5f * (2 / N);
+
+  // dfdv * (dvdm  + (x - m)) (use xMinusMean as storage for this expression)
+  xMinusMean.applyBroadcast(sd::broadcast::Add, &axes, dLdM, &xMinusMean);
+  xMinusMean.applyBroadcast(sd::broadcast::Multiply, &axes, dLdV, &xMinusMean);
+
+  // dLdI
+  *dLdI += xMinusMean;
+  if (applyScale) dLdI->applyBroadcast(sd::broadcast::Multiply, &axes, gamma, dLdI);
+
+  *dLdM = 0;  // put zeros so far
+  *dLdV = 0;  // put zeros so far
+
+  delete stdInv;
+  delete excludedAxes;
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(batchnorm_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, sd::DataType::ANY)
+      ->setAllowedInputTypes(1, sd::DataType::ANY)
+      ->setAllowedInputTypes(2, sd::DataType::ANY)
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedInputTypes(4, sd::DataType::ANY)
+      ->setAllowedInputTypes(5, sd::DataType::ANY)
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+DECLARE_SHAPE_FN(batchnorm_bp) {
+  sd::LongType * inShapeInfo = inputShape->at(0);
+  sd::LongType * meanShapeInfo = inputShape->at(1);
+
+  const bool applyScale = (bool)INT_ARG(0);
+  const bool applyOffset = (bool)INT_ARG(1);
+
+  DataType outType = DataTypeUtils::pickFloatingType(ArrayOptions::dataType(inShapeInfo));
+
+  auto shapes = SHAPELIST();
+
+  // dLdI shapeInfo
+  shapes->push_back(ConstantShapeHelper::getInstance().createShapeInfo(outType, inShapeInfo));
+
+  // dLdM shapeInfo
+  shapes->push_back(ConstantShapeHelper::getInstance().createShapeInfo(outType, meanShapeInfo));
+
+  // dLdV shapeInfo (same as dLdM)
+  shapes->push_back(shapes->at(shapes->size() - 1));
+
+  // dLdG shapeInfo (same as dLdM)
+  if (applyScale) shapes->push_back(shapes->at(shapes->size() - 1));
+
+  // dLdB shapeInfo (same as dLdM)
+  if (applyOffset) shapes->push_back(shapes->at(shapes->size() - 1));
+
+  return shapes;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/bias_add.cpp

@@ -0,0 +1,112 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_biasadd)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/addBias.h>
+
+namespace sd {
+namespace ops {
+
+////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(biasadd, 2, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto bias = INPUT_VARIABLE(1);
+
+  auto output = OUTPUT_VARIABLE(0);
+
+  const bool isNCHW = !block.getBArguments()->empty() ? B_ARG(0) : false;
+  const sd::LongType channelDim = isNCHW ? 1 : input->rankOf() - 1;  // second or last
+
+  REQUIRE_TRUE(bias->rankOf() == 1, 0, "BIASADD CUSTOM_OP: bias array should have rank = 1, but got %i instead !",
+               bias->rankOf());
+
+  REQUIRE_TRUE(bias->sizeAt(0) == input->sizeAt(channelDim), 0,
+               "BIASADD CUSTOM_OP: shapes of bias %s and input %s arrays are not suitable for broadcast operation "
+               "along channel dimension %i !",
+               ShapeUtils::shapeAsString(bias).c_str(), ShapeUtils::shapeAsString(input).c_str(), channelDim);
+
+  REQUIRE_TRUE(output->isSameShape(input), 0,
+               "BIASADD CUSTOM_OP: wrong shape of output array, expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(output).c_str());
+
+  helpers::addBias(block, *input, *bias, *output, isNCHW);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(bias_add, biasadd);
+
+////////////////////////////////////////////////////////////////////
+DECLARE_SHAPE_FN(biasadd) {
+  auto xShape = inputShape->at(0);
+  auto yShape = inputShape->at(1);
+
+  auto dtype = ArrayOptions::dataType(yShape);
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().castToDataType(xShape, dtype));
+  return ret;
+}
+
+DECLARE_TYPES(biasadd) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(biasadd_bp, 3, 2, false, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto bias = INPUT_VARIABLE(1);
+  auto gradO = INPUT_VARIABLE(2);
+
+  auto gradI = OUTPUT_VARIABLE(0);
+  auto gradB = OUTPUT_VARIABLE(1);
+
+  const bool isNCHW = !block.getBArguments()->empty() ? B_ARG(0) : false;
+  const int channelDim = isNCHW ? 1 : input->rankOf() - 1;  // second or last
+
+  gradI->assign(gradO);
+
+  std::vector<sd::LongType> channel;
+  channel.push_back(channelDim);
+  auto dims = ShapeUtils::evalDimsToExclude(gradO->rankOf(), 1,channel.data());
+  gradO->reduceAlongDimension(sd::reduce::Sum, gradB, dims);
+  delete dims;
+  return sd::Status::OK;
+}
+DECLARE_SYN(BiasAddGrad, biasadd_bp);
+
+////////////////////////////////////////////////////////////////////
+DECLARE_SHAPE_FN(biasadd_bp) {
+  auto input = inputShape->at(0);
+  auto bias = inputShape->at(1);
+  return SHAPELIST(CONSTANT(input), CONSTANT(bias));
+}
+
+DECLARE_TYPES(biasadd_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/col2im.cpp

@@ -0,0 +1,94 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 17.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_col2im)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/col2im.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(col2im, 1, 1, false, 0, 9) {
+  auto x = INPUT_VARIABLE(0);
+  auto z = OUTPUT_NULLIFIED(0);
+
+  REQUIRE_TRUE(x->rankOf() == 6, 0, "col2im input should be 6D, but got %i instead", x->rankOf());
+  REQUIRE_TRUE(z->rankOf() == 4, 0, "col2im output should be 4D, but got %i instead", z->rankOf());
+
+  LongType strideY = INT_ARG(0);
+  LongType strideX = INT_ARG(1);
+  LongType padHeight = INT_ARG(2);
+  LongType padWidth = INT_ARG(3);
+  LongType imgHeight = INT_ARG(4);
+  LongType imgWidth = INT_ARG(5);
+  LongType dY = INT_ARG(6);  // Dilation in height/y dimension
+  LongType dX = INT_ARG(7);  // Dilation in width/x dimension
+
+  LaunchContext* ctx = block.launchContext();
+  helpers::col2im(*ctx, x, z, strideY, strideX, padHeight, padWidth, imgHeight, imgWidth, dY, dX);
+
+  return Status::OK;
+}
+DECLARE_SHAPE_FN(col2im) {
+  auto inShape = inputShape->at(0);
+
+  LongType bS = shape::shapeOf(inShape)[0];
+  LongType iD = shape::shapeOf(inShape)[1];
+
+  LongType sY = INT_ARG(0);
+  LongType sX = INT_ARG(1);
+  LongType pY = INT_ARG(2);
+  LongType pX = INT_ARG(3);
+  LongType inY = INT_ARG(4);
+  LongType inX = INT_ARG(5);
+  LongType dY = INT_ARG(6);  // Dilation, height/y dimension
+  LongType dX = INT_ARG(7);  // Dilation, width/x dimension
+  bool isSameMode = INT_ARG(8) > 0;
+
+  LongType* zShape;
+  ALLOCATE(zShape, block.getWorkspace(), shape::shapeInfoLength(4), sd::LongType);
+
+  zShape[0] = 4;
+  zShape[1] = bS;
+  zShape[2] = iD;
+  zShape[3] = inY;
+  zShape[4] = inX;
+
+  zShape[shape::shapeInfoLength(zShape) - 2] = 1;
+  zShape[shape::shapeInfoLength(zShape) - 1] = 99;
+
+  ShapeUtils::updateStridesAndType(zShape, inShape, 'c');
+
+  return SHAPELIST(CONSTANT(zShape));
+}
+
+DECLARE_TYPES(col2im) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedOutputTypes(0, INHERIT)
+      ->setSameMode(true);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/conv1d.cpp

@@ -0,0 +1,394 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//  @author Yurii Shyrma
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_conv1d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/DeclarableOp.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(conv1d, 2, 1, false, 0, 5) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iW, iC] (NWC) or [bS, iC, iW] (NCW)
+  auto weights = INPUT_VARIABLE(1);                             // [kW, iC, oC], [oC, iC, kW], [oC, kW, iC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oW, oC] (NWC) or [bS, oC, oW] (NCW)sa
+
+
+  LongType kW = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) width
+  LongType sW = INT_ARG(1);                                                                          // strides width
+  LongType pW = INT_ARG(2);                                                                          // paddings width
+  LongType dW = INT_ARG(3);                                                                          // dilations width
+  LongType paddingMode = INT_ARG(4);                                                                 // 0-VALID, 1-SAME
+  /**
+   * TODO: fix java -> c++ NCW/NWC conversion.
+   */
+  LongType isNCW = block.getIArguments()->size() > 5 ? INT_ARG(5) : 1;  // INT_ARG(4): 1-NWC, 0-NCW
+  LongType originalNCW = isNCW;
+  //normally nchw is 0 and 1 being passed in, we're using it as a boolean here
+  //so we want it to be whether nchw is 0 or not.
+  isNCW = isNCW == 0;
+
+
+    LongType wFormat =
+        block.getIArguments()->size() > 6 ? INT_ARG(6) : 0;  // 0 - [kW, iC, oC], 1 - [oC, iC, kW], 2 - [oC, kW, iC]
+
+
+  LongType bS = input->sizeAt(0);             // batch size
+  LongType   iC = ConvolutionUtils::inChannels(weights->shapeInfo(), wFormat);
+  LongType    iW = ConvolutionUtils::inputWidth(input->shapeInfo(), isNCW);
+  LongType    oC = ConvolutionUtils::outChannels(weights->shapeInfo(), wFormat);
+  LongType   oW = ConvolutionUtils::calcOutDimConv(iW,kW,sW,pW,dW,paddingMode);  // batch size, input channels, input height/width, output channels, output height/width;
+
+
+  std::vector<LongType> reshapeForInput, reshapeForOutput;
+  if (!isNCW) {
+    reshapeForInput = {bS, 1, iW, iC};      // [bS, iW, iC] -> [bS, 1, iW, iC]
+    reshapeForOutput = {bS, 1, oW, oC};  // [bS, oW, oC] -> [bS, 1, oW, oC]
+  } else {
+    reshapeForInput = {bS, iC, 1, iW};      // [bS, iC, iW] -> [bS, iC, 1, iW]
+    reshapeForOutput = {bS,oC, 1, oW};  // [bS, oC, oW] -> [bS, oC, 1, oW]
+  }
+
+  auto inputReshaped = input->reshape(input->ordering(), reshapeForInput,false);
+  auto outputReshaped = output->reshape(output->ordering(), reshapeForOutput, false);
+  std::vector<LongType> weightsShape = {1, weights->sizeAt(0), weights->sizeAt(1), weights->sizeAt(2)};
+  auto weightsReshaped = weights->reshape(
+      weights->ordering(),
+      weightsShape,false);  // [kW, iC, oC] -> [1, kW, iC, oC]
+
+
+
+  conv2d conv2d;
+  Status ret = Status::OK;
+  if(bias == nullptr) {
+    //note this might look strange but we get a segfault otherwise.
+    //this problem was actually the source of a very strange JVM hang.
+    ret = conv2d.execute({inputReshaped, weightsReshaped}, {outputReshaped}, {},
+                         {1, kW, 1, sW, 0, pW, 1, dW, paddingMode, originalNCW}, {});
+
+    output->assign(outputReshaped);
+
+  } else {
+    ret = conv2d.execute({inputReshaped, weightsReshaped, bias}, {outputReshaped}, {},
+                         {1, kW, 1, sW, 0, pW, 1, dW, paddingMode, originalNCW}, {});
+
+    output->assign(outputReshaped);
+
+  }
+
+
+  return ret;
+}
+
+DECLARE_SHAPE_FN(conv1d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  LongType wFormat =
+      block.getIArguments()->size() > 6 ? INT_ARG(6) : 0;  // 0 - [kW, iC, oC], 1 - [oC, iC, kW], 2 - [oC, kW, iC]
+
+
+  LongType kW = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo,0));  // filter(kernel) width
+  LongType sW = INT_ARG(1);                                                                          // strides width
+  LongType pW = INT_ARG(2);                                                                          // paddings width
+  LongType dW = INT_ARG(3);                                                                          // dilations width
+  LongType paddingMode = INT_ARG(4);                                                                 // 0-VALID, 1-SAME
+  LongType isNCW = block.getIArguments()->size() > 5 ? INT_ARG(5) : 1;  // INT_ARG(4): 1-NWC, 0-NCW
+
+  //normally nchw is 0 and 1 being passed in, we're using it as a boolean here
+  //so we want it to be whether nchw is 0 or not.
+  isNCW = isNCW == 0;
+
+  const LongType rank = 3;  // 4
+
+  LongType bS = shape::sizeAt(inputShapeInfo, 0);             // batch size
+  LongType   iC = ConvolutionUtils::inChannels(weightsShapeInfo, wFormat);
+  LongType    iW = ConvolutionUtils::inputWidth(inputShapeInfo, isNCW);
+  LongType    oC = ConvolutionUtils::outChannels(weightsShapeInfo, wFormat);
+  LongType   oW = ConvolutionUtils::calcOutDimConv(iW,kW,sW,pW,dW,paddingMode);  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType* outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);
+
+  outputShapeInfo[0] = 3;
+  outputShapeInfo[1] = bS;
+  if (isNCW) {
+    outputShapeInfo[2] = oC;
+    outputShapeInfo[3] = oW;
+  } else {
+    outputShapeInfo[2] = oW;
+    outputShapeInfo[3] = oC;
+  }
+
+  sd::LongType * second = shape::calcStridesFortran(outputShapeInfo,shape::rank(outputShapeInfo));
+  shape::setStride(outputShapeInfo,second);
+  shape::setOrder(outputShapeInfo, 'f');
+  ArrayOptions::setDataType(outputShapeInfo, ArrayOptions::dataType(inputShapeInfo));
+  delete[] second;
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+
+}
+
+DECLARE_TYPES(conv1d) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, {ALL_FLOATS, ALL_INTS, QINT8, QINT16})
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(conv1d_bp, 3, 2, false, 0, 5) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iW, iC] (NWC) or [bS, iC, iW] (NCW)
+  auto weights = INPUT_VARIABLE(1);                             // [kW, iC, oC], [oC, iC, kW], [oC, kW, iC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto gradO = block.width() > 3 ? INPUT_VARIABLE(3)
+                                 : INPUT_VARIABLE(2);  // [bS, oW, oC] (NWC) or [bS, oC, oW] (NCW), epsilon_next
+
+  auto gradI = OUTPUT_NULLIFIED(0);                                // [bS, iW, iC] (NWC) or [bS, iC, iW] (NCW), epsilon
+  auto gradW = OUTPUT_NULLIFIED(1);                                // [kW, iC, oC], [oC, iC, kW], [oC, kW, iC]
+  auto gradB = block.width() > 3 ? OUTPUT_NULLIFIED(2) : nullptr;  // [oC]
+
+  LongType kW = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) width
+  LongType sW = INT_ARG(1);                                                          // strides width
+  LongType pW = INT_ARG(2);                                                          // paddings width
+  LongType dW = INT_ARG(3);                                                          // dilations width
+  LongType paddingMode = INT_ARG(4);                                                 // 0-VALID, 1-SAME, 2-CAUSAL
+  LongType isNCW = block.getIArguments()->size() > 5 ? !INT_ARG(5) : 1;              // INT_ARG(4): 1-NWC, 0-NCW
+  LongType wFormat =
+      block.getIArguments()->size() > 6 ? INT_ARG(6) : 0;  // 0 - [kW, iC, oC], 1 - [oC, iC, kW], 2 - [oC, kW, iC]
+
+  const LongType rank = 3;
+  REQUIRE_TRUE(input->rankOf() == rank, 0,
+               "CUSTOM CONV1D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == rank, 0,
+               "CUSTOM CONV1D_BP OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weights->rankOf());
+
+  LongType indIOioC, indIiW, indWoC(0 == wFormat ? 2 : 0);
+  if (!isNCW) {
+    indIOioC = 2;
+    indIiW = 1;
+  } else {
+    indIOioC = 1;
+    indIiW = 2;
+  }
+
+  const LongType bS = input->sizeAt(0);         // batch size
+  const LongType iW = input->sizeAt(indIiW);    // input width
+  const LongType iC = input->sizeAt(indIOioC);  // input channels
+  const LongType oC = weights->sizeAt(indWoC);  // output channels
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, 1, kW, 1, sW, 0, pW, 1, dW, 1, iW, paddingMode);
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoW, 0, indIOioC, indIiW});
+  std::vector<LongType> expectedWeightsShape =
+      0 == wFormat ? std::vector<LongType>({kW, iC, oC})
+                   : (1 == wFormat ? std::vector<LongType>({oC, iC, kW}) : std::vector<LongType>({oC, kW, iC}));
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM CONV1D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+  REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+               "CUSTOM CONV1D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, bias->rankOf(), bias->lengthOf());
+
+
+  std::vector<LongType> reshapeForInput, reshapeForGradO;
+  if (!isNCW) {
+    if(!gradO->isScalar()) {
+      reshapeForGradO = {gradO->sizeAt(0), 1, gradO->sizeAt(1), gradO->sizeAt(2)};  // [bS, oW, oC] -> [bS, 1, oW, oC]
+      reshapeForInput = {input->sizeAt(0), 1, input->sizeAt(1), input->sizeAt(2)};  // [bS, iW, iC] -> [bS, 1, iW, iC]
+    } else {
+      reshapeForGradO = {input->sizeAt(0),  input->sizeAt(1), input->sizeAt(2),1};  // [bS, oW, oC] -> [bS, 1, oW, oC]
+      reshapeForInput = {input->sizeAt(0),  input->sizeAt(1), input->sizeAt(2),1};  // [bS, iW, iC] -> [bS, 1, iW, iC]
+
+    }
+  } else {
+    if (!gradO->isScalar()) {
+      reshapeForGradO = {gradO->sizeAt(0), gradO->sizeAt(1), 1, gradO->sizeAt(2)};  // [bS, oC, oW] -> [bS, oC, 1, oW]
+      reshapeForInput = {input->sizeAt(0), input->sizeAt(1), 1, input->sizeAt(2)};  // [bS, iC, iW] -> [bS, iC, 1, iW]
+    } else {
+      reshapeForGradO = {input->sizeAt(0), 1, input->sizeAt(1), input->sizeAt(2)};  // [bS, oW, oC] -> [bS, 1, oW, oC]
+      reshapeForInput = {input->sizeAt(0), 1, input->sizeAt(1), input->sizeAt(2)};  // [bS, iW, iC] -> [bS, 1, iW, iC]
+    }
+  }
+
+  auto inputReshaped = input->reshape(input->ordering(), reshapeForInput,false);
+  auto gradIReshaped = !gradO->isScalar() ? gradI->reshape(gradI->ordering(), reshapeForInput, false) : gradI;
+  auto gradOReshaped = !gradO->isScalar() ?gradO->reshape(gradO->ordering(), reshapeForGradO,false) : gradO;
+  std::vector<LongType> weightsShape = {1, weights->sizeAt(0), weights->sizeAt(1), weights->sizeAt(2)};
+
+  auto weightsReshaped = weights->reshape(
+      weights->ordering(),
+      weightsShape,false);  // [kW, iC, oC] -> [1, kW, iC, oC]
+  auto gradWReshaped =
+      !gradO->isScalar() ?gradW->reshape(gradW->ordering(), weightsShape,
+                                                      false) : gradW;  // [kW, iC, oC] -> [1, kW, iC, oC]
+  Status ret = Status::OK;
+
+  conv2d_bp conv2dBP;
+  if(bias == nullptr) {
+    if(gradO->isScalar()) {
+      gradIReshaped->assign(gradO);
+      gradWReshaped->assign(gradO);
+    } else {
+      std::vector<NDArray *> inputs = {inputReshaped, weightsReshaped, gradOReshaped};
+      std::vector<NDArray *> outputs = {gradIReshaped, gradWReshaped};
+      //note this might look strange but we get a segfault otherwise.
+      //this problem was actually the source of a very strange JVM hang.
+      ret = conv2dBP.execute(inputs,
+                             outputs, {},
+                             {1, kW, 1, sW, 0, pW, 1, dW, paddingMode, !isNCW, wFormat}, {});
+
+    }
+
+
+  } else {
+    if(gradO->isScalar()) {
+      gradIReshaped->assign(gradO);
+      gradWReshaped->assign(gradO);
+      gradB->assign(gradO);
+    } else {
+      std::vector<NDArray *> inputs = {inputReshaped, weightsReshaped,bias, gradOReshaped};
+      std::vector<NDArray *> outputs = {gradIReshaped, gradWReshaped, gradB};
+
+      ret = conv2dBP.execute(inputs,
+                             outputs, {},
+                             {1, kW, 1, sW, 0, pW, 1, dW, paddingMode, !isNCW, wFormat}, {});
+    }
+
+  }
+
+  if(gradIReshaped->buffer() != gradI->buffer()) {
+    gradI->assign(gradIReshaped);
+  }
+
+
+
+  if(gradWReshaped->buffer() != gradW->buffer()) {
+    gradW->assign(gradWReshaped);
+  }
+
+  if(bias != nullptr) {
+    if(gradB->buffer() != gradB->buffer()) {
+      gradB->assign(gradB);
+    }
+  }
+
+
+
+  return ret;
+}
+
+DECLARE_SHAPE_FN(conv1d_bp) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iW, iC] (NWC) or [bS, iC, iW] (NCW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kW, iC, oC], [oC, iC, kW], [oC, kW, iC]
+  LongType const* biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;  // [oC]
+  LongType const* gradOShapeInfo =
+      block.width() > 3 ? inputShape->at(3)
+                        : inputShape->at(2);  // [bS, oW, oC] (NWC) or [bS, oC, oW] (NCW), epsilon_next
+
+  const LongType rank = 3;
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "CUSTOM CONV1D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo[0]);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM CONV1D_BP OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo[0]);
+
+  LongType kW = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<LongType>(0)));  // filter(kernel) width
+  LongType sW = INT_ARG(1);                                                                          // strides width
+  LongType pW = INT_ARG(2);                                                                          // paddings width
+  LongType dW = INT_ARG(3);                                                                          // dilations width
+  LongType paddingMode = INT_ARG(4);                                                                 // 0-VALID, 1-SAME
+  LongType isNCW = block.getIArguments()->size() > 5 ? !INT_ARG(5) : 1;  // INT_ARG(4): 1-NWC, 0-NCW
+  LongType wFormat =
+      block.getIArguments()->size() > 6 ? INT_ARG(6) : 0;  // 0 - [kW, iC, oC], 1 - [oC, iC, kW], 2 - [oC, kW, iC]
+
+  LongType indIOioC, indIiW, indWoC(0 == wFormat ? 2 : 0);
+  if (!isNCW) {
+    indIOioC = 2;
+    indIiW = 1;
+  } else {
+    indIOioC = 1;
+    indIiW = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iW = inputShapeInfo[indIiW + 1];    // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, 1, kW, 1, sW, 0, pW, 1, dW, 1, iW, paddingMode);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoW, 0, indIOioC, indIiW});
+  std::vector<LongType> expectedWeightsShape =
+      0 == wFormat ? std::vector<LongType>({kW, iC, oC})
+                   : (1 == wFormat ? std::vector<LongType>({oC, iC, kW}) : std::vector<LongType>({oC, kW, iC}));
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape), 0,
+               "CUSTOM CONV1D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+  REQUIRE_TRUE(biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+               "CUSTOM CONV1D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto gradIshapeInfo = ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  if (biasShapeInfo) {
+    auto gradBshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo), CONSTANT(gradBshapeInfo));
+  }
+
+  return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo));
+}
+
+DECLARE_TYPES(conv1d_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, {ALL_FLOATS, ALL_INTS, QINT8, QINT16})
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS})
+      ->setAllowedOutputTypes(1, {ALL_FLOATS});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/conv2d.cpp

@@ -0,0 +1,518 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 06.03.2018
+//
+
+#ifndef LIBND4J_CONVO_OPS_H
+#define LIBND4J_CONVO_OPS_H
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_conv2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/OpRegistrator.h>
+#include <ops/declarable/helpers/convolutions.h>
+#include <system/op_boilerplate.h>
+
+#include <memory>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(conv2d, 2, 1, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
+
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  //normally nchw is 0 and 1 being passed in, we're using it as a boolean here
+  //so we want it to be whether nchw is 0 or not.
+  isNCHW = isNCHW == 0;
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  ConvolutionUtils::conv2d(block, input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, isSameMode, isNCHW,
+                           wFormat);
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(conv2d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  // output [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
+
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int paddingMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? INT_ARG(9) : 0;  // INT_ARG(9): 0-NCHW, 1-NHWC
+  LongType wFormat = block.getIArguments()->size() > 10
+                     ? INT_ARG(10)
+                     : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+  //normally nchw is 0 and 1 being passed in, we're using it as a boolean here
+  //so we want it to be whether nchw is 0 or not.
+  isNCHW = isNCHW == 0;
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(ConvolutionUtils::sizeOfKh(weightsShapeInfo,wFormat));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(ConvolutionUtils::sizeOfKw(weightsShapeInfo,wFormat));  // filter(kernel) width
+  const int rank = 4;  // 4
+
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "CUSTOM CONV2D OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo[0]);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM CONV2D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo[0]);
+
+
+
+  LongType bS = shape::sizeAt(inputShapeInfo, 0);             // batch size
+  LongType   iC = ConvolutionUtils::inChannels(weightsShapeInfo, wFormat);
+  LongType   iH = ConvolutionUtils::inputHeight(inputShapeInfo, isNCHW);
+  LongType    iW = ConvolutionUtils::inputWidth(inputShapeInfo, isNCHW);
+  LongType    oC = ConvolutionUtils::outChannels(weightsShapeInfo, wFormat);
+  std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  if(!ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape)) {
+    std::string errorMessage;
+    errorMessage += "CUSTOM CONV2D OP: wrong shape of weights array, expected is ";
+    errorMessage += ShapeUtils::shapeAsString(expectedWeightsShape);
+    errorMessage += ", but got ";
+    errorMessage += ShapeUtils::shapeAsString(weightsShapeInfo);
+    errorMessage += " instead !";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+
+  if (biasShapeInfo) {
+    if(biasShapeInfo[0] > 2 || oC != shape::length(biasShapeInfo)) {
+      std::string errorMessage;
+      errorMessage += "CUSTOM CONV2D OP: wrong shape of array with biases, expected rank, length: <=2, ";
+      errorMessage += std::to_string(oC);
+      errorMessage += ", but got ";
+      errorMessage += std::to_string(biasShapeInfo[0]);
+      errorMessage += ", ";
+      errorMessage += std::to_string(shape::length(biasShapeInfo));
+      errorMessage += " instead !";
+      THROW_EXCEPTION(errorMessage.c_str());
+    }
+
+  }
+
+  LongType* outputShapeInfo = new LongType[shape::shapeInfoLength(rank)];
+
+  outputShapeInfo[0] = 4;
+  LongType    oH = ConvolutionUtils::calcOutDimConv(iH, kH, sH, pH, dH, paddingMode);
+  LongType   oW = ConvolutionUtils::calcOutDimConv(iW,kW,sW,pW,dW,paddingMode);  // batch size, input channels, input height/width, output channels, output height/width;
+
+
+  /**
+   * NOTE: THIS BLOCK OF LOGIC PROBABLY LOOKS STRANGE.
+   * THIS IS FOR COMPATIBILITY WITH THE CONV2D implementation in dl4j.
+   */
+  sd::LongType strideCalcShape[4];
+  strideCalcShape[0] = oW;
+  strideCalcShape[1] = oH;
+  strideCalcShape[2] = bS;
+  strideCalcShape[3] = oC;
+
+  sd::LongType *permute = new sd::LongType[4];
+  permute[0] = 2;
+  permute[1] = 3;
+  permute[2] = 1;
+  permute[3] = 0;
+
+  sd::LongType * second = shape::calcStridesFortran(strideCalcShape,shape::rank(outputShapeInfo));
+  shape::doPermuteSwap(4, second,permute);
+  shape::doPermuteSwap(4, strideCalcShape,permute);
+
+
+  if(!isNCHW) {
+    permute[0] = 0;
+    permute[1] = 2;
+    permute[2] = 3;
+    permute[3] = 1;
+    shape::doPermuteSwap(4, strideCalcShape,permute);
+    shape::doPermuteSwap(4, second,permute);
+
+
+    sd::LongType * second2 = shape::calcStridesFortran(strideCalcShape,shape::rank(outputShapeInfo));
+    shape::doPermuteSwap(4, second2,permute);
+
+    shape::setShape(outputShapeInfo, strideCalcShape);
+    shape::setStride(outputShapeInfo,second);
+    shape::setOrder(outputShapeInfo, 'f');
+    ArrayOptions::setExtra(outputShapeInfo,ArrayOptions::defaultFlag());
+    ArrayOptions::setDataType(outputShapeInfo,ArrayOptions::dataType(inputShapeInfo));
+    delete[] second2;
+
+  } else {
+    shape::setShape(outputShapeInfo, strideCalcShape);
+    shape::setStride(outputShapeInfo,second);
+    shape::setOrder(outputShapeInfo, 'f');
+    ArrayOptions::setExtra(outputShapeInfo,ArrayOptions::defaultFlag());
+    ArrayOptions::setDataType(outputShapeInfo,ArrayOptions::dataType(inputShapeInfo));
+  }
+
+
+
+  delete[] second;
+  delete[] permute;
+  auto ret = ConstantShapeHelper::getInstance().createFromExisting(outputShapeInfo);
+  return SHAPELIST(ret);
+}
+
+
+DECLARE_TYPES(conv2d) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_TYPES(conv2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(conv2d_bp, 3, 2, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto gradO = block.width() > 3
+               ? INPUT_VARIABLE(3)
+               : INPUT_VARIABLE(2);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+
+  auto gradI = OUTPUT_VARIABLE(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+  auto gradW = OUTPUT_VARIABLE(1);  // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW, 1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  isNCHW = isNCHW == 0;
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM CONV2D_BP OP: rank of input array must be equal to 4, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM CONV2D_BP OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(
+      gradO->rankOf() == 4, 0,
+      "CUSTOM CONV2D_BP OP: rank of output's gradients (next epsilon) array must be equal to 4, but got %i instead !",
+      gradO->rankOf());
+  int paddingMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+
+
+  ConvolutionUtils::conv2dBP(block, input, weights, bias, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW,
+                             isSameMode, isNCHW, wFormat);
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(conv2d_bp) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;  // [oC]
+  auto gradOShapeInfo = block.width() > 3
+                        ? inputShape->at(3)
+                        : inputShape->at(2);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+
+
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int paddingMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? INT_ARG(9) : 0;  // INT_ARG(9): 0-NCHW, 1-NHWC
+  LongType wFormat = block.getIArguments()->size() > 10
+                     ? INT_ARG(10)
+                     : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+  //normally nchw is 0 and 1 being passed in, we're using it as a boolean here
+  //so we want it to be whether nchw is 0 or not.
+  isNCHW = isNCHW == 0;
+  // output [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(ConvolutionUtils::sizeOfKh(weightsShapeInfo,wFormat));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(ConvolutionUtils::sizeOfKw(weightsShapeInfo,wFormat));  // filter(kernel) width
+
+  const LongType rank = 4;
+
+  LongType bS = shape::sizeAt(inputShapeInfo, 0);             // batch size
+  LongType   iC = ConvolutionUtils::inChannels(weightsShapeInfo, wFormat);
+  LongType   iH = ConvolutionUtils::inputHeight(inputShapeInfo, isNCHW);
+  LongType    iW = ConvolutionUtils::inputWidth(inputShapeInfo, isNCHW);
+  LongType    oC = ConvolutionUtils::outChannels(weightsShapeInfo, wFormat);
+  std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  if(!ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape)) {
+    std::string errorMessage;
+    errorMessage += "CUSTOM CONV2D OP: wrong shape of weights array, expected is ";
+    errorMessage += ShapeUtils::shapeAsString(expectedWeightsShape);
+    errorMessage += ", but got ";
+    errorMessage += ShapeUtils::shapeAsString(weightsShapeInfo);
+    errorMessage += " instead !";
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+
+  if (biasShapeInfo) {
+    if(biasShapeInfo[0] > 2 || oC != shape::length(biasShapeInfo)) {
+      std::string errorMessage;
+      errorMessage += "CUSTOM CONV2D OP: wrong shape of array with biases, expected rank, length: <=2, ";
+      errorMessage += std::to_string(oC);
+      errorMessage += ", but got ";
+      errorMessage += std::to_string(biasShapeInfo[0]);
+      errorMessage += ", ";
+      errorMessage += std::to_string(shape::length(biasShapeInfo));
+      errorMessage += " instead !";
+      THROW_EXCEPTION(errorMessage.c_str());
+    }
+
+  }
+
+  sd::LongType * strideCalcShapeGradI = new sd::LongType[shape::rank(inputShapeInfo)];
+  strideCalcShapeGradI[0] = iC;
+  strideCalcShapeGradI[1] = bS;
+  strideCalcShapeGradI[2] = iH;
+  strideCalcShapeGradI[3] = iW;
+
+  sd::LongType *strides = new sd::LongType[4];
+  sd::LongType *permute = new sd::LongType[4];
+  permute[0] = 1;
+  permute[1] = isNCHW ? 0 : 2;
+  permute[2] = isNCHW ? 2 : 3;
+  permute[3] = isNCHW ? 3 : 0;
+  shape::calcStrides(strideCalcShapeGradI,shape::rank(inputShapeInfo),strides);
+  shape::doPermuteSwap(4, strideCalcShapeGradI, permute);
+  shape::doPermuteSwap(4, strides, permute);
+  auto shapeDesc = ShapeBuilders::createShapeInfo(ArrayOptions::dataType(inputShapeInfo),
+                                                  'c',
+                                                  4,
+                                                  strideCalcShapeGradI,
+                                                  block.getWorkspace(),
+                                                  false);
+  shape::setStride(shapeDesc,strides);
+  auto gradIshapeInfo = ConstantShapeHelper::getInstance().createFromExisting(shapeDesc);
+  RELEASE(strides,block.getWorkspace());
+  RELEASE(strideCalcShapeGradI,block.getWorkspace());
+  RELEASE(permute,block.getWorkspace());
+  auto gradWshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  if (biasShapeInfo) {
+    auto gradBshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(gradIshapeInfo, CONSTANT(gradWshapeInfo), CONSTANT(gradBshapeInfo));
+  }
+
+  return SHAPELIST(gradIshapeInfo, CONSTANT(gradWshapeInfo));
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(conv2d_input_bp, 3, 1, false, 0, 9) {
+  auto gradIShape = INPUT_VARIABLE(0);  // [4]
+  auto weights = INPUT_VARIABLE(1);     // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto gradO = INPUT_VARIABLE(2);       // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW, 1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  const int rank = gradO->rankOf();
+
+  REQUIRE_TRUE(weights->rankOf() == rank, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(gradIShape->rankOf() == 1, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: rank of array with output shape must be equal to 1, but got %i instead !",
+               gradIShape->rankOf());
+  REQUIRE_TRUE(gradIShape->lengthOf() == rank, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: length of array with output shape must be equal to 4, but got %i instead !",
+               gradIShape->lengthOf());
+
+  // create empty conv2d input array
+  std::vector<LongType> gradIShapeAsVector(rank);
+  for (int i = 0; i < rank; ++i) gradIShapeAsVector[i] = gradIShape->e<LongType>(i);
+  NDArray input(gradO->ordering(), gradIShapeAsVector, gradO->dataType(), block.launchContext());
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWoC, indWkH, indOoH);
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "CUSTOM CONV2D_INPUT_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
+               "got %s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM CONV2D_INPUT_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+
+  ConvolutionUtils::conv2dBP(block, &input, weights, nullptr, gradO, gradI, nullptr, nullptr, kH, kW, sH, sW, pH, pW,
+                             dH, dW, isSameMode, isNCHW, wFormat);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(conv2d_input_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(conv2d_input_bp) {
+  auto gradIShapeShapeInfo = inputShape->at(0);  // [4]
+  auto weightsShapeInfo = inputShape->at(1);     // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto gradOShapeInfo = inputShape->at(2);       // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+
+  const LongType rank = 4;
+
+  REQUIRE_TRUE(gradIShapeShapeInfo[0] == 1, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: rank of array with output shape must be equal to %i, but got %i instead !",
+               1, gradIShapeShapeInfo[0]);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo[0]);
+  REQUIRE_TRUE(gradOShapeInfo[0] == rank, 0,
+               "CUSTOM CONV2D_INPUT_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but got "
+               "%i instead !",
+               rank, gradOShapeInfo[0]);
+
+  const LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  const LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  const LongType sH = INT_ARG(2);                                               // strides height
+  const LongType sW = INT_ARG(3);                                               // strides width
+  const LongType pH = INT_ARG(4);                                               // paddings height
+  const LongType pW = INT_ARG(5);                                               // paddings width
+  const LongType dH = INT_ARG(6);                                               // dilations height
+  const LongType dW = INT_ARG(7);                                               // dilations width
+  const int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW, 1-NHWC
+  const int wFormat = block.getIArguments()->size() > 10
+                      ? INT_ARG(10)
+                      : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  int indIOioC, indIiH, indWoC(0 == wFormat ? 3 : 0), indOoH;
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+    indOoH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+    indOoH = 2;
+  }
+
+  std::vector<LongType> gradIShape = INPUT_VARIABLE(0)->template asVectorT<LongType>();
+
+  const LongType bS = gradIShape[0];                 // batch size
+  const LongType iH = gradIShape[indIiH];            // input height
+  const LongType iW = gradIShape[indIiH + 1];        // input width
+  const LongType iC = gradIShape[indIOioC];          // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(gradOShapeInfo, expectedGradOShape), 0,
+               "CUSTOM CONV2D_INPUT_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
+               "got %s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(),
+               ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape), 0,
+               "CUSTOM CONV2D_INPUT_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+
+  LongType* gradIshapeInfo(nullptr);
+  ALLOCATE(gradIshapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);
+
+  gradIshapeInfo[0] = rank;
+  gradIshapeInfo[1] = bS;
+
+  if (isNCHW) {
+    gradIshapeInfo[2] = iC;
+    gradIshapeInfo[3] = iH;
+    gradIshapeInfo[4] = iW;
+  } else {
+    gradIshapeInfo[2] = iH;
+    gradIshapeInfo[3] = iW;
+    gradIshapeInfo[4] = iC;
+  }
+
+  ShapeUtils::updateStridesAndType(gradIshapeInfo, gradOShapeInfo, shape::order(gradOShapeInfo));
+
+  return SHAPELIST(CONSTANT(gradIshapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif
+
+#endif  // LIBND4J_CONVO_OPS_H

libnd4j/include/ops/declarable/generic/nn/convo/conv3d.cpp

@@ -0,0 +1,452 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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, created on 05.02.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_conv3dnew)
+
+#include <helpers/MmulHelper.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/addBias.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(conv3dnew, 2, 1, false, 0, 13) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);  // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto output = OUTPUT_VARIABLE(0);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW)
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0,
+               "CUSTOM CONV3D OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 5, 0,
+               "CUSTOM CONV3D OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                          // strides height
+  LongType sW = INT_ARG(5);                                                          // strides width
+  LongType pD = INT_ARG(6);                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                         // dilations width
+  int paddingMode = INT_ARG(12);                                                // 0-SAME,  1-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;          // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                ? INT_ARG(14)
+                : 0;  // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
+                                             indIOioC, indIOioD, indWiC, indWoC, indWkD);
+
+  REQUIRE_TRUE(paddingMode < 2, 0,
+               "CUSTOM CONV3D OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM CONV3D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+  REQUIRE_TRUE(
+      bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+      "CUSTOM CONV3D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !",
+      oC, bias->rankOf(), bias->lengthOf());
+
+  ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
+
+  sd_debug("MKL-DNN is not used for conv3dnew!\n", 0);
+
+  std::vector<LongType> permutForOutput;
+
+  std::vector<LongType> permuteDims = {0,4,1,2,3};
+  if (isNCDHW)
+    permutForOutput = {0, 2, 3, 4, 1};  // [bS, oC, oD, oH, oW] -> [bS, oD, oH, oW, oC]
+  else
+    input = input->permute(permuteDims, false, false);
+
+  std::vector<LongType> wAxes;
+  if (0 == wFormat)
+    wAxes = {3, 0, 1, 2};
+  else if (1 == wFormat)
+    wAxes = {1, 2, 3, 4};
+  else
+    wAxes = {4, 1, 2, 3};
+  std::vector<sd::LongType> colShape = {bS, iC, kD, kH, kW, oD, oH, oW};
+
+  NDArray columns(input->ordering(), colShape, input->dataType(), block.launchContext());
+  ConvolutionUtils::vol2col(block, input, &columns, sD, sH, sW, pD, pH, pW, dD, dH,
+                            dW);  // [bS, iC, iD, iH, iW] is convoluted to [bS, iC, kD, kH, kW, oD, oH, oW]
+  // [bS, iC, kD, kH, kW, oD, oH, oW] x [kD, kH, kW, iC, oC] = [bS, oD, oH, oW, oC]
+  // [bS, iC, kD, kH, kW, oD, oH, oW] x [oC, iC, kD, kH, kW] = [bS, oD, oH, oW, oC]
+  // [bS, iC, kD, kH, kW, oD, oH, oW] x [oC, kD, kH, kW, iC] = [bS, oD, oH, oW, oC]
+  std::vector<LongType> mulDims = {1,2,3,4};
+  MmulHelper::tensorDot(&columns, weights, output, mulDims, wAxes, permutForOutput);
+
+  if (bias)
+    helpers::addBias(block, *output, *bias, *output, isNCDHW);
+
+  if (!isNCDHW) delete input;
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(conv3dnew) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, sd::DataType::ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(conv3dnew) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(2)));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                                          // strides height
+  LongType sW = INT_ARG(5);                                                                          // strides width
+  LongType pD = INT_ARG(6);                                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                                         // dilations width
+  int paddingMode = INT_ARG(12);                                                                // 1-SAME,  0-VALID;
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;  // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                ? INT_ARG(14)
+                : 0;  // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
+
+  const int rank = 5;
+  REQUIRE_TRUE(paddingMode < 2, 0,
+               "CUSTOM CONV3D OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "CUSTOM CONV3D OP: rank of input array must be equal to %i, but got %i instead !", rank, inputShapeInfo);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM CONV3D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo);
+
+  LongType indIOioC, indIiD, indWoC(0 == wFormat ? 4 : 0);
+  if (!isNCDHW) {
+    indIOioC = 4;
+    indIiD = 1;
+  } else {
+    indIOioC = 1;
+    indIiD = 2;
+  }
+
+  LongType bS = inputShapeInfo[1];             // batch size
+  LongType iD = inputShapeInfo[indIiD + 1];    // input depth
+  LongType iH = inputShapeInfo[indIiD + 2];    // input height
+  LongType iW = inputShapeInfo[indIiD + 3];    // input width
+  LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape), 0,
+               "CUSTOM CONV3D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+  REQUIRE_TRUE(
+      biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+      "CUSTOM CONV3D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !",
+      oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  LongType oD, oH, oW;  // output depth, height, width
+  ConvolutionUtils::calcOutSizePool3D(oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH, iW,
+                                      paddingMode);
+
+  sd::LongType* outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inputShapeInfo), sd::LongType);
+
+  outputShapeInfo[0] = rank;
+  outputShapeInfo[1] = bS;
+  if (isNCDHW) {
+    outputShapeInfo[2] = oC;
+    outputShapeInfo[3] = oD;
+    outputShapeInfo[4] = oH;
+    outputShapeInfo[5] = oW;
+  } else {
+    outputShapeInfo[2] = oD;
+    outputShapeInfo[3] = oH;
+    outputShapeInfo[4] = oW;
+    outputShapeInfo[5] = oC;
+  }
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, weightsShapeInfo, shape::order(inputShapeInfo));
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(conv3dnew_bp, 3, 2, false, 0, 13) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);  // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto gradO = block.width() > 3
+               ? INPUT_VARIABLE(3)
+               : INPUT_VARIABLE(2);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+
+  auto gradI = OUTPUT_VARIABLE(0);  // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
+  auto gradW = OUTPUT_VARIABLE(1);  // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
+  auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0,
+               "CUSTOM CONV3D_BP OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 5, 0,
+               "CUSTOM CONV3D_BP OP: rank of weights array must be equal to 5, but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(
+      gradO->rankOf() == 5, 0,
+      "CUSTOM CONV3D_BP OP: rank of output gradients (next epsilon) array must be equal to 5, but got %i instead !",
+      gradO->rankOf());
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                          // strides height
+  LongType sW = INT_ARG(5);                                                          // strides width
+  LongType pD = INT_ARG(6);                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                         // dilations width
+  int paddingMode = INT_ARG(12);                                                // 1-SAME,  0-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;          // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                ? INT_ARG(14)
+                : 0;  // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
+                                             indIOioC, indIOioD, indWiC, indWoC, indWkD);
+
+  LongType trueoD, trueoH, trueoW;  // true output depth/height/width
+  ConvolutionUtils::calcOutSizePool3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
+                                      iW, paddingMode);
+
+  REQUIRE_TRUE(paddingMode < 2, 0,
+               "CUSTOM CONV3D_BP OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
+  std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
+      {bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
+  REQUIRE_TRUE(
+      gradO->isSameShape(expectedGradOShape), 0,
+      "CUSTOM CONV3D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM CONV3D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+  REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+               "CUSTOM CONV3D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, bias->rankOf(), bias->lengthOf());
+
+  ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
+
+  sd_debug("MKL-DNN is not used for conv3dnew_bp!\n", 0);
+
+  std::vector<LongType> gradOaxesForDot;
+
+  std::vector<LongType> permute = {0, 4, 1, 2, 3};
+  if (!isNCDHW) {
+    gradOaxesForDot = {0, 1, 2, 3};                        // bS, oD, oH, oW
+    input =input->permute(permute, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    gradI = gradI->permute(permute, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+  } else {
+    gradOaxesForDot = {0, 2, 3, 4};  // bS, oD, oH, oW
+  }
+
+  std::vector<LongType> wPermut, colPermut;
+
+  if (0 == wFormat) {
+    wPermut = {3, 0, 1, 2, 4};
+    colPermut = {2, 3, 4, 1, 0, 5, 6, 7};
+  } else if (1 == wFormat) {
+    wPermut = {1, 2, 3, 4, 0};
+    colPermut = {1, 2, 3, 4, 0, 5, 6, 7};
+  } else {
+    wPermut = {4, 1, 2, 3, 0};
+    colPermut = {2, 3, 4, 1, 0, 5, 6, 7};
+  }
+
+  std::vector<sd::LongType> colShape = {bS, iC, kD, kH, kW, oD, oH, oW};
+  // ----- calculation of gradW and gradB ----- //
+  NDArray columns(input->ordering(), colShape, input->dataType(), block.launchContext());
+  ConvolutionUtils::vol2col(block, input, &columns, sD, sH, sW, pD, pH, pW, dD, dH,
+                            dW);  // [bS, iC, iD, iH, iW] is convoluted to [bS, iC, kD, kH, kW, oD, oH, oW]
+
+  std::vector<LongType> mulDims = {0,5,6,7};
+  MmulHelper::tensorDot(
+      &columns, gradO, gradW, mulDims, gradOaxesForDot,
+      wPermut);  // [bS, iC, kD, kH, kW, oD, oH, oW] x [bS, oD, oH, oW, oC]/[bS, oC, oD, oH, oW] = [iC, kD, kH, kW, oC]
+
+  //----- calculation of gradO -----//
+  if (gradB) {
+    std::vector<LongType> bShape = { gradB->lengthOf()};
+    if (gradB->rankOf() == 2) gradB =gradB->reshape(gradB->ordering(),bShape, false);
+    gradO->reduceAlongDimension(reduce::Sum, gradB, &gradOaxesForDot);  // sum over bS oD oH oW
+    if (gradB != OUTPUT_VARIABLE(2)) delete gradB;
+  }
+
+  //----- calculation of gradI -----//
+  // [kD, kH, kW, iC, oC] x [bS, oD, oH, oW, oC]/[bS, oC, oD, oH, oW] = [kD, kH, kW, iC, bS, oD, oH, oW]
+  // [oC, iC, kD, kH, kW] x [bS, oD, oH, oW, oC]/[bS, oC, oD, oH, oW] = [kD, kH, kW, iC, bS, oD, oH, oW]
+  // [oC, kD, kH, kW, iC] x [bS, oD, oH, oW, oC]/[bS, oC, oD, oH, oW] = [kD, kH, kW, iC, bS, oD, oH, oW]
+  std::vector<LongType> firstDims = {indWoC};
+  std::vector<LongType> secondDims = {indIOioC};
+
+  MmulHelper::tensorDot(weights, gradO, &columns, firstDims, secondDims, colPermut);
+  ConvolutionUtils::col2vol(block, columns, *gradI, sD, sH, sW, pD, pH, pW, dD, dH,
+                            dW);  // columns [bS, iC, kD, kH, kW, oD, oH, oW] is de-convoluted to  [bS, iC, iD, iH, iW]
+
+  if (!isNCDHW) {
+    delete input;
+    delete gradI;
+  }
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(conv3dnew_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, sd::DataType::ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(conv3dnew_bp) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
+  sd::LongType const* biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;  // [oC]
+  sd::LongType const* gradOShapeInfo =
+      block.width() > 3
+      ? inputShape->at(3)
+      : inputShape->at(2);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<sd::LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(2)));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                                          // strides height
+  LongType sW = INT_ARG(5);                                                                          // strides width
+  LongType pD = INT_ARG(6);                                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                                         // dilations width
+  int paddingMode = INT_ARG(12);                                                                // 1-SAME,  0-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;  // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                ? INT_ARG(14)
+                : 0;  // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
+
+  const int rank = 5;
+  REQUIRE_TRUE(paddingMode < 2, 0,
+               "CUSTOM CONV3D OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "CUSTOM CONV3D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM CONV3D_BP OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo);
+  REQUIRE_TRUE(
+      gradOShapeInfo[0] == rank, 0,
+      "CUSTOM CONV3D_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but got %i instead !",
+      rank, gradOShapeInfo);
+
+  sd::LongType indIOioC, indIiD, indWoC(0 == wFormat ? 4 : 0);
+  if (!isNCDHW) {
+    indIOioC = 4;
+    indIiD = 1;
+  } else {
+    indIOioC = 1;
+    indIiD = 2;
+  }
+
+  LongType bS = inputShapeInfo[1];             // batch size
+  LongType iD = inputShapeInfo[indIiD + 1];    // input depth
+  LongType iH = inputShapeInfo[indIiD + 2];    // input height
+  LongType iW = inputShapeInfo[indIiD + 3];    // input width
+  LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  LongType trueoD, trueoH, trueoW;  // true output depth/height/width
+  ConvolutionUtils::calcOutSizePool3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
+                                      iW, paddingMode);
+
+  std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
+      {bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIiD, indIiD + 1, indIiD + 2});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(gradOShapeInfo, expectedGradOShape), 0,
+      "CUSTOM CONV3D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape), 0,
+               "CUSTOM CONV3D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+  REQUIRE_TRUE(biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+               "CUSTOM CONV3D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto gradIshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  if (biasShapeInfo) {
+    auto gradBshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo), CONSTANT(gradBshapeInfo));
+  }
+
+  return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo));
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/deconv2d.cpp

@@ -0,0 +1,407 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_deconv2d)
+
+#include <helpers/MmulHelper.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/addBias.h>
+#include <ops/declarable/helpers/col2im.h>
+#include <ops/declarable/helpers/convolutions.h>
+#include <ops/declarable/helpers/im2col.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(deconv2d, 2, 1, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM DECONV2D OP: rank of input array must be equal to 4, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM DECONV2D OP: rank of weights array must be equal to 4, but got %i instead !", weights->rankOf());
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                          // strides height
+  LongType sW = INT_ARG(3);                                                          // strides width
+  sd::LongType pH = INT_ARG(4);                                                          // paddings height
+  sd::LongType  pW = INT_ARG(5);                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;             // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWoC, indWiC, indWkH, indOoH);
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+  REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+               "CUSTOM DECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+               "instead !",
+               oC, bias->rankOf(), bias->lengthOf());
+
+  std::vector<LongType> outputPermute = {0,3,1,2};
+  if (!isNCHW) output = output->permute(outputPermute, false, false);  // [bS, oH, oW, oC] -> [bS, oC, oH, oW]
+
+  std::vector<LongType> colPermut;
+  if (1 == wFormat)
+    colPermut = {1, 2, 3, 0, 4, 5};
+  else
+    colPermut = {2, 3, 1, 0, 4, 5};
+
+  if (isSameMode)  // Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not
+    // deconv) forward pass
+    ConvolutionUtils::calcPadding2D(pH, pW, iH, iW, oH, oW, kH, kW, sH, sW, dH, dW);
+
+  std::vector<sd::LongType> colShape = {bS, oC, kH, kW, iH, iW};
+  NDArray columns(input->ordering(), colShape, input->dataType(), block.launchContext());
+
+  //----- calculation of output -----//
+  // NHWC: [kH, kW, oC, iC] x [bS, iH, iW, iC] = [kH, kW, oC, bS, iH, iW]
+  // NHWC: [iC, oC, kH, kW] x [bS, iH, iW, iC] = [oC, kH, kW, bS, iH, iW]
+  // NHWC: [iC, kH, kW, oC] x [bS, iH, iW, iC] = [kH, kW, oC, bS, iH, iW]
+  std::vector<LongType> firstDims = {indWiC};
+  std::vector<LongType> secondDims = {indIOioC};
+  sd::MmulHelper::tensorDot(weights, input, &columns, firstDims, secondDims, colPermut);
+  LaunchContext* ctx = block.launchContext();
+  helpers::col2im(*ctx, &columns, output, sH, sW, pH, pW, oH, oW, dH,
+                  dW);  // [bS, oC, kH, kW, iH, iW] is de-convoluted to [bS, oC, oH, oW]
+
+  //----- add biases if required -----//
+  if (bias)
+    helpers::addBias(block, *output, *bias, *output, true);
+
+  if (!isNCHW) delete output;
+
+  return sd::Status::OK;
+}
+DECLARE_TYPES(deconv2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(deconv2d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  const int rank = 4;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(inputShapeInfo));
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                                          // strides height
+  LongType sW = INT_ARG(3);                                                                          // strides width
+  LongType pH = INT_ARG(4);                                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
+
+  LongType indIOioC, indIiH, indWoC(0 == wFormat ? 2 : (1 == wFormat ? 1 : 3));
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iH = inputShapeInfo[indIiH + 1];    // input height
+  const LongType iW = inputShapeInfo[indIiH + 2];    // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
+  REQUIRE_TRUE(shape::shapeEquals(4, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+  REQUIRE_TRUE(shape::rank(biasShapeInfo) <= 2 && oC == shape::length(biasShapeInfo), 0,
+               "CUSTOM DECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+               "instead !",
+               oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  LongType oH, oW;  // output height, width
+  ConvolutionUtils::calcOutSizeDeconv2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  sd::LongType outputShape[4];
+
+  outputShape[0] = bS;
+  if (isNCHW) {
+    outputShape[1] = oC;
+    outputShape[2] = oH;
+    outputShape[3] = oW;
+  } else {
+    outputShape[1] = oH;
+    outputShape[2] = oW;
+    outputShape[3] = oC;
+  }
+
+  auto desc = new  ShapeDescriptor(ArrayOptions::dataType(weightsShapeInfo), shape::order(inputShapeInfo), outputShape, 4);
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
+}
+
+DECLARE_TYPES(deconv2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(deconv2d_bp, 3, 2, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto gradO = block.width() > 3
+               ? INPUT_VARIABLE(3)
+               : INPUT_VARIABLE(2);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
+
+  auto gradI = OUTPUT_VARIABLE(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW), gradI
+  auto gradW = OUTPUT_VARIABLE(1);  // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
+  auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM DECONV2D_BP OP: rank of input array must be equal to 4, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM DECONV2D_BP OP: rank of weights array must be equal to 4 , but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(
+      gradO->rankOf() == 4, 0,
+      "CUSTOM DECONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to 4, but got %i instead !",
+      gradO->rankOf());
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                          // strides height
+  LongType sW = INT_ARG(3);                                                          // strides width
+  sd::LongType pH = INT_ARG(4);                                                          // paddings height
+  sd::LongType pW = INT_ARG(5);                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;             // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWoC, indWiC, indWkH, indOoH);
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizeDeconv2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "CUSTOM DECONV2D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DECONV2D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+  REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+               "CUSTOM DECONV2D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, bias->rankOf(), bias->lengthOf());
+
+  if (isSameMode) {  // SAME
+    // Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not deconv) forward
+    // pass
+    ConvolutionUtils::calcPadding2D(pH, pW, iH, iW, oH, oW, kH, kW, sH, sW, dH, dW);
+  }
+
+  // ----- calculation of gradI -> pass it through conv2d_ff ----- //
+  sd::ops::conv2d conv2d;
+
+  const sd::Status status =
+      conv2d.execute({gradO, weights}, {gradI}, {}, {kH, kW, sH, sW, pH, pW, dH, dW, isSameMode, !isNCHW, wFormat}, {});
+  if (status != sd::Status::OK) return status;
+
+  // -----prepare permutation arrays and axes for dot product ----- //
+  std::vector<LongType> inputAxes;
+
+  if (!isNCHW) {
+    std::vector<LongType> permuteDims = {0,3,1,2};
+    gradO = gradO->permute(permuteDims, false, false);  // [bS, oH, oW, oC] -> [bS, oC, oH, oW]
+    inputAxes = {0, 1, 2};                              // bS, iH, iW
+  } else
+    inputAxes = {0, 2, 3};  // bS, iH, iW
+
+  std::vector<LongType> gradWAxes;  // empty for wFormat = 1
+  if (0 == wFormat)
+    gradWAxes = {3, 2, 0, 1};
+  else if (2 == wFormat)
+    gradWAxes = {0, 3, 1, 2};
+  std::vector<sd::LongType> colShape = {bS, oC, kH, kW, iH, iW};
+
+  // ----- calculation of gradW ----- //
+  NDArray columns(input->ordering(), colShape, input->dataType(), block.launchContext());
+
+  LaunchContext* ctx = block.launchContext();
+  NDArray *zero = NDArrayFactory::create(0.f, input->getContext());
+  helpers::im2col(
+      *ctx, *gradO, columns, kH, kW, sH, sW, pH, pW, dH, dW,
+     *zero );  // [bS, oC, oH, oW] is convoluted to [bS, oC, kH, kW, iH, iW]
+  std::vector<LongType> mulDims = {0,4,5};
+  MmulHelper::tensorDot(input, &columns, gradW, inputAxes, mulDims,
+                        gradWAxes);  // [bS, iC, iH, iW]/[bS, iH, iW, iC] x [bS, oC, kH, kW, iH, iW] = [iC, oC, kH, kW]
+
+  // ----- calculation of gradB ----- //
+  if (gradB) {
+    std::vector<LongType> bShape = {gradB->lengthOf()};
+    if (gradB->rankOf() == 2) gradB = gradB->reshape(gradB->ordering(), bShape, false);
+    std::vector<sd::LongType> axesForReduction = {0, 2, 3};  // bS, oH, oW
+    gradO->reduceAlongDimension(reduce::Sum, gradB, &axesForReduction);  // sum over bS, oH, oW
+    if (gradB != OUTPUT_VARIABLE(2)) delete gradB;
+  }
+
+  if (!isNCHW) delete gradO;
+
+  delete zero;
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(deconv2d_bp) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
+  sd::LongType const* biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;  // [oC]
+  auto gradOShapeInfo = block.width() > 3
+                        ? inputShape->at(3)
+                        : inputShape->at(2);  // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
+
+  const int rank = 4;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(inputShapeInfo));
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D_BP OP: rank of weights array must be equal to %i , but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+  REQUIRE_TRUE(
+      shape::rank(gradOShapeInfo) == rank, 0,
+      "CUSTOM DECONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but got %i instead !",
+      rank, shape::rank(gradOShapeInfo));
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                                          // strides height
+  LongType sW = INT_ARG(3);                                                                          // strides width
+  LongType pH = INT_ARG(4);                                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
+
+  LongType indIOioC, indIiH, indOoH, indWoC(0 == wFormat ? 2 : (1 == wFormat ? 1 : 3));
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+    indOoH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+    indOoH = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iH = inputShapeInfo[indIiH + 1];    // input height
+  const LongType iW = inputShapeInfo[indIiH + 2];    // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizeDeconv2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
+  REQUIRE_TRUE(
+      shape::shapeEquals(4, expectedGradOShape.data(), shape::rank(gradOShapeInfo), shape::shapeOf(gradOShapeInfo)), 0,
+      "CUSTOM DECONV2D_BP OP: wrong shape of output gradients next epsilon) array, expected is %s, but got %s instead "
+      "!",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  REQUIRE_TRUE(shape::shapeEquals(4, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DECONV2D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+  REQUIRE_TRUE(biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+               "CUSTOM DECONV2D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+               "%i instead !",
+               oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto gradIShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  auto shapes = SHAPELIST(CONSTANT(gradIShapeInfo), CONSTANT(gradWShapeInfo));
+
+  if (biasShapeInfo != nullptr) {
+    auto gradBShapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    shapes->push_back(CONSTANT(gradBShapeInfo));
+  }
+
+  return shapes;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/deconv2d_tf.cpp

@@ -0,0 +1,191 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+// @author Yurii Shyrma
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_deconv2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(deconv2d_tf, 3, 1, false, 0, 9) {
+  auto gradO = INPUT_VARIABLE(2);       // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto weights = INPUT_VARIABLE(1);     // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto gradIShape = INPUT_VARIABLE(0);  // [4] - shape of input of conv2d (that is shape of gradI)
+
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                          // strides height
+  LongType sW = INT_ARG(3);                                                          // strides width
+  LongType pH = INT_ARG(4);                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;             // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                ? INT_ARG(10)
+                : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  const LongType rank = gradO->rankOf();
+
+  REQUIRE_TRUE(weights->rankOf() == rank, 0,
+               "CUSTOM DECONV2D_TF OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(gradIShape->rankOf() == 1, 0,
+               "CUSTOM DECONV2D_TF OP: rank of array with output shape must be equal to 1, but got %i instead !",
+               gradIShape->rankOf());
+  REQUIRE_TRUE(gradIShape->lengthOf() == rank, 0,
+               "CUSTOM DECONV2D_TF OP: length of array with output shape must be equal to 4, but got %i instead !",
+               gradIShape->lengthOf());
+
+
+  auto nonConst = const_cast<NDArray *>(gradIShape);
+  std::vector<sd::LongType> gradIShapeVector = nonConst->template asVectorT<sd::LongType>();
+  // create empty conv2d input array
+  NDArray input(gradO->ordering(), gradIShapeVector, gradO->dataType(), block.launchContext());
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWoC, indWkH, indOoH);
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "CUSTOM DECONV2D_TF OP: wrong shape of input array, basing on array with output shape expected is %s, "
+               "but got %s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DECONV2D_TF OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+
+  ConvolutionUtils::conv2dBP(block, &input, weights, nullptr, gradO, gradI, nullptr, nullptr, kH, kW, sH, sW, pH, pW,
+                             dH, dW, isSameMode, isNCHW, wFormat);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(deconv2d_tf) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(deconv2d_tf) {
+  auto gradOShapeInfo = inputShape->at(2);       // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto weightsShapeInfo = inputShape->at(1);     // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
+  auto gradIShapeShapeInfo = inputShape->at(0);  // [4]
+
+  const int rank = 4;
+
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D_TF OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+  REQUIRE_TRUE(shape::rank(gradOShapeInfo) == rank, 0,
+               "CUSTOM DECONV2D_TF OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(gradOShapeInfo));
+  REQUIRE_TRUE(shape::rank(gradIShapeShapeInfo) == 1, 0,
+               "CUSTOM DECONV2D_TF OP: rank of array with output shape must be equal to %i, but got %i instead !", 1,
+               shape::rank(gradIShapeShapeInfo));
+
+  const LongType kH =
+      INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<int>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) height
+  const LongType kW =
+      INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<int>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) width
+  const LongType sH = INT_ARG(2);                                                               // strides height
+  const LongType sW = INT_ARG(3);                                                               // strides width
+  const LongType pH = INT_ARG(4);                                                               // paddings height
+  const LongType pW = INT_ARG(5);                                                               // paddings width
+  const LongType dH = INT_ARG(6);                                                               // dilations height
+  const LongType dW = INT_ARG(7);                                                               // dilations width
+  const int isSameMode = INT_ARG(8);                                                       // 0-VALID, 1-SAME
+  const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;                  // INT_ARG(9): 1-NHWC, 0-NCHW
+  const int wFormat = block.getIArguments()->size() > 10
+                      ? INT_ARG(10)
+                      : 0;  // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
+
+  LongType indIOioC, indIiH, indWoC(0 == wFormat ? 3 : 0), indOoH;
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+    indOoH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+    indOoH = 2;
+  }
+
+  std::vector<sd::LongType> gradIShape = INPUT_VARIABLE(0)->template asVectorT<sd::LongType>();
+
+  const LongType bS = gradIShape[0];                 // batch size
+  const LongType iH = gradIShape[indIiH];            // input height
+  const LongType iW = gradIShape[indIiH + 1];        // input width
+  const LongType iC = gradIShape[indIOioC];          // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+  const LongType oH = gradOShapeInfo[indOoH + 1];    // input height
+  const LongType oW = gradOShapeInfo[indOoH + 2];    // input width
+
+  LongType trueiH, trueiW;  // output height, width
+  ConvolutionUtils::calcOutSizeDeconv2D(trueiH, trueiW, kH, kW, sH, sW, pH, pW, dH, dW, oH, oW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, trueiH, trueiW, 0, indIOioC, indIiH, indIiH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
+  REQUIRE_TRUE(expectedGradIShape == gradIShape, 0,
+               "CUSTOM DECONV2D_TF OP: wrong shape of array with output shape, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradIShape).c_str());
+  REQUIRE_TRUE(shape::shapeEquals(4, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DECONV2D_TF OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+
+  sd::LongType shape[4];
+  shape[0] = bS;
+
+  if (isNCHW) {
+    shape[1] = iC;
+    shape[2] = iH;
+    shape[3] = iW;
+  } else {
+    shape[1] = iH;
+    shape[2] = iW;
+    shape[3] = iC;
+  }
+  auto ret = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(weightsShapeInfo),shape::order(gradOShapeInfo),4,shape,0);
+  return SHAPELIST(ret);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/deconv3d.cpp

@@ -0,0 +1,432 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com), created on 05.09.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_deconv3d)
+
+#include <helpers/MmulHelper.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/addBias.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(deconv3d, 2, 1, false, 0, 13) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);  // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  auto output = OUTPUT_VARIABLE(0);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW)
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0,
+               "CUSTOM DECONV3D OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 5, 0,
+               "CUSTOM DECONV3D OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                          // strides height
+  LongType sW = INT_ARG(5);                                                          // strides width
+  LongType pD = INT_ARG(6);                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                         // dilations width
+  int isSameMode = INT_ARG(12);                                                 // 0-SAME,  1-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;          // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                    ? INT_ARG(14)
+                    : 0;  // 0 - [kD, kH, kW, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
+
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
+                                             indIOioC, indIOioD, indWoC, indWiC, indWkD);
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DECONV3D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+                 "CUSTOM DECONV3D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+                 "instead !",
+                 oC, bias->rankOf(), bias->lengthOf());
+
+
+  std::vector<LongType> outputPerm = {0, 4, 1, 2, 3};
+  if (!isNCDHW) output = output->permute(outputPerm, false, false);  // [bS, oD, oH, oW, oC] -> [bS, oC, oD, oH, oW]
+
+  std::vector<LongType> colPermut;
+  if (1 == wFormat)
+    colPermut = {1, 2, 3, 4, 0, 5, 6, 7};
+  else
+    colPermut = {2, 3, 4, 1, 0, 5, 6, 7};
+
+  if (isSameMode)  // Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not
+                   // deconv) forward pass
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+
+  std::vector<sd::LongType> columnsShape = {bS, oC, kD, kH, kW, iD, iH, iW};
+  NDArray columns(input->ordering(),columnsShape, input->dataType(), block.launchContext());
+
+  //----- calculation of output -----//
+  // [kD, kH, kW, oC, iC] x [bS, iD, iH, iW, iC] = [kD, kH, kW, oC, bS, iD, iH, iW]
+  // [iC, oC, kD, kH, kW] x [bS, iD, iH, iW, iC] = [oC, kD, kH, kW, bS, iD, iH, iW]
+  // [iC, kD, kH, kW, oC] x [bS, iD, iH, iW, iC] = [kD, kH, kW, oC, bS, iD, iH, iW]
+  std::vector<LongType> indWiCShape = {indWiC};
+  std::vector<LongType> indIOioCShape = {indIOioC};
+  sd::MmulHelper::tensorDot(weights, input, &columns, indWiCShape, indIOioCShape,
+                            colPermut);  // [bS, oC, kD, kH, kW, iD, iH, iW] -> [kD, kH, kW, oC, bS, iD, iH, iW]
+
+  ConvolutionUtils::col2vol(block, columns, *output, sD, sH, sW, pD, pH, pW, dD, dH,
+                            dW);  // [bS, oC, kD, kH, kW, iD, iH, iW] is de-convoluted to [bS, oC, oD, oH, oW]
+
+  //----- add biases if required -----//
+  if (bias)
+    helpers::addBias(block, *output, *bias, *output, true);
+
+  //if (!isNCDHW) delete output;
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(deconv3d) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, sd::DataType::ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(deconv3d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NDCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  const sd::LongType rank = 5;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DECONV3D OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(inputShapeInfo));
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DECONV3D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(2)));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                                          // strides height
+  LongType sW = INT_ARG(5);                                                                          // strides width
+  LongType pD = INT_ARG(6);                                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                                         // dilations width
+  int isSameMode = INT_ARG(12);                                                                 // 0-SAME,  1-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;  // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                    ? INT_ARG(14)
+                    : 0;  // 0 - [kD, kH, kW, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
+
+  LongType indIOioC, indIiD, indWoC(0 == wFormat ? 3 : (1 == wFormat ? 1 : 4));
+  if (!isNCDHW) {
+    indIOioC = 4;
+    indIiD = 1;
+  } else {
+    indIOioC = 1;
+    indIiD = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iD = inputShapeInfo[indIiD + 1];    // input depth
+  const LongType iH = inputShapeInfo[indIiD + 2];    // input height
+  const LongType iW = inputShapeInfo[indIiD + 3];    // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
+  REQUIRE_TRUE(shape::shapeEquals(5, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DECONV3D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+    REQUIRE_TRUE(shape::rank(biasShapeInfo) <= 2 && oC == shape::length(biasShapeInfo), 0,
+                 "CUSTOM DECONV3D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+                 "instead !",
+                 oC, shape::rank(biasShapeInfo), shape::length(biasShapeInfo));
+
+  LongType oD, oH, oW;  // output depth, height, width
+  ConvolutionUtils::calcOutSizeDeconv3D(oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH, iW,
+                                        isSameMode);
+
+
+
+ std::vector<sd::LongType> outputShape;
+  if (isNCDHW) {
+    outputShape = {bS,oC,oD,oH,oW};
+  } else {
+    outputShape = {bS,oD,oH,oW,oC};
+  }
+
+  ShapeDescriptor *shapeDescriptor = new ShapeDescriptor(ArrayOptions::dataType(inputShapeInfo), shape::order(inputShapeInfo),
+                                                         outputShape);
+  auto outputShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(shapeDescriptor);
+  delete shapeDescriptor;
+  return SHAPELIST(outputShapeInfo);
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(deconv3d_bp, 3, 2, false, 0, 13) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);  // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+  auto gradO = block.width() > 3
+                   ? INPUT_VARIABLE(3)
+                   : INPUT_VARIABLE(2);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+
+  auto gradI = OUTPUT_VARIABLE(0);  // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), gradI
+  auto gradW = OUTPUT_VARIABLE(1);  // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
+  auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0,
+               "CUSTOM DECONV3D_BP OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 5, 0,
+               "CUSTOM DECONV3D_BP OP: rank of weights array must be equal to 5 , but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(
+      gradO->rankOf() == 5, 0,
+      "CUSTOM DECONV3D_BP OP: rank of output gradients (next epsilon) array must be equal to 5, but got %i instead !",
+      gradO->rankOf());
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                          // strides height
+  LongType sW = INT_ARG(5);                                                          // strides width
+  LongType pD = INT_ARG(6);                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                         // dilations width
+  int isSameMode = INT_ARG(12);                                                 // 0-SAME,  1-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;          // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                    ? INT_ARG(14)
+                    : 0;  // 0 - [kD, kH, kW, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
+                                             indIOioC, indIOioD, indWoC, indWiC, indWkD);
+
+  LongType trueoD, trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizeDeconv3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
+                                        iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
+      {bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "CUSTOM DECONV3D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DECONV3D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+                 "CUSTOM DECONV3D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+                 "%i instead !",
+                 oC, bias->rankOf(), bias->lengthOf());
+
+  if (isSameMode)  // Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not
+                   // deconv) forward pass
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+
+  // ----- calculation of gradI -> pass it through conv3d_ff ----- //
+  sd::ops::conv3dnew conv3d;
+  const sd::Status status =
+      conv3d.execute({gradO, weights}, {gradI}, {},
+                     {kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isSameMode, !isNCDHW, wFormat}, {});
+  if (status != sd::Status::OK) return status;
+
+  // -----prepare permutation arrays and axes for dot product ----- //
+  std::vector<LongType> inputAxesForDot;
+
+  if (!isNCDHW) {
+    std::vector<LongType> grad0Permute = {0,4,1,2,3};
+    gradO = gradO->permute(grad0Permute, false, false);  // [bS, oD, oH, oW, oC] -> [bS, oC, oD, oH, oW]
+    inputAxesForDot = {0, 1, 2, 3};                        // bS, iD, iH, iW
+  } else
+    inputAxesForDot = {0, 2, 3, 4};  // bS, iD, iH, iW
+
+  std::vector<LongType> gradWAxes;  // empty for wFormat = 1
+  if (0 == wFormat)
+    gradWAxes = {4, 3, 0, 1, 2};
+  else if (2 == wFormat)
+    gradWAxes = {0, 4, 1, 2, 3};
+
+  // ----- calculation of gradW ----- //
+  auto columns = NDArrayFactory::create(input->ordering(), {bS, oC, kD, kH, kW, iD, iH, iW}, input->dataType(),
+                                        block.launchContext());
+  ConvolutionUtils::vol2col(block, gradO, columns, sD, sH, sW, pD, pH, pW, dD, dH,
+                            dW);  // [bS, oC, oD, oH, oW] is deconvoluted to [bS, oC, kD, kH, kW, iD, iH, iW]
+
+  std::vector<LongType> mulDims = {0,5,6,7};
+  MmulHelper::tensorDot(input, columns, gradW, inputAxesForDot, mulDims,
+                        gradWAxes);  // [bS, iC, iD, iH, iW]/[bS, iD, iH, iW, iC] x [bS, oC, kD, kH, kW, iD, iH, iW] =
+                                     // [iC, oC, kD, kH, kW]
+
+  // ----- calculation of gradB ----- //
+  if (gradB) {
+    std::vector<LongType> biasShape = {gradB->lengthOf()};
+    if (gradB->rankOf() == 2) gradB = gradB->reshape(gradB->ordering(), biasShape, false);
+    std::vector<sd::LongType> dims = {{0, 2, 3, 4}};
+    gradO->reduceAlongDimension(reduce::Sum, gradB, &dims);  // sum over bS, oD, oH, oW
+    if (gradB != OUTPUT_VARIABLE(2)) delete gradB;
+  }
+
+  if (!isNCDHW) delete gradO;
+
+  delete columns;
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(deconv3d_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, sd::DataType::ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(deconv3d_bp) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
+  auto biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;  // [oC]
+  auto gradOShapeInfo =
+      block.width() > 3
+          ? inputShape->at(3)
+          : inputShape->at(2);  // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+
+  const int rank = 5;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DECONV3D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(inputShapeInfo));
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DECONV3D_BP OP: rank of weights array must be equal to %i , but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+  REQUIRE_TRUE(
+      shape::rank(gradOShapeInfo) == rank, 0,
+      "CUSTOM DECONV3D_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but got %i instead !",
+      rank, shape::rank(gradOShapeInfo));
+
+  LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) depth
+  LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) height
+  LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(2)));  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                                          // strides depth
+  LongType sH = INT_ARG(4);                                                                          // strides height
+  LongType sW = INT_ARG(5);                                                                          // strides width
+  LongType pD = INT_ARG(6);                                                                          // paddings depth
+  LongType pH = INT_ARG(7);                                                                          // paddings height
+  LongType pW = INT_ARG(8);                                                                          // paddings width
+  LongType dD = INT_ARG(9);                                                                          // dilations depth
+  LongType dH = INT_ARG(10);                                                                         // dilations height
+  LongType dW = INT_ARG(11);                                                                         // dilations width
+  int isSameMode = INT_ARG(12);                                                                 // 0-SAME,  1-VALID
+  int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1;  // INT_ARG(13): 1-NDHWC, 0-NCDHW
+  int wFormat = block.getIArguments()->size() > 14
+                    ? INT_ARG(14)
+                    : 0;  // 0 - [kD, kH, kW, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
+
+  LongType indIOioC, indIiD, indWoC(0 == wFormat ? 3 : (1 == wFormat ? 1 : 4));
+  if (!isNCDHW) {
+    indIOioC = 4;
+    indIiD = 1;
+  } else {
+    indIOioC = 1;
+    indIiD = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iD = inputShapeInfo[indIiD + 1];    // input depth
+  const LongType iH = inputShapeInfo[indIiD + 2];    // input height
+  const LongType iW = inputShapeInfo[indIiD + 3];    // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  LongType trueoD, trueoH, trueoW;  // true output depth, height, width
+  ConvolutionUtils::calcOutSizeDeconv3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
+                                        iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
+      {bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIiD, indIiD + 1, indIiD + 2});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
+  REQUIRE_TRUE(
+      shape::shapeEquals(5, expectedGradOShape.data(), shape::rank(gradOShapeInfo), shape::shapeOf(gradOShapeInfo)), 0,
+      "CUSTOM DECONV3D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead "
+      "!",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  REQUIRE_TRUE(shape::shapeEquals(5, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DECONV3D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+    REQUIRE_TRUE(biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+                 "CUSTOM DECONV3D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
+                 "%i instead !",
+                 oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto gradIShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  auto shapes = SHAPELIST(CONSTANT(gradIShapeInfo), CONSTANT(gradWShapeInfo));
+
+  if (biasShapeInfo != nullptr) {
+    auto gradBShapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    shapes->push_back(CONSTANT(gradBShapeInfo));
+  }
+
+  return shapes;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/depthwiseConv2d.cpp

@@ -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
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_depthwise_conv2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+#include <system/op_boilerplate.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(depthwise_conv2d, 2, 1, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC] = iC*mC
+
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oH, oW, iC*mC] (NHWC) or [bS, iC*mC, oH, oW] (NCHW)
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM DEPTHWISECONV2D OP: rank of input array must be equal to 4, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM DEPTHWISECONV2D OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                          // strides height
+  LongType sW = INT_ARG(3);                                                          // strides width
+  LongType pH = INT_ARG(4);                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;             // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  LongType bS, iC, iH, iW, mC, oC, oH, oW;  // batch size, input channels, input height/width, channels multiplier(oC =
+                                       // iC*mC), output channels, output height/width
+  LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWmC, indWkH, indOoH);
+  mC = weights->sizeAt(indWmC);  // channels multiplier
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DEPTHWISECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  REQUIRE_TRUE(
+      output->sizeAt(indIOioC) == iC * mC, 0,
+      "CUSTOM DEPTHWISECONV2D OP: the output_channels must be equal to input_channels * channels_multiplier = %i !",
+      iC * mC);
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+                 "CUSTOM DEPTHWISECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
+                 "%i, %i instead !",
+                 oC, bias->rankOf(), bias->lengthOf());
+
+  ConvolutionUtils::depthwiseConv2d(block, input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, isSameMode,
+                                    isNCHW, wFormat);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(depthwise_conv2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+DECLARE_SHAPE_FN(depthwise_conv2d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC] = iC*mC
+
+  const int rank = 4;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DEPTHWISECONV2D OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo[0]);
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DEPTHWISECONV2D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo[0]);
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                                          // strides height
+  LongType sW = INT_ARG(3);                                                                          // strides width
+  LongType pH = INT_ARG(4);                                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  int indIOioC, indIiH, indWmC(0 == wFormat ? 3 : 0);
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+  }
+
+  const LongType bS = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(0));           // batch size
+  const LongType iH = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIiH));      // input height
+  const LongType iW = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIiH + 1));  // input width
+  const LongType iC = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIOioC));    // input channels
+  const LongType mC = shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(indWmC));    // channels multiplier(oC = iC*mC)
+  const LongType oC = iC * mC;                                    // output channels
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(shape::shapeEquals(4, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "DEPTHWISECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+    REQUIRE_TRUE(shape::rank(biasShapeInfo) <= 2 && oC == shape::length(biasShapeInfo), 0,
+                 "DEPTHWISECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+                 "instead !",
+                 oC, shape::rank(biasShapeInfo), shape::length(biasShapeInfo));
+
+  LongType oH, oW;  // output height, width
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  sd::LongType* outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inputShapeInfo), sd::LongType);
+
+  outputShapeInfo[0] = rank;
+  outputShapeInfo[1] = bS;
+
+  if (isNCHW) {
+    outputShapeInfo[2] = oC;
+    outputShapeInfo[3] = oH;
+    outputShapeInfo[4] = oW;
+  } else {
+    outputShapeInfo[2] = oH;
+    outputShapeInfo[3] = oW;
+    outputShapeInfo[4] = oC;
+  }
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, weightsShapeInfo, shape::order(inputShapeInfo));
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+DECLARE_TYPES(depthwise_conv2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(depthwise_conv2d_bp, 3, 2, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
+  auto weights = INPUT_VARIABLE(1);                             // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr;  // [oC] = [iC*mC]
+  auto gradO = block.width() > 3
+                   ? INPUT_VARIABLE(3)
+                   : INPUT_VARIABLE(2);  // [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
+
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW), epsilon
+  auto gradW = OUTPUT_NULLIFIED(1);  // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  auto gradB = block.width() > 3 ? OUTPUT_NULLIFIED(2) : nullptr;  // [oC]
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of input array must be equal to 4, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+  REQUIRE_TRUE(gradO->rankOf() == 4, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to 4, but "
+               "got %i instead !",
+               gradO->rankOf());
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                          // strides height
+  LongType sW = INT_ARG(3);                                                          // strides width
+  LongType pH = INT_ARG(4);                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;             // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  LongType bS, iC, iH, iW, mC, oC, oH, oW;  // batch size, input channels, input height/width, channels multiplier(oC =
+                                       // iC*mC), output channels, output height/width
+  LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWmC, indWkH, indOoH);
+  mC = weights->sizeAt(indWmC);  // channels multiplier
+
+  LongType trueoH, trueoW;  // correct output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, "
+               "but got %s instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+                 "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
+                 "got %i, %i instead !",
+                 oC, bias->rankOf(), bias->lengthOf());
+
+  ConvolutionUtils::depthwiseConv2dBP(block, input, weights, bias, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW,
+                                      dH, dW, isSameMode, isNCHW, wFormat);
+
+  return sd::Status::OK;
+}
+
+//////////////////////////////////////////////////////////////////////
+DECLARE_SHAPE_FN(depthwise_conv2d_bp) {
+  auto inputShapeInfo = inputShape->at(0);
+  auto weightsShapeInfo = inputShape->at(1);
+  auto biasShapeInfo = block.width() > 3 ? inputShape->at(2) : nullptr;
+  auto gradOShapeInfo = block.width() > 3 ? inputShape->at(3) : inputShape->at(2);
+
+  const LongType rank = 4;
+  REQUIRE_TRUE(shape::rank(inputShapeInfo) == rank, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               shape::rank(inputShapeInfo));
+  REQUIRE_TRUE(shape::rank(weightsShapeInfo) == rank, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               shape::rank(weightsShapeInfo));
+  REQUIRE_TRUE(shape::rank(gradOShapeInfo) == rank, 0,
+               "CUSTOM DEPTHWISECONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but "
+               "got %i instead !",
+               rank, shape::rank(gradOShapeInfo));
+
+  LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(0)));  // filter(kernel) height
+  LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(1)));  // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                                          // strides height
+  LongType sW = INT_ARG(3);                                                                          // strides width
+  LongType pH = INT_ARG(4);                                                                          // paddings height
+  LongType pW = INT_ARG(5);                                                                          // paddings width
+  LongType dH = INT_ARG(6);                                                                          // dilations height
+  LongType dW = INT_ARG(7);                                                                          // dilations width
+  int isSameMode = INT_ARG(8);                                                                  // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  int indIOioC, indIiH, indWmC(0 == wFormat ? 3 : 0);
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+  }
+
+  const LongType bS = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(0));           // batch size
+  const LongType iH = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIiH));      // input height
+  const LongType iW = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIiH + 1));  // input width
+  const LongType iC = shape::sizeAt(inputShapeInfo, static_cast<sd::LongType>(indIOioC));    // input channels
+  const LongType mC = shape::sizeAt(weightsShapeInfo, static_cast<sd::LongType>(indWmC));    // channels multiplier(oC = iC*mC)
+  const LongType oC = iC * mC;                                    // output channels
+
+  LongType trueoH, trueoW;  // correct output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indIiH, indIiH + 1});
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(
+      shape::shapeEquals(4, expectedGradOShape.data(), shape::rank(gradOShapeInfo), shape::shapeOf(gradOShapeInfo)), 0,
+      "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s "
+      "instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  REQUIRE_TRUE(shape::shapeEquals(4, expectedWeightsShape.data(), shape::rank(weightsShapeInfo),
+                                  shape::shapeOf(weightsShapeInfo)),
+               0, "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+    REQUIRE_TRUE(shape::rank(biasShapeInfo) <= 2 && oC == shape::length(biasShapeInfo), 0,
+                 "CUSTOM DEPTHWISECONV2D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
+                 "got %i, %i instead !",
+                 oC, shape::rank(biasShapeInfo), shape::length(biasShapeInfo));
+
+  auto gradIshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  if (biasShapeInfo) {
+    sd::LongType* gradBshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo), CONSTANT(gradBshapeInfo));
+  }
+
+  return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWshapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/dilation2d.cpp

@@ -0,0 +1,133 @@
+/* ******************************************************************************
+ *
+ *
+ * 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>
+#if NOT_EXCLUDED(OP_dilation2d)
+
+#include <ops/declarable/headers/convo.h>
+#include <ops/declarable/helpers/dilation2d.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(dilation2d, 2, 1, false, 0, 1) {
+  auto input = INPUT_VARIABLE(0);
+  auto weights = INPUT_VARIABLE(1);
+
+  auto output = OUTPUT_VARIABLE(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "Dilation2D: input should be 4D");
+  REQUIRE_TRUE(weights->rankOf() == 3, 0, "Dilation2D: weights should be 3D");
+
+  const LongType bS = input->sizeAt(0);
+  const LongType iC = input->sizeAt(3);
+  const bool isSameShape = INT_ARG(0) == 1;
+
+  REQUIRE_TRUE(input->sizeAt(3) == weights->sizeAt(2), 0,
+               "Dilation2D: number of input channels doesn't match number of channels in weights: %i vs %i",
+               input->sizeAt(3), weights->sizeAt(2));
+
+  std::vector<sd::LongType> strides(4);
+  std::vector<sd::LongType> rates(4);
+
+  if (block.width() > 2) {
+    REQUIRE_TRUE(block.width() >= 4, 0, "Dilation2D: number of input arrays should be 4 at least");
+
+    auto r = INPUT_VARIABLE(2);
+    auto s = INPUT_VARIABLE(3);
+
+    strides = s->template asVectorT<sd::LongType>();
+    rates = r->template asVectorT<sd::LongType>();
+  } else {
+    REQUIRE_TRUE(block.numI() >= 9, 0, "Dilation2D: number of Int arguments should be 9 at least");
+
+    int e = 1;
+    for (int cnt = 0; cnt < 4; cnt++) rates[cnt] = INT_ARG(e++);
+
+    for (int cnt = 0; cnt < 4; cnt++) strides[cnt] = INT_ARG(e++);
+  }
+
+  sd::LongType sH = 0, sW = 0;
+  sd::LongType dH = 0, dW = 0;
+  sd::LongType pH = 0, pW = 0;
+  sd::LongType oH = 0, oW = 0;
+
+  helpers::dilation_hw(block.launchContext(), input->shapeInfo(), weights->shapeInfo(), strides, rates, isSameShape,
+                       &sH, &sW, &pH, &pW, &dH, &dW, &oH, &oW);
+
+  REQUIRE_TRUE(oH > 0 && oW > 0, 0, "Dilation2D: outY and outX should have positive values, but got [%i, %i] instead",
+               oH, oW);
+
+  helpers::dilation2d(block.launchContext(), input, weights, output, sH, sW, pH, pW, dH, dW);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(dilation2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dilation2d) {
+  auto input = inputShape->at(0);
+  auto weights = inputShape->at(1);
+
+  const int bS = shape::sizeAt(input, static_cast<sd::LongType>(0));
+  const int iC = shape::sizeAt(input, static_cast<sd::LongType>(3));
+  const bool isSameShape = INT_ARG(0) == 1;
+
+  std::vector<sd::LongType> strides(4);
+  std::vector<sd::LongType> rates(4);
+
+  if (block.width() > 2) {
+    auto r = INPUT_VARIABLE(2);
+    auto s = INPUT_VARIABLE(3);
+
+    strides = s->template asVectorT<sd::LongType>();
+    rates = r->template asVectorT<sd::LongType>();
+  } else {
+    if (block.numI() < 9) {
+      auto newShape = ConstantShapeHelper::getInstance().scalarShapeInfo(block.dataType());
+      return SHAPELIST(newShape);
+    }
+
+    int e = 1;
+    for (int cnt = 0; cnt < 4; cnt++) rates[cnt] = INT_ARG(e++);
+
+    for (int cnt = 0; cnt < 4; cnt++) strides[cnt] = INT_ARG(e++);
+  }
+
+  sd::LongType sH = 0, sW = 0;
+  sd::LongType dH = 0, dW = 0;
+  sd::LongType pH = 0, pW = 0;
+  sd::LongType oH = 0, oW = 0;
+
+  helpers::dilation_hw(block.launchContext(), input, weights, strides, rates, isSameShape, &sH, &sW, &pH, &pW, &dH, &dW,
+                       &oH, &oW);
+
+  std::array<sd::LongType, 4> shape = {{bS, oH, oW, iC}};
+  auto newShape =
+      ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(weights), 'c', 4, shape.data(),0);
+  return SHAPELIST(newShape);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/im2col.cpp

@@ -0,0 +1,157 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 17.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_im2col)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/col2im.h>
+#include <ops/declarable/helpers/convolutions.h>
+#include <ops/declarable/helpers/im2col.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(im2col, 1, 1, false, 0, 9) {
+  auto x = INPUT_VARIABLE(0);
+  auto z = OUTPUT_NULLIFIED(0);
+
+  REQUIRE_TRUE(x->rankOf() == 4, 0, "im2col input should be 4D, but got %i instead", x->rankOf());
+  REQUIRE_TRUE(z->rankOf() == 6, 0, "im2col output should be 6D, but got %i instead", z->rankOf());
+
+  LongType kernelHeight = INT_ARG(0);
+  LongType kernelWidth = INT_ARG(1);
+  LongType strideY = INT_ARG(2);
+  LongType strideX = INT_ARG(3);
+  LongType padHeight = INT_ARG(4);
+  LongType padWidth = INT_ARG(5);
+  LongType dY = INT_ARG(6);  // Dilation, height/y dimension
+  LongType dX = INT_ARG(7);  // Dilation, width/x dimension
+  bool isSameMode = INT_ARG(8) > 0;
+  double zeroPadVal = 0.0;
+  if (block.getTArguments()->size() > 0) zeroPadVal = T_ARG(0);
+
+  // FIXME: zeropad value is void
+  LaunchContext* ctx = block.launchContext();
+  NDArray *zero =  NDArrayFactory::create(zeroPadVal, block.launchContext());
+  sd::ops::helpers::im2col(*ctx, *x, *z, kernelHeight, kernelWidth, strideY, strideX, padHeight, padWidth, dY, dX,
+                          *zero);
+
+  delete zero;
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(im2col) {
+  auto inShape = inputShape->at(0);
+
+  LongType bS = shape::shapeOf(inShape)[0];
+  LongType iD = shape::shapeOf(inShape)[1];
+  LongType inY = shape::shapeOf(inShape)[2];
+  LongType inX = shape::shapeOf(inShape)[3];
+
+  LongType kY = INT_ARG(0);
+  LongType kX = INT_ARG(1);
+  LongType sY = INT_ARG(2);
+  LongType sX = INT_ARG(3);
+  sd::LongType pY = INT_ARG(4);
+  sd::LongType pX = INT_ARG(5);
+  LongType dY = INT_ARG(6);  // Dilation, height/y dimension
+  LongType dX = INT_ARG(7);  // Dilation, width/x dimension
+  int paddingMode = INT_ARG(8);
+  bool isSameMode = INT_ARG(8) == 1;
+  // output is always 6d for im2col
+  sd::LongType* zShape;
+  ALLOCATE(zShape, block.getWorkspace(), shape::shapeInfoLength(6), sd::LongType);
+
+  LongType oY = 0;
+  LongType oX = 0;
+
+  ConvolutionUtils::calcOutSizePool2D(oY, oX, kY, kX, sY, sX, pY, pX, dY, dX, inY, inX, paddingMode);
+
+  if (isSameMode) ConvolutionUtils::calcPadding2D(pY, pX, oY, oX, inY, inX, kY, kX, sY, sX, dY, dX);
+
+  zShape[0] = 6;
+  zShape[1] = bS;
+  zShape[2] = iD;
+  zShape[3] = kY;
+  zShape[4] = kX;
+  zShape[5] = oY;
+  zShape[6] = oX;
+
+  zShape[shape::shapeInfoLength(zShape) - 2] = 1;
+  zShape[shape::shapeInfoLength(zShape) - 1] = 99;
+
+  ShapeUtils::updateStridesAndType(zShape, inShape, 'c');
+
+  return SHAPELIST(CONSTANT(zShape));
+}
+
+CUSTOM_OP_IMPL(im2col_bp, 2, 1, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);
+  auto gradAtOutput = INPUT_VARIABLE(1);
+  auto z = OUTPUT_NULLIFIED(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "im2col_bp input should be 4D, but got %i instead", input->rankOf());
+  REQUIRE_TRUE(gradAtOutput->rankOf() == 6, 0,
+               "im2col_bp gradient at output (input idx 1) should be 6D, but got %i instead", gradAtOutput->rankOf());
+  REQUIRE_TRUE(z->rankOf() == 4, 0, "im2col_bp output (grad at input) should be 4D, but got %i instead", z->rankOf());
+
+  LongType kernelHeight = INT_ARG(0);
+  LongType kernelWidth = INT_ARG(1);
+  LongType strideY = INT_ARG(2);
+  LongType strideX = INT_ARG(3);
+  LongType pH = INT_ARG(4);
+  LongType pW = INT_ARG(5);
+  LongType dY = INT_ARG(6);  // Dilation, height/y dimension
+  LongType dX = INT_ARG(7);  // Dilation, width/x dimension
+
+  // Assuming NCHW format here
+  int imgH = input->sizeAt(2);
+  int imgW = input->sizeAt(3);
+
+  LaunchContext* ctx = block.launchContext();
+  // FIXME:: all helpers should accept NDArray
+  ops::helpers::col2im(*ctx, gradAtOutput, z, strideY, strideX, pH, pW, imgH, imgW, dY, dX);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(im2col) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, DataType::ANY)
+      ->setAllowedOutputTypes(0, DataType::INHERIT)
+      ->setSameMode(true);
+}
+
+DECLARE_TYPES(im2col_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, DataType::ANY)
+      ->setAllowedOutputTypes(0, DataType::INHERIT)
+      ->setSameMode(true);
+}
+
+DECLARE_SHAPE_FN(im2col_bp) {
+  return SHAPELIST(CONSTANT(inputShape->at(0)));
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/ismax.cpp

@@ -0,0 +1,54 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com, created on 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_ismax)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/ismax.h>
+
+namespace sd {
+namespace ops {
+
+CONFIGURABLE_OP_IMPL(ismax, 1, 1, true, 0, -2) {
+  auto x = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+  auto dimensions = *(block.getIArguments());  // argI
+  int one = 1;
+  if (x->isScalar())
+    z->assign(one);
+  else
+    helpers::ismax(block.launchContext(), x, z, dimensions);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(IsMax, ismax);
+
+DECLARE_TYPES(ismax) {
+  getOpDescriptor()->setAllowedInputTypes(0, DataType::ANY)->setAllowedOutputTypes(0, DataType::ANY);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/pointwiseConv2d.cpp

@@ -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 Yurii Shyrma, created on 20.03.2018
+//
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+#if NOT_EXCLUDED(op_pointwise_conv2d)
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(pointwise_conv2d, 2, 1, false, 0, 0) {
+  auto input = INPUT_VARIABLE(0);                               // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weights = INPUT_VARIABLE(1);                             // [1, 1, iC, oC], [oC, iC, 1, 1], [oC, 1, 1, iC]
+  auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;  // [oC]
+
+  auto output = OUTPUT_VARIABLE(0);  // [bS, iH, iW, oC] (NHWC) or [bS, oC, iH, iW] (NCHW)
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0,
+               "CUSTOM POINTWISECONV2D OP: rank of input array must be equal to 4, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(weights->rankOf() == 4, 0,
+               "CUSTOM POINTWISECONV2D OP: rank of weights array must be equal to 4, but got %i instead !",
+               weights->rankOf());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2, 0,
+                 "CUSTOM POINTWISECONV2D OP: rank of biases array must be equal <= 2, but got %i instead !",
+                 bias->rankOf());
+
+  LongType kH = 1;                                                        // filter(kernel) height
+  LongType kW = 1;                                                        // filter(kernel) width
+  LongType sH = 1;                                                        // strides height
+  LongType sW = 1;                                                        // strides width
+  LongType pH = 0;                                                        // paddings height
+  LongType pW = 0;                                                        // paddings width
+  LongType dH = 1;                                                        // dilations height
+  LongType dW = 1;                                                        // dilations width
+  int isNCHW = block.getIArguments()->size() > 0 ? !INT_ARG(0) : 1;  // INT_ARG(0): 0-NCHW, 1-NHWC
+  int wFormat =
+      block.getIArguments()->size() > 1 ? INT_ARG(1) : 0;  // 0 - [1, 1, iC, oC], 1 - [oC, iC, 1, 1], 2 - [oC, 1, 1, iC]
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWoC, indWkH, indOoH);
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC, oC);
+  REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
+               "CUSTOM POINTWISECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
+                 "CUSTOM POINTWISECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
+                 "%i, %i instead !",
+                 oC, bias->rankOf(), bias->lengthOf());
+
+  ConvolutionUtils::conv2d(block, input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, 1 /*isSameMode*/,
+                           isNCHW, wFormat);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(pointwise_conv2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(pointwise_conv2d) {
+  auto inputShapeInfo = inputShape->at(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto weightsShapeInfo = inputShape->at(1);  // [1, 1, iC, oC], [oC, iC, 1, 1], [oC, 1, 1, iC]
+  auto biasShapeInfo = block.width() > 2 ? inputShape->at(2) : nullptr;  // [oC]
+
+  const int rank = 4;
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "CUSTOM POINTWISECONV2D OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo[0]);
+  REQUIRE_TRUE(weightsShapeInfo[0] == rank, 0,
+               "CUSTOM POINTWISECONV2D OP: rank of weights array must be equal to %i, but got %i instead !", rank,
+               weightsShapeInfo[0]);
+
+  int isNCHW = block.getIArguments()->size() > 0 ? !INT_ARG(0) : 1;  // INT_ARG(0): 0-NCHW, 1-NHWC
+  int wFormat =
+      block.getIArguments()->size() > 1 ? INT_ARG(1) : 0;  // 0 - [1, 1, iC, oC], 1 - [oC, iC, 1, 1], 2 - [oC, 1, 1, iC]
+
+  LongType indIOioC, indWoC(0 == wFormat ? 3 : 0);
+  if (!isNCHW)
+    indIOioC = 3;
+  else
+    indIOioC = 1;
+
+  const LongType bS = inputShapeInfo[1];             // batch size
+  const LongType iC = inputShapeInfo[indIOioC + 1];  // input channels
+  const LongType oC = weightsShapeInfo[indWoC + 1];  // output channels
+
+  std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC, oC);
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsShapeInfo, expectedWeightsShape), 0,
+               "POINTWISECONV2D OP: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsShape).c_str(),
+               ShapeUtils::shapeAsString(weightsShapeInfo).c_str());
+  if (biasShapeInfo)
+    REQUIRE_TRUE(biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+                 "POINTWISECONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
+                 "instead !",
+                 oC, biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto outputShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, weightsShapeInfo, true, block.getWorkspace());
+
+  // do not forget to put oC instead of iC in outputShapeInfo
+  outputShapeInfo[indIOioC + 1] = oC;
+
+  shape::updateStrides(outputShapeInfo, shape::order(inputShapeInfo), false);
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/sconv2d.cpp

@@ -0,0 +1,489 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 29/10/17.
+// @author Yurii Shyrma, changed on 20.03.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_sconv2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+#include <memory>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(sconv2d, 2, 1, false, 0, 9) {
+  NDArray *input = INPUT_VARIABLE(0);         // [bS, iH, iW, iC]  (NHWC) or [bS, iC, iH, iW]  (NCHW)
+  NDArray *weightsDepth = INPUT_VARIABLE(1);  // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  NDArray *weightsPoint = nullptr;            // [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
+  NDArray *bias = nullptr;                    // [oC], if weightsPoint=nullptr then oC = iC*mC
+
+  NDArray *output = OUTPUT_NULLIFIED(0);  // [bS, oH, oW, oC]  (NHWC) or [bS, oC, oH, oW]  (NCHW)
+
+  if (block.width() == 3) {
+    if ((INPUT_VARIABLE(2))->rankOf() == 4)
+      weightsPoint = INPUT_VARIABLE(2);
+    else
+      bias = INPUT_VARIABLE(2);
+  } else if (block.width() == 4) {
+    weightsPoint = INPUT_VARIABLE(2);
+    bias = INPUT_VARIABLE(3);
+  }
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, " SCONV2D OP: rank of input array must be equal to 4, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(weightsDepth->rankOf() == 4, 0,
+               " SCONV2D OP: rank of weightsDepth array must be equal to 4, but got %i instead !",
+               weightsDepth->rankOf());
+  if (weightsPoint)
+    REQUIRE_TRUE(weightsPoint->rankOf() == 4, 0,
+                 " SCONV2D OP: rank of weightsPoint array must be equal to 4, but got %i instead !",
+                 weightsPoint->rankOf());
+  if (bias)
+    REQUIRE_TRUE(bias->rankOf() == 1 || bias->rankOf() == 2, 0,
+                 " SCONV2D OP: rank of biases array must be equal to 1 or 2, but got %i instead !", bias->rankOf());
+  ;
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  LongType bS, iC, iH, iW, mC, oC, oH,
+      oW;  // batch size, input channels, input height/width, channels multiplier, output channels, output height/width
+  LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWmC, indWkH, indOoH);
+  mC = weightsDepth->sizeAt(indWmC);  // channels multiplier
+
+  std::vector<sd::LongType> expectedWeightsDShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(weightsDepth->isSameShape(expectedWeightsDShape), 0,
+               " SCONV2D OP: wrong shape of weightsDepth array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsDShape).c_str(),
+               ShapeUtils::shapeAsString(weightsDepth).c_str());
+  if (weightsPoint) {
+    std::vector<sd::LongType> expectedWeightsPShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC * mC, oC);
+    REQUIRE_TRUE(weightsPoint->isSameShape(expectedWeightsPShape), 0,
+                 " SCONV2D OP: wrong shape of weightsPoint array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedWeightsPShape).c_str(),
+                 ShapeUtils::shapeAsString(weightsPoint).c_str());
+  }
+  if (bias)
+    REQUIRE_TRUE(oC == bias->lengthOf(), 0,
+                 " SCONV2D OP: length of bias array must be equal to outChannels, but got %i instead",
+                 bias->lengthOf());
+
+  if (iC == 1) {
+    sd_debug("SCONV2D OP: for input_channels = 1 this op is equivalent to standard conv2d\n", "");
+    ConvolutionUtils::conv2d(block, input, weightsDepth, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, isSameMode,
+                             isNCHW, wFormat);
+    return sd::Status::OK;
+  }
+
+  ConvolutionUtils::sconv2d(block, input, weightsDepth, weightsPoint, bias, output, kH, kW, sH, sW, pH, pW, dH, dW,
+                            isSameMode, isNCHW, wFormat);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(sconv2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(sconv2d) {
+  auto inputShapeInfo = inputShape->at(0);          // [bS, iH, iW, iC]  (NHWC) or [bS, iC, iH, iW]  (NCHW)
+  auto weightsDShapeInfo = inputShape->at(1);       // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  sd::LongType const *weightsPShapeInfo = nullptr;  // [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
+  sd::LongType const *biasShapeInfo = nullptr;      // [oC], oC = iC*mC if weightsPoint=nullptr
+
+  if (block.width() == 3)
+    if (inputShape->at(2)[0] == 4)
+      weightsPShapeInfo = inputShape->at(2);
+    else
+      biasShapeInfo = inputShape->at(2);
+  else if (block.width() == 4) {
+    weightsPShapeInfo = inputShape->at(2);
+    biasShapeInfo = inputShape->at(3);
+  }
+
+  const LongType rank = 4;
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               "SCONV2D OP: rank of input array must be equal to %i, but got %i instead !", rank, inputShapeInfo[0]);
+  REQUIRE_TRUE(weightsDShapeInfo[0] == rank, 0,
+               "SCONV2D OP: rank of weightsDepth array must be equal to %i, but got %i instead !", rank,
+               weightsDShapeInfo[0]);
+  if (weightsPShapeInfo)
+    REQUIRE_TRUE(weightsPShapeInfo[0] == rank, 0,
+                 "SCONV2D OP: rank of weightsPoint array must be equal to %i, but got %i instead !", rank,
+                 weightsPShapeInfo[0]);
+  if (biasShapeInfo)
+    REQUIRE_TRUE(biasShapeInfo[0] <= 2, 0, "SCONV2D OP: rank of biases array must be <= 2, but got %i instead !",
+                 biasShapeInfo[0]);
+
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 1-NHWC, 0-NCHW
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  LongType indIOioC, indIiH, indWmC(0 == wFormat ? 3 : 0);
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];                                            // batch size
+  const LongType iH = inputShapeInfo[indIiH + 1];                                   // input height
+  const LongType iW = inputShapeInfo[indIiH + 2];                                   // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];                                 // input channels
+  const LongType mC = weightsDShapeInfo[indWmC + 1];                                // channel multiplier
+  const LongType oC = weightsPShapeInfo ? weightsPShapeInfo[indWmC + 1] : iC * mC;  // output channels (oC or iC*mC)
+
+  std::vector<sd::LongType> expectedWeightsDShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsDShapeInfo, expectedWeightsDShape), 0,
+               "SCONV2D OP: wrong shape of depth weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsDShape).c_str(),
+               ShapeUtils::shapeAsString(weightsDShapeInfo).c_str());
+  if (weightsPShapeInfo) {
+    std::vector<sd::LongType> expectedWeightsPShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC * mC, oC);
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsPShapeInfo, expectedWeightsPShape), 0,
+                 "SCONV2D OP: wrong shape of point array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedWeightsPShape).c_str(),
+                 ShapeUtils::shapeAsString(weightsPShapeInfo).c_str());
+  }
+  if (biasShapeInfo)
+    REQUIRE_TRUE(
+        biasShapeInfo[0] <= 2 && oC == shape::length(biasShapeInfo), 0,
+        "SCONV2D OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !", oC,
+        biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  LongType oH, oW;  // output height, width
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  sd::LongType *outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inputShapeInfo), sd::LongType);
+
+  outputShapeInfo[0] = 4;
+  outputShapeInfo[1] = bS;
+
+  if (isNCHW) {
+    outputShapeInfo[2] = oC;
+    outputShapeInfo[3] = oH;
+    outputShapeInfo[4] = oW;
+  } else {
+    outputShapeInfo[2] = oH;
+    outputShapeInfo[3] = oW;
+    outputShapeInfo[4] = oC;
+  }
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, weightsDShapeInfo, shape::order(inputShapeInfo));
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+DECLARE_TYPES(sconv2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sconv2d_bp, 3, 2, false, 0, 9) {
+  NDArray *input = INPUT_VARIABLE(0);         // [bS, iH, iW, iC]  (NHWC) or [bS, iC, iH, iW]  (NCHW)
+  NDArray *gradO = INPUT_VARIABLE(1);         // [bS, oH, oW, oC]  (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  NDArray *weightsDepth = INPUT_VARIABLE(2);  // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  NDArray *weightsPoint = nullptr;            // [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
+  NDArray *bias = nullptr;                    // [oC], oC = iC*mC if weightsPoint=nullptr
+
+  NDArray *gradI = OUTPUT_NULLIFIED(0);   // [bS, iH, iW, iC]  (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+  NDArray *gradWD = OUTPUT_NULLIFIED(1);  // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  NDArray *gradWP = nullptr;              // [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
+  NDArray *gradB = nullptr;               // [oC]
+
+  if (block.width() == 4) {
+    if ((INPUT_VARIABLE(3))->rankOf() == 4) {
+      weightsPoint = INPUT_VARIABLE(3);
+      gradWP = OUTPUT_NULLIFIED(2);
+    } else {
+      bias = INPUT_VARIABLE(3);
+      gradB = OUTPUT_NULLIFIED(2);
+    }
+  } else if (block.width() == 5) {
+    weightsPoint = INPUT_VARIABLE(3);
+    bias = INPUT_VARIABLE(4);
+    gradWP = OUTPUT_NULLIFIED(2);
+    gradB = OUTPUT_NULLIFIED(3);
+  }
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, " SCONV2D_BP OP: rank of input array must be equal to 4, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(gradO->rankOf() == 4, 0,
+               " SCONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to 4, but got %i instead !",
+               gradO->rankOf());
+  REQUIRE_TRUE(weightsDepth->rankOf() == 4, 0,
+               " SCONV2D_BP OP: rank of weightsDepth array must be equal to 4 !, but got %i instead !",
+               weightsDepth->rankOf());
+  if (weightsPoint) {
+    REQUIRE_TRUE(weightsPoint->rankOf() == 4, 0,
+                 " SCONV2D_BP OP: rank of weightsPoint array must be equal to 4, but got %i instead !",
+                 weightsPoint->rankOf());
+    REQUIRE_TRUE(gradWP->rankOf() == 4, 0,
+                 " SCONV2D_BP OP: rank of weightsPoint gradients array must be equal to 4, but got %i instead !",
+                 gradWP->rankOf());
+  }
+  if (bias) {
+    REQUIRE_TRUE(bias->rankOf() == 1 || bias->rankOf() == 2, 0,
+                 " SCONV2D_BP OP: rank of biases array must be equal to 1 or 2, but got %i instead !", bias->rankOf());
+    REQUIRE_TRUE(gradB->rankOf() == 1 || gradB->rankOf() == 2, 0,
+                 " SCONV2D_BP OP: rank of  biases gradientsarray must be equal to 1 or 2, but got %i instead !",
+                 gradB->rankOf());
+  }
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  LongType bS, iC, iH, iW, mC, oC, oH,
+      oW;  // batch size, input channels, input height/width, channels multiplier, output channels, output height/width
+  LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
+                                             indIiH, indWiC, indWmC, indWkH, indOoH);
+  mC = weightsDepth->sizeAt(indWmC);  // channels multiplier
+
+  std::vector<sd::LongType> expectedWeightsDShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(weightsDepth->isSameShape(expectedWeightsDShape), 0,
+               " SCONV2D_BP OP: wrong shape of weightsDepth array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsDShape).c_str(),
+               ShapeUtils::shapeAsString(weightsDepth).c_str());
+  REQUIRE_TRUE(gradWD->isSameShape(expectedWeightsDShape), 0,
+               " SCONV2D_BP OP: wrong shape of gradWD array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsDShape).c_str(), ShapeUtils::shapeAsString(gradWD).c_str());
+  if (weightsPoint) {
+    std::vector<sd::LongType> expectedWeightsPShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC * mC, oC);
+    REQUIRE_TRUE(weightsPoint->isSameShape(expectedWeightsPShape), 0,
+                 " SCONV2D_BP OP: wrong shape of weightsPoint array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedWeightsPShape).c_str(),
+                 ShapeUtils::shapeAsString(weightsPoint).c_str());
+    REQUIRE_TRUE(gradWP->isSameShape(expectedWeightsPShape), 0,
+                 " SCONV2D_BP OP: wrong shape of gradWP array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedWeightsPShape).c_str(), ShapeUtils::shapeAsString(gradWP).c_str());
+  }
+  if (bias) {
+    REQUIRE_TRUE(oC == bias->lengthOf(), 0,
+                 " SCONV2D_BP OP: length of bias array must be equal to outChannels, but got %i instead",
+                 bias->lengthOf());
+    REQUIRE_TRUE(oC == gradB->lengthOf(), 0,
+                 " SCONV2D_BP OP: length of biases gradients array must be equal to outChannels, but got %i instead",
+                 gradB->lengthOf());
+  }
+
+
+
+  // ----- if weightsPoint is present, perform pointwise backprop first and calculate gradWP at this step ----- //
+  if (weightsPoint) {
+    auto resultFFShape =
+        isNCHW ? std::vector<sd::LongType>({bS, mC * iC, oH, oW}) : std::vector<sd::LongType>({bS, oH, oW, mC * iC});
+    auto resultFF = NDArrayFactory::create_(input->ordering(), resultFFShape, input->dataType(), block.launchContext());
+    ConvolutionUtils::sconv2d(block, input, weightsDepth, nullptr, nullptr, resultFF, kH, kW, sH, sW, pH, pW, dH, dW,
+                              isSameMode, isNCHW, wFormat);
+
+    auto gradIDepthShape =
+        ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC * mC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
+    auto gradIDepth =
+        NDArrayFactory::create_(resultFF->ordering(), gradIDepthShape, resultFF->dataType(),
+                                block.launchContext());  // [bS, oH, oW, iC*mC]  (NHWC) or [bS, iC*mC, oH, oW] (NCHW)
+
+    ConvolutionUtils::conv2dBP(block, resultFF, weightsPoint, bias, gradO, gradIDepth, gradWP, gradB, 1, 1, 1, 1, 0, 0,
+                               1, 1, isSameMode, isNCHW, wFormat);  // in this case oH=iH and oW=iW
+
+    gradO = gradIDepth;
+    bias = gradB = nullptr;  // if pointwise backprop was done then don't calculate gradB at depthwise_conv2d_bp step
+
+    delete resultFF;
+  }
+
+  // ----- apply depthwise_conv2d_bp ----- //
+  ConvolutionUtils::depthwiseConv2dBP(block, input, weightsDepth, bias, gradO, gradI, gradWD, gradB, kH, kW, sH, sW, pH,
+                                      pW, dH, dW, isSameMode, isNCHW, wFormat);
+
+  if (weightsPoint) delete gradO;
+
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(sconv2d_bp) {
+  auto inputShapeInfo = inputShape->at(0);          // [bS, iH, iW, iC]  (NHWC) or [bS, iC, iH, iW]  (NCHW)
+  auto gradOShapeInfo = inputShape->at(1);          // [bS, oH, oW, oC]  (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto weightsDShapeInfo = inputShape->at(2);       // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
+  sd::LongType const *weightsPShapeInfo = nullptr;  // [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
+  sd::LongType const *biasShapeInfo = nullptr;      // [oC], oC = iC*mC if weightsPoint=nullptr
+
+  if (block.width() == 4) {
+    if (inputShape->at(3)[0] == 4)
+      weightsPShapeInfo = inputShape->at(3);
+    else
+      biasShapeInfo = inputShape->at(3);
+  } else if (block.width() == 5) {
+    weightsPShapeInfo = inputShape->at(3);
+    biasShapeInfo = inputShape->at(4);
+  }
+
+  const LongType rank = 4;
+  REQUIRE_TRUE(inputShapeInfo[0] == rank, 0,
+               " SCONV2D_BP OP: rank of input array must be equal to %i, but got %i instead !", rank,
+               inputShapeInfo[0]);
+  REQUIRE_TRUE(
+      gradOShapeInfo[0] == rank, 0,
+      " SCONV2D_BP OP: rank of output gradients (next epsilon) array must be equal to %i, but got %i instead !", rank,
+      gradOShapeInfo[0]);
+  REQUIRE_TRUE(weightsDShapeInfo[0] == rank, 0,
+               " SCONV2D_BP OP: rank of weightsDepth array must be equal to %i, but got %i instead !", rank,
+               weightsDShapeInfo[0]);
+  if (weightsPShapeInfo)
+    REQUIRE_TRUE(weightsPShapeInfo[0] == rank, 0,
+                 " SCONV2D_BP OP: rank of weightsPoint array must be equal to %i, but got %i instead !", rank,
+                 weightsPShapeInfo[0]);
+  if (biasShapeInfo)
+    REQUIRE_TRUE(biasShapeInfo[0] == 1 || biasShapeInfo[0] == 2, 0,
+                 " SCONV2D_BP OP: rank of biases array must be 1 or 2, but got %i instead !", biasShapeInfo[0]);
+
+
+  LongType kH = INT_ARG(0);                                               // filter(kernel) height
+  LongType kW = INT_ARG(1);                                               // filter(kernel) width
+  LongType sH = INT_ARG(2);                                               // strides height
+  LongType sW = INT_ARG(3);                                               // strides width
+  LongType pH = INT_ARG(4);                                               // paddings height
+  LongType pW = INT_ARG(5);                                               // paddings width
+  LongType dH = INT_ARG(6);                                               // dilations height
+  LongType dW = INT_ARG(7);                                               // dilations width
+  int isSameMode = INT_ARG(8);                                       // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1;  // INT_ARG(9): 0-NCHW,  1-NHWC
+  int wFormat = block.getIArguments()->size() > 10
+                    ? INT_ARG(10)
+                    : 0;  // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
+
+  int indIOioC, indIiH, indWmC(0 == wFormat ? 3 : 0);
+  if (!isNCHW) {
+    indIOioC = 3;
+    indIiH = 1;
+  } else {
+    indIOioC = 1;
+    indIiH = 2;
+  }
+
+  const LongType bS = inputShapeInfo[1];                                            // batch size
+  const LongType iH = inputShapeInfo[indIiH + 1];                                   // input height
+  const LongType iW = inputShapeInfo[indIiH + 2];                                   // input width
+  const LongType iC = inputShapeInfo[indIOioC + 1];                                 // input channels
+  const LongType mC = weightsDShapeInfo[indWmC + 1];                                // channel multiplier
+  const LongType oC = weightsPShapeInfo ? weightsPShapeInfo[indWmC + 1] : iC * mC;  // output channels (oC or iC*mC)
+
+  LongType trueoH, trueoW;  // true output height, width
+  ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  std::vector<sd::LongType> expectedGradOShapeInfo =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indIiH, indIiH + 1});
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(gradOShapeInfo, expectedGradOShapeInfo), 0,
+      "SCONV2D_BP OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShapeInfo).c_str(), ShapeUtils::shapeAsString(gradOShapeInfo).c_str());
+  std::vector<sd::LongType> expectedWeightsDShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsDShapeInfo, expectedWeightsDShape), 0,
+               "SCONV2D_BP OP: wrong shape of depth weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWeightsDShape).c_str(),
+               ShapeUtils::shapeAsString(weightsDShapeInfo).c_str());
+  if (weightsPShapeInfo) {
+    std::vector<sd::LongType> expectedWeightsPShape = ConvolutionUtils::expectWeightsShape(wFormat, 1, 1, iC * mC, oC);
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(weightsPShapeInfo, expectedWeightsPShape), 0,
+                 "SCONV2D_BP OP: wrong shape of point array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedWeightsPShape).c_str(),
+                 ShapeUtils::shapeAsString(weightsPShapeInfo).c_str());
+  }
+  if (biasShapeInfo)
+    REQUIRE_TRUE(
+        (biasShapeInfo[0] == 1 || biasShapeInfo[0] == 2) && oC == shape::length(biasShapeInfo), 0,
+        "SCONV2D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !", oC,
+        biasShapeInfo[0], shape::length(biasShapeInfo));
+
+  auto gradIshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+  auto gradWDshapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(weightsDShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+
+  sd::LongType *gradWPshapeInfo(nullptr), *gradBshapeInfo(nullptr);
+
+  if (weightsPShapeInfo && biasShapeInfo) {
+    gradWPshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(weightsPShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    gradBshapeInfo = ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWDshapeInfo), CONSTANT(gradWPshapeInfo),
+                     CONSTANT(gradBshapeInfo));
+  }
+
+  if (weightsPShapeInfo && !biasShapeInfo) {
+    gradWPshapeInfo =
+        ShapeBuilders::copyShapeInfoAndType(weightsPShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWDshapeInfo), CONSTANT(gradWPshapeInfo));
+  }
+
+  if (!weightsPShapeInfo && biasShapeInfo) {
+    gradBshapeInfo = ShapeBuilders::copyShapeInfoAndType(biasShapeInfo, gradOShapeInfo, false, block.getWorkspace());
+    return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWDshapeInfo), CONSTANT(gradBshapeInfo));
+  }
+
+  return SHAPELIST(CONSTANT(gradIshapeInfo), CONSTANT(gradWDshapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/upsampling2d.cpp

@@ -0,0 +1,126 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119, created on 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com), changed on 03.05.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_upsampling2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(upsampling2d, 1, 1, false, 0, 2) {
+  auto input = INPUT_VARIABLE(0);  // [bS, iC, iH, iW] (NCHW) or [bS, iH, iW, iC] (NHWC)
+  auto output =
+      OUTPUT_NULLIFIED(0);  // [bS, iC, factorH*iH, factorW*iW ] (NCHW) or [bS, factorH*iH, factorW*iW, iC] (NHWC)
+
+  const int factorH = INT_ARG(0);
+  const int factorW = INT_ARG(1);
+  const int isNCHW = block.getIArguments()->size() > 2 ? INT_ARG(2) : 0;  // INT_ARG(2): 0-NCHW,  1-NHWC
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "UPSAMPLING2D op: input should be 4D, but got %i instead!", input->rankOf());
+  REQUIRE_TRUE(output->rankOf() == 4, 0, "UPSAMPLING2D op: output should be 4D, but got %i instead!", output->rankOf());
+
+  ConvolutionUtils::upsampling2d(block, *input, *output, factorH, factorW, (bool)isNCHW);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(upsampling, upsampling2d);
+
+DECLARE_TYPES(upsampling2d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(upsampling2d) {
+  auto inputShapeInfo = inputShape->at(0);
+
+  REQUIRE_TRUE(inputShapeInfo[0] == 4, 0, "UPSAMPLING2D op: input should be 4D, but got %i instead!",
+               inputShapeInfo[0]);
+
+  const LongType factorH = INT_ARG(0);
+  const LongType factorW = INT_ARG(1);
+  const int isNCHW = block.getIArguments()->size() > 2 ? INT_ARG(2) : 0;  // INT_ARG(2): 0-NCHW,  1-NHWC
+
+  sd::LongType *outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inputShapeInfo[0]), sd::LongType);
+
+  outputShapeInfo[0] = inputShapeInfo[0];
+  outputShapeInfo[1] = inputShapeInfo[1];
+
+  if (isNCHW) {
+    outputShapeInfo[2] = inputShapeInfo[2];
+    outputShapeInfo[3] = inputShapeInfo[3] * factorH;
+    outputShapeInfo[4] = inputShapeInfo[4] * factorW;
+  } else {
+    outputShapeInfo[2] = inputShapeInfo[2] * factorH;
+    outputShapeInfo[3] = inputShapeInfo[3] * factorW;
+    outputShapeInfo[4] = inputShapeInfo[4];
+  }
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, inputShapeInfo, shape::order(inputShapeInfo));
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+DECLARE_TYPES(upsampling2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(upsampling2d_bp, 2, 1, false, 0, 0) {
+  // NDArray<T>* input = INPUT_VARIABLE(0);             // [bS, iC, iH, iW] (NCHW) or [bS, iH, iW, iC] (NHWC)
+  auto gradO =
+      INPUT_VARIABLE(1);  // [bS, iC, factorH*iH, factorW*iW ] (NCHW) or [bS, factorH*iH, factorW*iW, iC] (NHWC)
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iC, iH, iW] (NCHW) or [bS, iH, iW, iC] (NHWC)
+
+  const int isNCHW = block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // INT_ARG(0): 0-NCHW,  1-NHWC
+
+  REQUIRE_TRUE(gradO->rankOf() == 4, 0, "UPSAMPLING2D_BP op: output's gradient array must be 4D, but got %i instead!",
+               gradO->rankOf());
+  REQUIRE_TRUE(gradI->rankOf() == 4, 0, "UPSAMPLING2D_BP op: input's gradient array must be 4D, but got %i instead!",
+               gradI->rankOf());
+
+  ConvolutionUtils::upsampling2dBP(block, *gradO, *gradI, (bool)isNCHW);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(upsampling_bp, upsampling2d_bp);
+
+DECLARE_SHAPE_FN(upsampling2d_bp) {
+  REQUIRE_TRUE(inputShape->at(0)[0] == 4, 0, "UPSAMPLING2D_BP op: input array must be 4D, but got %i instead!",
+               inputShape->at(0)[0]);
+  REQUIRE_TRUE(inputShape->at(1)[0] == 4, 0,
+               "UPSAMPLING2D_BP op: output's gradient array must be 4D, but got %i instead!", inputShape->at(1)[0]);
+
+  auto gradIShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShape->at(0), inputShape->at(1), false, block.getWorkspace());
+
+  return SHAPELIST(CONSTANT(gradIShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/convo/upsampling3d.cpp

@@ -0,0 +1,129 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 04.05.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_upsampling3d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(upsampling3d, 1, 1, false, 0, 3) {
+  auto input = INPUT_VARIABLE(0);     // [bS, iC, iD, iH, iW] (NCDHW) or [bS, iD, iH, iW, iC] (NDHWC)
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, iC, factorD*iD, factorH*iH, factorW*iW ] (NCDHW) or [bS, factorD*iD,
+                                      // factorH*iH, factorW*iW, iC] (NDHWC)
+
+  const LongType factorD = INT_ARG(0);
+  const LongType factorH = INT_ARG(1);
+  const LongType factorW = INT_ARG(2);
+  const int isNCDHW = block.getIArguments()->size() > 3 ? INT_ARG(3) : 0;  // INT_ARG(3): 0-NCDHW,  1-NDHWC
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0, "UPSAMPLING3D op: input should be 5D, but got %i instead!", input->rankOf());
+  REQUIRE_TRUE(output->rankOf() == 5, 0, "UPSAMPLING3D op: output should be 5D, but got %i instead!", output->rankOf());
+
+  ConvolutionUtils::upsampling3d(block, *input, *output, factorD, factorH, factorW, (bool)isNCDHW);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(upsampling3d) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(upsampling3d) {
+  auto inputShapeInfo = inputShape->at(0);
+
+  REQUIRE_TRUE(inputShapeInfo[0] == 5, 0, "UPSAMPLING2D op: input should be 5D, but got %i instead!",
+               inputShapeInfo[0]);
+
+  const LongType factorD = INT_ARG(0);
+  const LongType factorH = INT_ARG(1);
+  const LongType factorW = INT_ARG(2);
+  const int isNCDHW = block.getIArguments()->size() > 3 ? INT_ARG(3) : 0;  // INT_ARG(3): 0-NCHW,  1-NHWC
+
+  sd::LongType *outputShapeInfo = nullptr;
+  ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inputShapeInfo[0]), sd::LongType);
+
+  outputShapeInfo[0] = inputShapeInfo[0];
+  outputShapeInfo[1] = inputShapeInfo[1];
+
+  if (isNCDHW) {
+    outputShapeInfo[2] = inputShapeInfo[2];
+    outputShapeInfo[3] = inputShapeInfo[3] * factorD;
+    outputShapeInfo[4] = inputShapeInfo[4] * factorH;
+    outputShapeInfo[5] = inputShapeInfo[5] * factorW;
+  } else {
+    outputShapeInfo[2] = inputShapeInfo[2] * factorD;
+    outputShapeInfo[3] = inputShapeInfo[3] * factorH;
+    outputShapeInfo[4] = inputShapeInfo[4] * factorW;
+    outputShapeInfo[5] = inputShapeInfo[5];
+  }
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, inputShapeInfo, shape::order(inputShapeInfo));
+
+  return SHAPELIST(CONSTANT(outputShapeInfo));
+}
+
+DECLARE_TYPES(upsampling3d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(upsampling3d_bp, 2, 1, false, 0, 0) {
+  // NDArray<T>* input = INPUT_VARIABLE(0);             // [bS, iC, iD, iH, iW] (NCDHW) or [bS, iD, iH, iW, iC] (NDHWC)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, iC, factorD*iD, factorH*iH, factorW*iW ] (NCDHW) or [bS, factorD*iD,
+                                     // factorH*iH, factorW*iW, iC] (NDHWC)
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iC, iD, iH, iW] (NCDHW) or [bS, iD, iH, iW, iC] (NDHWC)
+
+  const int isNCDHW = block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // INT_ARG(0): 0-NCHW,  1-NHWC
+
+  // REQUIRE_TRUE(input->rankOf() == 5, 0, "UPSAMPLING3D_BP op: input array must be 4D, but got %i instead!",
+  // input->rankOf());
+  REQUIRE_TRUE(gradO->rankOf() == 5, 0, "UPSAMPLING3D_BP op: output's gradient array must be 4D, but got %i instead!",
+               gradO->rankOf());
+  REQUIRE_TRUE(gradI->rankOf() == 5, 0, "UPSAMPLING3D_BP op: input's gradient array must be 4D, but got %i instead!",
+               gradI->rankOf());
+
+  ConvolutionUtils::upsampling3dBP(block, *gradO, *gradI, (bool)isNCDHW);
+
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(upsampling3d_bp) {
+  REQUIRE_TRUE(inputShape->at(0)[0] == 5, 0, "UPSAMPLING3D_BP op: input array must be 4D, but got %i instead!",
+               inputShape->at(0)[0]);
+  REQUIRE_TRUE(inputShape->at(1)[0] == 5, 0,
+               "UPSAMPLING3D_BP op: output's gradient array must be 4D, but got %i instead!", inputShape->at(1)[0]);
+
+  auto gradIShapeInfo =
+      ShapeBuilders::copyShapeInfoAndType(inputShape->at(0), inputShape->at(1), false, block.getWorkspace());
+
+  return SHAPELIST(CONSTANT(gradIShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/dot_product_attention.cpp

@@ -0,0 +1,244 @@
+/*
+*  ******************************************************************************
+*  *
+*  *
+*  * This program and the accompanying materials are made available under the
+*  * terms of the Apache License, Version 2.0 which is available at
+*  * https://www.apache.org/licenses/LICENSE-2.0.
+*  *
+*  * See the NOTICE file distributed with this work for additional
+*  * information regarding copyright ownership.
+*  * Unless required by applicable law or agreed to in writing, software
+*  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+*  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+*  * License for the specific language governing permissions and limitations
+*  * under the License.
+*  *
+*  * SPDX-License-Identifier: Apache-2.0
+*  *****************************************************************************
+*/
+
+//
+// @author Paul Dubs
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_dot_product_attention)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/reverse.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(dot_product_attention, 3, -1, false, 0, 2) {
+  auto queries = INPUT_VARIABLE(0);
+  auto keys = INPUT_VARIABLE(1);
+  auto values = INPUT_VARIABLE(2);
+  auto mask = block.width() > 3 ? INPUT_VARIABLE(3) : nullptr;
+
+  auto output = OUTPUT_VARIABLE(0);
+  NDArray *weights;
+  bool outputWeights = INT_ARG(1);
+  if (outputWeights) {
+    weights = OUTPUT_VARIABLE(1);
+  } else {
+    auto weightShape = ShapeUtils::evalShapeForMatmul(keys->shapeInfo(), queries->shapeInfo(), true, false);
+    weights = new NDArray('c', weightShape, values->dataType(), block.launchContext());
+  }
+
+  int normalization = INT_ARG(0);
+
+  REQUIRE_TRUE(queries->rankOf() == keys->rankOf() && keys->rankOf() == values->rankOf(), 0,
+               "dot_product_attention: Queries, Keys and Values must have same rank. "
+               "But got queries = %s, keys = %s, values = %s",
+               ShapeUtils::shapeAsString(queries).c_str(), ShapeUtils::shapeAsString(keys).c_str(),
+               ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(queries->rankOf() == 3 || queries->rankOf() == 4, 0,
+               "dot_product_attention: Queries, Keys and Values must be rank 3 arrays for single headed attention "
+               "or rank 4 arrays for multi headed attention. But got rank = %i",
+               queries->rankOf());
+
+  REQUIRE_TRUE(queries->sizeAt(0) == keys->sizeAt(0) && keys->sizeAt(0) == values->sizeAt(0), 0,
+               "dot_product_attention: Queries, Keys and Values must have the same mini batch size. "
+               "But got queries = %i, keys = %i, values = %i",
+               queries->sizeAt(0), keys->sizeAt(0), values->sizeAt(0));
+
+  REQUIRE_TRUE(queries->sizeAt(-2) == keys->sizeAt(-2), 0,
+               "dot_product_attention: Queries and Keys must have the same feature size. "
+               "But got queries = %i, keys = %i",
+               queries->sizeAt(-2), keys->sizeAt(-2));
+
+  REQUIRE_TRUE(keys->sizeAt(-1) == values->sizeAt(-1), 0,
+               "dot_product_attention: Keys and Values must have the same timestep length. "
+               "But got keys = %i, values = %i",
+               keys->sizeAt(-1), values->sizeAt(-1));
+
+  sd::ops::matmul mmul;
+  mmul.execute({keys, queries}, {weights}, {}, {1}, {});
+  if (normalization) {
+    *weights /= sqrt((double)keys->sizeAt(-2));
+  }
+
+  if (mask != nullptr) {
+    NDArray *reshapedMask;
+    if (weights->rankOf() == 4) {
+      std::vector<sd::LongType> shape = {mask->sizeAt(0), 1, mask->sizeAt(1), 1};
+      reshapedMask = mask->reshape(mask->ordering(),shape);
+    } else {
+      std::vector<sd::LongType> shape = {mask->sizeAt(0), mask->sizeAt(1), 1};
+      reshapedMask = mask->reshape(mask->ordering(),shape);
+    }
+
+    // the mask is 0 for positions we want to skip, and 1 for positions we want to keep. By subtracting 1 from
+    // it we get -1 for those we want to skip and 0 for those we want to keep. Multiplying it by 1e9 then
+    // turns all of those we want to skip into very large negative values. By adding this to the weights
+    // before going through the softmax, we effectively push all masked positions to zero after softmax.
+    //
+    // we are using 1e9 to mean effectively infinity
+    auto applyMask = *reshapedMask  * 1e9;
+    *weights -= *applyMask;
+    delete reshapedMask;
+    delete applyMask;
+
+  }
+
+  int softmaxDim = -2;
+  sd::ops::softmax softmax;
+  softmax.execute({weights}, std::vector<NDArray *>{weights}, {}, {softmaxDim}, {}, {}, true);
+
+  mmul.execute({values, weights}, {output}, {}, {}, {});
+
+  if (!outputWeights) {
+    delete weights;
+  }
+
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(dot_product_attention) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dot_product_attention) {
+  auto query_shape = inputShape->at(0);
+  auto keys_shape = inputShape->at(1);
+  auto values_shape = inputShape->at(2);
+
+  auto weights_shape = ConstantShapeHelper::getInstance().createShapeInfo(
+      sd::ArrayOptions::dataType(values_shape), 'c',
+      ShapeUtils::evalShapeForMatmul(keys_shape, query_shape, true, false));
+  auto output_shape = ConstantShapeHelper::getInstance().createShapeInfo(
+      sd::ArrayOptions::dataType(values_shape), 'c',
+      ShapeUtils::evalShapeForMatmul(values_shape, weights_shape, false, false));
+
+  if (INT_ARG(1)) {
+    return SHAPELIST(output_shape, weights_shape);
+  } else {
+    return SHAPELIST(output_shape);
+  }
+}
+
+CUSTOM_OP_IMPL(dot_product_attention_bp, 4, 3, false, 0, 1) {
+  auto queries = INPUT_VARIABLE(0);
+  auto keys = INPUT_VARIABLE(1);
+  auto values = INPUT_VARIABLE(2);
+  auto eps = INPUT_VARIABLE(3);
+  auto mask = block.width() > 4 ? INPUT_VARIABLE(4) : nullptr;
+
+  auto dLdq = OUTPUT_VARIABLE(0);
+  auto dLdk = OUTPUT_VARIABLE(1);
+  auto dLdv = OUTPUT_VARIABLE(2);
+
+  int normalization = INT_ARG(0);
+
+  REQUIRE_TRUE(queries->rankOf() == keys->rankOf() && keys->rankOf() == values->rankOf(), 0,
+               "dot_product_attention: Queries, Keys and Values must have same rank. "
+               "But got queries = %s, keys = %s, values = %s",
+               ShapeUtils::shapeAsString(queries).c_str(), ShapeUtils::shapeAsString(keys).c_str(),
+               ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(queries->rankOf() == 3 || queries->rankOf() == 4, 0,
+               "dot_product_attention: Queries, Keys and Values must be rank 3 arrays for single headed attention "
+               "or rank 4 arrays for multi headed attention. But got rank = %i",
+               queries->rankOf());
+
+  REQUIRE_TRUE(queries->sizeAt(0) == keys->sizeAt(0) && keys->sizeAt(0) == values->sizeAt(0), 0,
+               "dot_product_attention: Queries, Keys and Values must have the same mini batch size. "
+               "But got queries = %i, keys = %i, values = %i",
+               queries->sizeAt(0), keys->sizeAt(0), values->sizeAt(0));
+
+  REQUIRE_TRUE(queries->sizeAt(-2) == keys->sizeAt(-2), 0,
+               "dot_product_attention: Queries and Keys must have the same feature size. "
+               "But got queries = %i, keys = %i",
+               queries->sizeAt(-2), keys->sizeAt(-2));
+
+  REQUIRE_TRUE(keys->sizeAt(-1) == values->sizeAt(-1), 0,
+               "dot_product_attention: Keys and Values must have the same timestep length. "
+               "But got keys = %i, values = %i",
+               keys->sizeAt(-1), values->sizeAt(-1));
+
+  double factor;
+  if (normalization) factor = sqrt((double)keys->sizeAt(-2));
+
+  auto weightShape = ShapeUtils::evalShapeForMatmul(keys->shapeInfo(), queries->shapeInfo(), true, false);
+
+  sd::ops::matmul mmul;
+  NDArray preSoftmax('c', weightShape, values->dataType(), block.launchContext());
+  mmul.execute({keys, queries}, {&preSoftmax}, {}, {1}, {});
+
+  if (normalization) preSoftmax /= factor;
+  NDArray *reshapedMask;
+  if (mask != nullptr && !mask->isEmpty()) {
+    if (preSoftmax.rankOf() == 4) {
+      std::vector<sd::LongType> shape = {mask->sizeAt(0), 1, mask->sizeAt(1), 1};
+      reshapedMask = mask->reshape(mask->ordering(), shape);
+    } else {
+      std::vector<sd::LongType> shape = {mask->sizeAt(0), mask->sizeAt(1), 1};
+      reshapedMask = mask->reshape(mask->ordering(), shape);
+    }
+
+    *reshapedMask  *= 1e9;
+    preSoftmax -= *reshapedMask;
+  }
+
+  int softmaxDim = -2;
+
+  NDArray weights('c', weightShape, values->dataType(), block.launchContext());
+  sd::ops::softmax softmax;
+  softmax.execute({&preSoftmax}, {&weights}, {}, {softmaxDim}, {});
+  sd::ops::matmul_bp mmul_bp;
+  NDArray dLdw(weights.shapeInfo(), block.workspace());
+  mmul_bp.execute({values, &weights, eps}, {dLdv, &dLdw}, {}, {}, {});
+
+  NDArray dLds(preSoftmax.shapeInfo(), block.workspace());
+  sd::ops::softmax_bp softmax_bp;
+  softmax_bp.execute({&preSoftmax, &dLdw,&weights}, {&dLds}, {}, {softmaxDim}, {});
+
+  if (normalization) dLds /= factor;
+  if(mask != nullptr) {
+    dLds *= *reshapedMask;
+  }
+  mmul_bp.execute({keys, queries, &dLds}, std::vector<NDArray *>{dLdk, dLdq}, {}, {1}, {});
+
+  delete reshapedMask;
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(dot_product_attention_bp) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dot_product_attention_bp) {
+  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/dot_product_attention_v2.cpp

@@ -0,0 +1,292 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// @author Paul Dubs
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_dot_product_attention_v2)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/reverse.h>
+#include <helpers/AttentionHelper.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(dot_product_attention_v2, -2, -1, false, -2, -2) {
+  auto queries = INPUT_VARIABLE(0);
+  auto values = INPUT_VARIABLE(1);
+
+  REQUIRE_TRUE(queries->rankOf() == values->rankOf(),0,"Queries and values must be same ranks!");
+
+  bool reshapedQ = false;
+  if(queries->rankOf() == 2) {
+    reshapedQ = true;
+    std::vector<sd::LongType> qShape = {1,queries->sizeAt(0), queries->sizeAt(-1)};
+    std::vector<sd::LongType> vShape = {1,values->sizeAt(0), values->sizeAt(-1)};
+    queries = queries->reshape('c',qShape);
+    values = values->reshape('c', vShape);
+  }
+
+
+
+
+
+  auto keys = block.width() > 2  ? INPUT_VARIABLE(2) : values;
+  if(reshapedQ && block.width() > 2) {
+    std::vector<sd::LongType> shape = {1,keys->sizeAt(0), keys->sizeAt(-1)};
+    keys = keys->reshape('c', shape);
+  }
+
+  auto qMask = block.width() > 3 ? INPUT_VARIABLE(3) : nullptr;
+  auto vMask = block.width() > 4 ? INPUT_VARIABLE(4) : nullptr;
+
+  //reshape to handle different shapes
+  if(qMask != nullptr && reshapedQ) {
+    std::vector<sd::LongType> qShape = {1,qMask->sizeAt(0), qMask->sizeAt(-1)};
+    qMask = qMask->reshape('c', qShape);
+  }
+
+  if(vMask != nullptr && reshapedQ) {
+    std::vector<sd::LongType> vShape = {1,vMask->sizeAt(0), vMask->sizeAt(-1)};
+    vMask = vMask->reshape('c',vShape);
+  }
+
+
+  auto dropout = block.numT() > 1 ? T_ARG(1) : 0.0;
+
+  auto scale = block.numT() > 0 ? T_ARG(0) : 1.0;
+
+  auto useCausalMask = block.numB() > 0 ? B_ARG(0) : false;
+  auto training = block.numB() > 1 ? B_ARG(1) : false;
+
+
+
+  std::vector<sd::NDArray*> inputs = {queries,values,keys};
+  std::vector<sd::NDArray *> masks2 = {qMask,vMask};
+
+
+
+
+  //TODO: handle reshape for rank 2 for each variable here
+  auto applyScoresOut = OUTPUT_VARIABLE(0);
+  auto attentionScores = OUTPUT_VARIABLE(1);
+  auto attentionLogits = OUTPUT_VARIABLE(2);
+  auto dropoutMask = dropout > 0.0 ? OUTPUT_VARIABLE(3) : nullptr;
+  if(reshapedQ) {
+    applyScoresOut->reshapei('c', {1,applyScoresOut->sizeAt(0), applyScoresOut->sizeAt(1)});
+    attentionLogits->reshapei('c', {1,attentionLogits->sizeAt(0), attentionLogits->sizeAt(1)});
+    attentionScores->reshapei('c', {1,attentionScores->sizeAt(0), attentionScores->sizeAt(1)});
+
+  }
+  AttentionHelper::doAttention(inputs, masks2, training, useCausalMask, dropout, scale, attentionScores,
+                               block.randomSeed(), applyScoresOut, attentionLogits, dropoutMask);
+
+
+  //delete extra memory
+  if(reshapedQ) {
+    delete queries;
+    delete values;
+    if(block.width() > 2) {
+      delete keys;
+    }
+    if(qMask != nullptr)
+      delete qMask;
+
+    if(vMask != nullptr)
+      delete vMask;
+
+    applyScoresOut->reshapei('c', {applyScoresOut->sizeAt(1), applyScoresOut->sizeAt(-1)});
+    attentionLogits->reshapei('c', {attentionLogits->sizeAt(1), attentionLogits->sizeAt(-1)});
+    attentionScores->reshapei('c', {attentionScores->sizeAt(1), attentionScores->sizeAt(-1)});
+
+  }
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(dot_product_attention_v2) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dot_product_attention_v2) {
+  auto firstInputType = INPUT_VARIABLE(0)->dataType();
+  auto queries = INPUT_VARIABLE(0);
+  auto values = INPUT_VARIABLE(1);
+  auto keys = block.width() > 2  ? INPUT_VARIABLE(2) : values;
+
+  int batchSize = queries->sizeAt(0);
+  int tq = queries->rankOf() == 3 ? queries->sizeAt(-2) : queries->sizeAt(-1);
+  int tv = queries->rankOf() == 3 ? values->sizeAt(-2) : values->sizeAt(-1);
+  int dim = queries->rankOf() == 3 ? values->sizeAt(-1) : 1;
+
+  auto dropout = block.numT() > 1 ? block.getTArguments()->at(1) : 0.0;
+  //inputs: batchSize,Tq,dim batchSize,Tq,Tv
+  //outputs: batchSize,Tq, dim batchSize,Tq,Tv
+  std::vector<sd::LongType> outShape;
+  std::vector<sd::LongType> scoresShape1;
+
+
+  if(queries->rankOf() == 3) {
+    outShape.push_back(batchSize);
+    outShape.push_back(tq);
+
+    scoresShape1.push_back(batchSize);
+    scoresShape1.push_back(tq);
+
+    outShape.push_back(dim);
+    scoresShape1.push_back(dim);
+  } else {
+    outShape.push_back(batchSize);
+    outShape.push_back(tv);
+
+    scoresShape1.push_back(batchSize);
+    scoresShape1.push_back(dim);
+
+  }
+
+
+
+
+  auto constOutputScores = ConstantShapeHelper::getInstance().bufferForShapeInfo(firstInputType, 'c', outShape)->primary();
+  auto attentionScoresShape = ConstantShapeHelper::getInstance().bufferForShapeInfo(firstInputType, 'c', scoresShape1)->primary();
+  auto attentionLogitsShape = ConstantShapeHelper::getInstance().bufferForShapeInfo(firstInputType, 'c', scoresShape1)->primary();
+
+  if(dropout > 0) {
+    return SHAPELIST(constOutputScores, attentionScoresShape, attentionLogitsShape, attentionScoresShape);
+  } else {
+    return SHAPELIST(constOutputScores, attentionScoresShape, attentionLogitsShape);
+  }
+
+
+
+}
+
+CUSTOM_OP_IMPL(dot_product_attention_v2_bp, -2, 3, false, 0, -2) {
+  auto queries = INPUT_VARIABLE(0);
+  auto values = INPUT_VARIABLE(1);
+  auto keys = INPUT_VARIABLE(2);
+  bool reshapedQ = false;
+  if(queries->rankOf() == 2) {
+    reshapedQ = true;
+    std::vector<sd::LongType> qShape = {1,queries->sizeAt(0), queries->sizeAt(-1)};
+    std::vector<sd::LongType> vShape = {1,values->sizeAt(0), values->sizeAt(-1)};
+    std::vector<sd::LongType> kShape = {1,keys->sizeAt(0), keys->sizeAt(-1)};
+    queries = queries->reshape('c', qShape);
+    values = values->reshape('c', vShape);
+    keys = keys->reshape('c',kShape);
+  }
+
+
+
+  auto attentionScoresOut = INPUT_VARIABLE(3);
+  auto attentionScoresWeights = INPUT_VARIABLE(4);
+  auto attentionScoreLogits = INPUT_VARIABLE(5);
+  if(reshapedQ) {
+    attentionScoresOut->reshapei('c', {1,attentionScoresOut->sizeAt(0), attentionScoresOut->sizeAt(1)});
+    attentionScoreLogits->reshapei('c', {1,attentionScoreLogits->sizeAt(0), attentionScoreLogits->sizeAt(1)});
+    attentionScoresWeights->reshapei('c', {1,attentionScoresWeights->sizeAt(0), attentionScoresWeights->sizeAt(1)});
+
+  }
+
+
+  auto eps = INPUT_VARIABLE(6);
+  if(reshapedQ) {
+    eps->reshapei('c', {1,eps->sizeAt(0), eps->sizeAt(1)});
+  }
+  auto dropoutMask = block.width() > 7 ? INPUT_VARIABLE(7) : nullptr;
+
+  auto qMask = block.width() > 8 ? INPUT_VARIABLE(8) : nullptr;
+  auto vMask = block.width() > 9 ? INPUT_VARIABLE(9) : nullptr;
+
+  if(qMask != nullptr && qMask->rankOf() == 2) {
+    std::vector<sd::LongType> qShape = {1,qMask->sizeAt(0), qMask->sizeAt(-1)};
+    qMask =qMask->reshape('c', qShape);
+  }
+
+  if(vMask != nullptr && vMask->rankOf() == 2) {
+    std::vector<sd::LongType> vShape = {1,vMask->sizeAt(0), vMask->sizeAt(-1)};
+    vMask = vMask->reshape('c', vShape);
+  }
+
+  auto dLdq = OUTPUT_VARIABLE(0);
+  auto dLdv = OUTPUT_VARIABLE(1);
+  auto dLdk = OUTPUT_VARIABLE(2);
+  if(reshapedQ) {
+    dLdq->reshapei('c', {1,dLdq->sizeAt(0), dLdq->sizeAt(1)});
+    dLdv->reshapei('c', {1,dLdv->sizeAt(0), dLdv->sizeAt(1)});
+    dLdk->reshapei('c', {1,dLdk->sizeAt(0), dLdk->sizeAt(1)});
+  }
+  auto scale = block.numT() > 1 ? T_ARG(0) : 1.0;
+  auto dropout = block.numT() > 0 ? T_ARG(1) : 0.0;
+
+
+  auto useCausalMask = block.numB() > 0 ? B_ARG(0) : false;
+  auto training = block.numB() > 1 ? B_ARG(1) : false;
+
+
+
+  std::vector<sd::NDArray*> inputs = {queries,values,keys,attentionScoresOut,attentionScoresWeights,attentionScoreLogits,eps};
+  if(dropoutMask != nullptr) {
+    inputs.push_back(dropoutMask);
+  }
+
+  std::vector<sd::NDArray *> masks2 = {qMask,vMask};
+  std::vector<sd::NDArray *> outputs = {dLdq,dLdv,dLdk};
+
+  int seed = block.randomSeed();
+  AttentionHelper::dotProductAttentionBpHelper(queries, keys, values, scale, dLdq, dLdk, dLdv, eps, seed, qMask, vMask,
+                                               useCausalMask, dropout, training, attentionScoresWeights,
+                                               attentionScoreLogits, dropoutMask);
+
+  if(reshapedQ) {
+    delete queries;
+    delete values;
+    delete keys;
+    if(qMask != nullptr)
+      delete qMask;
+    if(vMask != nullptr)
+      delete vMask;
+
+    dLdq->reshapei('c', {dLdq->sizeAt(1), dLdq->sizeAt(2)});
+    dLdv->reshapei('c', {dLdv->sizeAt(1), dLdv->sizeAt(2)});
+    dLdk->reshapei('c', {dLdk->sizeAt(1), dLdk->sizeAt(2)});
+    eps->reshapei('c', {eps->sizeAt(1), eps->sizeAt(2)});
+  }
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(dot_product_attention_v2_bp) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dot_product_attention_v2_bp) {
+  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/embedding_lookup.cpp

@@ -0,0 +1,113 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by GS <sgazeos@gmail.com>
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_embedding_lookup)
+
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+
+#include <numeric>
+#include <vector>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(embedding_lookup, 2, 1, false, 0, 1) {
+  auto input = INPUT_VARIABLE(0);    // lookup param
+  auto indices = INPUT_VARIABLE(1);  // indices, as is
+  auto output = OUTPUT_VARIABLE(0);  //
+
+  if (block.width() > 2) {  // multiple input
+    indices = INPUT_VARIABLE(block.width() - 1);
+    std::vector<sd::LongType> dims(input->rankOf());
+    int i = output->rankOf() - input->rankOf();
+    for (auto& v : dims) {
+      v = i++;
+    }
+
+    ResultSet outputView = output->allTensorsAlongDimension(dims);
+    REQUIRE_TRUE(static_cast<sd::LongType >(block.width()) > output->sizeAt(0), 0,
+                 "embedding_lookup: input list should be greater then %i, but %i given.", output->sizeAt(0),
+                 block.width());
+    for (sd::LongType e = 0; e < indices->lengthOf(); ++e) {
+      sd::LongType thisIndex = (*indices).e<sd::LongType>(e);
+      input = INPUT_VARIABLE(thisIndex);  // lookup param
+
+      outputView.at(e)->assign(input);
+    }
+  } else {
+    int indexRank = indices->rankOf();
+    REQUIRE_TRUE(indexRank > 0, 0,
+                 "embedded_lookup: input array of indexes can't be single scalar, the requirement is: rank > 0 !");
+
+    int inputRank = input->rankOf();
+    int lastIndDim = indices->lengthOf();
+    int partition_mode = INT_ARG(0);  // partition_mode == 0 - i.e. 'mod' , 1 - 'div'
+
+    sd::ops::gather op;
+
+    auto result(op.evaluate({input, indices}, {0}));
+    REQUIRE_TRUE(result.status() == sd::Status::OK, 0, "embedding_lookup: cannot retrieve results from gather op.");
+    REQUIRE_TRUE(result.at(0)->isSameShape(output), 0, "embedding_lookup: wrong shape of return from gather op.");
+    output->assign(result.at(0));
+  }
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(embedding_lookup) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes(sd::DataType::ANY);
+}
+
+DECLARE_SHAPE_FN(embedding_lookup) {
+  auto inShapeInfo = inputShape->at(0);
+  auto indicesShapeInfo = inputShape->at(1);
+  int inRank = shape::rank(inShapeInfo);
+  if (inputShape->size() == 2u) {
+    int outRank = inRank;
+
+    std::vector<sd::LongType> shapeInfo(outRank);
+
+    shapeInfo[0] = indicesShapeInfo[1];  // vector - how many elements
+    for (sd::LongType e = 1; e < outRank; e++) shapeInfo[e] = shape::sizeAt(inShapeInfo, e);
+
+    auto outShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(inShapeInfo),
+                                                                           shape::order(inShapeInfo), shapeInfo);
+    return SHAPELIST(outShapeInfo);
+  }
+
+  int outRank = inRank + 1;
+  std::vector<sd::LongType> shapeInfo(outRank);
+  auto indices = INPUT_VARIABLE(block.width() - 1);
+  shapeInfo[0] = indices->lengthOf();  // vector - how many elements
+  for (sd::LongType e = 1; e < outRank; e++) shapeInfo[e] = shape::sizeAt(inShapeInfo, e);
+
+  auto outShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(inShapeInfo),
+                                                                         shape::order(inShapeInfo), shapeInfo);
+  return SHAPELIST(outShapeInfo);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/fusedBatchNorm.cpp

@@ -0,0 +1,196 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 29/10/17.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_fused_batch_norm)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+
+DECLARE_TYPES(fused_batch_norm) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+CUSTOM_OP_IMPL(fused_batch_norm, 3, 3, false, 0, 2) {
+  auto x = INPUT_VARIABLE(0);       // [bS,iH,iW,iD] (NHWC) or [bS,iD,iH,iW] (NCHW)
+  auto scale = INPUT_VARIABLE(1);   // [iD]
+  auto offset = INPUT_VARIABLE(2);  // [iD]
+
+  auto y = OUTPUT_VARIABLE(0);          // [bS,iH,iW,iD] (NHWC) or [bS,iD,iH,iW] (NCHW)
+  auto batchMean = OUTPUT_VARIABLE(1);  // [iD]
+  auto batchVar = OUTPUT_VARIABLE(2);   // [iD]
+
+  const bool dataFormat = (bool)INT_ARG(0);  // 0->NHWC, 1->NCHW
+  const bool isTraining = (bool)INT_ARG(1);
+  sd_debug("CUSTOM_OP fused_batch_norm: data format, is NCHW: %d, isTraining: %d\n", dataFormat, isTraining);
+
+  REQUIRE_TRUE(x->rankOf() == 4, 0,
+               "CUSTOM_OP fused_batch_norm: the rank of input x array must be equal to 4, but got %i instead !",
+               x->rankOf());
+
+  int  iD;         // input height, input width, input depth(number of channels)
+  if (dataFormat) {
+    iD = x->sizeAt(1);
+  } else {
+    iD = x->sizeAt(3);
+  }
+
+  auto* xCast = x->cast(sd::DataType::FLOAT32);
+  if (dataFormat) {
+    std::vector<LongType> permute = {0,2,3,1};
+    auto* xCastPermuted = xCast->permute(permute, false, false);
+    delete xCast;
+    xCast = xCastPermuted;
+  }
+  REQUIRE_TRUE(scale->rankOf() == 1 && scale->sizeAt(0) == iD, 0,
+               "CUSTOM_OP fused_batch_norm: wrong shape of input scale array, expected is [%i], but got %s instead", iD,
+               ShapeUtils::shapeAsString(scale).c_str());
+  REQUIRE_TRUE(offset->rankOf() == 1 && offset->sizeAt(0) == iD, 0,
+               "CUSTOM_OP fused_batch_norm: wrong shape of input offset array, expected is [%i], but got %s instead",
+               iD, ShapeUtils::shapeAsString(offset).c_str());
+
+  NDArray *mean(nullptr), *variance(nullptr);
+  if (!isTraining) {
+    mean = INPUT_VARIABLE(3);
+    variance = INPUT_VARIABLE(4);
+    REQUIRE_TRUE(mean->rankOf() == 1 && mean->sizeAt(0) == iD, 0,
+                 "CUSTOM_OP fused_batch_norm: wrong shape of input mean array, expected is [%i], but got %s instead",
+                 iD, ShapeUtils::shapeAsString(mean).c_str());
+    REQUIRE_TRUE(
+        variance->rankOf() == 1 && variance->sizeAt(0) == iD, 0,
+        "CUSTOM_OP fused_batch_norm: wrong shape of input variance array, expected is [%i], but got %s instead", iD,
+        ShapeUtils::shapeAsString(variance).c_str());
+  } else {
+    // REQUIRE_TRUE(block.width() == 3, 0, "CUSTOM_OP fused_batch_norm: when isTraining=true then number of input arrays
+    // must be equal to 3, but got %i instead !", block.width());
+    std::vector<sd::LongType> shape = {iD};
+    mean = NDArrayFactory::create_(scale->ordering(), shape, scale->dataType(), block.launchContext());
+    variance = NDArrayFactory::create_(scale->ordering(), shape, scale->dataType(), block.launchContext());
+  }
+
+  float epsilon;
+  if (block.getTArguments()->size() > 0) {
+    epsilon = (float)(T_ARG(0) > 1.001e-5 ? T_ARG(0) : 1.001e-5);
+  } else {
+    epsilon = 0.001f;
+  }
+
+  const int restSize = x->lengthOf() / iD;
+
+  auto* xAffected = NDArrayFactory::create(x->ordering(), {restSize, iD}, mean->dataType(), block.launchContext());
+  xAffected->assign(xCast);
+
+  const int restSizeMinusOne = (restSize > 1) ? (restSize - 1) : 1;
+  const float restSizeInv = 1.0f / restSize;
+  const float restSizeAdjust = (float)restSize / restSizeMinusOne;
+
+  if (isTraining) {
+    std::vector<sd::LongType > dim = {0};
+    auto* sum = xAffected->reduceAlongDimension(reduce::Sum, &dim);
+    *sum *= restSizeInv;
+    mean->assign(sum);
+    delete sum;
+    *batchMean = *mean;
+  } else
+    *batchMean = 0.;
+
+  auto* xCentered = (*xAffected) - (*mean);
+  *xAffected -= *mean;
+
+  if (isTraining) {
+    int power = 2;
+    xAffected->applyScalar(scalar::Pow, power, xAffected);
+    std::vector<sd::LongType > dim = {0};
+
+    auto* sum = xAffected->reduceAlongDimension(reduce::Sum, &dim);
+    *sum *= restSizeInv;
+    variance->assign(sum);
+    delete sum;
+    auto* varOutput = (*variance) * restSizeAdjust;
+    batchVar->assign(varOutput);
+    delete varOutput;
+  } else
+    *batchVar = 0.;
+
+  // Break down: ((*variance + epsilon).transform(transform::RSqrt) * (*scale)).cast(xAffected->dataType())
+  auto* variancePlusEps = (*variance) + epsilon;
+  variancePlusEps->applyTransform(transform::RSqrt, variancePlusEps);
+  auto* scaledVariance = (*variancePlusEps) * (*scale);
+  delete variancePlusEps;
+  auto* scaledVarianceCast = scaledVariance->cast(xAffected->dataType());
+  delete scaledVariance;
+  
+  auto* xScaled1 = (*xCentered) * (*scaledVarianceCast);
+  delete xCentered;
+  delete scaledVarianceCast;
+  
+  auto* xShifted1 = (*xScaled1) + (*offset);
+  delete xScaled1;
+  
+  if (dataFormat) {
+    // need to reshape from matrix to 4d then permute the ordering due to NWHC  ordering
+    auto* newShapePtr = xCast->getShapeAsVector();
+    std::vector<LongType> newShape = *newShapePtr;
+    delete newShapePtr;
+    auto* reshaped = xShifted1->reshape(xCast->ordering(), newShape, false);
+    delete xShifted1;
+    reshaped->permutei({0, 3, 1, 2}, 0, false);
+    y->assign(reshaped);
+    delete reshaped;
+
+  } else {  // NWHC case
+    y->assign(xShifted1);
+    delete xShifted1;
+  }
+
+  if (isTraining) {
+    delete mean;
+    delete variance;
+  }
+
+  if(xCast != x) {
+    delete xCast;
+  }
+  delete xAffected;
+
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(fused_batch_norm) {
+  auto xShapeInfo = inputShape->at(0);
+  auto scaleShapeInfo = inputShape->at(1);
+
+  const bool dataFormat = (bool)INT_ARG(0);  // 0->NHWC, 1->NCHW
+  const int iD = dataFormat ? xShapeInfo[2] : xShapeInfo[4];
+
+  REQUIRE_TRUE(scaleShapeInfo[0] == 1 && scaleShapeInfo[1] == iD, 0,
+               "CUSTOM_OP fused_batch_norm: wrong shape of input scale array, expected is [%i], but got %s instead", iD,
+               ShapeUtils::shapeAsString(scaleShapeInfo).c_str());
+  return SHAPELIST(CONSTANT(xShapeInfo), CONSTANT(scaleShapeInfo), CONSTANT(scaleShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/layer_norm.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 Paul Dubs
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_layer_norm)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/addBias.h>
+#include <ops/declarable/helpers/reverse.h>
+
+namespace sd {
+namespace ops {
+
+CONFIGURABLE_OP_IMPL(layer_norm, 2, 1, false, 0, -1) {
+  auto input = INPUT_VARIABLE(0);
+  auto gain = INPUT_VARIABLE(1);
+  auto output = OUTPUT_VARIABLE(0);
+
+  std::vector<sd::LongType> axis = *block.getIArguments();
+
+  const bool isNCHW = block.getBArguments()->size() > 0 ? B_ARG(0) : true;  // 0-NCHW,  1-NHWC
+  const int dimC = isNCHW ? 1 : input->rankOf() - 1;
+
+  REQUIRE_TRUE(gain->rankOf() == 1 && gain->sizeAt(0) == input->sizeAt(dimC), 0,
+               "LAYER_NORM OP: wrong shape of gain array, expected is {%i}, but got %s instead !", input->sizeAt(dimC),
+               ShapeUtils::shapeAsString(gain).c_str());
+
+  NDArray *bias = nullptr;
+  if (block.width() > 2) {
+    bias = INPUT_VARIABLE(2);
+    REQUIRE_TRUE(bias->rankOf() == 1 && bias->sizeAt(0) == input->sizeAt(dimC), 0,
+                 "LAYER_NORM OP: wrong shape of bias array, expected is {%i}, but got %s instead !",
+                 input->sizeAt(dimC), ShapeUtils::shapeAsString(bias).c_str());
+  }
+
+  std::vector<sd::LongType> longAxis = ArrayUtils::toLongVector(axis);
+
+  sd::ops::standardize standardizeOp;
+  std::vector<NDArray *> inputs = {input};
+  std::vector<NDArray *> outputs = {output};
+  std::vector<double> targs = {};
+  std::vector<bool> bargs = {};
+  auto status = standardizeOp.execute(inputs, outputs, targs, longAxis, bargs);
+  if (status != sd::Status::OK) {
+    std::string errorMessage;
+    errorMessage += "LAYER_NORM OP: standardize operation failed with status ";
+    errorMessage += std::to_string(static_cast<int>(status));
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  std::vector<sd::LongType> dimcVec = {dimC};
+  output->applyBroadcast(sd::broadcast::Multiply, &dimcVec, gain, output);
+  if (bias != nullptr) {
+    helpers::addBias(block, *output, *bias, *output, isNCHW);
+  }
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(layer_norm) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+CUSTOM_OP_IMPL(layer_norm_bp, 3, -1, false, 0, -1) {
+  auto input = INPUT_VARIABLE(0);
+  auto gain = INPUT_VARIABLE(1);
+  auto bias = block.width() == 4 ? INPUT_VARIABLE(2) : nullptr;
+  auto eps = block.width() == 4 ? INPUT_VARIABLE(3) : INPUT_VARIABLE(2);
+
+  auto dLdx = OUTPUT_VARIABLE(0);
+  auto dLdg = OUTPUT_VARIABLE(1);
+  auto dLdb = block.width() == 4 ? OUTPUT_VARIABLE(2) : nullptr;
+
+  const bool isNCHW = block.getBArguments()->size() > 0 ? B_ARG(0) : true;  //  0-NCHW,  1-NHWC
+  const int dimC = isNCHW ? 1 : input->rankOf() - 1;
+
+  REQUIRE_TRUE(gain->rankOf() == 1 && gain->sizeAt(0) == input->sizeAt(dimC), 0,
+               "LAYER_NORM_BP OP: wrong shape of gain array, expected is {%i}, but got %s instead !",
+               input->sizeAt(dimC), ShapeUtils::shapeAsString(gain).c_str());
+
+  std::vector<sd::LongType> axis = *block.getIArguments();
+
+  std::vector<sd::LongType> longAxis = ArrayUtils::toLongVector(axis);
+
+  if (bias != nullptr) {
+    REQUIRE_TRUE(bias->rankOf() == 1 && bias->sizeAt(0) == input->sizeAt(dimC), 0,
+                 "LAYER_NORM_BP OP: wrong shape of bias array, expected is {%i}, but got %s instead !",
+                 input->sizeAt(dimC), ShapeUtils::shapeAsString(bias).c_str());
+    std::vector<sd::LongType> dimCVector = {dimC};
+    auto vec = ShapeUtils::evalDimsToExclude(input->rankOf(),1,dimCVector.data());
+    eps->reduceAlongDimension(sd::reduce::Sum, dLdb, vec);
+    delete vec;
+  }
+
+  NDArray standardized(input->shapeInfo(), false, block.launchContext());
+
+  sd::ops::standardize standardizeOp;
+  std::vector<NDArray *> inputs = {input};
+  std::vector<NDArray *> outputs = {&standardized};
+  std::vector<double> targs = {};
+  std::vector<bool> bargs = {};
+
+  auto status = standardizeOp.execute(inputs, outputs, targs, longAxis, bargs);
+  if (status != sd::Status::OK) {
+    std::string errorMessage;
+    errorMessage += "LAYER_NORM_BP OP: standardize operation failed with status ";
+    errorMessage += std::to_string(static_cast<int>(status));
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+  standardized.applyPairwiseTransform(sd::pairwise::Multiply, eps, &standardized);
+  std::vector<sd::LongType> dimCVector = {dimC};
+  auto vec = ShapeUtils::evalDimsToExclude(input->rankOf(),1,dimCVector.data());
+  standardized.reduceAlongDimension(sd::reduce::Sum, dLdg, vec);
+  delete vec;
+
+  sd::ops::standardize_bp standardizeBp;
+  std::vector<sd::LongType> dimvC = {dimC};
+  eps->applyBroadcast(sd::broadcast::Multiply, &dimvC, gain, dLdx);
+
+  auto dLdx_tmp = dLdx->dup();
+  std::vector<NDArray *> standardizeBpArgs = {input, dLdx_tmp};
+  std::vector<NDArray *> standardizeBpOut = {dLdx};
+  status = standardizeBp.execute(standardizeBpArgs, standardizeBpOut, targs, longAxis, bargs);
+  if (status != sd::Status::OK) {
+    delete dLdx_tmp;
+    std::string errorMessage;
+    errorMessage += "LAYER_NORM_BP OP: standardize_bp operation failed with status ";
+    errorMessage += std::to_string(static_cast<int>(status));
+    THROW_EXCEPTION(errorMessage.c_str());
+  }
+
+  delete dLdx_tmp;
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(layer_norm_bp) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(layer_norm_bp) {
+  if (inputShape->size() > 3) {
+    return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)));
+  }
+  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/logSoftmax.cpp

@@ -0,0 +1,93 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 01.02.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_log_softmax)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/activations.h>
+
+namespace sd {
+namespace ops {
+
+DECLARE_TYPES(log_softmax) { getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS})->setSameMode(true); }
+
+CONFIGURABLE_OP_IMPL(log_softmax, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  const int rank = input->rankOf();
+  const int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
+
+  REQUIRE_TRUE(dim < rank, 0,
+               "LOG_SOFTMAX OP: the value of input integer parameter (dimension) must be less than input array rank "
+               "%i, but got dimension = %i instead !",
+               rank, dim);
+
+  helpers::logSoftmax(block.launchContext(), input, output, dim);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(log_softmax_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, DataType::ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(log_softmax_bp, 3, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto gradO = INPUT_VARIABLE(1);
+  auto softmaxOut = INPUT_VARIABLE(2);
+  auto gradI = OUTPUT_VARIABLE(0);
+  gradI->assign(softmaxOut);
+  const int rank = input->rankOf();
+  const int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
+
+  REQUIRE_TRUE(dim < rank, 0,
+               "LOG_SOFTMAX_BP OP: the value of input integer parameter (dimension) must be less than input array rank "
+               "%i, but got dimension = %i instead !",
+               rank, dim);
+
+  helpers::softmax(block.launchContext(), input, gradI, dim);
+
+  std::vector<sd::LongType> dimVec;
+  dimVec.push_back(dim);
+  auto* sumGradOj = gradO->reduceAlongDimension(reduce::Sum, &dimVec, true);
+  
+  // Break down: *gradO - *gradI * sumGradOj
+  auto* product = (*gradI) * (*sumGradOj);
+  auto* assign = (*gradO) - (*product);
+  delete product;
+  delete sumGradOj;
+  
+  gradI->assign(assign);
+  delete assign;
+
+  return sd::Status::OK;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/lrn.cpp

@@ -0,0 +1,79 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119 on 29/10/17
+// @author GS <sgazeos@gmail.com> 2/16/18
+// @author Yurii Shyrma (iuriish@yahoo.com) -> back prop author
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lrn)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lrn.h>
+
+namespace sd {
+namespace ops {
+
+DECLARE_TYPES(lrn) { getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS}); }
+
+CONFIGURABLE_OP_IMPL(lrn, 1, 1, true, 3, 1) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "lrn: Input rank of 4 expected, but got %i instead", input->rankOf());
+
+  double alpha = T_ARG(1);
+  double beta = T_ARG(2);
+  double bias = T_ARG(0);
+  int depth = INT_ARG(0);
+
+  return helpers::lrnFunctor(block, input, output, depth, bias, alpha, beta);
+}
+
+DECLARE_TYPES(lrn_bp) {
+  getOpDescriptor()->setAllowedInputTypes(sd::DataType::ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+CONFIGURABLE_OP_IMPL(lrn_bp, 2, 1, true, 3, 1) {
+  auto input = INPUT_VARIABLE(0);
+  auto gradO = INPUT_VARIABLE(1);
+  auto gradI = OUTPUT_VARIABLE(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "lrn_bp: Input rank of 4 expected, but got %i instead", input->rankOf());
+  REQUIRE_TRUE(input->isSameShape(gradO), 0,
+               "lrn_bp: Both input and grad_output should have the same shape, but got %s and %s correspondingly !",
+               ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+
+  // FIXME: double/float?
+  float bias = T_ARG(0);
+  float alpha = T_ARG(1);
+  float beta = T_ARG(2);
+  int depth = INT_ARG(0);
+
+  helpers::lrnBP(block, *input, *gradO, *gradI, depth, bias, alpha, beta);
+
+  return sd::Status::OK;
+}
+DECLARE_SYN(local_response_normalization, lrn);
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/multi_head_dot_product_attention.cpp

@@ -0,0 +1,303 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// @author Paul Dubs
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_multi_head_dot_product_attention)
+
+#include <helpers/AttentionHelper.h>
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(multi_head_dot_product_attention, 7, -1, false, 0, 2) {
+  auto queries = INPUT_VARIABLE(0);  //[batch, nIn, timeSteps]
+  auto keys = INPUT_VARIABLE(1);     //[batch, nIn, timeSteps]
+  auto values = INPUT_VARIABLE(2);   //[batch, nIn, timeSteps]
+  auto Wq = INPUT_VARIABLE(3);       //[nHeads, headSize, nIn]
+  auto Wk = INPUT_VARIABLE(4);       //[nHeads, headSize, nIn]
+  auto Wv = INPUT_VARIABLE(5);       //[nHeads, headSize, nIn]
+  auto Wo = INPUT_VARIABLE(6);       //[nHeads * headSize, nOut]
+  auto mask = block.width() > 7 ? INPUT_VARIABLE(7) : nullptr;
+
+  auto output = OUTPUT_VARIABLE(0);
+  int normalization = INT_ARG(0);
+  int weights = INT_ARG(1);
+
+  auto numHeads = Wk->sizeAt(0);
+  auto miniBatchSize = queries->sizeAt(0);
+  auto queryCount = queries->sizeAt(2);
+  auto projectedValuesSize = Wv->sizeAt(1);
+  auto outSize = Wo->sizeAt(1);
+
+  REQUIRE_TRUE(queries->rankOf() == keys->rankOf() && keys->rankOf() == values->rankOf(), 0,
+               "multi_head_dot_product_attention: Queries, Keys and Values must have same rank. "
+               "But got queries = %s, keys = %s, values = %s",
+               ShapeUtils::shapeAsString(queries).c_str(), ShapeUtils::shapeAsString(keys).c_str(),
+               ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(queries->rankOf() == 3, 0,
+               "multi_head_dot_product_attention: Queries, Keys and Values must be rank 3 arrays"
+               "But got rank = %i",
+               queries->rankOf());
+
+  REQUIRE_TRUE(Wq->rankOf() == Wk->rankOf() && Wk->rankOf() == Wv->rankOf(), 0,
+               "multi_head_dot_product_attention: Input projections weights must have the same rank. "
+               "But got Wq = %s, Wk = %s, Wv = %s",
+               ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(Wk).c_str(),
+               ShapeUtils::shapeAsString(Wv).c_str());
+
+  REQUIRE_TRUE(Wq->sizeAt(0) == Wk->sizeAt(0) && Wk->sizeAt(0) == Wv->sizeAt(0), 0,
+               "multi_head_dot_product_attention: Projections weights must have the same number of attention heads. "
+               "But got Wq = %s, Wk = %s, Wv = %s",
+               ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(Wk).c_str(),
+               ShapeUtils::shapeAsString(Wv).c_str());
+
+  REQUIRE_TRUE(Wo->rankOf() == 2, 0,
+               "multi_head_dot_product_attention: Output projection weights must have rank 2. "
+               "But got Wo = %s",
+               ShapeUtils::shapeAsString(Wo).c_str());
+
+  REQUIRE_TRUE(Wq->sizeAt(2) == queries->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Query projection matrix Wq has incompatible size to queries matrix."
+               "Expected Wq[2] = queries[1] = %i, but got Wq = %s, queries = %s ",
+               queries->sizeAt(1), ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(queries).c_str());
+
+  REQUIRE_TRUE(Wk->sizeAt(2) == keys->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Key projection matrix Wk has incompatible size to keys matrix."
+               "Expected Wk[2] = keys[1] = %i, but got Wk = %s, keys = %s ",
+               keys->sizeAt(1), ShapeUtils::shapeAsString(Wk).c_str(), ShapeUtils::shapeAsString(keys).c_str());
+
+  REQUIRE_TRUE(Wv->sizeAt(2) == values->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Value projection matrix Wv has incompatible size to values matrix."
+               "Expected Wv[2] = values[1] = %i, but got Wv = %s, values = %s ",
+               values->sizeAt(1), ShapeUtils::shapeAsString(Wv).c_str(), ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(
+      Wo->sizeAt(0) == (Wv->sizeAt(1) * Wv->sizeAt(0)), 0,
+      "multi_head_dot_product_attention: Output projection matrix Wo has incompatible size to attention result."
+      "Expected Wo[0] = Wv[0] * Wv[1] = %i, but got Wo = %s, Wv = %",
+      (Wv->sizeAt(1) * Wv->sizeAt(0)), ShapeUtils::shapeAsString(Wo).c_str(), ShapeUtils::shapeAsString(Wv).c_str());
+
+  // Project queries, keys, values
+  auto projectedQueries = AttentionHelper::multiHeadProject(
+      queries, Wq, block.launchContext());  //[minibatch, numHeads, projectedSize, seqLength]
+  auto projectedKeys = AttentionHelper::multiHeadProject(
+      keys, Wk, block.launchContext());  //[minibatch, numHeads, projectedSize, seqLength]
+  auto projectedValues = AttentionHelper::multiHeadProject(
+      values, Wv, block.launchContext());  //[minibatch, numHeads, projectedSize, seqLength]
+
+  std::vector<sd::LongType> shape = {projectedQueries.sizeAt(0), projectedValues.sizeAt(1), projectedValues.sizeAt(2), projectedQueries.sizeAt(3)};
+  // Apply Attention
+  // attnResults = [minibatch, numHeads, projectedSize, seqLenth
+  NDArray attnResults(
+      'c',
+      shape,
+      projectedValues.dataType(), block.launchContext());
+  sd::ops::dot_product_attention attention;
+  attention.execute({&projectedQueries, &projectedKeys, &projectedValues, mask},
+                    {&attnResults, weights ? OUTPUT_VARIABLE(1) : nullptr}, {}, {normalization, weights}, {});
+
+  // Project attention results
+  attnResults.permutei({0, 3, 1, 2}, 0, false);
+  attnResults.reshapei(attnResults.ordering(), {miniBatchSize * queryCount, numHeads * projectedValuesSize});
+
+  sd::ops::matmul mmul;
+  std::vector<sd::LongType> projShape ={attnResults.sizeAt(0), Wo->sizeAt(1)};
+  NDArray projRes('c', projShape, values->dataType(), block.launchContext());
+  mmul.execute({&attnResults, Wo}, {&projRes}, {}, {}, {});
+  projRes.reshapei(projRes.ordering(), {miniBatchSize, queryCount, outSize});
+  projRes.permutei({0, 2, 1}, 0, false);
+
+  // FIXME: bad for performance
+  output->assign(&projRes);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(multi_head_dot_product_attention) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(multi_head_dot_product_attention) {
+  auto queryShape = inputShape->at(0);
+  auto keysShape = inputShape->at(1);
+  auto valuesShape = inputShape->at(2);
+  auto WkShape = inputShape->at(3);
+  auto WoShape = inputShape->at(6);
+
+  auto batchSize = shape::sizeAt(queryShape, static_cast<sd::LongType>(0));
+  auto outSize = shape::sizeAt(WoShape, static_cast<sd::LongType>(1));
+  auto queryCount = shape::sizeAt(queryShape, static_cast<sd::LongType>(2));
+  auto numHeads = shape::sizeAt(WkShape, static_cast<sd::LongType>(0));
+  auto timeSteps = shape::sizeAt(keysShape, static_cast<sd::LongType>(2));
+
+  auto weightsShape = ConstantShapeHelper::getInstance().createShapeInfo(sd::ArrayOptions::dataType(valuesShape), 'c',
+                                                                         {batchSize, numHeads, timeSteps, queryCount});
+  auto outputShape = ConstantShapeHelper::getInstance().createShapeInfo(sd::ArrayOptions::dataType(valuesShape), 'c',
+                                                                        {batchSize, outSize, queryCount});
+
+  if (INT_ARG(1)) {
+    return SHAPELIST(outputShape, weightsShape);
+  } else {
+    return SHAPELIST(outputShape);
+  }
+}
+
+CUSTOM_OP_IMPL(multi_head_dot_product_attention_bp, 8, 7, false, 0, 1) {
+  auto queries = INPUT_VARIABLE(0);
+  auto keys = INPUT_VARIABLE(1);
+  auto values = INPUT_VARIABLE(2);
+  auto Wq = INPUT_VARIABLE(3);
+  auto Wk = INPUT_VARIABLE(4);
+  auto Wv = INPUT_VARIABLE(5);
+  auto Wo = INPUT_VARIABLE(6);
+  auto eps = INPUT_VARIABLE(7);
+  auto mask = block.width() > 8 ? INPUT_VARIABLE(8) : nullptr;
+
+  auto dLdq = OUTPUT_VARIABLE(0);
+  auto dLdk = OUTPUT_VARIABLE(1);
+  auto dLdv = OUTPUT_VARIABLE(2);
+  auto dLdWq = OUTPUT_VARIABLE(3);
+  auto dLdWk = OUTPUT_VARIABLE(4);
+  auto dLdWv = OUTPUT_VARIABLE(5);
+  auto dLdWo = OUTPUT_VARIABLE(6);
+
+  int normalization = INT_ARG(0);
+
+  auto numHeads = Wk->sizeAt(0);
+  auto miniBatchSize = queries->sizeAt(0);
+  auto queryCount = queries->sizeAt(2);
+  auto outSize = Wo->sizeAt(1);
+  auto projectedValuesSize = Wv->sizeAt(1);
+
+  REQUIRE_TRUE(queries->rankOf() == keys->rankOf() && keys->rankOf() == values->rankOf(), 0,
+               "multi_head_dot_product_attention: Queries, Keys and Values must have same rank. "
+               "But got queries = %s, keys = %s, values = %s",
+               ShapeUtils::shapeAsString(queries).c_str(), ShapeUtils::shapeAsString(keys).c_str(),
+               ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(queries->rankOf() == 3, 0,
+               "multi_head_dot_product_attention: Queries, Keys and Values must be rank 3 arrays"
+               "But got rank = %i",
+               queries->rankOf());
+
+  REQUIRE_TRUE(Wq->rankOf() == Wk->rankOf() && Wk->rankOf() == Wv->rankOf(), 0,
+               "multi_head_dot_product_attention: Input projections weights must have the same rank. "
+               "But got Wq = %s, Wk = %s, Wv = %s",
+               ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(Wk).c_str(),
+               ShapeUtils::shapeAsString(Wv).c_str());
+
+  REQUIRE_TRUE(Wq->sizeAt(0) == Wk->sizeAt(0) && Wk->sizeAt(0) == Wv->sizeAt(0), 0,
+               "multi_head_dot_product_attention: Projections weights must have the same number of attention heads. "
+               "But got Wq = %s, Wk = %s, Wv = %s",
+               ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(Wk).c_str(),
+               ShapeUtils::shapeAsString(Wv).c_str());
+
+  REQUIRE_TRUE(Wo->rankOf() == 2, 0,
+               "multi_head_dot_product_attention: Output projection weights must have rank 2. "
+               "But got Wo = %s",
+               ShapeUtils::shapeAsString(Wo).c_str());
+
+  REQUIRE_TRUE(Wq->sizeAt(2) == queries->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Query projection matrix Wq has incompatible size to queries matrix."
+               "Expected Wq[2] = queries[1] = %i, but got Wq = %s, queries = %s ",
+               queries->sizeAt(1), ShapeUtils::shapeAsString(Wq).c_str(), ShapeUtils::shapeAsString(queries).c_str());
+
+  REQUIRE_TRUE(Wk->sizeAt(2) == keys->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Key projection matrix Wk has incompatible size to keys matrix."
+               "Expected Wk[2] = keys[1] = %i, but got Wk = %s, keys = %s ",
+               keys->sizeAt(1), ShapeUtils::shapeAsString(Wk).c_str(), ShapeUtils::shapeAsString(keys).c_str());
+
+  REQUIRE_TRUE(Wv->sizeAt(2) == values->sizeAt(1), 0,
+               "multi_head_dot_product_attention: Value projection matrix Wv has incompatible size to values matrix."
+               "Expected Wv[2] = values[1] = %i, but got Wv = %s, values = %s ",
+               values->sizeAt(1), ShapeUtils::shapeAsString(Wv).c_str(), ShapeUtils::shapeAsString(values).c_str());
+
+  REQUIRE_TRUE(
+      Wo->sizeAt(0) == (Wv->sizeAt(1) * Wv->sizeAt(0)), 0,
+      "multi_head_dot_product_attention: Output projection matrix Wo has incompatible size to attention result."
+      "Expected Wo[0] = Wv[0] * Wv[1] = %i, but got Wo = %s, Wv = %",
+      (Wv->sizeAt(1) * Wv->sizeAt(0)), ShapeUtils::shapeAsString(Wo).c_str(), ShapeUtils::shapeAsString(Wv).c_str());
+
+  // Project queries, keys, values
+  auto projectedQueries = AttentionHelper::multiHeadProject(queries, Wq, block.launchContext());
+  auto projectedKeys = AttentionHelper::multiHeadProject(keys, Wk, block.launchContext());
+  auto projectedValues = AttentionHelper::multiHeadProject(values, Wv, block.launchContext());
+  std::vector<sd::LongType> shape = {projectedQueries.sizeAt(0), projectedValues.sizeAt(1), projectedValues.sizeAt(2), projectedQueries.sizeAt(3)};
+  // Apply Attention
+  NDArray attnResults(
+      'c',
+      shape,
+      projectedValues.dataType(), block.launchContext());
+  sd::ops::dot_product_attention attention;
+  attention.execute({&projectedQueries, &projectedKeys, &projectedValues, mask}, {&attnResults}, {}, {normalization, 0},
+                    {});
+
+  // Project attention results
+  attnResults.permutei({0, 3, 1, 2}, 0, false);
+  attnResults.reshapei(attnResults.ordering(), {miniBatchSize * queryCount, numHeads * projectedValuesSize});
+
+  std::vector<sd::LongType> perm = {0,2,1};
+  // dLdWo
+  auto epsPerm = eps->permute(perm, false, false);
+  std::vector<sd::LongType> epsShape =  {miniBatchSize * queryCount, outSize};
+  auto epsPostReshape = epsPerm->reshape(eps->ordering(), epsShape);
+  sd::ops::matmul_bp matmulBp;
+  NDArray dLdPreWo(attnResults.shapeInfo(), false, block.launchContext());
+  matmulBp.execute({&attnResults, Wo, epsPostReshape}, std::vector<NDArray *>{&dLdPreWo, dLdWo}, {}, {}, {});
+  delete epsPostReshape;
+  // dLdAttn
+  dLdPreWo.reshapei({miniBatchSize, queryCount, numHeads, projectedValues.sizeAt(2)});
+  dLdPreWo.permutei({0, 2, 3, 1}, 0, false);
+
+  sd::ops::dot_product_attention_bp attentionBp;
+  NDArray dLdProjectedQueries(projectedQueries.shapeInfo(), false, block.launchContext());
+  NDArray dLdProjectedKeys(projectedKeys.shapeInfo(), false, block.launchContext());
+  NDArray dLdProjectedValues(projectedValues.shapeInfo(), false, block.launchContext());
+  attentionBp.execute({&projectedQueries, &projectedKeys, &projectedValues, &dLdPreWo, mask},
+                      {&dLdProjectedQueries, &dLdProjectedKeys, &dLdProjectedValues}, {}, {normalization}, {});
+
+  AttentionHelper::multiHeadProjectBp(queries, Wq, &dLdProjectedQueries, dLdq, dLdWq, block.launchContext());
+  AttentionHelper::multiHeadProjectBp(keys, Wk, &dLdProjectedKeys, dLdk, dLdWk, block.launchContext());
+  AttentionHelper::multiHeadProjectBp(values, Wv, &dLdProjectedValues, dLdv, dLdWv, block.launchContext());
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(multi_head_dot_product_attention_bp) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(multi_head_dot_product_attention_bp) {
+  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)), CONSTANT(inputShape->at(3)),
+                   CONSTANT(inputShape->at(4)), CONSTANT(inputShape->at(5)), CONSTANT(inputShape->at(6)));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/avgpool2d.cpp

@@ -0,0 +1,227 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com), changed on 14.05.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_avgpool2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+CUSTOM_OP_IMPL(avgpool2d, 1, 1, false, 0, 10) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_NULLIFIED(0);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
+  // mode;
+
+  const LongType kH = INT_ARG(0);
+  const LongType kW = INT_ARG(1);
+  const LongType sH = INT_ARG(2);
+  const LongType sW = INT_ARG(3);
+  LongType pH = INT_ARG(4);
+  LongType pW = INT_ARG(5);
+  const LongType dH = INT_ARG(6);
+  const LongType dW = INT_ARG(7);
+  const auto isSameMode = static_cast<bool>(INT_ARG(8));
+  const auto extraParam0 = INT_ARG(9);
+  const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 0-NCHW, 1-NHWC
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D op: input should have rank of 4, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType oH = 0;
+  LongType oW = 0;
+
+  const LongType iH = static_cast<LongType>(isNCHW ? input->sizeAt(2) : input->sizeAt(1));
+  const LongType iW = static_cast<LongType>(isNCHW ? input->sizeAt(3) : input->sizeAt(2));
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm = {0,3,1,2};
+    input = input->permute(perm, false, false);    // [bS, iH, iW, iC] -> [bS, iC, iH, iW] - permute() already returns NDArray*
+    output = output->permute(perm, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW] - permute() already returns NDArray*
+  }
+
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  if (isSameMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
+  // poolingMode; 9 - divisor;
+  ConvolutionUtils::pooling2d(block, *input, *output, kH, kW, sH, sW, pH, pW, dH, dW, AVG_POOL,
+                              extraParam0);
+
+  if (!isNCHW) {
+     delete input;
+      delete output;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SYN(AvgPool2D, avgpool2d);
+DECLARE_SYN(AvgPool, avgpool2d);
+DECLARE_SYN(avgpool, avgpool2d);
+
+DECLARE_TYPES(avgpool2d) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(avgpool2d) {
+  auto inShape = inputShape->at(0);
+  auto shapeOf = shape::shapeOf(inShape);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
+  // mode;
+  const LongType kH = INT_ARG(0);
+  const LongType kW = INT_ARG(1);
+  const LongType sH = INT_ARG(2);
+  const LongType sW = INT_ARG(3);
+  const LongType pH = INT_ARG(4);
+  const LongType pW = INT_ARG(5);
+  const LongType dH = INT_ARG(6);
+  const LongType dW = INT_ARG(7);
+  const int isSameMode = INT_ARG(8);
+
+  const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 0-NCHW, 1-NHWC
+
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  const LongType bS = shapeOf[0];
+  const LongType iD = isNCHW ? shapeOf[1] : shapeOf[3];
+  const LongType iH = isNCHW ? shapeOf[2] : shapeOf[1];
+  const LongType iW = isNCHW ? shapeOf[3] : shapeOf[2];
+
+  const char order = shape::order(inShape);  // output order must be equal to input order
+
+  // calculate output Height/Width
+  LongType oH, oW;
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  // allocate memory for new shape
+  LongType *newShape = new LongType[4];
+  if (isNCHW) {
+    newShape[0] = bS;
+    newShape[1] = iD;
+    newShape[2] = oH;
+    newShape[3] = oW;
+  } else {
+    newShape[0] = bS;
+    newShape[1] = oH;
+    newShape[2] = oW;
+    newShape[3] = iD;
+  }
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape),
+                                                                             shape::order(inShape),
+                                                                             4,
+                                                                             newShape)->primary());
+  delete[] newShape;
+  return ret;
+}
+
+DECLARE_TYPES(avgpool2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(avgpool2d_bp, 2, 1, false, 0, 10) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+
+  LongType kH = INT_ARG(0);          // filter(kernel) height
+  LongType kW = INT_ARG(1);          // filter(kernel) width
+  LongType sH = INT_ARG(2);          // strides height
+  LongType sW = INT_ARG(3);          // strides width
+  LongType pH = INT_ARG(4);          // paddings height
+  LongType pW = INT_ARG(5);          // paddings width
+  LongType dH = INT_ARG(6);          // dilations height
+  LongType dW = INT_ARG(7);          // dilations width
+  int isSameMode = INT_ARG(8);  // 0-VALID, 1-SAME
+  int extraParam0 = INT_ARG(9);
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 0-NCHW, 1-NHWC
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D_BP op: input should have rank of 4, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D_BP op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
+                                             indWiC, indWoC, indWkH, indOoH);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
+  std::vector<LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
+  REQUIRE_TRUE(
+      gradO->isSameShape(expectedGradOShape), 0,
+      "AVGPOOL2D_BP op: wrong shape of output's gradients array (next epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(
+      gradI->isSameShape(expectedGradIShape), 0,
+      "AVGPOOL2D_BP op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm = {0,3,1,2};
+    input = input->permute(perm, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW] - permute() already returns NDArray*
+    gradI = gradI->permute(perm, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW] - permute() already returns NDArray*
+    gradO = gradO->permute(perm, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW] - permute() already returns NDArray*
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
+  // poolingMode; 9 - divisor;
+  ConvolutionUtils::pooling2dBP(block, *input, *gradO, *gradI, kH, kW, sH, sW, pH, pW, dH, dW, 1, extraParam0);
+
+  if (!isNCHW) {
+    delete input;
+    delete gradI;
+    delete gradO;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(avgpool2d_bp) {
+  REQUIRE_TRUE(inputShape->at(0)[0] == 4, 0, "AVGPOOL2D_BP op: input array must be 4D, but got %i instead!",
+               inputShape->at(0)[0]);
+  REQUIRE_TRUE(inputShape->at(1)[0] == 4, 0,
+               "AVGPOOL2D_BP op: output's gradient array (next epsilon) must be 4D, but got %i instead!",
+               inputShape->at(1)[0]);
+
+  auto desc = new ShapeDescriptor(inputShape->at(0), ArrayOptions::dataType(inputShape->at(1)), false);
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/avgpool3d.cpp

@@ -0,0 +1,241 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 01.03.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_avgpool3dnew)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(avgpool3dnew, 1, 1, false, 0, 14) {
+  auto input = INPUT_VARIABLE(0);     // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
+
+  LongType kD = INT_ARG(0);           // filter(kernel) depth
+  LongType kH = INT_ARG(1);           // filter(kernel) height
+  LongType kW = INT_ARG(2);           // filter(kernel) width
+  LongType sD = INT_ARG(3);           // strides depth
+  LongType sH = INT_ARG(4);           // strides height
+  LongType sW = INT_ARG(5);           // strides width
+  LongType pD = INT_ARG(6);           // paddings depth
+  LongType pH = INT_ARG(7);           // paddings height
+  LongType pW = INT_ARG(8);           // paddings width
+  LongType dD = INT_ARG(9);           // dilations depth
+  LongType dH = INT_ARG(10);          // dilations height
+  LongType dW = INT_ARG(11);          // dilations width
+  int isSameMode = INT_ARG(12);  // 1-SAME,  0-VALID
+  int extraParam0 = INT_ARG(13);
+  int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 0-NCDHW, 1-NDHWC
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0, "AVGPOOL3DNEW OP: rank of input array must be equal to 5, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "AVGPOOL3DNEW OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
+                                             indIOioD, indWiC, indWoC, indWkD);
+
+  std::vector<LongType> expectedOutputShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
+               "AVGPOOL3DNEW OP: wrong shape of output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
+
+  if (!isNCDHW) {
+    std::vector<sd::LongType> perm = {0, 4, 1, 2, 3};
+    input = input->permute(perm, false, false);    // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    output =output->permute(perm, false, false);  // [bS, oD, oH, oW, iC] -> [bS, iC, oD, oH, oW]
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+
+  // T extraParams[] = {};
+  ConvolutionUtils::pooling3d(block, *input, *output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, 1, extraParam0);
+
+  if (!isNCDHW) {
+     delete input;
+     delete output;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(avgpool3dnew) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(avgpool3dnew) {
+  LongType kD = INT_ARG(0);                                                  // filter(kernel) depth
+  LongType kH = INT_ARG(1);                                                  // filter(kernel) height
+  LongType kW = INT_ARG(2);                                                  // filter(kernel) width
+  LongType sD = INT_ARG(3);                                                  // strides depth
+  LongType sH = INT_ARG(4);                                                  // strides height
+  LongType sW = INT_ARG(5);                                                  // strides width
+  LongType pD = INT_ARG(6);                                                  // paddings depth
+  LongType pH = INT_ARG(7);                                                  // paddings height
+  LongType pW = INT_ARG(8);                                                  // paddings width
+  LongType dD = INT_ARG(9);                                                  // dilations depth
+  LongType dH = INT_ARG(10);                                                 // dilations height
+  LongType dW = INT_ARG(11);                                                 // dilations width
+  int isSameMode = INT_ARG(12);                                         // 1-SAME,  0-VALID
+  int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 0-NCDHW, 1-NDHWC
+
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "AVGPOOL3DNEW op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  auto inputShapeInfo = inputShape->at(0);
+
+  LongType idxID, idxIC;
+  if (isNCDHW) {
+    idxID = 2;
+    idxIC = 1;
+  } else {
+    idxID = 1;
+    idxIC = 4;
+  }
+
+  LongType bS = inputShapeInfo[1];          // batch size
+  LongType iC = inputShapeInfo[idxIC + 1];  // input channels
+  LongType iD = inputShapeInfo[idxID + 1];  // input depth
+  LongType iH = inputShapeInfo[idxID + 2];  // input height
+  LongType iW = inputShapeInfo[idxID + 3];  // input width
+
+  LongType oD, oH, oW;  // output depth, height, width
+  ConvolutionUtils::calcOutSizePool3D(oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH, iW,
+                                      isSameMode);
+
+  LongType outputShape[5];
+
+  outputShape[0] = bS;
+
+  if (isNCDHW) {
+    outputShape[1] = iC;
+    outputShape[2] = oD;
+    outputShape[3] = oH;
+    outputShape[4] = oW;
+  } else {
+    outputShape[1] = oD;
+    outputShape[2] = oH;
+    outputShape[3] = oW;
+    outputShape[4] = iC;
+  }
+  // TF DOC: A Tensor. Has the same type as input.
+  // TF DOC: A Tensor. Has the same type as input.
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inputShapeInfo),
+                                                                             shape::order(inputShapeInfo),
+                                                                             5,
+                                                                             outputShape)->primary());
+  return ret;
+}
+
+DECLARE_TYPES(avgpool3dnew_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(avgpool3dnew_bp, 2, 1, false, 0, 14) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
+
+  const LongType kD = INT_ARG(0);            // filter(kernel) depth
+  const LongType kH = INT_ARG(1);            // filter(kernel) height
+  const LongType kW = INT_ARG(2);            // filter(kernel) width
+  const LongType sD = INT_ARG(3);            // strides depth
+  const LongType sH = INT_ARG(4);            // strides height
+  const LongType sW = INT_ARG(5);            // strides width
+  LongType pD = INT_ARG(6);                  // paddings depth
+  LongType pH = INT_ARG(7);                  // paddings height
+  LongType pW = INT_ARG(8);                  // paddings width
+  const LongType dD = INT_ARG(9);            // dilations depth
+  const LongType dH = INT_ARG(10);           // dilations height
+  const LongType dW = INT_ARG(11);           // dilations width
+  const int isSameMode = INT_ARG(12);   // 1-SAME,  0-VALID
+  const int extraParam0 = INT_ARG(13);  // define what divisor to use while averaging
+  const int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 0-NCDHW, 1-NDHWC
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0, "AVGPOOL3DNEW_BP op: input should have rank of 5, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "AVGPOOL3DNEW_BP op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
+                                             indIOioD, indWiC, indWoC, indWkD);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  std::vector<LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "AVGPOOL3DNEW_BP op: wrong shape of output's gradients array (next epsilon), expected is %s, but got %s "
+               "instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(
+      gradI->isSameShape(expectedGradIShape), 0,
+      "AVGPOOL3DNEW_BP op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
+
+  if (!isNCDHW) {
+    std::vector<sd::LongType> perm = {0, 4, 1, 2, 3};
+    input = input->permute(perm, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    gradI = gradI->permute(perm, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    gradO =gradO->permute(perm, false, false);  // [bS, oD, oH, oW, iC] -> [bS, iC, oD, oH, oW]
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
+  // poolingMode; 9 - divisor;
+  ConvolutionUtils::pooling3dBP(block, *input, *gradO, *gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, 1,
+                                extraParam0);
+
+  if (!isNCDHW) {
+    delete input;
+    delete gradI;
+    delete gradO;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(avgpool3dnew_bp) {
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().castToDataType(inputShape->at(0),
+                                                                         ArrayOptions::dataType(inputShape->at(1))));
+  return ret;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/maxpool2d.cpp

@@ -0,0 +1,225 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com), changed on 09.05.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_maxpool2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+// maxpool2d corresponds to poolingMode=0
+CUSTOM_OP_IMPL(maxpool2d, 1, 1, false, 0, 9) {
+  auto input = INPUT_VARIABLE(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D OP: input array should have rank of 4, but got %i instead",
+               input->rankOf());
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
+  // mode;
+  auto output = OUTPUT_NULLIFIED(0);
+
+  const LongType kH = INT_ARG(0);
+  const LongType kW = INT_ARG(1);
+  const LongType sH = INT_ARG(2);
+  const LongType sW = INT_ARG(3);
+  LongType pH = INT_ARG(4);
+  LongType pW = INT_ARG(5);
+  const LongType dH = INT_ARG(6);
+  const LongType dW = INT_ARG(7);
+  const bool isSameMode = INT_ARG(8);
+
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType oH = 0;
+  LongType oW = 0;
+
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 1-NHWC, 0-NCHW
+
+  const LongType iH = isNCHW ? input->sizeAt(2) : input->sizeAt(1);
+  const LongType iW = isNCHW ? input->sizeAt(3) : input->sizeAt(2);
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm = {0, 3, 1, 2};
+    input = input->permute(perm, false, false);    // [bS, iH, iW, iC] -> [bS, iC, iH, iW] - permute() already returns NDArray*
+    output = output->permute(perm, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW] - permute() already returns NDArray*
+  }
+
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+  if (isSameMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width;
+  // poolingMode; 9 - divisor;
+  ConvolutionUtils::pooling2d(block, *input, *output, kH, kW, sH, sW, pH, pW, dH, dW, MAX_POOL, 1);
+
+  if (!isNCHW) {
+    delete input;
+    delete output;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SYN(MaxPool2D, maxpool2d);
+DECLARE_SYN(MaxPool, maxpool2d);
+DECLARE_SYN(maxpool, maxpool2d);
+
+DECLARE_TYPES(maxpool2d) { getOpDescriptor()->setAllowedInputTypes(ANY)->setSameMode(true); }
+
+DECLARE_SHAPE_FN(maxpool2d) {
+  // NDArray<T> *x = block.getVariables().at(0)->getNDArray();
+  auto inShape = inputShape->at(0);
+  auto shapeOf = shape::shapeOf(inShape);
+  // 0 - number of dimensions; 1,2 - kernel Height/Width; 3,4 - stride Height/Width; 5,6 - pad Height/Width; 7,8 -
+  // dilation Height/Width; 9,10 - input Height/Width; 11 - batch size; 12 - input depth; 13 - same mode;
+  LongType kH = INT_ARG(0);
+  LongType kW = INT_ARG(1);
+  LongType sH = INT_ARG(2);
+  LongType sW = INT_ARG(3);
+  LongType pH = INT_ARG(4);
+  LongType pW = INT_ARG(5);
+  LongType dH = INT_ARG(6);
+  LongType dW = INT_ARG(7);
+  int isSameMode = INT_ARG(8);
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 1-NHWC, 0-NCHW
+
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType bS = shapeOf[0];
+  LongType iC = isNCHW ? shapeOf[1] : shapeOf[3];
+  LongType iH = isNCHW ? shapeOf[2] : shapeOf[1];
+  LongType iW = isNCHW ? shapeOf[3] : shapeOf[2];
+
+  char order = shape::order(inShape);  // output order must be equal to input order
+
+  // calculate output Height/Width
+  LongType oH, oW;
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+
+  // allocate memory for new shape
+  LongType newShape[4];
+
+  newShape[0] = bS;
+  if (isNCHW) {
+    newShape[1] = iC;
+    newShape[2] = oH;
+    newShape[3] = oW;
+  } else {
+    newShape[1] = oH;
+    newShape[2] = oW;
+    newShape[3] = iC;
+  }
+
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape),
+                                                                             order,
+                                                                             4,
+                                                                             newShape)->primary());
+  return ret;
+}
+
+DECLARE_TYPES(maxpool2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(maxpool2d_bp, 2, 1, false, 0, 10) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+
+  LongType kH = INT_ARG(0);                                                 // filter(kernel) height
+  LongType kW = INT_ARG(1);                                                 // filter(kernel) width
+  LongType sH = INT_ARG(2);                                                 // strides height
+  LongType sW = INT_ARG(3);                                                 // strides width
+  LongType pH = INT_ARG(4);                                                 // paddings height
+  LongType pW = INT_ARG(5);                                                 // paddings width
+  LongType dH = INT_ARG(6);                                                 // dilations height
+  LongType dW = INT_ARG(7);                                                 // dilations width
+  int isSameMode = INT_ARG(8);                                         // 0-VALID, 1-SAME
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // INT_ARG(10): 1-NHWC, 0-NCHW
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D_BP op: input should have rank of 4, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D_BP op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
+                                             indWiC, indWoC, indWkH, indOoH);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
+  std::vector<LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
+  REQUIRE_TRUE(
+      gradO->isSameShape(expectedGradOShape), 0,
+      "MAXPOOL2D_BP op: wrong shape of output's gradients array (next epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(
+      gradI->isSameShape(expectedGradIShape), 0,
+      "MAXPOOL2D_BP op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm = {0, 3, 1, 2};
+    input = input->permute(perm, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW]
+    gradI = gradI->permute(perm, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW]
+    gradO = gradO->permute(perm, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW]
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
+
+
+  ConvolutionUtils::pooling2dBP(block, *input, *gradO, *gradI, kH, kW, sH, sW, pH, pW, dH, dW, 0., 1.);
+
+  if (!isNCHW) {
+    delete input;
+    delete gradI;
+    delete gradO;
+  }
+
+  return Status::OK;
+}
+DECLARE_SYN(MaxPool2D_bp, maxpool2d_bp);
+DECLARE_SYN(MaxPool_bp, maxpool2d_bp);
+
+DECLARE_SHAPE_FN(maxpool2d_bp) {
+  REQUIRE_TRUE(inputShape->at(0)[0] == 4, 0, "MAXPOOL2D_BP op: input array must be 4D, but got %i instead!",
+               inputShape->at(0)[0]);
+  REQUIRE_TRUE(inputShape->at(1)[0] == 4, 0,
+               "MAXPOOL2D_BP op: output's gradient array (next epsilon) must be 4D, but got %i instead!",
+               inputShape->at(1)[0]);
+
+  auto desc = new ShapeDescriptor(inputShape->at(0), ArrayOptions::dataType(inputShape->at(1)), false);
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/maxpool3d.cpp

@@ -0,0 +1,241 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 19.02.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_maxpool3dnew)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(maxpool3dnew, 1, 1, false, 0, 14) {
+  auto input = INPUT_VARIABLE(0);     // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto output = OUTPUT_NULLIFIED(0);  // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
+
+  LongType kD = INT_ARG(0);            // filter(kernel) depth
+  LongType kH = INT_ARG(1);            // filter(kernel) height
+  LongType kW = INT_ARG(2);            // filter(kernel) width
+  LongType sD = INT_ARG(3);            // strides depth
+  LongType sH = INT_ARG(4);            // strides height
+  LongType sW = INT_ARG(5);            // strides width
+  LongType pD = INT_ARG(6);            // paddings depth
+  LongType pH = INT_ARG(7);            // paddings height
+  LongType pW = INT_ARG(8);            // paddings width
+  LongType dD = INT_ARG(9);            // dilations depth
+  LongType dH = INT_ARG(10);           // dilations height
+  LongType dW = INT_ARG(11);           // dilations width
+  int isSameMode = INT_ARG(12);   // 1-SAME,  0-VALID
+  int extraParam0 = INT_ARG(13);  // unnecessary for max case, required only for avg and pnorm cases
+  int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 1-NDHWC, 0-NCDHW
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0, "MAXPOOL3DNEW OP: rank of input array must be equal to 5, but got %i instead !",
+               input->rankOf());
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "MAXPOOL3DNEW op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
+                                             indIOioD, indWiC, indWoC, indWkD);
+
+  std::vector<LongType> expectedOutputShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
+               "MAXPOOL3D op: wrong shape of output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
+  // REQUIRE_TRUE(iD   >= kD && iH   >= kH && iW   >= kW, 0, "MAXPOOL3D OP: the input depth/height/width must be greater
+  // or equal to kernel(filter) depth/height/width, but got [%i, %i, %i] and [%i, %i, %i] correspondingly !", iD,iH,iW,
+  // kD,kH,kW); REQUIRE_TRUE(kD/2 >= pD && kH/2 >= pH && kW/2 >= pW, 0, "MAXPOOL3D OP: pad depth/height/width must not
+  // be greater than half of kernel depth/height/width, but got [%i, %i, %i] and [%i, %i, %i] correspondingly !",
+  // pD,pH,pW, kD,kH,kW);
+
+  if (!isNCDHW) {
+    std::vector<sd::LongType> perm = {0, 4, 1, 2, 3};
+    input = input->permute(perm, false, false);    // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    output = output->permute(perm, false, false);  // [bS, oD, oH, oW, iC] -> [bS, iC, oD, oH, oW]
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+
+  ConvolutionUtils::pooling3d(block, *input, *output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, 0, 1);
+
+  if (!isNCDHW) {
+    delete input;
+    delete output;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(maxpool3dnew) { getOpDescriptor()->setAllowedInputTypes(ANY)->setSameMode(true); }
+
+DECLARE_SHAPE_FN(maxpool3dnew) {
+  LongType kD = INT_ARG(0);           // filter(kernel) depth
+  LongType kH = INT_ARG(1);           // filter(kernel) height
+  LongType kW = INT_ARG(2);           // filter(kernel) width
+  LongType sD = INT_ARG(3);           // strides depth
+  LongType sH = INT_ARG(4);           // strides height
+  LongType sW = INT_ARG(5);           // strides width
+  LongType pD = INT_ARG(6);           // paddings depth
+  LongType pH = INT_ARG(7);           // paddings height
+  LongType pW = INT_ARG(8);           // paddings width
+  LongType dD = INT_ARG(9);           // dilations depth
+  LongType dH = INT_ARG(10);          // dilations height
+  LongType dW = INT_ARG(11);          // dilations width
+  int isSameMode = INT_ARG(12);  // 1-SAME,  0-VALID
+  // int extraParam0 = INT_ARG(13);
+  int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 1-NDHWC, 0-NCDHW
+
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "MAXPOOL3DNEW op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  auto inputShapeInfo = inputShape->at(0);
+
+  LongType idxID, idxIC;
+  if (isNCDHW) {
+    idxID = 2;
+    idxIC = 1;
+  } else {
+    idxID = 1;
+    idxIC = 4;
+  }
+
+  LongType bS = inputShapeInfo[1];          // batch size
+  LongType iC = inputShapeInfo[idxIC + 1];  // input channels
+  LongType iD = inputShapeInfo[idxID + 1];  // input depth
+  LongType iH = inputShapeInfo[idxID + 2];  // input height
+  LongType iW = inputShapeInfo[idxID + 3];  // input width
+
+  LongType oD, oH, oW;  // output depth, height, width
+  ConvolutionUtils::calcOutSizePool3D(oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH, iW,
+                                      isSameMode);
+
+  LongType outputShape[5];
+
+  outputShape[0] = bS;
+  if (isNCDHW) {
+    outputShape[1] = iC;
+    outputShape[2] = oD;
+    outputShape[3] = oH;
+    outputShape[4] = oW;
+  } else {
+    outputShape[1] = oD;
+    outputShape[2] = oH;
+    outputShape[3] = oW;
+    outputShape[4] = iC;
+  }
+
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inputShapeInfo),
+                                                                             shape::order(inputShapeInfo),
+                                                                             5,
+                                                                             outputShape)->primary());
+  return ret;
+}
+
+DECLARE_TYPES(maxpool3dnew_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(maxpool3dnew_bp, 2, 1, false, 0, 14) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
+
+  const LongType kD = INT_ARG(0);           // filter(kernel) depth
+  const LongType kH = INT_ARG(1);           // filter(kernel) height
+  const LongType kW = INT_ARG(2);           // filter(kernel) width
+  const LongType sD = INT_ARG(3);           // strides depth
+  const LongType sH = INT_ARG(4);           // strides height
+  const LongType sW = INT_ARG(5);           // strides width
+  LongType pD = INT_ARG(6);                 // paddings depth
+  LongType pH = INT_ARG(7);                 // paddings height
+  LongType pW = INT_ARG(8);                 // paddings width
+  const LongType dD = INT_ARG(9);           // dilations depth
+  const LongType dH = INT_ARG(10);          // dilations height
+  const LongType dW = INT_ARG(11);          // dilations width
+  const int isSameMode = INT_ARG(12);  // 1-SAME,  0-VALID
+  int extraParam0 = INT_ARG(13);       // unnecessary for max case, required only for avg and pnorm cases
+  int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1;  // 1-NDHWC, 0-NCDHW
+
+  REQUIRE_TRUE(input->rankOf() == 5, 0, "MAXPOOL3DNEW_BP op: input should have rank of 5, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
+               "MAXPOOL3DNEW_BP op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
+
+  LongType bS, iC, iD, iH, iW, oC, oD, oH,
+      oW;  // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
+  LongType indIOioC, indIOioD, indWoC, indWiC, indWkD;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
+                                             indIOioD, indWiC, indWoC, indWkD);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  std::vector<LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
+  REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
+               "MAXPOOL3DNEW_BP op: wrong shape of output's gradients array (next epsilon), expected is %s, but got %s "
+               "instead !",
+               ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(
+      gradI->isSameShape(expectedGradIShape), 0,
+      "MAXPOOL3DNEW_BP op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
+
+  if (!isNCDHW) {
+    std::vector<sd::LongType> perm = {0, 4, 1, 2, 3};
+    input = input->permute(perm, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    gradI = gradI->permute(perm, false, false);  // [bS, iD, iH, iW, iC] -> [bS, iC, iD, iH, iW]
+    gradO = gradO->permute(perm, false, false);  // [bS, oD, oH, oW, iC] -> [bS, iC, oD, oH, oW]
+  }
+
+  if (isSameMode)  // SAME
+    ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
+  // [bS, iC, kD, kH, kW, oD, oH, oW] is de-convoluted to  [bS, iC, iD, iH, iW]
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
+  // poolingMode; 9 - unnecessary;
+  ConvolutionUtils::pooling3dBP(block, *input, *gradO, *gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, 0, 1);
+
+  if (!isNCDHW) {
+    delete input;
+    delete gradI;
+    delete gradO;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(maxpool3dnew_bp) {
+  auto desc = new ShapeDescriptor(inputShape->at(0), ArrayOptions::dataType(inputShape->at(1)), false);
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/maxpool_with_argmax.cpp

@@ -0,0 +1,67 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by GS <sgazeos@gmail.com> at 2/20/18
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_max_pool_with_argmax)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+#include <ops/declarable/helpers/max_pooling.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(max_pool_with_argmax, 1, 2, false, 0, 9) {
+  auto x = INPUT_VARIABLE(0);
+  auto z = OUTPUT_NULLIFIED(0);
+  auto indices = OUTPUT_NULLIFIED(1);
+
+  REQUIRE_TRUE(x->rankOf() == 4, 0, "max_pool_with_argmax: Input should have rank of 4, but got %i instead",
+               x->rankOf());
+
+  auto argI = *(block.getIArguments());
+
+  helpers::maxPoolingFunctor(block.launchContext(), block, x, z, argI, indices);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(max_pool_with_argmax) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(ANY)
+      ->setAllowedOutputTypes(0, {ALL_FLOATS, ALL_INTS})
+      ->setAllowedOutputTypes(1, {ALL_INDICES});
+}
+
+DECLARE_SHAPE_FN(max_pool_with_argmax) {
+  auto in = inputShape->at(0);
+  auto dtype = block.numD() ? D_ARG(0) : INT64;
+  // First shape info uses original data type from 'in'
+  auto valuesShape = ConstantShapeHelper::getInstance().bufferForShapeInfo(in)->primary();
+  // Second one needs to be cast to dtype
+  auto indicesShape = ConstantShapeHelper::getInstance().castToDataType(valuesShape, dtype);
+
+  return SHAPELIST(valuesShape, indicesShape);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/pooling/pnormpool2d.cpp

@@ -0,0 +1,223 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 29/10/17.
+// @author Yurii Shyrma (iuriish@yahoo.com), changed on 14.05.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_pnormpool2d)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/convolutions.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(pnormpool2d, 1, 1, false, 0, 10) {
+  REQUIRE_OK(this->validateInputLengthMatch(block));
+  REQUIRE_OK(this->validateInputDimensionsMatch(block));
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_NULLIFIED(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "PNORMPOOL2D op: input should have rank of 4, but got %i instead",
+               input->rankOf());
+
+  LongType kY = INT_ARG(0);
+  LongType kX = INT_ARG(1);
+  LongType sY = INT_ARG(2);
+  LongType sX = INT_ARG(3);
+  LongType pY = INT_ARG(4);
+  LongType pX = INT_ARG(5);
+  LongType dY = INT_ARG(6);
+  LongType dX = INT_ARG(7);
+  bool isSameMode = static_cast<bool>(INT_ARG(8));
+  auto extraParam0 = INT_ARG(9);
+
+  REQUIRE_TRUE(dY != 0 && dX != 0, 0, "PNORMPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dY, dX);
+
+  LongType oY = 0;
+  LongType oX = 0;
+
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // 1-NHWC, 0-NCHW
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm = {0, 3, 1, 2};
+    input = input->permute(perm, false, false);    // [bS, iH, iW, iC] -> [bS, iC, iH, iW]
+    output = output->permute(perm, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW]
+  }
+
+  const LongType inY = static_cast<LongType>(input->sizeAt(2));
+  const LongType inX = static_cast<LongType>(input->sizeAt(3));
+
+  ConvolutionUtils::calcOutSizePool2D(oY, oX, kY, kX, sY, sX, pY, pX, dY, dX, inY, inX, isSameMode);
+
+  if (isSameMode) ConvolutionUtils::calcPadding2D(pY, pX, oY, oX, inY, inX, kY, kX, sY, sX, dY, dX);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
+  // poolingMode; 9 - divisor;
+  ConvolutionUtils::pooling2d(block, *input, *output, kY, kX, sY, sX, pY, pX, dY, dX, PNORM_POOL,
+                              extraParam0);
+
+  if (!isNCHW) {
+    delete input;
+    delete output;
+  }
+
+  return Status::OK;
+}
+DECLARE_SYN(PnormPool2D, pnormpool2d);
+DECLARE_SYN(PnormPool, pnormpool2d);
+DECLARE_SYN(pnormpool, pnormpool2d);
+
+DECLARE_TYPES(pnormpool2d) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(pnormpool2d) {
+  auto inShape = inputShape->at(0);
+  auto shapeOf = shape::shapeOf(inShape);
+
+  // 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
+  // mode;
+  std::vector<LongType> argI = *(block.getIArguments());
+  LongType kH = INT_ARG(0);
+  LongType kW = INT_ARG(1);
+  LongType sH = INT_ARG(2);
+  LongType sW = INT_ARG(3);
+  LongType pH = INT_ARG(4);
+  LongType pW = INT_ARG(5);
+  LongType dH = INT_ARG(6);
+  LongType dW = INT_ARG(7);
+  int isSameMode = INT_ARG(8);
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // 1-NHWC, 0-NCHW
+
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "PNORMPOOL2D op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType bS = shapeOf[0];
+  LongType iC = isNCHW ? shapeOf[1] : shapeOf[3];
+  LongType iH = isNCHW ? shapeOf[2] : shapeOf[1];
+  LongType iW = isNCHW ? shapeOf[3] : shapeOf[2];
+  char order = shape::order(inShape);  // output order must be equal to input order
+
+  // calculate output Height/Width
+  LongType oH, oW;
+  ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
+  // allocate memory for new shape
+  LongType newShape[4];
+
+  newShape[0] = bS;
+  if (isNCHW) {
+    newShape[1] = iC;
+    newShape[2] = oH;
+    newShape[3] = oW;
+  } else {
+    newShape[1] = oH;
+    newShape[2] = oW;
+    newShape[3] = iC;
+  }
+
+  auto ret = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape),
+                                                                             order,
+                                                                             4,
+                                                                             newShape)->primary());
+  return ret;
+}
+
+DECLARE_TYPES(pnormpool2d_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(pnormpool2d_bp, 2, 1, false, 1, 10) {
+  auto input = INPUT_VARIABLE(0);    // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
+  auto gradO = INPUT_VARIABLE(1);    // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
+  auto gradI = OUTPUT_NULLIFIED(0);  // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
+
+  LongType kH = INT_ARG(0);          // filter(kernel) height
+  LongType kW = INT_ARG(1);          // filter(kernel) width
+  LongType sH = INT_ARG(2);          // strides height
+  LongType sW = INT_ARG(3);          // strides width
+  LongType pH = INT_ARG(4);          // paddings height
+  LongType pW = INT_ARG(5);          // paddings width
+  LongType dH = INT_ARG(6);          // dilations height
+  LongType dW = INT_ARG(7);          // dilations width
+  int isSameMode = INT_ARG(8);  // 0-VALID, 1-SAME
+  int pnorm = INT_ARG(9);
+  int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1;  // 1-NHWC, 0-NCHW
+
+  // FIXME: double?
+  double eps = T_ARG(0);
+
+  REQUIRE_TRUE(input->rankOf() == 4, 0, "PNORMPOOL2D_BP op: input should have rank of 4, but got %i instead",
+               input->rankOf());
+  REQUIRE_TRUE(dH != 0 && dW != 0, 0, "PNORMPOOL2D_BP op: dilation must not be zero, but got instead {%i, %i}", dH, dW);
+
+  LongType bS, iC, iH, iW, oC, oH,
+      oW;  // batch size, input channels, input height/width, output channels, output height/width;
+  LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;  // corresponding indexes
+  ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
+                                             indWiC, indWoC, indWkH, indOoH);
+
+  std::vector<LongType> expectedGradOShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
+  std::vector<LongType> expectedGradIShape =
+      ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
+  REQUIRE_TRUE(
+      gradO->isSameShape(expectedGradOShape), 0,
+      "PNORMPOOL2D_BP op: wrong shape of output's gradients array (next epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
+  REQUIRE_TRUE(
+      gradI->isSameShape(expectedGradIShape), 0,
+      "PNORMPOOL2D_BP op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
+
+  if (!isNCHW) {
+    std::vector<sd::LongType> perm2 = {0, 3, 1, 2};
+    input = input->permute(perm2, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW]
+    gradI = gradI->permute(perm2, false, false);  // [bS, iH, iW, iC] -> [bS, iC, iH, iW]
+    gradO = gradO->permute(perm2, false, false);  // [bS, oH, oW, iC] -> [bS, iC, oH, oW]
+  }
+
+
+  ConvolutionUtils::pooling2dBP(block, *input, *gradO, *gradI, kH, kW, sH, sW, pH, pW, dH, dW, 2, pnorm);
+
+  if (!isNCHW) {
+     delete input;
+     delete gradI;
+      delete gradO;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(pnormpool2d_bp) {
+  REQUIRE_TRUE(inputShape->at(0)[0] == 4, 0, "PNORMPOOL2D_BP op: input array must be 4D, but got %i instead!",
+               inputShape->at(0)[0]);
+  REQUIRE_TRUE(inputShape->at(1)[0] == 4, 0,
+               "PNORMPOOL2D_BP op: output's gradient array (next epsilon) must be 4D, but got %i instead!",
+               inputShape->at(1)[0]);
+
+  auto desc = new ShapeDescriptor(inputShape->at(0), ArrayOptions::dataType(inputShape->at(1)), false);
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/dynamicBidirectionalRNN.cpp

@@ -0,0 +1,295 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 05.04.2018
+//
+
+#include <ops/declarable/CustomOperations.h>
+#if NOT_EXCLUDED(OP_dynamic_bidirectional_rnn)
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(dynamic_bidirectional_rnn, 7, 4, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);
+  // input [time x bS x inSize] or [bS x time x inSize], shape depends on timeMajor parameter
+  auto WxFW = INPUT_VARIABLE(1);  // input-to-hidden  weights for forward  RNN, [inSize  x numUnitsFW]
+  auto WhFW = INPUT_VARIABLE(2);  // hidden-to-hidden weights for forward  RNN, [numUnitsFW x numUnitsFW]
+  auto bFW = INPUT_VARIABLE(3);   // biases for forward  RNN, [2*numUnitsFW]
+  auto WxBW = INPUT_VARIABLE(4);  // input-to-hidden  weights for backward RNN, [inSize  x numUnitsBW]
+  auto WhBW = INPUT_VARIABLE(5);  // hidden-to-hidden weights for backward RNN, [numUnitsBW x numUnitsBW]
+  auto bBW = INPUT_VARIABLE(6);   // biases for backward RNN, [2*v]
+
+  NDArray* h0FW = nullptr;  // initial cell output for forward  RNN (at time step = 0) [bS x numUnitsFW]
+  NDArray* h0BW = nullptr;  // initial cell output for backward RNN (at time step = 0) [bS x numUnitsBW]
+  NDArray* maxTimeStep =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  const int timeMajor =
+      block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // if non zero then [time, bS, ...], else [bS, time, ...]
+
+  switch (block.width()) {
+    case 8:
+      maxTimeStep = INPUT_VARIABLE(7);
+      break;
+    case 9:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      break;
+    case 10:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      maxTimeStep = INPUT_VARIABLE(9);
+      break;
+  }
+
+  auto hFW = OUTPUT_VARIABLE(0);  // cell outputs for forward RNN  [time x bS x numUnitsFW] or [bS x time x numUnitsFW],
+                                  // shape depends on timeMajor parameter
+  auto hBW = OUTPUT_VARIABLE(1);  // cell outputs for backward RNN [time x bS x numUnitsBW] or [bS x time x numUnitsBW],
+                                  // shape depends on timeMajor parameter
+  auto hFWFinal = OUTPUT_VARIABLE(2);  // final cell out for forward  RNN [bS x numUnitsFW]
+  auto hBWFinal = OUTPUT_VARIABLE(3);  // final cell out for backward RNN [bS x numUnitsBF]
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input array must have rank = 3, but got %i instead !",
+               x->rankOf());
+  REQUIRE_TRUE(WxFW->rankOf() == 2, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for forward  RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxFW->rankOf());
+  REQUIRE_TRUE(WxBW->rankOf() == 2, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for backward RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxBW->rankOf());
+
+  const int inRank = x->rankOf();
+  int time = timeMajor ? x->sizeAt(0) : x->sizeAt(1);
+  const int bS = timeMajor ? x->sizeAt(1) : x->sizeAt(0);
+  const int numUnitsFW = WxFW->sizeAt(1);
+  const int numUnitsBW = WxBW->sizeAt(1);
+
+  std::vector<LongType> expectedWhFWshape = {numUnitsFW, numUnitsFW};
+  std::vector<LongType> expectedWhBWshape = {numUnitsBW, numUnitsBW};
+  std::vector<LongType> expectedbFWshape = {2 * numUnitsFW};
+  std::vector<LongType> expectedbBWshape = {2 * numUnitsBW};
+  REQUIRE_TRUE(WhFW->isSameShape(expectedWhFWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for forward "
+               " RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhFWshape).c_str(), ShapeUtils::shapeAsString(WhFW).c_str());
+  REQUIRE_TRUE(WhBW->isSameShape(expectedWhBWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for "
+               "backward RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhBWshape).c_str(), ShapeUtils::shapeAsString(WhBW).c_str());
+  REQUIRE_TRUE(bFW->isSameShape(expectedbFWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for forward  RNN), expected "
+               "is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbFWshape).c_str(), ShapeUtils::shapeAsString(bFW).c_str());
+  REQUIRE_TRUE(bBW->isSameShape(expectedbBWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for backward RNN), expected "
+               "is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbBWshape).c_str(), ShapeUtils::shapeAsString(bBW).c_str());
+  if (h0FW) {
+    std::vector<LongType> expectedh0FWshape = {bS, numUnitsFW};
+    REQUIRE_TRUE(h0FW->isSameShape(expectedh0FWshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for forward  "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0FWshape).c_str(), ShapeUtils::shapeAsString(h0FW).c_str());
+  }
+  if (h0BW) {
+    std::vector<LongType> expectedh0BWshape = {bS, numUnitsBW};
+    REQUIRE_TRUE(h0BW->isSameShape(expectedh0BWshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for backward "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0BWshape).c_str(), ShapeUtils::shapeAsString(h0BW).c_str());
+  }
+  if (maxTimeStep) {
+    std::vector<LongType> expectedmaxTimeStepshape = {bS};
+    REQUIRE_TRUE(maxTimeStep->isSameShape(expectedmaxTimeStepshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of maxTimeStep array, expected is [%i], but "
+                 "got %s instead !",
+                 bS, ShapeUtils::shapeAsString(maxTimeStep).c_str());
+  }
+
+  // forward steps
+  dynamic_rnn dynamicRnn;
+  auto resultsFW = dynamicRnn.evaluate({x, WxFW, WhFW, bFW, h0FW, maxTimeStep}, {timeMajor});
+  hFW->assign(resultsFW.at(0));  // [time x bS x numUnitsFW] or [bS x time x numUnitsFW]
+  hFWFinal->assign(resultsFW.at(1));
+
+  auto seqLen = maxTimeStep;
+  if (seqLen == nullptr) {
+    // FIXME: which datatype should be used here?
+    std::vector<sd::LongType> shape = {bS};
+    seqLen = new NDArray(x->ordering(),shape, INT64, block.launchContext());
+    seqLen->assign(time);  // set each element of seqLen to be equal to time
+  }
+
+  // reverse x
+  reverse_sequence reverse;
+  auto resultsIn = timeMajor ? reverse.evaluate({x, seqLen}, {0, 1}) : reverse.evaluate({x, seqLen}, {1, 0});
+  REQUIRE_TRUE(resultsIn.status() == sd::Status::OK, 0,
+               "dynamic_bidirectional_rnn: there is a problem with reverse on the sequence.");
+  auto revInput = resultsIn.at(0);
+
+  // backward steps
+  auto resultsBW = dynamicRnn.evaluate({revInput, WxBW, WhBW, bBW, h0BW, maxTimeStep}, {timeMajor});
+  auto hBWtemp = resultsBW.at(0);  // [time x bS x numUnitsBW] or [ bS x time xnumUnitsBW]
+  hBWFinal->assign(resultsBW.at(1));
+
+  // reverse hBWtemp
+  auto resultsOut =
+      timeMajor ? reverse.evaluate({hBWtemp, seqLen}, {0, 1}) : reverse.evaluate({hBWtemp, seqLen}, {1, 0});
+  hBW->assign(resultsOut.at(0));
+
+  if (seqLen != maxTimeStep) delete seqLen;
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(dynamic_bidirectional_rnn) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dynamic_bidirectional_rnn) {
+  auto x =
+      INPUT_VARIABLE(0);  // input [time x bS x inSize] or [bS x time x inSize], shape depends on timeMajor parameter
+  auto WxFW = INPUT_VARIABLE(1);  // input-to-hidden  weights for forward  RNN, [inSize  x numUnitsFW]
+  auto WhFW = INPUT_VARIABLE(2);  // hidden-to-hidden weights for forward  RNN, [numUnitsFW x numUnitsFW]
+  auto bFW = INPUT_VARIABLE(3);   // biases for forward  RNN, [2*numUnitsFW]
+  auto WxBW = INPUT_VARIABLE(4);  // input-to-hidden  weights for backward RNN, [inSize  x numUnitsBW]
+  auto WhBW = INPUT_VARIABLE(5);  // hidden-to-hidden weights for backward RNN, [numUnitsBW x numUnitsBW]
+  auto bBW = INPUT_VARIABLE(6);   // biases for backward RNN, [2*numUnitsBW]
+
+  NDArray* h0FW = nullptr;  // initial cell output for forward  RNN (at time step = 0) [bS x numUnitsFW]
+  NDArray* h0BW = nullptr;  // initial cell output for backward RNN (at time step = 0) [bS x numUnitsBW]
+  NDArray* maxTimeStep =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  const int timeMajor =
+      block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // if true then [time, bS, ...], else [bS, time, ...]
+
+  switch (block.width()) {
+    case 8:
+      maxTimeStep = INPUT_VARIABLE(7);
+      break;
+    case 9:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      break;
+    case 10:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      maxTimeStep = INPUT_VARIABLE(9);
+      break;
+  }
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input array must have rank = 3, but got %i instead !",
+               x->rankOf());
+  REQUIRE_TRUE(WxFW->rankOf() == 2, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for forward  RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxFW->rankOf());
+  REQUIRE_TRUE(WxBW->rankOf() == 2, 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for backward RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxBW->rankOf());
+
+  const int inRank = x->rankOf();
+  const int time = timeMajor ? x->sizeAt(0) : x->sizeAt(1);
+  const int bS = timeMajor ? x->sizeAt(1) : x->sizeAt(0);
+  const int numUnitsFW = WxFW->sizeAt(1);
+  const int numUnitsBW = WxBW->sizeAt(1);
+
+  std::vector<LongType> expectedWhFWshape = {numUnitsFW, numUnitsFW};
+  std::vector<LongType> expectedWhBWshape = {numUnitsBW, numUnitsBW};
+  std::vector<LongType> expectedbFWshape = {2 * numUnitsFW};
+  std::vector<LongType> expectedbBWshape = {2 * numUnitsBW};
+
+  REQUIRE_TRUE(WhFW->isSameShape(expectedWhFWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for forward "
+               " RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhFWshape).c_str(), ShapeUtils::shapeAsString(WhFW).c_str());
+  REQUIRE_TRUE(WhBW->isSameShape(expectedWhBWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for "
+               "backward RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhBWshape).c_str(), ShapeUtils::shapeAsString(WhBW).c_str());
+  REQUIRE_TRUE(bFW->isSameShape(expectedbFWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for forward  RNN), expected "
+               "is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbFWshape).c_str(), ShapeUtils::shapeAsString(bFW).c_str());
+  REQUIRE_TRUE(bBW->isSameShape(expectedbBWshape), 0,
+               "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for backward RNN), expected "
+               "is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbBWshape).c_str(), ShapeUtils::shapeAsString(bBW).c_str());
+  if (h0FW) {
+    std::vector<LongType> expectedh0FWshape = {bS, numUnitsFW};
+    REQUIRE_TRUE(h0FW->isSameShape(expectedh0FWshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for forward  "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0FWshape).c_str(), ShapeUtils::shapeAsString(h0FW).c_str());
+  }
+  if (h0BW) {
+    std::vector<LongType> expectedh0BWshape = {bS, numUnitsBW};
+    REQUIRE_TRUE(h0BW->isSameShape(expectedh0BWshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for backward "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0BWshape).c_str(), ShapeUtils::shapeAsString(h0BW).c_str());
+  }
+  if (maxTimeStep) {
+    std::vector<LongType> expectedmaxTimeStepshape = {bS};
+    REQUIRE_TRUE(maxTimeStep->isSameShape(expectedmaxTimeStepshape), 0,
+                 "DYNAMIC_BIDIRECTIONAL_RNN custom operation: wrong shape of maxTimeStep array, expected is [%i], but "
+                 "got %s instead !",
+                 bS, ShapeUtils::shapeAsString(maxTimeStep).c_str());
+  }
+
+  // evaluate output shapeInfos
+  LongType *hFWShapeInfo(nullptr), *hBWShapeInfo(nullptr), *hFWFinalPrevShapeInfo(nullptr),
+      *hBWFinalPrevShapeInfo(nullptr);
+  ALLOCATE(hFWShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank), sd::LongType);
+  ALLOCATE(hBWShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank), sd::LongType);
+  ALLOCATE(hFWFinalPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+  ALLOCATE(hBWFinalPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+
+  hFWShapeInfo[0] = hBWShapeInfo[0] = inRank;
+  hFWShapeInfo[1] = hBWShapeInfo[1] = timeMajor ? time : bS;
+  hFWShapeInfo[2] = hBWShapeInfo[2] = timeMajor ? bS : time;
+  hFWShapeInfo[3] = numUnitsFW;
+  hBWShapeInfo[3] = numUnitsBW;
+  hFWFinalPrevShapeInfo[0] = hBWFinalPrevShapeInfo[0] = inRank - 1;
+  hFWFinalPrevShapeInfo[1] = hBWFinalPrevShapeInfo[1] = bS;
+  hFWFinalPrevShapeInfo[2] = numUnitsFW;
+  hBWFinalPrevShapeInfo[2] = numUnitsBW;
+
+  ShapeUtils::updateStridesAndType(hFWShapeInfo, x->shapeInfo(), x->ordering());
+  ShapeUtils::updateStridesAndType(hBWShapeInfo, x->shapeInfo(), x->ordering());
+  ShapeUtils::updateStridesAndType(hFWFinalPrevShapeInfo, x->shapeInfo(), x->ordering());
+  ShapeUtils::updateStridesAndType(hBWFinalPrevShapeInfo, x->shapeInfo(), x->ordering());
+
+  return SHAPELIST(CONSTANT(hFWShapeInfo), CONSTANT(hBWShapeInfo), CONSTANT(hFWFinalPrevShapeInfo),
+                   CONSTANT(hBWFinalPrevShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/dynamicRNN.cpp

@@ -0,0 +1,202 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 05.04.2018
+//
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/rnn.h>
+
+#if NOT_EXCLUDED(OP_dynamic_rnn)
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(dynamic_rnn, 4, 2, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // input [time x bS x inSize] or [bS x time x inSize], depends on timeMajor parameter
+  auto Wx = INPUT_VARIABLE(1);  // input-to-hidden  weights, [inSize  x numUnits]
+  auto Wh = INPUT_VARIABLE(2);  // hidden-to-hidden weights, [numUnits x numUnits]
+  auto b = INPUT_VARIABLE(3);   // biases for, [2*numUnits]
+
+  NDArray* h0 = nullptr;  // initial cell output (at time step = 0) [bS x numUnits]
+  NDArray* maxTimeStep =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  const int timeMajor =
+      block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // if true then [time, bS, ...], else [bS, time, ...]
+
+  if (block.width() == 5) {
+    if ((*INPUT_VARIABLE(4)).rankOf() == 2)
+      h0 = INPUT_VARIABLE(4);
+    else
+      maxTimeStep = INPUT_VARIABLE(4);
+  } else if (block.width() == 6) {
+    h0 = INPUT_VARIABLE(4);
+    maxTimeStep = INPUT_VARIABLE(5);
+  }
+
+  auto h = OUTPUT_VARIABLE(0);  // cell outputs [time x bS x numUnits] or [bS x time x numUnits], depends on timeMajor parameter
+  auto hFinal = OUTPUT_VARIABLE(1);  // at the end it will store cell final non-zero output [bS x numUnits]
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0,
+               "DYNAMIC_RNN custom operation: input array x must have rank = 3, but got %i instead !", x->rankOf());
+  REQUIRE_TRUE(Wx->rankOf() == 2, 0,
+               "DYNAMIC_RNN custom operation: input-to-hidden weights array must have rank = 2, but got %i instead !",
+               Wx->rankOf());
+
+  const int inRank = x->rankOf();
+  const int time = timeMajor ? x->sizeAt(0) : x->sizeAt(1);
+  const int bS = timeMajor ? x->sizeAt(1) : x->sizeAt(0);
+  const int numUnits = Wx->sizeAt(1);
+
+  std::vector<LongType> expectedWhShape = {numUnits, numUnits};
+  std::vector<LongType> expectedBShape = {2 * numUnits};
+  REQUIRE_TRUE(Wh->isSameShape(expectedWhShape), 0,
+               "DYNAMIC_RNN custom operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(expectedWhShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(expectedBShape), 0,
+               "DYNAMIC_RNN custom operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedBShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  if (h0) {
+    std::vector<LongType> expectedh0Shape = {bS, numUnits};
+    REQUIRE_TRUE(
+        h0->isSameShape(expectedh0Shape), 0,
+        "DYNAMIC_RNN custom operation: wrong shape of initial cell output array, expected is %s but got %s instead !",
+        ShapeUtils::shapeAsString(expectedh0Shape).c_str(), ShapeUtils::shapeAsString(h0).c_str());
+  }
+  if (maxTimeStep) {
+    std::vector<LongType> expectedmaxTimeStepShape = {bS};
+    REQUIRE_TRUE(maxTimeStep->isSameShape(expectedmaxTimeStepShape), 0,
+                 "DYNAMIC_RNN custom operation: wrong shape of maxTimeStep array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedmaxTimeStepShape).c_str(),
+                 ShapeUtils::shapeAsString(maxTimeStep).c_str());
+  }
+
+  if (timeMajor == false) {
+    std::vector<sd::LongType> perm = {1, 0, 2};
+    x = x->permute(perm, false, false);  // [bS x time x inSize]   -> [time x bS x inSize]
+    h = h->permute(perm, false, false);  // [bS x time x numUnits] -> [time x bS x numUnits]
+  }
+
+  helpers::rnnTimeLoop(block.launchContext(), x, Wx, Wh, b, h0, maxTimeStep, h, hFinal);
+
+  if (timeMajor == false) {
+    delete x;
+    delete h;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(dynamic_rnn) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedInputTypes(4, {ALL_FLOATS, ALL_INTS})
+      ->setAllowedInputTypes(5, {ALL_FLOATS, ALL_INTS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS})
+      ->setAllowedOutputTypes(1, {ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(dynamic_rnn) {
+  auto xShapeInfo =
+      inputShape->at(0);  // input [time x bS x inSize] or [bS x time x inSize], depends on timeMajor parameter
+  auto WxShapeInfo = inputShape->at(1);  // input-to-hidden  weights, [inSize  x numUnits]
+  auto WhShapeInfo = inputShape->at(2);  // hidden-to-hidden weights, [numUnits x numUnits]
+  auto bShapeInfo = inputShape->at(3);   // biases for, [2*numUnits]
+
+  LongType const* h0ShapeInfo = nullptr;  // initial cell output (at time step = 0) [bS x numUnits]
+  LongType const* maxTimeStepShapeInfo =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  const int timeMajor =
+      block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;  // if true then [time, bS, ...], else [bS, time, ...]
+
+  if (block.width() == 5) {
+    if (inputShape->at(4)[0] == 2)
+      h0ShapeInfo = inputShape->at(4);
+    else
+      maxTimeStepShapeInfo = inputShape->at(4);
+  } else if (block.width() == 6) {
+    h0ShapeInfo = inputShape->at(4);
+    maxTimeStepShapeInfo = inputShape->at(5);
+  }
+
+  REQUIRE_TRUE(xShapeInfo[0] == 3, 0,
+               "DYNAMIC_RNN custom operation: input array x must have rank = 3, but got %i instead !", xShapeInfo[0]);
+  REQUIRE_TRUE(WxShapeInfo[0] == 2, 0,
+               "DYNAMIC_RNN custom operation: input-to-hidden weights array must have rank = 2, but got %i instead !",
+               WxShapeInfo[0]);
+
+  const int inRank = xShapeInfo[0];
+  const int time = timeMajor ? xShapeInfo[1] : xShapeInfo[2];
+  const int bS = timeMajor ? xShapeInfo[2] : xShapeInfo[1];
+  const int numUnits = WxShapeInfo[2];
+
+  std::vector<LongType> expectedWhShape = {numUnits, numUnits};
+  std::vector<LongType> expectedBShape = {2 * numUnits};
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhShapeInfo, expectedWhShape), 0,
+               "DYNAMIC_RNN custom operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(expectedWhShape).c_str(), ShapeUtils::shapeAsString(WhShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, expectedBShape), 0,
+               "DYNAMIC_RNN custom operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedBShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  if (h0ShapeInfo) {
+    std::vector<LongType> expectedh0Shape = {bS, numUnits};
+    REQUIRE_TRUE(
+        ShapeUtils::areShapesEqual(h0ShapeInfo, expectedh0Shape), 0,
+        "DYNAMIC_RNN custom operation: wrong shape of initial cell output array, expected is %s but got %s instead !",
+        ShapeUtils::shapeAsString(expectedh0Shape).c_str(), ShapeUtils::shapeAsString(h0ShapeInfo).c_str());
+  }
+  if (maxTimeStepShapeInfo) {
+    std::vector<LongType> expectedmaxTimeStepShape = {bS};
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(maxTimeStepShapeInfo, expectedmaxTimeStepShape), 0,
+                 "DYNAMIC_RNN custom operation: wrong shape of maxTimeStep array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedmaxTimeStepShape).c_str(),
+                 ShapeUtils::shapeAsString(maxTimeStepShapeInfo).c_str());
+  }
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *hPrevShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank), sd::LongType);
+  ALLOCATE(hPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+
+  hShapeInfo[0] = inRank;
+  hPrevShapeInfo[0] = inRank - 1;
+  hShapeInfo[1] = timeMajor ? time : bS;
+  hShapeInfo[2] = timeMajor ? bS : time;
+  hPrevShapeInfo[1] = bS;
+  hShapeInfo[3] = hPrevShapeInfo[2] = numUnits;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, WhShapeInfo, shape::order(xShapeInfo));
+  ShapeUtils::updateStridesAndType(hPrevShapeInfo, WhShapeInfo, shape::order(xShapeInfo));
+
+  return SHAPELIST(CONSTANT(hShapeInfo), CONSTANT(hPrevShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/gru.cpp

@@ -0,0 +1,219 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_gru)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/gru.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(gru, 5, 1, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // input [time, bS, nIn]
+  auto hI = INPUT_VARIABLE(1);  // initial cell output (at time step = 0) [bS, nOut]
+
+  auto Wx = INPUT_VARIABLE(2);  // input-to-hidden  weights, [nIn, 3*nOut]
+  auto Wh = INPUT_VARIABLE(3);  // hidden-to-hidden weights, [nOut, 3*nOut]
+  auto b = INPUT_VARIABLE(4);   // biases, [3*nOut]
+
+  auto h = OUTPUT_VARIABLE(0);  // cell outputs [time, bS, nOut], that is per each time step
+
+
+  auto linearBeforeReset = block.numB() > 0 ? B_ARG(0) : false;
+
+  const int bS = x->sizeAt(1);
+  const int nIn = x->sizeAt(2);
+  const int nOut = hI->sizeAt(1);
+
+  const std::vector<LongType> h0CorrectShape = {bS, nOut};
+  const std::vector<LongType> wxCorrectShape = {nIn, 3 * nOut};
+  const std::vector<LongType> whCorrectShape = {nOut, 3 * nOut};
+  const std::vector<LongType> bCorrectShape = {3 * nOut};
+
+  REQUIRE_TRUE(hI->isSameShape(h0CorrectShape), 0,
+               "GRU operation: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(h0CorrectShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(wxCorrectShape), 0,
+               "GRU operation: wrong shape of input-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wxCorrectShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(whCorrectShape), 0,
+               "GRU operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(whCorrectShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "GRU operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+
+  helpers::gruTimeLoop(block.launchContext(), x, hI, Wx, Wh, b, h, linearBeforeReset);
+
+  return Status::OK;
+}
+
+//////////////////////////////////////////////////////////////////////////
+DECLARE_TYPES(gru) { getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS}); }
+
+//////////////////////////////////////////////////////////////////////////
+DECLARE_SHAPE_FN(gru) {
+  auto x = INPUT_VARIABLE(0);   // input [time, bS, nIn]
+  auto hI = INPUT_VARIABLE(1);  // initial cell output (at time step = 0) [bS, nOut]
+
+  auto Wx = INPUT_VARIABLE(2);  // input-to-hidden  weights, [nIn, 3*nOut]
+  auto Wh = INPUT_VARIABLE(3);  // hidden-to-hidden weights, [nOut, 3*nOut]
+  auto b = INPUT_VARIABLE(4);   // biases, [3*nOut]
+
+  const int time = x->sizeAt(0);
+  const int bS = x->sizeAt(1);
+  const int nIn = x->sizeAt(2);
+  const int nOut = hI->sizeAt(1);
+  const std::vector<LongType> h0CorrectShape = {bS, nOut};
+  const std::vector<LongType> wxCorrectShape = {nIn, 3 * nOut};
+  const std::vector<LongType> whCorrectShape = {nOut, 3 * nOut};
+  const std::vector<LongType> bCorrectShape = {3 * nOut};
+ 
+  REQUIRE_TRUE(hI->isSameShape(h0CorrectShape), 0,
+               "GRU operation: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(h0CorrectShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(wxCorrectShape), 0,
+               "GRU operation: wrong shape of input-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wxCorrectShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(whCorrectShape), 0,
+               "GRU operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(whCorrectShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "GRU operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+
+  auto hShapeInfo =
+      ConstantShapeHelper::getInstance().createShapeInfo(hI->dataType(), hI->ordering(), {time, bS, nOut});
+
+  return SHAPELIST(hShapeInfo);
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(gru_bp, 6, 5, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // input [time, bS, nIn]
+  auto hI = INPUT_VARIABLE(1);  // initial cell output (at time step = 0) [bS, nOut]
+
+  auto Wx = INPUT_VARIABLE(2);  // input-to-hidden  weights, [nIn, 3*nOut]
+  auto Wh = INPUT_VARIABLE(3);  // hidden-to-hidden weights, [nOut, 3*nOut]
+  auto b = INPUT_VARIABLE(4);   // biases, [3*nOut]
+
+  auto dLdh = INPUT_VARIABLE(5);  // gradient vs. ff output, [time, bS, nOut]
+
+  auto dLdx = OUTPUT_VARIABLE(0);    // gradient vs. ff input, [time, bS, nIn]
+  auto dLdhI = OUTPUT_NULLIFIED(1);  // gradient vs. initial cell output, [bS, nOut]
+  auto dLdWx = OUTPUT_NULLIFIED(2);  // gradient vs. input-to-hidden  weights, [nIn, 3*nOut]
+  auto dLdWh = OUTPUT_NULLIFIED(3);  // gradient vs. hidden-to-hidden weights, [nOut, 3*nOut]
+  auto dLdb = OUTPUT_NULLIFIED(4);   // gradient vs. biases [3*nOut]
+
+  const int time = x->sizeAt(0);
+  const int bS = x->sizeAt(1);
+  const int nIn = x->sizeAt(2);
+  const int nOut = hI->sizeAt(1);
+
+  const std::vector<LongType> h0CorrectShape = {bS, nOut};
+  const std::vector<LongType> wxCorrectShape = {nIn, 3 * nOut};
+  const std::vector<LongType> whCorrectShape = {nOut, 3 * nOut};
+  const std::vector<LongType> bCorrectShape = {3 * nOut};
+  const std::vector<LongType> hCorrectShape = {time, bS, nOut};
+
+  REQUIRE_TRUE(hI->isSameShape(h0CorrectShape), 0,
+               "GRU_BP operation: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(h0CorrectShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(wxCorrectShape), 0,
+               "GRU_BP operation: wrong shape of input-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wxCorrectShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(whCorrectShape), 0,
+               "GRU_BP operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(whCorrectShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "GRU_BP operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(dLdh->isSameShape(hCorrectShape), 0,
+               "GRU_BP operation: wrong shape of gradient vs. ff output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(hCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdh).c_str());
+
+  helpers::gruTimeLoopBp(block.launchContext(), x, hI, Wx, Wh, b, dLdh, dLdx, dLdhI, dLdWx, dLdWh, dLdb);
+
+  return Status::OK;
+}
+
+//////////////////////////////////////////////////////////////////////////
+DECLARE_TYPES(gru_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+DECLARE_SHAPE_FN(gru_bp) {
+  auto x = INPUT_VARIABLE(0);   // input [time, bS, nIn]
+  auto hI = INPUT_VARIABLE(1);  // initial cell output (at time step = 0) [bS, nOut]
+
+  auto Wx = INPUT_VARIABLE(2);  // input-to-hidden  weights, [nIn, 3*nOut]
+  auto Wh = INPUT_VARIABLE(3);  // hidden-to-hidden weights, [nOut, 3*nOut]
+  auto b = INPUT_VARIABLE(4);   // biases, [3*nOut]
+
+  auto dLdh = INPUT_VARIABLE(5);  // gradient vs. ff output, [time, bS, nOut]
+
+  const int time = x->sizeAt(0);
+  const int bS = x->sizeAt(1);
+  const int nIn = x->sizeAt(2);
+  const int nOut = hI->sizeAt(1);
+
+  const std::vector<LongType> h0CorrectShape = {bS, nOut};
+  const std::vector<LongType> wxCorrectShape = {nIn, 3 * nOut};
+  const std::vector<LongType> whCorrectShape = {nOut, 3 * nOut};
+  const std::vector<LongType> bCorrectShape = {3 * nOut};
+  const std::vector<LongType> hCorrectShape = {time, bS, nOut};
+
+  REQUIRE_TRUE(hI->isSameShape(h0CorrectShape), 0,
+               "GRU_BP operation: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(h0CorrectShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(wxCorrectShape), 0,
+               "GRU_BP operation: wrong shape of input-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wxCorrectShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(whCorrectShape), 0,
+               "GRU_BP operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(whCorrectShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "GRU_BP operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(dLdh->isSameShape(hCorrectShape), 0,
+               "GRU_BP operation: wrong shape of gradient vs. ff output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(hCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdh).c_str());
+
+  auto dLdxShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(dLdh->dataType(), x->shapeInfo());
+  auto dLdhIShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(dLdh->dataType(), hI->shapeInfo());
+  auto dLdWxShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(dLdh->dataType(), Wx->shapeInfo());
+  auto dLdWhShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(dLdh->dataType(), Wh->shapeInfo());
+  auto dLdbShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(dLdh->dataType(), b->shapeInfo());
+
+  return SHAPELIST(dLdxShapeInfo, dLdhIShapeInfo, dLdWxShapeInfo, dLdWhShapeInfo, dLdbShapeInfo);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/gruCell.cpp

@@ -0,0 +1,281 @@
+/* ******************************************************************************
+ *
+ *
+ * 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)
+// @author Alex Black
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_gruCell)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/gru.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(gruCell, 6, 4, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);  // input [bS, nIn], nIn - input size
+  auto hLast =
+      INPUT_VARIABLE(1);  // previous cell output [bS, nU],  that is at previous time step t-1, nU - number of units
+  auto Wru = INPUT_VARIABLE(2);  // RU weights - [nIn+nU, 2*nU] - reset and update gates (input/recurrent weights)
+  auto Wc = INPUT_VARIABLE(3);   // C weights - [nIn+nU, nU] - cell gate (input/recurrent weights)
+  auto bru = INPUT_VARIABLE(4);  // reset and update biases, [2*nU] - reset and update gates
+  auto bc = INPUT_VARIABLE(5);   // cell biases, [nU]
+
+  auto r = OUTPUT_VARIABLE(0);  // Reset gate output [bS, nU]
+  auto u = OUTPUT_VARIABLE(1);  // Update gate output [bS, nU]
+  auto c = OUTPUT_VARIABLE(2);  // Cell gate output [bS, nU]
+  auto h = OUTPUT_VARIABLE(3);  // current cell output [bS, nU]
+
+  REQUIRE_TRUE(x->rankOf() == 2 && hLast->rankOf() == 2, 0,
+               "gruCell: Input ranks must be 2 for inputs 0 and 1 (x, hLast) - got %i, %i", x->rankOf(),
+               hLast->rankOf());
+
+  const int rank = x->rankOf();
+  const auto bS = x->sizeAt(0);
+  const auto nIn = x->sizeAt(1);
+  const auto nU = hLast->sizeAt(1);
+
+  REQUIRE_TRUE(x->sizeAt(0) == hLast->sizeAt(0), 0,
+               "gruCell: Input minibatch sizes (dimension 0) must be same for x and hLast");
+  REQUIRE_TRUE(Wru->rankOf() == 2 && Wc->rankOf() == 2, 0,
+               "gruCell: weight arrays (Wru, Wc) arrays must be 2, got %i and %i", Wru->rankOf(), Wc->rankOf());
+  REQUIRE_TRUE(Wru->sizeAt(0) == (nIn + nU) && Wc->sizeAt(0) == (nIn + nU), 0,
+               "gruCell: Weights size(0) must be equal to nIn + nU, got %i", Wru->sizeAt(0));
+  REQUIRE_TRUE(Wru->sizeAt(1) == (2 * nU), 0,
+               "gruCell: Weights (reset and update) size(1) must be equal to 2*nU, got %i", Wru->sizeAt(1));
+  REQUIRE_TRUE(Wc->sizeAt(1) == nU, 0, "gruCell: Weights (cell) size(1) must be equal to nU, got %i", Wc->sizeAt(1));
+  REQUIRE_TRUE(bru->rankOf() == 1 && bru->sizeAt(0) == (2 * nU), 0,
+               "gruCell: reset/update biases must be rank 1, size 2*nU");
+  REQUIRE_TRUE(bc->rankOf() == 1 && bc->sizeAt(0) == nU, 0, "gruCell: cell biases must be rank 1, size nU");
+  REQUIRE_TRUE(r->rankOf() == 2 && u->rankOf() == 2 && c->rankOf() == 2 && h->rankOf() == 2 && r->sizeAt(0) == bS &&
+                   u->sizeAt(0) == bS && c->sizeAt(0) == bS && h->sizeAt(0) == bS && r->sizeAt(1) == nU &&
+                   u->sizeAt(1) == nU && c->sizeAt(1) == nU && h->sizeAt(1) == nU,
+               0, "gruCell: Output arrays must all be rank 2 with size(0) == batchSize and size(1) == nU");
+
+  helpers::gruCell(block.launchContext(), x, hLast, Wru, Wc, bru, bc, r, u, c, h);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(gruCell) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedInputTypes(4, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(gruCell) {
+  auto x = inputShape->at(0);      // input [bS x nIn]
+  auto hLast = inputShape->at(1);  // previous cell output [bS x nU],  that is at previous time step t-1
+  auto Wru = inputShape->at(2);    // RU weights - [(nIn+nU), 2*nU] - reset and update gates (input/recurrent weights)
+  auto Wc = inputShape->at(3);     // C weights - [(nIn+nU), nU] - cell gate (input/recurrent weights)
+  auto bru = inputShape->at(4);    // reset and update biases, [2*nU] - reset and update gates
+  auto bc = inputShape->at(5);     // cell biases, [nU]
+
+  REQUIRE_TRUE(shape::rank(x) == 2 && shape::rank(hLast) == 2, 0,
+               "gruCell: Input ranks must be 2 for inputs 0 and 1 (x, hLast) - got %i, %i", shape::rank(x),
+               shape::rank(hLast));
+
+  const int rank = x[0];
+  const auto bS = x[1];
+  const auto nIn = x[2];
+  const auto nU = hLast[2];
+
+  REQUIRE_TRUE(x[1] == hLast[1], 0, "gruCell: Input minibatch sizes (dimension 0) must be same for x and hLast");
+  REQUIRE_TRUE(shape::rank(Wru) == 2 && shape::rank(Wc) == 2, 0,
+               "gruCell: weight arrays (Wru, Wc) arrays must be 2, got %i and %i", shape::rank(Wru), shape::rank(Wc));
+  REQUIRE_TRUE(Wru[1] == (nIn + nU) && Wc[1] == (nIn + nU), 0,
+               "gruCell: Weights size(0) must be equal to nIn + nU, got %i and %i", Wru[1], Wc[1]);
+  REQUIRE_TRUE(Wru[2] == (2 * nU), 0, "gruCell: Weights (reset and update) size(1) must be equal to 2*nU, got %i",
+               Wru[2]);
+  REQUIRE_TRUE(Wc[2] == nU, 0, "gruCell: Weights (cell) size(1) must be equal to nU, got %i", Wc[2]);
+  REQUIRE_TRUE(shape::rank(bru) == 1 && bru[1] == (2 * nU), 0,
+               "gruCell: reset/update biases must be rank 1, size 2*nU");
+  REQUIRE_TRUE(shape::rank(bc) == 1 && bc[1] == nU, 0, "gruCell: cell biases must be rank 1, size nU");
+
+  LongType *s0(nullptr);
+  ALLOCATE(s0, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x nU]
+
+  s0[0] = rank;
+  s0[1] = bS;
+  s0[2] = nU;
+
+  ShapeUtils::updateStridesAndType(s0, x, shape::order(hLast));
+  auto ts0 = ConstantShapeHelper::getInstance().createFromExisting(s0);
+
+  // 4 output shapes, all [bs, nU]
+  return SHAPELIST(ts0, ts0, ts0, ts0);
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(gruCell_bp, 10, 6, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);     // input [bS x iS]
+  auto hi = INPUT_VARIABLE(1);    // previous cell output [bS x nU]
+  auto W = INPUT_VARIABLE(2);     // weights, [iS+nU x 2*nU]
+  auto Wc = INPUT_VARIABLE(3);    // c weights, [iS+nU x nU]
+  auto b = INPUT_VARIABLE(4);     // biases, [2*nU]
+  auto bc = INPUT_VARIABLE(5);    // biases, [nU]
+  auto dLdr = INPUT_VARIABLE(6);  // gradient wrt reset gate, [bS, nU]
+  auto dLdu = INPUT_VARIABLE(7);  // gradient wrt update gate, [bS, nU]
+  auto dLdc = INPUT_VARIABLE(8);  // gradient wrt cell state, [bS, nU]
+  auto dLdh = INPUT_VARIABLE(9);  // gradient wrt current cell output, [bS, nU]
+
+  auto dLdx = OUTPUT_VARIABLE(0);   // gradient wrt x,  [bS, iS]
+  auto dLdhi = OUTPUT_VARIABLE(1);  // gradient wrt hi, [bS, nU]
+  auto dLdW = OUTPUT_VARIABLE(2);   // gradient wrt W,  [iS+nU x 2*nU]
+  auto dLdWc = OUTPUT_VARIABLE(3);  // gradient wrt Wc, [iS+nU x nU]
+  auto dLdb = OUTPUT_VARIABLE(4);   // gradient wrt biases, [2*nU]
+  auto dLdbc = OUTPUT_VARIABLE(5);  // gradient wrt c biases, [nU]
+
+  const LongType bS = x->sizeAt(0);
+  const LongType iS = x->sizeAt(1);
+  const LongType nU = hi->sizeAt(1);
+
+  REQUIRE_TRUE(x->rankOf() == 2, 0, "GRU_CELL_BP: rank of input array x must be 2, but got %i instead", x->rankOf());
+
+  const std::vector<LongType> hiCorrectShape = {bS, nU};
+  const std::vector<LongType> wCorrectShape = {iS + nU, 2 * nU};
+  const std::vector<LongType> wcCorrectShape = {iS + nU, nU};
+  const std::vector<LongType> bCorrectShape = {2 * nU};
+  const std::vector<LongType> bcCorrectShape = {nU};
+
+  REQUIRE_TRUE(hi->isSameShape(hiCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(hi).c_str());
+  REQUIRE_TRUE(W->isSameShape(wCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(W).c_str());
+  REQUIRE_TRUE(Wc->isSameShape(wcCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of c weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wcCorrectShape).c_str(), ShapeUtils::shapeAsString(Wc).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(bc->isSameShape(bcCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of c biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bcCorrectShape).c_str(), ShapeUtils::shapeAsString(bc).c_str());
+  REQUIRE_TRUE(
+      dLdr->isSameShape(hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdr array (gradient wrt reset gate), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdr).c_str());
+  REQUIRE_TRUE(
+      dLdu->isSameShape(hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdu array (gradient wrt update gate), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdu).c_str());
+  REQUIRE_TRUE(
+      dLdc->isSameShape(hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdc array (gradient wrt cell state), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdc).c_str());
+  REQUIRE_TRUE(dLdh->isSameShape(hiCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of dLdh array (gradient wrt current cell output), expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdh).c_str());
+
+  helpers::gruCellBp(block.launchContext(), x, hi, W, Wc, b, bc, dLdr, dLdu, dLdc, dLdh, dLdx, dLdhi, dLdW, dLdWc, dLdb,
+                     dLdbc);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(gruCell_bp) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedInputTypes(2, {ALL_FLOATS})
+      ->setAllowedInputTypes(3, {ALL_FLOATS})
+      ->setAllowedInputTypes(4, {ALL_FLOATS})
+      ->setAllowedInputTypes(5, {ALL_FLOATS})
+      ->setAllowedInputTypes(6, {ALL_FLOATS})
+      ->setAllowedInputTypes(7, {ALL_FLOATS})
+      ->setAllowedInputTypes(8, {ALL_FLOATS})
+      ->setAllowedInputTypes(9, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(gruCell_bp) {
+  auto xShapeInfo = inputShape->at(0);     // [bS x iS]
+  auto hiShapeInfo = inputShape->at(1);    // [bS x nU]
+  auto wShapeInfo = inputShape->at(2);     // [iS+nU x 2*nU]
+  auto wcShapeInfo = inputShape->at(3);    // [iS+nU x nU]
+  auto bShapeInfo = inputShape->at(4);     // [2*nU]
+  auto bcShapeInfo = inputShape->at(5);    // [nU]
+  auto dLdrShapeInfo = inputShape->at(6);  // [bS, nU]
+  auto dLduShapeInfo = inputShape->at(7);  // [bS, nU]
+  auto dLdcShapeInfo = inputShape->at(8);  // [bS, nU]
+  auto dLdhShapeInfo = inputShape->at(9);  // [bS, nU]
+
+  const int rank = xShapeInfo[0];  // = 2
+  const LongType bS = xShapeInfo[1];
+  const LongType iS = xShapeInfo[2];
+  const LongType nU = hiShapeInfo[2];
+
+  REQUIRE_TRUE(xShapeInfo[0] == 2, 0, "GRU_CELL_BP: rank of input array x must be 2, but got %i instead",
+               xShapeInfo[0]);
+
+  const std::vector<LongType> hiCorrectShape = {bS, nU};
+  const std::vector<LongType> wCorrectShape = {iS + nU, 2 * nU};
+  const std::vector<LongType> wcCorrectShape = {iS + nU, nU};
+  const std::vector<LongType> bCorrectShape = {2 * nU};
+  const std::vector<LongType> bcCorrectShape = {nU};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(hiShapeInfo, hiCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of previous cell output array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(hiShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wcShapeInfo, wcCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of c weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wcCorrectShape).c_str(), ShapeUtils::shapeAsString(wcShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bcShapeInfo, bcCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of c biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bcCorrectShape).c_str(), ShapeUtils::shapeAsString(bcShapeInfo).c_str());
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(dLdrShapeInfo, hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdr array (gradient wrt reset gate), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdrShapeInfo).c_str());
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(dLduShapeInfo, hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdu array (gradient wrt update gate), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLduShapeInfo).c_str());
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(dLdcShapeInfo, hiCorrectShape), 0,
+      "GRU_CELL_BP op: wrong shape of dLdc array (gradient wrt cell state), expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdcShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(dLdhShapeInfo, hiCorrectShape), 0,
+               "GRU_CELL_BP op: wrong shape of dLdh array (gradient wrt current cell output), expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(hiCorrectShape).c_str(), ShapeUtils::shapeAsString(dLdhShapeInfo).c_str());
+  return SHAPELIST(CONSTANT(xShapeInfo), CONSTANT(hiShapeInfo), CONSTANT(wShapeInfo), CONSTANT(wcShapeInfo),
+                   CONSTANT(bShapeInfo), CONSTANT(bcShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstm.cpp

@@ -0,0 +1,182 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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, created on 15.02.2018
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstm)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstm.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstm, 8, 2, false, 3, 2) {
+  auto x = INPUT_VARIABLE(0);   // input [time x bS x inSize]
+  auto h0 = INPUT_VARIABLE(1);  // initial cell output (at time step = 0) [bS x numProj], in case of projection=false ->
+                                // numProj == numUnits !!!
+  auto c0 = INPUT_VARIABLE(2);  // initial cell state  (at time step = 0) [bS x numUnits],
+
+  auto Wx = INPUT_VARIABLE(3);  // input-to-hidden  weights, [inSize  x 4*numUnits]
+  auto Wh = INPUT_VARIABLE(4);  // hidden-to-hidden weights, [numProj x 4*numUnits]
+  auto Wc = INPUT_VARIABLE(5);  // diagonal weights for peephole connections [3*numUnits]
+  auto Wp = INPUT_VARIABLE(6);  // projection weights [numUnits x numProj]
+  auto b = INPUT_VARIABLE(7);   // biases, [4*numUnits]
+
+  auto h = OUTPUT_VARIABLE(0);  // cell outputs [time x bS x numProj], that is per each time step
+  auto c = OUTPUT_VARIABLE(1);  // cell states  [time x bS x numUnits] that is per each time step
+
+  const int peephole = INT_ARG(0);  // if 1, provide peephole connections
+  const int projection = INT_ARG(1);
+   // if 1, then projection is performed, if false then numProj==numUnits is mandatory!!!!
+
+  // FIXME: double
+  const double clippingCellValue = T_ARG(0);
+   // clipping value for ct, if it is not equal to zero, then cell state is clipped
+  const double clippingProjValue =  T_ARG(1);
+   // clipping value for projected ht, if it is not equal to zero, then projected cell output is clipped
+  const double forgetBias = T_ARG(2);
+
+  const int rank = x->rankOf();
+  const int time = x->sizeAt(0);
+  const int bS = x->sizeAt(1);
+  const int inSize = x->sizeAt(2);
+  const int numProj = h0->sizeAt(1);
+  const int numUnits = c0->sizeAt(1);
+
+  // input shapes validation
+  const std::vector<LongType> correctH0Shape = {bS, numProj};
+  const std::vector<LongType> correctC0Shape = {bS, numUnits};
+  const std::vector<LongType> correctWxShape = {inSize, 4 * numUnits};
+  const std::vector<LongType> correctWhShape = {numProj, 4 * numUnits};
+  const std::vector<LongType> correctWcShape = {3 * numUnits};
+  const std::vector<LongType> correctWpShape = {numUnits, numProj};
+  const std::vector<LongType> correctBShape = {4 * numUnits};
+
+  REQUIRE_TRUE(h0->isSameShape(correctH0Shape), 0,
+               "LSTM operation: wrong shape of initial cell output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctH0Shape).c_str(), ShapeUtils::shapeAsString(h0).c_str());
+  REQUIRE_TRUE(c0->isSameShape(correctC0Shape), 0,
+               "LSTM operation: wrong shape of initial cell state,  expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctC0Shape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(correctWxShape), 0,
+               "LSTM operation: wrong shape of input-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWxShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(correctWhShape), 0,
+               "LSTM operation: wrong shape of hidden-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWhShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(
+      Wc->isSameShape(correctWcShape), 0,
+      "LSTM operation: wrong shape of diagonal weights for peephole connections, expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(correctWcShape).c_str(), ShapeUtils::shapeAsString(Wc).c_str());
+  REQUIRE_TRUE(Wp->isSameShape(correctWpShape), 0,
+               "LSTM operation: wrong shape of projection weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWpShape).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+  REQUIRE_TRUE(b->isSameShape(correctBShape), 0,
+               "LSTM operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(!(!projection && numUnits != numProj), 0,
+               "LSTM operation: projection option is switched of, and in this case output dimensionality for the "
+               "projection matrices (numProj) must be equal to number of units in lstmCell !");
+
+  helpers::lstmTimeLoop(block.launchContext(), x, h0, c0, Wx, Wh, Wc, Wp, b, h, c,
+                        {(double)peephole, (double)projection, clippingCellValue, clippingProjValue, forgetBias});
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstm) { getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS}); }
+
+DECLARE_SHAPE_FN(lstm) {
+  auto xShapeInfo = inputShape->at(0);   // input [time x bS x inSize]
+  auto h0ShapeInfo = inputShape->at(1);  // initial cell output (at time step = 0) [bS x numProj], in case of
+                                         // projection=false -> numProj == numUnits !!!
+  auto c0ShapeInfo = inputShape->at(2);  // initial cell state  (at time step = 0) [bS x numUnits],
+
+  auto WxShapeInfo = inputShape->at(3);  // input-to-hidden  weights, [inSize  x 4*numUnits]
+  auto WhShapeInfo = inputShape->at(4);  // hidden-to-hidden weights, [numProj x 4*numUnits]
+  auto WcShapeInfo = inputShape->at(5);  // diagonal weights for peephole connections [3*numUnits]
+  auto WpShapeInfo = inputShape->at(6);  // projection weights [numUnits x numProj]
+  auto bShapeInfo = inputShape->at(7);   // biases, [4*numUnits]
+
+  const int rank = xShapeInfo[0];
+  const int time = xShapeInfo[1];
+  const int bS = xShapeInfo[2];
+  const int inSize = xShapeInfo[3];
+  const int numProj = h0ShapeInfo[2];
+  const int numUnits = c0ShapeInfo[2];
+
+  // input shapes validation
+  const std::vector<LongType> correctH0Shape = {bS, numProj};
+  const std::vector<LongType> correctC0Shape = {bS, numUnits};
+  const std::vector<LongType> correctWxShape = {inSize, 4 * numUnits};
+  const std::vector<LongType> correctWhShape = {numProj, 4 * numUnits};
+  const std::vector<LongType> correctWcShape = {3 * numUnits};
+  const std::vector<LongType> correctWpShape = {numUnits, numProj};
+  const std::vector<LongType> correctBShape = {4 * numUnits};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(h0ShapeInfo, correctH0Shape), 0,
+               "LSTM operation: wrong shape of initial cell output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctH0Shape).c_str(), ShapeUtils::shapeAsString(h0ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, correctC0Shape), 0,
+               "LSTM operation: wrong shape of initial cell state,  expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctC0Shape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WxShapeInfo, correctWxShape), 0,
+               "LSTM operation: wrong shape of input-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWxShape).c_str(), ShapeUtils::shapeAsString(WxShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhShapeInfo, correctWhShape), 0,
+               "LSTM operation: wrong shape of hidden-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWhShape).c_str(), ShapeUtils::shapeAsString(WhShapeInfo).c_str());
+  REQUIRE_TRUE(
+      ShapeUtils::areShapesEqual(WcShapeInfo, correctWcShape), 0,
+      "LSTM operation: wrong shape of diagonal weights for peephole connections, expected is %s, but got %s instead !",
+      ShapeUtils::shapeAsString(correctWcShape).c_str(), ShapeUtils::shapeAsString(WcShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WpShapeInfo, correctWpShape), 0,
+               "LSTM operation: wrong shape of projection weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWpShape).c_str(), ShapeUtils::shapeAsString(WpShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, correctBShape), 0,
+               "LSTM operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *cShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [time x bS x numProj]
+  ALLOCATE(cShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [time x bS x numUnits]
+
+  hShapeInfo[0] = cShapeInfo[0] = rank;
+  hShapeInfo[1] = cShapeInfo[1] = time;
+  hShapeInfo[2] = cShapeInfo[2] = bS;
+  hShapeInfo[3] = numProj;
+  cShapeInfo[3] = numUnits;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, xShapeInfo, shape::order(h0ShapeInfo));
+  ShapeUtils::updateStridesAndType(cShapeInfo, xShapeInfo, shape::order(c0ShapeInfo));
+
+  return SHAPELIST(CONSTANT(hShapeInfo), CONSTANT(cShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstmBlock.cpp

@@ -0,0 +1,132 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+// lstmBlock: Full LSTM layer in one op
+// @author Alex Black
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstmBlock)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstmBlock.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmBlock, 9, 7, false, 2, 2) {
+  auto maxTSLength = INPUT_VARIABLE(0);
+  auto x = INPUT_VARIABLE(1);      // input [seqLen, bS, nIn] at time t
+  auto cLast = INPUT_VARIABLE(2);  // previous cell state  [bS, nOut], time t-1
+  auto yLast = INPUT_VARIABLE(3);  // previous output [bS, nOut], time t-1
+
+  auto W = INPUT_VARIABLE(4);
+   // Weights - concatenated (input-to-hidden, hidden-to-hidden weights)  weights, [(nIn+nOut), 4*nOut]
+  auto Wci = INPUT_VARIABLE(5);  // weights - cell peephole (t-1) connections to input modulation gate, [nOut]
+  auto Wcf = INPUT_VARIABLE(6);  // weights - cell peephole (t-1) connections to forget gate, [nOut]
+  auto Wco = INPUT_VARIABLE(7);  // weights - cell peephole (t) connections to output gate, [nOut]
+  auto b = INPUT_VARIABLE(8);    // biases, [4*nOut]
+
+  auto i = OUTPUT_VARIABLE(0);  // Output - input modulation gate activations [seqLen, bS, nOut]
+  auto c = OUTPUT_VARIABLE(1);  // Activations, cell state (pre tanh) [seqLen, bs, nOut]
+  auto f = OUTPUT_VARIABLE(2);  // Output - forget gate activations [seqLen, bs, nOut]
+  auto o = OUTPUT_VARIABLE(3);  // Output - output gate activations [seqLen, bs, nOut]
+  auto z = OUTPUT_VARIABLE(4);  // Output - input gate activations [seqLen, bs, nOut]
+  auto h = OUTPUT_VARIABLE(5);  // Cell state, post tanh [seqLen, bs, nOut]
+  auto y = OUTPUT_VARIABLE(6);  // current cell output [seqLen, bS, numProj], time t
+
+  const int peephole = INT_ARG(0);    // if 1, provide peephole connections
+  const int dataFormat = INT_ARG(1);  // 0=TNS=[seqLen,bS,nIn]; 1=NST=[bS,nIn,seqLen]; 2=NTS=[bS,seqLen,nIn]
+  const double forgetBias = T_ARG(0);
+  const double clippingCellValue = T_ARG(1);
+  // clipping value for ct, if it is not equal to zero, then cell state is clipped
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0, "lstmBlock: Input array 1 (x) rank must be got input with rank %i", x->rankOf());
+  REQUIRE_TRUE(cLast->rankOf() == 2 && yLast->rankOf() == 2, 0,
+               "lstmBlock: Input ranks must be 2 for inputs 2/3 (cLast, yLast) - got %i, %i", cLast->rankOf(),
+               yLast->rankOf());
+  REQUIRE_TRUE(W->rankOf() == 2, 0, "lstmBlock: Weights array rank must be 2");
+  REQUIRE_TRUE(b->rankOf() == 1, 0, "lstmBlock: Biases must be rank 1");
+  REQUIRE_TRUE(i->rankOf() == 3 && c->rankOf() == 3 && f->rankOf() == 3 && o->rankOf() == 3 && z->rankOf() == 3 &&
+                   h->rankOf() == 3 && y->rankOf() == 3,
+               0, "lstmBlock: Output arrays must all be rank 3");
+
+  helpers::lstmBlockTimeLoop(maxTSLength, x, cLast, yLast, W, Wci, Wcf, Wco, b, i, c, f, o, z, h, y,
+                             {(double)peephole, forgetBias, clippingCellValue}, dataFormat);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmBlock) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmBlock) {
+  auto x = inputShape->at(1);
+  auto cLast = inputShape->at(2);
+  auto yLast = inputShape->at(3);
+  auto W = inputShape->at(4);
+  auto b = inputShape->at(8);
+
+  REQUIRE_TRUE(shape::rank(x) == 3, 0, "lstmBlock: Input array 1 (x) rank must be got input with rank %i",
+               shape::rank(x));
+  REQUIRE_TRUE(shape::rank(cLast) == 2 && shape::rank(yLast) == 2, 0,
+               "lstmBlock: Input ranks must be 2 for inputs 2/3 (cLast, yLast) - got %i, %i", shape::rank(cLast),
+               shape::rank(yLast));
+  REQUIRE_TRUE(shape::rank(W) == 2, 0, "lstmBlock: Weights array rank must be 2");
+  REQUIRE_TRUE(shape::rank(b) == 1, 0, "lstmBlock: Biases must be rank 1");
+
+  const int dataFormat = INT_ARG(1);  // 0=TNS=[seqLen,bS,size]; 1=NST=[bS,size,seqLen]; 2=NTS=[bS,seqLen,size]
+  int bs;
+  int t;
+  int nOut = cLast[2];  // rank, bs, nOut, ...]
+
+  LongType *s(nullptr);
+  ALLOCATE(s, block.getWorkspace(), shape::shapeInfoLength(3), sd::LongType);  // [time, bS, nOut]
+  s[0] = 3;
+  if (dataFormat == 0) {
+    //[rank, seqLen, bs, nIn, ...]
+    s[1] = x[1];  // seqLen
+    s[2] = x[2];  // bS
+    s[3] = nOut;
+  } else if (dataFormat == 1) {
+    //[rank, bs, nIn, seqLen, ...]
+    s[1] = x[1];  // bS
+    s[2] = nOut;
+    s[3] = x[3];  // seqLen
+  } else {
+    //[rank, bs, seqLen, nIn, ...]
+    s[1] = x[1];  // bS
+    s[2] = x[2];  // seqLen
+    s[3] = nOut;
+  }
+  ShapeUtils::updateStridesAndType(s, x, 'c');
+
+  auto s1 = CONSTANT(s);
+
+  // 7 outputs, all same shape/type
+  return SHAPELIST(s1, s1, s1, s1, s1, s1, s1);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstmBlockCell.cpp

@@ -0,0 +1,144 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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 Alex Black
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstmBlockCell)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstmBlock.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmBlockCell, 8, 7, false, 2, 1) {
+  // Notation: mostly following https://arxiv.org/pdf/1503.04069.pdf
+  auto xt = INPUT_VARIABLE(0);     // input [bS, inSize] at time t
+  auto cLast = INPUT_VARIABLE(1);  // previous cell state  [bS, numUnits], time t-1
+  auto yLast = INPUT_VARIABLE(2);  // previous output [bS, numUnits], time t-1
+
+  auto W = INPUT_VARIABLE(3);    // Weights - concatenated (input-to-hidden, hidden-to-hidden weights)  weights,
+                                 // [(inSize+numUnits), 4*numUnits]
+  auto Wci = INPUT_VARIABLE(4);  // weights - cell peephole (t-1) connections to input modulation gate, [numUnits]
+  auto Wcf = INPUT_VARIABLE(5);  // weights - cell peephole (t-1) connections to forget gate, [numUnits]
+  auto Wco = INPUT_VARIABLE(6);  // weights - cell peephole (t) connections to output gate, [numUnits]
+  auto b = INPUT_VARIABLE(7);    // biases, [4*numUnits]
+
+  auto i = OUTPUT_VARIABLE(0);  // Output - input modulation gate activations [bS, numUnits]
+  auto c = OUTPUT_VARIABLE(1);  // Activations, cell state (pre tanh) [bs, numUnits]
+  auto f = OUTPUT_VARIABLE(2);  // Output - forget gate activations [bs, numUnits]
+  auto o = OUTPUT_VARIABLE(3);  // Output - output gate activations [bs, numUnits]
+  auto z = OUTPUT_VARIABLE(4);  // Output - input gate activations [bs, numUnits]
+  auto h = OUTPUT_VARIABLE(5);  // Cell state, post tanh [bs, numUnits]
+  auto y = OUTPUT_VARIABLE(6);  // current cell output [bS, numProj], time t
+
+  const int peephole = INT_ARG(0);  // if 1, provide peephole connections
+
+  const double forgetBias = T_ARG(0);
+  const double clippingCellValue = T_ARG(1);
+  // clipping value for ct, if it is not equal to zero, then cell state is clipped
+
+  REQUIRE_TRUE(xt->rankOf() == 2 && cLast->rankOf() == 2 && yLast->rankOf() == 2, 0,
+               "lstmBlockCell: Input ranks must be 2 for inputs 0/1/2 (x, cLast, outLast) - got %i, %i, %i",
+               xt->rankOf(), cLast->rankOf(), yLast->rankOf());
+  const int rank = xt->rankOf();
+  const int bS = xt->sizeAt(0);
+  const int inSize = xt->sizeAt(1);
+  const int numUnits = cLast->sizeAt(1);
+
+  REQUIRE_TRUE(xt->sizeAt(0) == yLast->sizeAt(0) && xt->sizeAt(0) == cLast->sizeAt(0), 0,
+               "lstmBlockCell: Input minibatch sizes (dimension 0) must be same for xt, cLast, yLast");
+  REQUIRE_TRUE(W->rankOf() == 2, 0, "lstmBlockCell: Weights array rank must be 2");
+  REQUIRE_TRUE(W->sizeAt(0) == (inSize + numUnits), 0,
+               "lstmBlockCell: Weights size(0) must be equal to inSize + numUnits, got %i", W->sizeAt(0));
+  REQUIRE_TRUE(W->sizeAt(1) == (4 * numUnits), 0, "lstmBlockCell: Weights size(1) must be equal to 4*numUnits, got %i",
+               W->sizeAt(1));
+  REQUIRE_TRUE(b->rankOf() == 1 && b->sizeAt(0) == (4 * numUnits), 0,
+               "lstmBlockCell: Biases must be rank 1, size 4*numUnits");
+  REQUIRE_TRUE(i->rankOf() == 2 && c->rankOf() == 2 && f->rankOf() == 2 && o->rankOf() == 2 && z->rankOf() == 2 &&
+                   h->rankOf() == 2 && y->rankOf() == 2 && i->sizeAt(0) == bS && c->sizeAt(0) == bS &&
+                   f->sizeAt(0) == bS && o->sizeAt(0) == bS && z->sizeAt(0) == bS && h->sizeAt(0) == bS &&
+                   y->sizeAt(0) == bS && i->sizeAt(1) == numUnits && c->sizeAt(1) == numUnits &&
+                   f->sizeAt(1) == numUnits && o->sizeAt(1) == numUnits && z->sizeAt(1) == numUnits &&
+                   h->sizeAt(1) == numUnits && y->sizeAt(1) == numUnits,
+               0, "lstmBlockCell: Output arrays must all be rank 2 with size(0) == batchSize and size(1) == numUnits");
+
+  // calculations
+  helpers::lstmBlockCell(xt, cLast, yLast, W, Wci, Wcf, Wco, b, i, c, f, o, z, h, y,
+                         {(double)peephole, forgetBias, clippingCellValue});
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmBlockCell) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmBlockCell) {
+  auto xt = inputShape->at(0);     // input [bS, inSize] at time t
+  auto cLast = inputShape->at(1);  // previous cell state  [bS, numUnits], time t-1
+  auto yLast = inputShape->at(2);  // previous output [bS, numUnits], time t-1
+
+  auto W = inputShape->at(3);    // Weights - concatenated (input-to-hidden, hidden-to-hidden weights)  weights,
+                                 // [(inSize+numUnits), 4*numUnits]
+  auto Wci = inputShape->at(4);  // weights - cell peephole (t-1) connections to input modulation gate, [numUnits]
+  auto Wcf = inputShape->at(5);  // weights - cell peephole (t-1) connections to forget gate, [numUnits]
+  auto Wco = inputShape->at(6);  // weights - cell peephole (t) connections to output gate, [numUnits]
+  auto b = inputShape->at(7);    // biases, [4*numUnits]
+
+  REQUIRE_TRUE(shape::rank(xt) == 2 && shape::rank(cLast) == 2 && shape::rank(yLast) == 2, 0,
+               "lstmBlockCell: Input ranks must be 2 for inputs 0/1/2 (x, cLast, outLast) - got %i, %i, %i",
+               shape::rank(xt), shape::rank(cLast), shape::rank(yLast));
+  const int inSize = xt[2];
+  const int numUnits = cLast[2];  //[rank, bS, nOut, ...]
+  REQUIRE_TRUE(xt[1] == yLast[1] && xt[1] == cLast[1], 0,
+               "lstmBlockCell: Input minibatch sizes (dimension 0) must be same for xt, cLast, yLast");
+  REQUIRE_TRUE(shape::rank(W) == 2, 0, "lstmBlockCell: Weights array rank must be rank 2, got %i", shape::rank(W));
+  REQUIRE_TRUE(W[1] == (inSize + numUnits), 0,
+               "lstmBlockCell: Weights size(0) must be equal to inSize + numUnits, got %i", W[1]);
+  REQUIRE_TRUE(W[2] == (4 * numUnits), 0, "lstmBlockCell: Weights size(1) must be equal to 4*numUnits, got %i", W[2]);
+  REQUIRE_TRUE(shape::rank(b) == 1 && b[1] == (4 * numUnits), 0,
+               "lstmBlockCell: Biases must be rank 1, size 4*numUnits");
+
+  // evaluate output shapeInfos
+  const int bS = xt[1];
+  LongType *s(nullptr);
+  ALLOCATE(s, block.getWorkspace(), shape::shapeInfoLength(2), sd::LongType);  // [bS, numUnits]
+
+  s[0] = 2;
+  s[1] = bS;
+  s[2] = numUnits;
+
+  ShapeUtils::updateStridesAndType(s, xt, 'c');
+
+  auto s1 = CONSTANT(s);
+
+  // 7 outputs, all same shape: z, i, f, o, h, c, y
+  return SHAPELIST(s1, s1, s1, s1, s1, s1, s1);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstmCell.cpp

@@ -0,0 +1,184 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 30.11.2017
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstmCell)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstm.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmCell, 8, 2, false, 3, 2) {
+  auto xt = INPUT_VARIABLE(0);    // input [bS x inSize]
+  auto ht_1 = INPUT_VARIABLE(1);  // previous cell output [bS x numProj],  that is at previous time step t-1, in case of
+  // projection=false -> numProj=numUnits!!!
+  auto ct_1 = INPUT_VARIABLE(2);  // previous cell state  [bS x numUnits], that is at previous time step t-1
+
+  auto Wx = INPUT_VARIABLE(3);  // input-to-hidden  weights, [inSize  x 4*numUnits]
+  auto Wh = INPUT_VARIABLE(4);  // hidden-to-hidden weights, [numProj x 4*numUnits]
+  auto Wc = INPUT_VARIABLE(5);  // diagonal weights for peephole connections [3*numUnits]
+  auto Wp = INPUT_VARIABLE(6);  // projection weights [numUnits x numProj]
+  auto b = INPUT_VARIABLE(7);   // biases, [4*numUnits]
+
+  auto ht = OUTPUT_VARIABLE(0);  // current cell output [bS x numProj], that is at current time step t
+  auto ct = OUTPUT_VARIABLE(1);  // current cell state  [bS x numUnits], that is at current time step t
+
+  const int peephole = INT_ARG(0);  // if 1, provide peephole connections
+  const int projection = INT_ARG(1);
+  // if 1, then projection is performed, if false then numProj==numUnits is mandatory!!!!
+
+  // FIXME: double?
+  const double clippingCellValue = T_ARG(0);
+  // clipping value for ct, if it is not equal to zero, then cell state is clipped
+  const double clippingProjValue = T_ARG(1);
+  // clipping value for projected ht, if it is not equal to zero, then projected cell output is clipped
+  const double forgetBias = T_ARG(2);
+
+  const int rank = xt->rankOf();
+  const int bS = xt->sizeAt(0);
+  const int inSize = xt->sizeAt(1);
+  const int numProj = ht_1->sizeAt(1);
+  const int numUnits = ct_1->sizeAt(1);
+
+  // input shapes validation
+  const std::vector<LongType> correctHt_1Shape = {bS, numProj};
+  const std::vector<LongType> correctCt_1Shape = {bS, numUnits};
+  const std::vector<LongType> correctWxShape = {inSize, 4 * numUnits};
+  const std::vector<LongType> correctWhShape = {numProj, 4 * numUnits};
+  const std::vector<LongType> correctWcShape = {3 * numUnits};
+  const std::vector<LongType> correctWpShape = {numUnits, numProj};
+  const std::vector<LongType> correctBShape = {4 * numUnits};
+
+  REQUIRE_TRUE(ht_1->isSameShape(correctHt_1Shape), 0,
+               "LSTMCELL operation: wrong shape of initial cell output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctHt_1Shape).c_str(), ShapeUtils::shapeAsString(ht_1).c_str());
+  REQUIRE_TRUE(ct_1->isSameShape(correctCt_1Shape), 0,
+               "LSTMCELL operation: wrong shape of initial cell state,  expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctCt_1Shape).c_str(), ShapeUtils::shapeAsString(ct_1).c_str());
+  REQUIRE_TRUE(Wx->isSameShape(correctWxShape), 0,
+               "LSTMCELL operation: wrong shape of input-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWxShape).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  REQUIRE_TRUE(Wh->isSameShape(correctWhShape), 0,
+               "LSTMCELL operation: wrong shape of hidden-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWhShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(Wc->isSameShape(correctWcShape), 0,
+               "LSTMCELL operation: wrong shape of diagonal weights for peephole connections, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(correctWcShape).c_str(), ShapeUtils::shapeAsString(Wc).c_str());
+  REQUIRE_TRUE(Wp->isSameShape(correctWpShape), 0,
+               "LSTMCELL operation: wrong shape of projection weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWpShape).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+  REQUIRE_TRUE(b->isSameShape(correctBShape), 0,
+               "LSTMCELL operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(!(!projection && numUnits != numProj), 0,
+               "LSTMCELL operation: projection option is switched of, and in this case output dimensionality for the "
+               "projection matrices (numProj) must be equal to number of units in lstmCell !");
+
+  // calculations
+  helpers::lstmCell(block.launchContext(), xt, ht_1, ct_1, Wx, Wh, Wc, Wp, b, ht, ct,
+                    {(double)peephole, (double)projection, clippingCellValue, clippingProjValue, forgetBias});
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmCell) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmCell) {
+  auto xtShapeInfo = inputShape->at(0);    // input [bS x inSize]
+  auto ht_1ShapeInfo = inputShape->at(1);  // previous cell output [bS x numProj],  that is at previous time step t-1,
+  // in case of projection=false -> numProj=numUnits!!!
+  auto ct_1ShapeInfo = inputShape->at(2);  // previous cell state  [bS x numUnits], that is at previous time step t-1
+
+  auto WxShapeInfo = inputShape->at(3);  // input-to-hidden  weights, [inSize  x 4*numUnits]
+  auto WhShapeInfo = inputShape->at(4);  // hidden-to-hidden weights, [numProj x 4*numUnits]
+  auto WcShapeInfo = inputShape->at(5);  // diagonal weights for peephole connections [3*numUnits]
+  auto WpShapeInfo = inputShape->at(6);  // projection weights [numUnits x numProj]
+  auto bShapeInfo = inputShape->at(7);   // biases, [4*numUnits]
+
+  const int rank = shape::rank(xtShapeInfo);
+  const auto bS = xtShapeInfo[1];
+  const auto inSize = xtShapeInfo[2];
+  const auto numProj = ht_1ShapeInfo[2];
+  const auto numUnits = ct_1ShapeInfo[2];
+
+  // input shapes validation
+  const std::vector<LongType> correctHt_1Shape = {bS, numProj};
+  const std::vector<LongType> correctCt_1Shape = {bS, numUnits};
+  const std::vector<LongType> correctWxShape = {inSize, 4 * numUnits};
+  const std::vector<LongType> correctWhShape = {numProj, 4 * numUnits};
+  const std::vector<LongType> correctWcShape = {3 * numUnits};
+  const std::vector<LongType> correctWpShape = {numUnits, numProj};
+  const std::vector<LongType> correctBShape = {4 * numUnits};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(ht_1ShapeInfo, correctHt_1Shape), 0,
+               "LSTMCELL operation: wrong shape of initial cell output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctHt_1Shape).c_str(), ShapeUtils::shapeAsString(ht_1ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(ct_1ShapeInfo, correctCt_1Shape), 0,
+               "LSTMCELL operation: wrong shape of initial cell state,  expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctCt_1Shape).c_str(), ShapeUtils::shapeAsString(ct_1ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WxShapeInfo, correctWxShape), 0,
+               "LSTMCELL operation: wrong shape of input-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWxShape).c_str(), ShapeUtils::shapeAsString(WxShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhShapeInfo, correctWhShape), 0,
+               "LSTMCELL operation: wrong shape of hidden-to-hidden weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWhShape).c_str(), ShapeUtils::shapeAsString(WhShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WcShapeInfo, correctWcShape), 0,
+               "LSTMCELL operation: wrong shape of diagonal weights for peephole connections, expected is %s, but got "
+               "%s instead !",
+               ShapeUtils::shapeAsString(correctWcShape).c_str(), ShapeUtils::shapeAsString(WcShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WpShapeInfo, correctWpShape), 0,
+               "LSTMCELL operation: wrong shape of projection weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWpShape).c_str(), ShapeUtils::shapeAsString(WpShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, correctBShape), 0,
+               "LSTMCELL operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *cShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x numProj]
+  ALLOCATE(cShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x numUnits]
+
+  hShapeInfo[0] = cShapeInfo[0] = rank;
+  hShapeInfo[1] = cShapeInfo[1] = bS;
+  hShapeInfo[2] = numProj;
+  cShapeInfo[2] = numUnits;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, xtShapeInfo, shape::order(ht_1ShapeInfo));
+  ShapeUtils::updateStridesAndType(cShapeInfo, xtShapeInfo, shape::order(ct_1ShapeInfo));
+
+  auto result = SHAPELIST(ConstantShapeHelper::getInstance().bufferForShapeInfo(hShapeInfo)->primary(),
+                          ConstantShapeHelper::getInstance().bufferForShapeInfo(cShapeInfo)->primary());
+  RELEASE(hShapeInfo, block.workspace());
+  RELEASE(cShapeInfo, block.workspace());
+  return result;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstmLayer.cpp

@@ -0,0 +1,967 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * This program and the accompanying materials are made available under the
+ *  * terms of the Apache License, Version 2.0 which is available at
+ *  * https://www.apache.org/licenses/LICENSE-2.0.
+ *  *
+ *  * See the NOTICE file distributed with this work for additional
+ *  * information regarding copyright ownership.
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  * License for the specific language governing permissions and limitations
+ *  * under the License.
+ *  *
+ *  * SPDX-License-Identifier: Apache-2.0
+ *  *****************************************************************************
+ */
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstmLayer)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstmLayer.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmLayer, 3, 1, false, 1, 5) {
+  // equations (no peephole connections)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = ft ◦ ct-1 + it ◦ c't
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // equations (peephole connections are present)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  Wpi ◦ ct-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  Wpf ◦ ct-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = clip(ft ◦ ct-1 + it ◦ c't)
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  Wpo ◦ ct  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // notations:
+  // bS - batch size
+  // sL - sequence length, number of time steps
+  // nIn - input size
+  // nOut - output size (hidden size)
+
+  //     INPUTS:
+
+  // *******
+  // input x:
+  // 1) [sL, bS, nIn]  when dataFormat == 0
+  // 2) [bS, sL, nIn]  when dataFormat == 1
+  // 3) [bS, nIn, sL]  when dataFormat == 2
+
+  // *******
+  // input weights Wx:
+  // 1) [nIn, 4*nOut]    when directionMode <  2
+  // 2) [2, nIn, 4*nOut] when directionMode >= 2
+
+  // *******
+  // recurrent weights Wr:
+  // 1) [nOut, 4*nOut]    when directionMode <  2
+  // 2) [2, nOut, 4*nOut] when directionMode >= 2
+
+  // *******
+  // peephole weights Wp, optional:
+  // 1) [3*nOut]    when directionMode <  2
+  // 2) [2, 3*nOut] when directionMode >= 2
+
+  // *******
+  // biases b, optional:
+  // 1) [4*nOut]    when directionMode <  2
+  // 2) [2, 4*nOut] when directionMode >= 2
+
+  // *******
+  // sequence length array seqLen, optional:
+  // 1) [bS]
+
+  // *******
+  // initial output hI, optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // initial cell state cI (same shape as in hI), optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  //     OUTPUTS:
+
+  // *******
+  // output h, optional:
+  // 1) [sL, bS, nOut]    when directionMode <= 2 && dataFormat == 0
+  // 2) [bS, sL, nOut]    when directionMode <= 2 && dataFormat == 1
+  // 3) [bS, nOut, sL]    when directionMode <= 2 && dataFormat == 2
+  // 4) [sL, bS, 2*nOut]  when directionMode == 3 && dataFormat == 0
+  // 5) [bS, sL, 2*nOut]  when directionMode == 3 && dataFormat == 1
+  // 6) [bS, 2*nOut, sL]  when directionMode == 3 && dataFormat == 2
+  // 7) [sL, 2, bS, nOut] when directionMode == 4 && dataFormat == 3
+
+  // *******
+  // output at last step hL, optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // cell state at last step cL (same shape as in hL), optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // !!! dimension 4*nOut implies order it, ft, c't, ot
+  // !!! dimension 3*nOut implies order it, ft, ot
+
+  const auto dataFormat = INT_ARG(0);  // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
+  // for bidirectional: 3 = [sL, bS, nIn] && [sL, 2, bS, nOut] (for ONNX)
+  const auto directionMode = INT_ARG(1);
+  // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
+  // extra output dim (in conjunction with format dataFormat = 3)
+
+  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
+  // relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
+  const auto gateAct = INT_ARG(2);  // activation for input (i), forget (f) and output (o) gates
+  const auto cellAct = INT_ARG(3);  // activation for cell state (c)
+  const auto outAct = INT_ARG(4);   // activation for output (h)
+
+  const auto hasBiases = B_ARG(0);   // indicates whether biases array is provided
+  const auto hasSeqLen = B_ARG(1);   // indicates whether seqLen array is provided
+  const auto hasInitH = B_ARG(2);    // indicates whether initial output is provided
+  const auto hasInitC = B_ARG(3);    // indicates whether initial cell state is provided
+  const auto hasPH = B_ARG(4);       // indicates whether peephole connections are present
+  const auto retFullSeq = B_ARG(5);  // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
+  const auto retLastH = B_ARG(6);    // indicates whether to return output at last time step only
+  const auto retLastC = B_ARG(7);    // indicates whether to return cells state at last time step only
+
+  const auto gateActHasAlpha = gateAct == 3 || gateAct == 4 || gateAct == 5 || gateAct == 6 || gateAct == 8;
+  const auto cellActHasAlpha = cellAct == 3 || cellAct == 4 || cellAct == 5 || cellAct == 6 || cellAct == 8;
+  const auto outActHasAlpha = outAct == 3 || outAct == 4 || outAct == 5 || outAct == 6 || outAct == 8;
+  const auto gateActHasBeta = gateAct == 3 || gateAct == 6;
+  const auto cellActHasBeta = cellAct == 3 || cellAct == 6;
+  const auto outActHasBeta = outAct == 3 || outAct == 6;
+
+  LongType count = 1;
+  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping
+  const auto gateAlpha = gateActHasAlpha ? T_ARG(count++) : 0;
+  const auto gateBeta = gateActHasBeta ? T_ARG(count++) : 0;
+  const auto cellAlpha = cellActHasAlpha ? T_ARG(count++) : 0;
+  const auto cellBeta = cellActHasBeta ? T_ARG(count++) : 0;
+  const auto outAlpha = outActHasAlpha ? T_ARG(count++) : 0;
+  const auto outBeta = outActHasBeta ? T_ARG(count++) : 0;
+
+  const auto x = INPUT_VARIABLE(0);   // input
+  const auto Wx = INPUT_VARIABLE(1);  // input weights
+  const auto Wr = INPUT_VARIABLE(2);  // recurrent weights
+
+  count = 3;
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;       // biases
+  const auto seqLen = hasSeqLen ? INPUT_VARIABLE(count++) : nullptr;  // seqLen vector
+  const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr;       // initial output
+  const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr;       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr;          // peephole weights
+
+  REQUIRE_TRUE(dataFormat < 3 || (dataFormat == 3 && directionMode == 4), 0,
+               "LSTM_LAYER operation: if argument dataFormat = 3, then directionMode = 4, but got dataFormat = %i and "
+               "directionMode = %i instead !",
+               dataFormat, directionMode);
+  REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER operation: cell clipping value should be nonnegative (>=0) !");
+  REQUIRE_TRUE(retFullSeq || retLastH || retLastC, 0,
+               "LSTM_LAYER operation: please specify what output arrays to produce !");
+
+  count = 0;
+  auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr;  // output
+  auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr;   // output at last step
+  auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr;   // cell state at last step
+
+  // evaluate dimensions
+  const LongType sL = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
+  const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
+  const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
+  const LongType nOut = Wx->sizeAt(-1) / 4;
+
+  // inputs validations
+  if (directionMode < 2) {  // no bidirectional
+
+    // Wx validation
+    if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+    // Wr validation
+    if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+    // biases validation
+    if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+    // initial output validation
+    if (hI != nullptr && (hI->rankOf() != 2 || hI->sizeAt(0) != bS || hI->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+    // initial cell  validation
+    if (cI != nullptr && (cI->rankOf() != 2 || cI->sizeAt(0) != bS || cI->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+    // peephole weights validation
+    if (Wp != nullptr && (Wp->rankOf() != 1 || Wp->sizeAt(0) != 3 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+  } else {  // bidirectional
+    // Wx validation
+    if (Wx->rankOf() != 3 || Wx->sizeAt(0) != 2 || Wx->sizeAt(1) != nIn)
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+    // Wr validation
+    if (Wr->rankOf() != 3 || Wr->sizeAt(0) != 2 || Wr->sizeAt(1) != nOut || Wr->sizeAt(2) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+    // biases validation
+    if (b != nullptr && (b->rankOf() != 2 || b->sizeAt(0) != 2 || b->sizeAt(1) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+    // initial output validation
+    if (hI != nullptr && (hI->rankOf() != 3 || hI->sizeAt(0) != 2 || hI->sizeAt(1) != bS || hI->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+    // initial cell  validation
+    if (cI != nullptr && (cI->rankOf() != 3 || cI->sizeAt(0) != 2 || cI->sizeAt(1) != bS || cI->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+    // peephole weights validation
+    if (Wp != nullptr && (Wp->rankOf() != 2 || Wp->sizeAt(0) != 2 || Wp->sizeAt(1) != 3 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, 3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+  }
+
+  std::vector<float> params = {
+      static_cast<float>(dataFormat), static_cast<float>(directionMode), static_cast<float>(cellClip),
+      static_cast<float>(gateAct),    static_cast<float>(gateAlpha),     static_cast<float>(gateBeta),
+      static_cast<float>(cellAct),    static_cast<float>(cellAlpha),     static_cast<float>(cellBeta),
+      static_cast<float>(outAct),     static_cast<float>(outAlpha),      static_cast<float>(outBeta)};
+
+  if (directionMode == 0) {  // forward
+
+    helpers::lstmLayerTimeLoop(x, Wx, Wr, b, seqLen, hI, cI, Wp, params, true, h, hL, cL);
+  } else if (directionMode == 1) {  // backward
+
+    helpers::lstmLayerTimeLoop(x, Wx, Wr, b, seqLen, hI, cI, Wp, params, false, h, hL, cL);
+  } else {  // bidirectional
+
+    NDArray *WxFwd = (*Wx)({0, 1, 0, 0, 0, 0});
+    NDArray *WxBwd = (*Wx)({1, 2, 0, 0, 0, 0});
+    NDArray *WrFwd = (*Wr)({0, 1, 0, 0, 0, 0});
+    NDArray *WrBwd = (*Wr)({1, 2, 0, 0, 0, 0});
+
+
+    NDArray *WpFwd(nullptr), *WpBwd(nullptr), *bFwd(nullptr), *bBwd(nullptr), *hIFwd(nullptr), *hIBwd(nullptr),
+        *cIFwd(nullptr), *cIBwd(nullptr), *hLFwd(nullptr), *hLBwd(nullptr), *cLFwd(nullptr), *cLBwd(nullptr),
+        *hFwd(nullptr), *hBwd(nullptr);
+
+    if (Wp) {
+      WpFwd = (*Wp)({0, 1, 0, 0});
+      WpBwd = (*Wp)({1, 2, 0, 0});
+    }
+    if (b) {
+      bFwd = (*b)({0, 1, 0, 0});
+      bBwd = (*b)({1, 2, 0, 0});
+    }
+    if (hI) {
+      hIFwd = (*hI)({0, 1, 0, 0, 0, 0});
+      hIBwd = (*hI)({1, 2, 0, 0, 0, 0});
+    }
+    if (cI) {
+      cIFwd =(*cI)({0, 1, 0, 0, 0, 0});
+      cIBwd = (*cI)({1, 2, 0, 0, 0, 0});
+    }
+    if (hL) {
+      hLFwd = (*hL)({0, 1, 0, 0, 0, 0});
+      hLBwd = (*hL)({1, 2, 0, 0, 0, 0});
+    }
+    if (cL) {
+      cLFwd = (*cL)({0, 1, 0, 0, 0, 0});
+      cLBwd = (*cL)({1, 2, 0, 0, 0, 0});
+    }
+
+    if (h) {
+      if (directionMode == 2) {  // sum
+        hFwd = h;
+        hBwd = new NDArray(h, false, h->getContext());
+      } else if (directionMode == 3) {  // concat
+        hFwd = dataFormat <= 1 ? (*h)({0, 0, 0, 0, 0, nOut}) : (*h)({0, 0, 0, nOut, 0, 0});
+        hBwd = dataFormat <= 1 ? (*h)({0, 0, 0, 0, nOut, 2 * nOut}) : (*h)({0, 0, nOut, 2 * nOut, 0, 0});
+      } else {  // directionMode == 4
+        hFwd = (*h)({0, 0, 0, 1, 0, 0, 0, 0});
+        hBwd = (*h)({0, 0, 1, 2, 0, 0, 0, 0});
+      }
+    }
+
+    // FIXME - following two calls are independent and may run in different streams
+    helpers::lstmLayerTimeLoop(x, WxFwd, WrFwd, bFwd, seqLen, hIFwd, cIFwd, WpFwd, params, true, hFwd, hLFwd, cLFwd);
+    helpers::lstmLayerTimeLoop(x, WxBwd, WrBwd, bBwd, seqLen, hIBwd, cIBwd, WpBwd, params, false, hBwd, hLBwd, cLBwd);
+
+    if (h && directionMode == 2) *h += *hBwd;
+
+    delete WpFwd;
+    delete WpBwd;
+    delete bFwd;
+    delete bBwd;
+    delete hIFwd;
+    delete hIBwd;
+    delete cIFwd;
+    delete cIBwd;
+    delete hLFwd;
+    delete hLBwd;
+    delete cLFwd;
+    delete cLBwd;
+    delete hBwd;
+    delete WxFwd;
+    delete WxBwd;
+    delete WrFwd;
+    delete WrBwd;
+    if (hFwd != h) delete hFwd;
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmLayer) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmLayer) {
+  const auto dataFormat = INT_ARG(0);  // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
+  // for bidirectional: 3 = [sL, 2, bS, nIn] (for ONNX)
+  const auto directionMode = INT_ARG(1);  // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional
+  // concat, 4 = bidirectional extra output dim
+
+  const auto retFullSeq = B_ARG(5);  // indicates whether to return whole h {h_0, h_1, ... , h_sL-1}, if true, format
+  // would be [sL,bS,nOut] (exact shape depends on dataFormat argument)
+  const auto retLastH = B_ARG(6);    // indicates whether to return output at last time step only, in this case shape
+  // would be [bS, nOut] (exact shape depends on dataFormat argument)
+  const auto retLastC = B_ARG(7);  // indicates whether to return cells state at last time step only, in this case shape
+  // would be [bS, nOut] (exact shape depends on dataFormat argument)
+
+  const auto x = INPUT_VARIABLE(0);   // input
+  const auto Wx = INPUT_VARIABLE(1);  // input weights
+  const auto Wr = INPUT_VARIABLE(2);  // recurrent weights
+
+  // evaluate dimensions
+  const LongType sL = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
+  const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
+  const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
+  const LongType nOut = Wx->sizeAt(-1) / 4;
+
+  DataType type;
+  if (x->isR())
+    type = x->dataType();
+  else
+    type = FLOAT32;
+
+  auto shapes = SHAPELIST();
+
+  // evaluate h shape (output)
+  if (retFullSeq) {
+    std::vector<LongType> hShape;
+
+    if (directionMode <= 2) {  // single direction or bidirectional with sum
+      if (dataFormat == 0)
+        hShape = {sL, bS, nOut};
+      else if (dataFormat == 1)
+        hShape = {bS, sL, nOut};
+      else if (dataFormat == 2)
+        hShape = {bS, nOut, sL};
+    } else if (directionMode == 3) {  // bidirectional with concat
+
+      if (dataFormat == 0)
+        hShape = {sL, bS, 2 * nOut};
+      else if (dataFormat == 1)
+        hShape = {bS, sL, 2 * nOut};
+      else if (dataFormat == 2)
+        hShape = {bS, 2 * nOut, sL};
+    } else {  // bidirectional with extra output dimension equal to 2
+      hShape = {sL, 2, bS, nOut};
+    }
+
+    shapes->push_back(ConstantShapeHelper::getInstance().createShapeInfo(type, x->ordering(), hShape));
+  }
+
+  // evaluate hL shape (output at last step)
+  if (retLastH) {
+    std::vector<LongType> hLShape;
+
+    if (directionMode < 2)
+      hLShape = {bS, nOut};
+    else
+      hLShape = {2, bS, nOut};
+
+    shapes->push_back(ConstantShapeHelper::getInstance().createShapeInfo(type, x->ordering(), hLShape));
+
+    if (retLastC)  // cL and hL have same shapes
+      shapes->push_back(shapes->at(shapes->size() - 1));
+  }
+
+  // evaluate cL shape (cell state at last step)
+  if (retLastC && !retLastH) {
+    std::vector<LongType> cLShape;
+
+    if (directionMode < 2)
+      cLShape = {bS, nOut};
+    else
+      cLShape = {2, bS, nOut};
+
+    shapes->push_back(ConstantShapeHelper::getInstance().createShapeInfo(type, x->ordering(), cLShape));
+  }
+
+  return shapes;
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmLayer_bp, 4, 1, false, 1, 5) {
+  // equations (no peephole connections)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = ft ◦ ct-1 + it ◦ c't
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // equations (peephole connections are present)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  Wpi ◦ ct-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  Wpf ◦ ct-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = clip(ft ◦ ct-1 + it ◦ c't)
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  Wpo ◦ ct  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // notations:
+  // bS - batch size
+  // sL - sequence length, number of time steps
+  // nIn - input size
+  // nOut - output size (hidden size)
+
+  //     INPUTS:
+
+  // *******
+  // input x:
+  // 1) [sL, bS, nIn]  when dataFormat == 0
+  // 2) [bS, sL, nIn]  when dataFormat == 1
+  // 3) [bS, nIn, sL]  when dataFormat == 2
+
+  // *******
+  // input weights Wx:
+  // 1) [nIn, 4*nOut]    when directionMode <  2
+  // 2) [2, nIn, 4*nOut] when directionMode >= 2
+
+  // *******
+  // recurrent weights Wr:
+  // 1) [nOut, 4*nOut]    when directionMode <  2
+  // 2) [2, nOut, 4*nOut] when directionMode >= 2
+
+  // *******
+  // peephole weights Wp, optional:
+  // 1) [3*nOut]    when directionMode <  2
+  // 2) [2, 3*nOut] when directionMode >= 2
+
+  // *******
+  // biases b, optional:
+  // 1) [4*nOut]    when directionMode <  2
+  // 2) [2, 4*nOut] when directionMode >= 2
+
+  // *******
+  // sequence length array seqLen, optional:
+  // 1) [bS]
+
+  // *******
+  // initial output hI, optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // initial cell state cI (same shape as in hI), optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs. output dLdh, optional:
+  // 1) [sL, bS, nOut]    when directionMode <= 2 && dataFormat == 0
+  // 2) [bS, sL, nOut]    when directionMode <= 2 && dataFormat == 1
+  // 3) [bS, nOut, sL]    when directionMode <= 2 && dataFormat == 2
+  // 4) [sL, bS, 2*nOut]  when directionMode == 3 && dataFormat == 0
+  // 5) [bS, sL, 2*nOut]  when directionMode == 3 && dataFormat == 1
+  // 6) [bS, 2*nOut, sL]  when directionMode == 3 && dataFormat == 2
+  // 7) [sL, 2, bS, nOut] when directionMode == 4 && dataFormat == 3
+
+  // *******
+  // gradient vs output at last time step dLdhL, optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs cell state at last time step dLdcL(same shape as in dLdhL), optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  //     OUTPUTS:
+
+  // *******
+  // gradient vs. input dLdx:
+  // 1) [sL, bS, nIn]  when dataFormat == 0
+  // 2) [bS, sL, nIn]  when dataFormat == 1
+  // 3) [bS, nIn, sL]  when dataFormat == 2
+
+  // *******
+  // gradient vs. input weights dLdWx:
+  // 1) [nIn, 4*nOut]    when directionMode <  2
+  // 2) [2, nIn, 4*nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs. recurrent weights dLdWr:
+  // 1) [nOut, 4*nOut]    when directionMode <  2
+  // 2) [2, nOut, 4*nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs. peephole weights dLdWp, optional:
+  // 1) [3*nOut]    when directionMode <  2
+  // 2) [2, 3*nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs. biases dLdb, optional:
+  // 1) [4*nOut]    when directionMode <  2
+  // 2) [2, 4*nOut] when directionMode >= 2
+
+  // gradient vs. sequence length array dLdsL, optional (do not calculate it!!!):
+  // 1) [bS] always
+
+  // *******
+  // gradient vs. initial output dLdhI, optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // *******
+  // gradient vs. initial cell state dLdcI (same shape as in dLdhI), optional:
+  // 1) [bS, nOut]    when directionMode <  2
+  // 2) [2, bS, nOut] when directionMode >= 2
+
+  // !!! dimension 4*nOut implies order it, ft, c't, ot
+  // !!! dimension 3*nOut implies order it, ft, ot
+
+  const auto dataFormat = INT_ARG(0);  // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
+  // for bidirectional: 3 = [sL, bS, nIn] && [sL, 2, bS, nOut] (for ONNX)
+  const auto directionMode = INT_ARG(1);
+  // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
+  // extra output dim (in conjunction with format dataFormat = 3)
+
+  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
+  // relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
+  const auto gateAct = INT_ARG(2);  // activation for input (i), forget (f) and output (o) gates
+  const auto cellAct = INT_ARG(3);  // activation for cell state (c)
+  const auto outAct = INT_ARG(4);   // activation for output (h)
+
+  const auto hasBiases = B_ARG(0);   // indicates whether biases array is provided
+  const auto hasSeqLen = B_ARG(1);   // indicates whether seqLen array is provided
+  const auto hasInitH = B_ARG(2);    // indicates whether initial output is provided
+  const auto hasInitC = B_ARG(3);    // indicates whether initial cell state is provided
+  const auto hasPH = B_ARG(4);       // indicates whether peephole connections are present
+  const auto retFullSeq = B_ARG(5);  // indicates whether gradient vs. outputs is given for whole time sequence dLdh
+  // {dLdh_0, dLdh_1, ... , dLdh_sL-1}
+  const auto retLastH = B_ARG(6);    // indicates whether gradient vs. output at last time step (dLdhL) is given
+  const auto retLastC = B_ARG(7);    // indicates whether gradient vs. cell state at last time step (dLdcL) is given
+
+  const auto gateActHasAlpha = gateAct == 3 || gateAct == 4 || gateAct == 5 || gateAct == 6 || gateAct == 8;
+  const auto cellActHasAlpha = cellAct == 3 || cellAct == 4 || cellAct == 5 || cellAct == 6 || cellAct == 8;
+  const auto outActHasAlpha = outAct == 3 || outAct == 4 || outAct == 5 || outAct == 6 || outAct == 8;
+  const auto gateActHasBeta = gateAct == 3 || gateAct == 6;
+  const auto cellActHasBeta = cellAct == 3 || cellAct == 6;
+  const auto outActHasBeta = outAct == 3 || outAct == 6;
+
+  LongType count = 1;
+  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping
+  const auto gateAlpha = gateActHasAlpha ? T_ARG(count++) : 0;
+  const auto gateBeta = gateActHasBeta ? T_ARG(count++) : 0;
+  const auto cellAlpha = cellActHasAlpha ? T_ARG(count++) : 0;
+  const auto cellBeta = cellActHasBeta ? T_ARG(count++) : 0;
+  const auto outAlpha = outActHasAlpha ? T_ARG(count++) : 0;
+  const auto outBeta = outActHasBeta ? T_ARG(count++) : 0;
+
+  REQUIRE_TRUE(dataFormat < 3 || (dataFormat == 3 && directionMode == 4), 0,
+               "LSTM_LAYER_BP operation: if argument dataFormat = 3, then directionMode = 4, but got dataFormat = %i "
+               "and directionMode = %i instead !",
+               dataFormat, directionMode);
+  REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER_BP operation: cell clipping value should be nonnegative (>=0) !");
+  REQUIRE_TRUE(
+      retFullSeq || retLastH || retLastC, 0,
+      "LSTM_LAYER_BP operation: please specify at least one of three input gradient arrays: dLdh, dLdhL or dLdcL !");
+
+  const auto x = INPUT_VARIABLE(0);   // input
+  const auto Wx = INPUT_VARIABLE(1);  // input weights
+  const auto Wr = INPUT_VARIABLE(2);  // recurrent weights
+
+  // evaluate dimensions
+  const LongType sL = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
+  const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
+  const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
+  const LongType nOut = Wx->sizeAt(-1) / 4;
+
+  // continue with input
+  count = 3;
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;       // biases
+  const auto seqLen = hasSeqLen ? INPUT_VARIABLE(count++) : nullptr;  // seqLen vector
+  const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr;       // initial output
+  const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr;       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr;          // peephole weights
+
+  NDArray *dLdh = nullptr;
+  NDArray *dLdhL = nullptr;
+  NDArray *dLdcL = nullptr;
+  std::unique_ptr<NDArray> temp_dLdh, temp_dLdhL, temp_dLdcL;
+  std::vector<LongType> expdLdhShape;
+  // gradient vs. output
+  if (retFullSeq) {
+    int factor = directionMode <= 2 ? 1 : 2;
+    if (dataFormat == 0)
+      expdLdhShape = std::vector<LongType>{sL, bS, factor * nOut};
+    else if (dataFormat == 1)
+      expdLdhShape = std::vector<LongType>{bS, sL, factor * nOut};
+    else if (dataFormat == 2)
+      expdLdhShape = std::vector<LongType>{bS, factor * nOut, sL};
+    else
+      expdLdhShape = std::vector<LongType>{sL, 2, bS, nOut};
+
+    dLdh = INPUT_VARIABLE(count++);
+    if (dLdh->isScalar()) {
+      temp_dLdh.reset(NDArrayFactory::valueOf(expdLdhShape, *dLdh, x->ordering()));
+      dLdh = temp_dLdh.get();
+    }
+  }
+  // gradient vs. output at last time step
+  if (retLastH) {
+    dLdhL = INPUT_VARIABLE(count++);
+    if (dLdhL->isScalar()) {
+      std::vector<sd::LongType> shape =  directionMode < 2 ? std::vector<LongType>{bS, nOut} : std::vector<LongType>{2, bS, nOut};
+      temp_dLdhL.reset(NDArrayFactory::valueOf(
+          shape, *dLdhL,
+          x->ordering()));
+      // refresh
+      dLdhL = temp_dLdhL.get();
+    }
+  }
+  // gradient vs. cell state at last time step
+  if (retLastC) {
+    dLdcL = INPUT_VARIABLE(count++);
+    if (dLdcL->isScalar()) {
+      std::vector<sd::LongType> shape = directionMode < 2 ? std::vector<LongType>{bS, nOut} : std::vector<LongType>{2, bS, nOut};
+      temp_dLdcL.reset(NDArrayFactory::valueOf(
+          shape, *dLdcL,
+          x->ordering()));
+      // refresh
+      dLdcL = temp_dLdcL.get();
+    }
+  }
+
+  count = 3;
+  auto dLdx = OUTPUT_VARIABLE(0);                               // gradient vs. input
+  auto dLdWx = OUTPUT_NULLIFIED(1);                             // gradient vs. input weights
+  auto dLdWr = OUTPUT_NULLIFIED(2);                             // gradient vs. recurrent weights
+  auto dLdb = hasBiases ? OUTPUT_NULLIFIED(count++) : nullptr;  // gradient vs. biases
+  auto dLdsL = hasSeqLen ? INPUT_VARIABLE(count++) : nullptr;   // gradient vs. seqLen vector, we don't calculate it !!!
+  auto dLdhI = hasInitH ? OUTPUT_NULLIFIED(count++) : nullptr;  // gradient vs. initial output
+  auto dLdcI = hasInitC ? OUTPUT_NULLIFIED(count++) : nullptr;  // gradient vs. initial cell state
+  auto dLdWp = hasPH ? OUTPUT_NULLIFIED(count) : nullptr;       // gradient vs. peephole weights
+
+  // inputs validations
+  if (directionMode < 2) {  // no bidirectional
+
+    // Wx validation
+    if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+    // Wr validation
+    if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+    // biases validation
+    if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_BP operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+    // initial output validation
+    if (hI != nullptr && (hI->rankOf() != 2 || hI->sizeAt(0) != bS || hI->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of initial output, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+    // initial cell  validation
+    if (cI != nullptr && (cI->rankOf() != 2 || cI->sizeAt(0) != bS || cI->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+    // peephole weights validation
+    if (Wp != nullptr && (Wp->rankOf() != 1 || Wp->sizeAt(0) != 3 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_BP operation: wrong peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+    // gradient vs. output at last time step validation
+    if (dLdhL != nullptr && (dLdhL->rankOf() != 2 || dLdhL->sizeAt(0) != bS || dLdhL->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of gradient vs. output at last time step, expected is %s, but "
+                 "got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(dLdhL).c_str());
+    // gradient vs. cell state at last time step validation
+    if (dLdcL != nullptr && (dLdcL->rankOf() != 2 || dLdcL->sizeAt(0) != bS || dLdcL->sizeAt(1) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of gradient vs. cell state at last time, expected is %s, but "
+                 "got %s instead !",
+                 ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(dLdcL).c_str());
+  } else {  // bidirectional
+    // Wx validation
+    if (Wx->rankOf() != 3 || Wx->sizeAt(0) != 2 || Wx->sizeAt(1) != nIn)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+    // Wr validation
+    if (Wr->rankOf() != 3 || Wr->sizeAt(0) != 2 || Wr->sizeAt(1) != nOut || Wr->sizeAt(2) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+    // biases validation
+    if (b != nullptr && (b->rankOf() != 2 || b->sizeAt(0) != 2 || b->sizeAt(1) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_BP operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+    // initial output validation
+    if (hI != nullptr && (hI->rankOf() != 3 || hI->sizeAt(0) != 2 || hI->sizeAt(1) != bS || hI->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of initial output, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+    // initial cell  validation
+    if (cI != nullptr && (cI->rankOf() != 3 || cI->sizeAt(0) != 2 || cI->sizeAt(1) != bS || cI->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+    // peephole weights validation
+    if (Wp != nullptr && (Wp->rankOf() != 2 || Wp->sizeAt(0) != 2 || Wp->sizeAt(1) != 3 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_BP operation: wrong peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({2, 3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+    // gradient vs. output at last time step validation
+    if (dLdhL != nullptr &&
+        (dLdhL->rankOf() != 3 || dLdhL->sizeAt(0) != 2 || dLdhL->sizeAt(1) != bS || dLdhL->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of gradient vs. output at last time step, expected is %s, but "
+                 "got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(dLdhL).c_str());
+    // gradient vs. cell state at last time step validation
+    if (dLdcL != nullptr &&
+        (dLdcL->rankOf() != 3 || dLdcL->sizeAt(0) != 2 || dLdcL->sizeAt(1) != bS || dLdcL->sizeAt(2) != nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_BP operation: wrong shape of gradient vs. cell state at last time, expected is %s, but "
+                 "got %s instead !",
+                 ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(dLdcL).c_str());
+  }
+
+  // gradient vs. output  validation
+  if (dLdh) {
+    REQUIRE_TRUE(
+        dLdh->isSameShape(expdLdhShape), 0,
+        "LSTM_LAYER_CELL_BP operation: wrong shape of gradient vs. output, expected is %s, but got %s instead !",
+        ShapeUtils::shapeAsString(expdLdhShape).c_str(), ShapeUtils::shapeAsString(dLdh).c_str());
+  }
+
+  std::vector<float> params = {
+      static_cast<float>(dataFormat), static_cast<float>(directionMode), static_cast<float>(cellClip),
+      static_cast<float>(gateAct),    static_cast<float>(gateAlpha),     static_cast<float>(gateBeta),
+      static_cast<float>(cellAct),    static_cast<float>(cellAlpha),     static_cast<float>(cellBeta),
+      static_cast<float>(outAct),     static_cast<float>(outAlpha),      static_cast<float>(outBeta)};
+
+  if (directionMode == 0) {  // forward
+
+    helpers::lstmLayerTimeLoopBp(x, Wx, Wr, b, seqLen, hI, cI, Wp, dLdh, dLdhL, dLdcL, params, true, dLdx, dLdWx, dLdWr,
+                                 dLdb, dLdhI, dLdcI, dLdWp);
+
+  } else if (directionMode == 1) {  // backward
+
+    helpers::lstmLayerTimeLoopBp(x, Wx, Wr, b, seqLen, hI, cI, Wp, dLdh, dLdhL, dLdcL, params, false, dLdx, dLdWx,
+                                 dLdWr, dLdb, dLdhI, dLdcI, dLdWp);
+
+  } else {  // bidirectional
+
+    NDArray *WxFwd = (*Wx)({0, 1, 0, 0, 0, 0});
+    NDArray *WxBwd = (*Wx)({1, 2, 0, 0, 0, 0});
+    NDArray *dLdWxFwd = (*dLdWx)({0, 1, 0, 0, 0, 0});
+    NDArray *dLdWxBwd = (*dLdWx)({1, 2, 0, 0, 0, 0});
+
+    NDArray *WrFwd = (*Wr)({0, 1, 0, 0, 0, 0});
+    NDArray *WrBwd = (*Wr)({1, 2, 0, 0, 0, 0});
+    NDArray *dLdWrFwd = (*dLdWr)({0, 1, 0, 0, 0, 0});
+    NDArray *dLdWrBwd = (*dLdWr)({1, 2, 0, 0, 0, 0});
+
+
+
+    NDArray *WpFwd(nullptr), *WpBwd(nullptr), *bFwd(nullptr), *bBwd(nullptr), *hIFwd(nullptr), *hIBwd(nullptr),
+        *cIFwd(nullptr), *cIBwd(nullptr), *dLdhFwd(nullptr), *dLdhBwd(nullptr), *dLdhLFwd(nullptr), *dLdhLBwd(nullptr),
+        *dLdcLFwd(nullptr), *dLdcLBwd(nullptr), *dLdWpFwd(nullptr), *dLdWpBwd(nullptr), *dLdbFwd(nullptr),
+        *dLdbBwd(nullptr), *dLdhIFwd(nullptr), *dLdhIBwd(nullptr), *dLdcIFwd(nullptr), *dLdcIBwd(nullptr);
+
+    if (Wp) {
+      WpFwd = (*Wp)({0, 1, 0, 0});
+      WpBwd = (*Wp)({1, 2, 0, 0});
+      dLdWpFwd = (*dLdWp)({0, 1, 0, 0});
+      dLdWpBwd = (*dLdWp)({1, 2, 0, 0});
+    }
+    if (b) {
+      bFwd = (*b)({0, 1, 0, 0});
+      bBwd = (*b)({1, 2, 0, 0});
+      dLdbFwd = (*dLdb)({0, 1, 0, 0});
+      dLdbBwd = (*dLdb)({1, 2, 0, 0});
+    }
+    if (hI) {
+      hIFwd = (*hI)({0, 1, 0, 0, 0, 0});
+      hIBwd = (*hI)({1, 2, 0, 0, 0, 0});
+      dLdhIFwd = (*dLdhI)({0, 1, 0, 0, 0, 0});
+      dLdhIBwd = (*dLdhI)({1, 2, 0, 0, 0, 0});
+    }
+    if (cI) {
+      cIFwd = (*cI)({0, 1, 0, 0, 0, 0});
+      cIBwd = (*cI)({1, 2, 0, 0, 0, 0});
+      dLdcIFwd = (*dLdcI)({0, 1, 0, 0, 0, 0});
+      dLdcIBwd = (*dLdcI)({1, 2, 0, 0, 0, 0});
+    }
+    if (dLdhL) {
+      dLdhLFwd = (*dLdhL)({0, 1, 0, 0, 0, 0});
+      dLdhLBwd = (*dLdhL)({1, 2, 0, 0, 0, 0});
+    }
+    if (dLdcL) {
+      dLdcLFwd = (*dLdcL)({0, 1, 0, 0, 0, 0});
+      dLdcLBwd = (*dLdcL)({1, 2, 0, 0, 0, 0});
+    }
+
+    if (dLdh) {
+      if (directionMode == 2) {  // sum
+        dLdhFwd = dLdh;
+        dLdhBwd = dLdh;
+      } else if (directionMode == 3) {  // concat
+        dLdhFwd = dataFormat <= 1 ? (*dLdh)({0, 0, 0, 0, 0, nOut}) : (*dLdh)({0, 0, 0, nOut, 0, 0});
+        dLdhBwd = dataFormat <= 1 ? (*dLdh)({0, 0, 0, 0, nOut, 2 * nOut})
+                                              : (*dLdh)({0, 0, nOut, 2 * nOut, 0, 0});
+      } else {  // directionMode == 4
+        dLdhFwd = (*dLdh)({0, 0, 0, 1, 0, 0, 0, 0});
+        dLdhBwd = (*dLdh)({0, 0, 1, 2, 0, 0, 0, 0});
+      }
+    }
+
+
+
+    NDArray *dLdxBwd = dLdx->ulike();
+
+    // FIXME - following two calls are independent and may run in different streams
+    helpers::lstmLayerTimeLoopBp(x, WxFwd, WrFwd, bFwd, seqLen, hIFwd, cIFwd, WpFwd, dLdhFwd, dLdhLFwd, dLdcLFwd,
+                                 params, true, dLdx,dLdWxFwd, dLdWrFwd, dLdbFwd, dLdhIFwd, dLdcIFwd, dLdWpFwd);
+    helpers::lstmLayerTimeLoopBp(x, WxBwd, WrBwd, bBwd, seqLen, hIBwd, cIBwd, WpBwd, dLdhBwd, dLdhLBwd, dLdcLBwd,
+                                 params, false, dLdxBwd, dLdWxBwd, dLdWrBwd, dLdbBwd, dLdhIBwd, dLdcIBwd, dLdWpBwd);
+
+    *dLdx += *dLdxBwd;
+
+    delete WpFwd;
+    delete WpBwd;
+    delete bFwd;
+    delete bBwd;
+    delete hIFwd;
+    delete hIBwd;
+    delete cIFwd;
+    delete cIBwd;
+    delete dLdhLFwd;
+    delete dLdhLBwd;
+    delete dLdcLFwd;
+    delete dLdcLBwd;
+    delete dLdWpFwd;
+    delete dLdWpBwd;
+    delete dLdbFwd;
+    delete dLdbBwd;
+    delete dLdhIFwd;
+    delete dLdhIBwd;
+    delete dLdcIFwd;
+    delete dLdcIBwd;
+    delete WxFwd;
+    delete WxBwd;
+    delete dLdWxFwd;
+    delete dLdWxBwd;
+
+    delete dLdWrBwd;
+    delete WrFwd;
+    delete WrBwd;
+    delete dLdWrFwd;
+    if (!(dLdh && directionMode == 2)) {
+      delete dLdhFwd;
+      delete dLdhBwd;
+    }
+
+    if(directionMode > 2) {
+      delete dLdhFwd;
+      delete dLdhBwd;
+    }
+
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmLayer_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmLayer_bp) {
+  const auto hasBiases = B_ARG(0);  // indicates whether biases array is provided
+  const auto hasSeqLen = B_ARG(1);  // indicates whether seqLen array is provided
+  const auto hasInitH = B_ARG(2);   // indicates whether initial output is provided
+  const auto hasInitC = B_ARG(3);   // indicates whether initial cell state is provided
+  const auto hasPH = B_ARG(4);      // indicates whether peephole connections are present
+
+  int count = 3;
+  const auto x = INPUT_VARIABLE(0);                                   // input
+  const auto Wx = INPUT_VARIABLE(1);                                  // input weights
+  const auto Wr = INPUT_VARIABLE(2);                                  // recurrent weights
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;       // biases
+  const auto seqLen = hasSeqLen ? INPUT_VARIABLE(count++) : nullptr;  // seqLen vector
+  const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr;       // initial output
+  const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr;       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr;          // peephole weights
+
+  auto outShapes = SHAPELIST(x->shapeInfo(), Wx->shapeInfo(), Wr->shapeInfo());
+
+  if (b != nullptr) {
+    outShapes->push_back(b->shapeInfo());
+  }
+  if (seqLen != nullptr) {
+    outShapes->push_back(seqLen->shapeInfo());
+  }
+  if (hI != nullptr) {
+    outShapes->push_back(hI->shapeInfo());
+  }
+  if (cI != nullptr) {
+    outShapes->push_back(cI->shapeInfo());
+  }
+
+  if (Wp != nullptr) {
+    outShapes->push_back(Wp->shapeInfo());
+  }
+
+  return outShapes;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/lstmLayerCell.cpp

@@ -0,0 +1,362 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lstmLayerCell)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/lstmLayer.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmLayerCell, 5, 2, false, 1, 3) {
+  // equations (no peephole connections)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = ft ◦ ct-1 + it ◦ c't
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // equations (peephole connections are present)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  Wpi ◦ ct-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  Wpf ◦ ct-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = clip(ft ◦ ct-1 + it ◦ c't)
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  Wpo ◦ ct  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // notations:
+  // bS - batch size
+  // nIn - input size
+  // nOut - output size (hidden size)
+
+  // INPUTS:
+  // input x:                          [bS, nIn] or [nIn]
+  // input weights Wx:                 [nIn, 4*nOut]
+  // recurrent weights Wr:             [nOut, 4*nOut]
+  // initial (previous) output hI:     [bS, nOut] or [nOut]
+  // initial (previous) cell state cI: [bS, nOut] or [nOut]
+  // biases b (optional):              [4*nOut]
+  // peephole weights Wp (optional):   [3*nOut]
+
+  // OUTPUTS:
+  // current output h:     [bS, nOut] or [nOut]
+  // current cell state c: [bS, nOut] or [nOut]
+
+  // !!! dimension 4*nOut implies order it, ft, c't, ot
+  // !!! dimension 3*nOut implies order it, ft, ot
+
+  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
+  // relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
+  const auto gateAct = INT_ARG(0);  // activation for input (i), forget (f) and output (o) gates
+  const auto cellAct = INT_ARG(1);  // activation for cell state (c)
+  const auto outAct = INT_ARG(2);   // activation for output (h)
+
+  const auto hasBiases = B_ARG(0);  // indicates whether biases array is provided
+  const auto hasPH = B_ARG(1);      // indicates whether peephole connections are present
+
+  const auto gateActHasAlpha = gateAct == 3 || gateAct == 4 || gateAct == 5 || gateAct == 6 || gateAct == 8;
+  const auto cellActHasAlpha = cellAct == 3 || cellAct == 4 || cellAct == 5 || cellAct == 6 || cellAct == 8;
+  const auto outActHasAlpha = outAct == 3 || outAct == 4 || outAct == 5 || outAct == 6 || outAct == 8;
+  const auto gateActHasBeta = gateAct == 3 || gateAct == 6;
+  const auto cellActHasBeta = cellAct == 3 || cellAct == 6;
+  const auto outActHasBeta = outAct == 3 || outAct == 6;
+
+  LongType count = 1;
+  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping
+  const auto gateAlpha = gateActHasAlpha ? T_ARG(count++) : 0;
+  const auto gateBeta = gateActHasBeta ? T_ARG(count++) : 0;
+  const auto cellAlpha = cellActHasAlpha ? T_ARG(count++) : 0;
+  const auto cellBeta = cellActHasBeta ? T_ARG(count++) : 0;
+  const auto outAlpha = outActHasAlpha ? T_ARG(count++) : 0;
+  const auto outBeta = outActHasBeta ? T_ARG(count++) : 0;
+
+  count = 3;
+  const auto x = INPUT_VARIABLE(0);                              // input
+  const auto Wx = INPUT_VARIABLE(1);                             // input weights
+  const auto Wr = INPUT_VARIABLE(2);                             // recurrent weights
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;  // biases
+  const auto hI = INPUT_VARIABLE(count++);                       // initial output
+  const auto cI = INPUT_VARIABLE(count++);                       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count) : nullptr;       // peephole weights
+
+  REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER_CELL operation: cell clipping value should be nonnegative (>=0) !");
+
+  auto h = OUTPUT_VARIABLE(0);
+  auto c = OUTPUT_VARIABLE(1);
+
+  // evaluate dimensions
+  const LongType bS = x->rankOf() == 1 ? 0 : x->sizeAt(0);
+  const LongType nIn = x->sizeAt(-1);
+  const LongType nOut = Wx->sizeAt(-1) / 4;
+
+  // inputs validations
+  // Wx validation
+  if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  // Wr validation
+  if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+  // initial output/cell validation
+  std::vector<LongType> exphIcIShape =
+      x->rankOf() == 1 ? std::vector<LongType>{nOut} : std::vector<LongType>{bS, nOut};
+  REQUIRE_TRUE(hI->isSameShape(exphIcIShape), 0,
+               "LSTM_LAYER_CELL operation: wrong shape of initial output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(exphIcIShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(cI->isSameShape(exphIcIShape), 0,
+               "LSTM_LAYER_CELL operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(exphIcIShape).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+  // biases validation
+  if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_CELL operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  // peephole weights validation
+  if (Wp != nullptr && (Wp->rankOf() != 1 || Wp->sizeAt(0) != 3 * nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL operation: wrong shape of peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+
+  std::vector<float> params = {
+      static_cast<float>(0) /*ignore*/, static_cast<float>(0) /*ignore*/, static_cast<float>(cellClip),
+      static_cast<float>(gateAct),      static_cast<float>(gateAlpha),    static_cast<float>(gateBeta),
+      static_cast<float>(cellAct),      static_cast<float>(cellAlpha),    static_cast<float>(cellBeta),
+      static_cast<float>(outAct),       static_cast<float>(outAlpha),     static_cast<float>(outBeta)};
+
+  helpers::lstmLayerCell(x, Wx, Wr, b, hI, cI, Wp, params, h, c);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmLayerCell) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmLayerCell) {
+  const auto hasBiases = B_ARG(0);  // indicates whether biases array is provided
+
+  LongType count = hasBiases ? 4 : 3;
+  const auto hI = INPUT_VARIABLE(count++);  // initial output
+  const auto cI = INPUT_VARIABLE(count);    // initial cell state
+
+  return new ShapeList({hI->shapeInfo(), cI->shapeInfo()});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(lstmLayerCellBp, 7, 5, false, 1, 3) {
+  // equations (no peephole connections)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = ft ◦ ct-1 + it ◦ c't
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // equations (peephole connections are present)
+  // it  = σ(Wxi * xt  +  Wri * ht-1  +  Wpi ◦ ct-1  +  bi)
+  // ft  = σ(Wxf * xt  +  Wrf * ht-1  +  Wpf ◦ ct-1  +  bf)
+  // c't = tanh(Wxc * xt  +  Wrc * ht-1  +  bc)
+  // ct  = clip(ft ◦ ct-1 + it ◦ c't)
+  // ot  = σ(Wxo * xt  +  Wro * ht-1  +  Wpo ◦ ct  +  bo)
+  // ht  = ot ◦ tanh(ct)
+
+  // notations:
+  // bS - batch size
+  // nIn - input size
+  // nOut - output size (hidden size)
+
+  // INPUTS:
+  // input x:                          [bS, nIn] or [nIn]
+  // input weights Wx:                 [nIn, 4*nOut]
+  // recurrent weights Wr:             [nOut, 4*nOut]
+  // initial (previous) output hI:     [bS, nOut] or [nOut]
+  // initial (previous) cell state cI: [bS, nOut] or [nOut]
+  // gradient wrt output dLdh:         [bS, nOut] or [nOut]
+  // gradient wrt cell state dLdc:     [bS, nOut] or [nOut]
+  // peephole weights Wp (optional):   [3*nOut]
+  // biases b (optional):              [4*nOut]
+
+  // OUTPUTS:
+  // gradient wrt x dLdx:              [bS, nIn] or [nIn]
+  // gradient wrt Wx dLdWx:            [nIn, 4*nOut]
+  // gradient wrt Wr dLdWr:            [nOut, 4*nOut]
+  // gradient wrt hI dLdhI:            [bS, nOut] or [nOut]
+  // gradient wrt cI dLdcI:            [bS, nOut] or [nOut]
+  // gradient wrt b dLdb (optional):   [4*nOut]
+  // gradient wrt Wp dLdWp (optional): [3*nOut]
+
+  // !!! dimension 4*nOut implies order it, ft, c't, ot
+  // !!! dimension 3*nOut implies order it, ft, ot
+
+  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
+  // relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
+  const auto gateAct = INT_ARG(0);  // activation for input (i), forget (f) and output (o) gates
+  const auto cellAct = INT_ARG(1);  // activation for cell state (c)
+  const auto outAct = INT_ARG(2);   // activation for output (h)
+
+  const auto hasBiases = B_ARG(0);  // indicates whether biases array is provided
+  const auto hasPH = B_ARG(1);      // indicates whether peephole connections are present
+
+  const auto gateActHasAlpha = gateAct == 3 || gateAct == 4 || gateAct == 5 || gateAct == 6 || gateAct == 8;
+  const auto cellActHasAlpha = cellAct == 3 || cellAct == 4 || cellAct == 5 || cellAct == 6 || cellAct == 8;
+  const auto outActHasAlpha = outAct == 3 || outAct == 4 || outAct == 5 || outAct == 6 || outAct == 8;
+  const auto gateActHasBeta = gateAct == 3 || gateAct == 6;
+  const auto cellActHasBeta = cellAct == 3 || cellAct == 6;
+  const auto outActHasBeta = outAct == 3 || outAct == 6;
+
+  LongType count = 1;
+  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping
+  const auto gateAlpha = gateActHasAlpha ? T_ARG(count++) : 0;
+  const auto gateBeta = gateActHasBeta ? T_ARG(count++) : 0;
+  const auto cellAlpha = cellActHasAlpha ? T_ARG(count++) : 0;
+  const auto cellBeta = cellActHasBeta ? T_ARG(count++) : 0;
+  const auto outAlpha = outActHasAlpha ? T_ARG(count++) : 0;
+  const auto outBeta = outActHasBeta ? T_ARG(count++) : 0;
+
+  count = 3;
+  const auto x = INPUT_VARIABLE(0);                              // input
+  const auto Wx = INPUT_VARIABLE(1);                             // input weights
+  const auto Wr = INPUT_VARIABLE(2);                             // recurrent weights
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;  // biases
+  const auto hI = INPUT_VARIABLE(count++);                       // initial output
+  const auto cI = INPUT_VARIABLE(count++);                       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr;     // peephole weights
+  const auto dLdh = INPUT_VARIABLE(count);                       // gradient wrt output
+
+  REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER_CELL_BP operation: cell clipping value should be nonnegative (>=0) !");
+
+  count = 3;
+  auto dLdx = OUTPUT_VARIABLE(0);
+  auto dLdWx = OUTPUT_VARIABLE(1);
+  auto dLdWr = OUTPUT_VARIABLE(2);
+  auto dLdb = hasBiases ? OUTPUT_VARIABLE(count++) : nullptr;
+  auto dLdhI = OUTPUT_VARIABLE(count++);
+  auto dLdcI = OUTPUT_VARIABLE(count++);
+  auto dLdWp = hasPH ? OUTPUT_VARIABLE(count) : nullptr;
+
+  // evaluate dimensions
+  const LongType bS = x->rankOf() == 1 ? 0 : x->sizeAt(0);
+  const LongType nIn = x->sizeAt(-1);
+  const LongType nOut = Wx->sizeAt(-1) / 4;
+
+  // inputs validations
+  // Wx validation
+  if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL_BP operation: wrong shape of input weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
+  // Wr validation
+  if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL_BP operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
+  // initial output/cell validation
+  std::vector<LongType> exphIcIShape =
+      x->rankOf() == 1 ? std::vector<LongType>{nOut} : std::vector<LongType>{bS, nOut};
+  REQUIRE_TRUE(hI->isSameShape(exphIcIShape), 0,
+               "LSTM_LAYER_CELL_BP operation: wrong shape of initial output, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(exphIcIShape).c_str(), ShapeUtils::shapeAsString(hI).c_str());
+  REQUIRE_TRUE(cI->isSameShape(exphIcIShape), 0,
+               "LSTM_LAYER_CELL_BP operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(exphIcIShape).c_str(), ShapeUtils::shapeAsString(cI).c_str());
+  REQUIRE_TRUE(dLdh->isSameShape(exphIcIShape), 0,
+               "LSTM_LAYER_CELL_BP operation: wrong shape of dLdh gradient, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(exphIcIShape).c_str(), ShapeUtils::shapeAsString(dLdh).c_str());
+  // biases validation
+  if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
+    REQUIRE_TRUE(false, 0, "LSTM_LAYER_CELL_BP operation: wrong shape of biases, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  if (dLdb != nullptr && (dLdb->rankOf() != 1 || dLdb->sizeAt(0) != 4 * nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL_BP operation: wrong shape of dLdb gradient, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(dLdb).c_str());
+  // peephole weights validation
+  if (Wp != nullptr && (Wp->rankOf() != 1 || Wp->sizeAt(0) != 3 * nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL_BP operation: wrong shape of peephole weights, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({3 * nOut}).c_str(), ShapeUtils::shapeAsString(Wp).c_str());
+  if (dLdWp != nullptr && (dLdWp->rankOf() != 1 || dLdWp->sizeAt(0) != 3 * nOut))
+    REQUIRE_TRUE(false, 0,
+                 "LSTM_LAYER_CELL_BP operation: wrong shape of dLdWp gradient, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({3 * nOut}).c_str(), ShapeUtils::shapeAsString(dLdWp).c_str());
+
+  std::vector<float> params = {
+      static_cast<float>(0) /*ignore*/, static_cast<float>(0) /*ignore*/, static_cast<float>(cellClip),
+      static_cast<float>(gateAct),      static_cast<float>(gateAlpha),    static_cast<float>(gateBeta),
+      static_cast<float>(cellAct),      static_cast<float>(cellAlpha),    static_cast<float>(cellBeta),
+      static_cast<float>(outAct),       static_cast<float>(outAlpha),     static_cast<float>(outBeta)};
+
+  std::vector<LongType> zShape =
+      x->rankOf() == 1 ? std::vector<LongType>({4 * nOut}) : std::vector<LongType>({bS, 4 * nOut});
+
+  NDArray z(x->ordering(), zShape, x->dataType(), block.launchContext());
+  NDArray *a = z.ulike();
+  NDArray *h = cI->ulike();
+  NDArray *c = cI->ulike();
+
+  helpers::lstmLayerCell(x, Wx, Wr, b, hI, cI, Wp, params, &z, a, h, c);
+
+  helpers::lstmLayerCellBp(x, Wx, Wr, b, hI, cI, Wp, dLdh, nullptr, nullptr, &z, a, c, params, dLdx, dLdWx, dLdWr,
+                           dLdhI, dLdcI, dLdb, dLdWp);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(lstmLayerCellBp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(lstmLayerCellBp) {
+  const auto hasBiases = B_ARG(0);  // indicates whether biases array is provided
+  const auto hasPH = B_ARG(1);      // indicates whether peephole connections are present
+
+  LongType count = 3;
+  const auto x = INPUT_VARIABLE(0);                              // input
+  const auto Wx = INPUT_VARIABLE(1);                             // input weights
+  const auto Wr = INPUT_VARIABLE(2);                             // recurrent weights
+  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;  // biases
+  const auto hI = INPUT_VARIABLE(count++);                       // initial output
+  const auto cI = INPUT_VARIABLE(count++);                       // initial cell state
+  const auto Wp = hasPH ? INPUT_VARIABLE(count) : nullptr;       // peephole weights
+
+  auto shapes = SHAPELIST(x->shapeInfo(), Wx->shapeInfo(), Wr->shapeInfo());
+
+  if (b != nullptr) shapes->push_back(b->shapeInfo());
+
+  shapes->push_back(hI->shapeInfo());
+  shapes->push_back(cI->shapeInfo());
+
+  if (Wp != nullptr) shapes->push_back(Wp->shapeInfo());
+
+  return shapes;
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/sru.cpp

@@ -0,0 +1,663 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// implementation of operations for Simple Recurrent Unit: arXiv:1709.02755v2 [cs.CL] 12 Sep 2017
+//
+//@author Yurii Shyrma
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_sru)
+
+#include <helpers/MmulHelper.h>
+#include <helpers/PointersManager.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/sru.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sru, 5, 2, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);  // X, input 3d tensor [bS x inSize x time], time - number of time steps, bS - batch size,
+  // inSize - number of features
+  auto w = INPUT_VARIABLE(1);  // W, 2d tensor of weights [3*inSize x inSize]
+  auto b = INPUT_VARIABLE(2);  // B, row of biases with twice length [2*inSize]
+  auto c0 = INPUT_VARIABLE(3);  // C_{0}, 2d tensor of initial state [bS x inSize] at time t=0
+  auto mask = block.width() > 4 ? INPUT_VARIABLE(4) : nullptr;  // optional,  2d tensor of dropout mask [bS x inSize]
+
+  auto h = OUTPUT_VARIABLE(0);  // cell outputs, [bS x inSize x time]
+  auto c = OUTPUT_VARIABLE(1);  // cell states,  [bS x inSize x time]
+
+  const int rank = x->rankOf();  // = 3
+  const auto bS = x->sizeAt(0);
+  const auto inSize = x->sizeAt(1);
+  const auto time = x->sizeAt(2);
+
+  // input shapes validation
+  REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
+               "SRU operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               w->rankOf());
+  REQUIRE_TRUE(b->rankOf() == 1, 0, "SRU operation: wrong rank of biases  array, expected is %i, but got %i instead !",
+               1, b->rankOf());
+  REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
+               "SRU operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0->rankOf());
+  if (mask)
+  REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
+               "SRU operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               mask->rankOf());
+
+  const std::vector<LongType> wCorrectShape = {3 * inSize, inSize};
+  const std::vector<LongType> bCorrectShape = {2 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, inSize};
+
+  REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
+               "SRU operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "SRU operation: wrong shape of biases  array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
+               "SRU operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
+  if (mask)
+  REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
+               "SRU operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
+
+  //  xm = x * mask
+  auto xm = x;
+  if (mask) {
+    xm = new NDArray(x->shapeInfo(), true, block.launchContext());
+    std::vector<LongType> dims = {0, 1};
+    x->applyBroadcast(broadcast::Multiply,&dims , mask, xm);
+  }
+
+  // time loop
+  helpers::sruTimeLoop(block.launchContext(), xm, c0, w, b, h, c);
+
+  if (mask) delete xm;
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(sru) { getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS}); }
+
+DECLARE_SHAPE_FN(sru) {
+  auto xShapeInfo = inputShape->at(0);   // X, input 3d tensor [bS x inSize x time], time - number of time steps, bS -
+  // batch size, inSize - number of features
+  auto wShapeInfo = inputShape->at(1);   // W, 2d tensor of weights [3*inSize x inSize]
+  auto bShapeInfo = inputShape->at(2);   // B, row of biases with twice length [2*inSize]
+  auto c0ShapeInfo = inputShape->at(3);  // C_{0}, 2d tensor of initial state [bS x inSize] at time t=0
+  auto maskShapeInfo =
+      block.width() > 4 ? inputShape->at(4) : nullptr;  // optional,  2d tensor of dropout mask [bS x inSize]
+
+  const int rank = xShapeInfo[0];  // = 3
+  const int bS = xShapeInfo[1];
+  const int inSize = xShapeInfo[2];
+  const int time = xShapeInfo[3];
+
+  // input shapes validation
+  REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
+               "SRU operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               wShapeInfo[0]);
+  REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
+               "SRU operation: wrong rank of biases  array, expected is %i, but got %i instead !", 1, bShapeInfo[0]);
+  REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
+               "SRU operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0ShapeInfo[0]);
+  if (maskShapeInfo)
+  REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
+               "SRU operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               maskShapeInfo[0]);
+
+  const std::vector<LongType> wCorrectShape = {3 * inSize, inSize};
+  const std::vector<LongType> bCorrectShape = {2 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, inSize};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
+               "SRU operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
+               "SRU operation: wrong shape of biases  array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
+               "SRU operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
+  if (maskShapeInfo)
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
+               "SRU operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
+
+  LongType * newShapeInfo1 = nullptr;
+  ALLOCATE(newShapeInfo1, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x inSize x time]
+
+  newShapeInfo1[0] = rank;
+  newShapeInfo1[1] = bS;
+  newShapeInfo1[2] = inSize;
+  newShapeInfo1[3] = time;
+
+  ShapeUtils::updateStridesAndType(newShapeInfo1, xShapeInfo, shape::order(xShapeInfo));
+  auto result = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo1)->primary();
+  RELEASE(newShapeInfo1, block.getWorkspace());
+  return SHAPELIST(result, result);
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sru_bp, 8, 4, true, 0, 0) {
+  auto x = INPUT_VARIABLE(0);
+  // X, input 3d tensor [bS x K x N], N - number of time steps, bS - batch size, K - number of features
+  auto w = INPUT_VARIABLE(1);         // W, 2d tensor of weights [3K x K]
+  auto b = INPUT_VARIABLE(2);         // B, row of biases with twice length [1 x 2*K]
+  auto c0 = INPUT_VARIABLE(3);        // C_{0}, 2d tensor of initial state [bS x K] at time t=0
+  auto c = INPUT_VARIABLE(4);         // C, [bS x K x N]
+  auto inGradCt = INPUT_VARIABLE(5);  // [bS x K]
+  auto inGradH = INPUT_VARIABLE(6);   // [bS x K x N]
+  NDArray* mask = nullptr;            // optional,  2d tensor of dropout mask [bS x K]
+
+  bool applyMask = false;
+  if (block.width() > 7) {
+    mask = INPUT_VARIABLE(7);
+    applyMask = true;
+  }
+
+  auto gradX = OUTPUT_VARIABLE(0);     // [bS x K x N]
+  auto gradW = OUTPUT_VARIABLE(1);     // [bS x 3K x K]
+  auto gradB = OUTPUT_VARIABLE(2);     // [1 x 2K]
+  auto gradInit = OUTPUT_VARIABLE(3);  // [bS x K]
+
+  const int bS = x->shapeOf()[0];
+  const int K = x->shapeOf()[1];
+  const int N = x->shapeOf()[2];  // N - number of time steps
+
+  std::vector<sd::LongType> gradBiasShape = {bS, 2 * K, N};
+  std::vector<sd::LongType> gradUShape = {bS, 3 * K, N};
+  std::vector<sd::LongType> gradHXShape = {bS, K, N};
+  std::vector<sd::LongType> gctShape = {bS, K};
+  std::vector<sd::LongType> gradTanhShape = {bS, K};
+  std::vector<sd::LongType> gradCtShape = {bS, K};
+  std::vector<sd::LongType> ftMinusShape = {bS, K};
+  std::vector<sd::LongType> rtMinusShape = {bS, K};
+  std::vector<sd::LongType> temp1Shape = {bS, K};
+  std::vector<sd::LongType> temp2Shape = {bS, K};
+
+  auto gradBias = NDArrayFactory::create_(x->ordering(), gradBiasShape, gradX->dataType(), block.launchContext());
+  auto gradU = NDArrayFactory::create_(x->ordering(), gradUShape, gradX->dataType(), block.launchContext());
+  auto gradHX = NDArrayFactory::create_(x->ordering(), gradHXShape, gradX->dataType(), block.launchContext());
+  auto gct = NDArrayFactory::create_(c->ordering(), gctShape, gradX->dataType(), block.launchContext());
+  auto gradTanh = NDArrayFactory::create_(c->ordering(), gradTanhShape, gradX->dataType(), block.launchContext());
+  auto gradCt = NDArrayFactory::create_(c->ordering(), gradCtShape, gradX->dataType(), block.launchContext());
+  auto ftMinus = NDArrayFactory::create_(c->ordering(), ftMinusShape, gradX->dataType(), block.launchContext());
+  auto rtMinus = NDArrayFactory::create_(c->ordering(), rtMinusShape, gradX->dataType(), block.launchContext());
+  auto temp1 = NDArrayFactory::create_(c->ordering(), temp1Shape, gradX->dataType(), block.launchContext());
+  auto temp2 = NDArrayFactory::create_(c->ordering(), temp2Shape, gradX->dataType(), block.launchContext());
+  std::vector<LongType> axes = {0, 1};
+  //  x = x * mask
+  if (applyMask) x->applyBroadcast(broadcast::Multiply, &axes, mask, x);  // apply mask
+  // multiplication matrix wi = matmul(w,x), U = WX
+  auto wi = MmulHelper::mmul(w, x, nullptr, 1., 0.);  // U [bS x 3K x N]
+
+  auto wiZ = (*wi)({0, 0, 0, K, 0, 0}, true);                // [bS x K x N]
+  auto wiF = (*wi)({0, 0, K, 2 * K, 0, 0}, true);            // forget gate [bS x K x N]
+  auto wiR = (*wi)({0, 0, 2 * K, 3 * K, 0, 0}, true);        // reset gate [bS x K x N]
+  auto bF = (*b)({0, 0, 0, K}, true);                        // biases for forget gate [1 x K]
+  auto bR = (*b)({0, 0, K, 2 * K}, true);                    // biases for reset gate [1 x K]
+  auto gradBF = (*gradBias)({0, 0, 0, K, 0, 0}, true);       // [bS x K x N]
+  auto gradBR = (*gradBias)({0, 0, K, 2 * K, 0, 0}, true);   // [bS x K x N]
+  auto gradUZ = (*gradU)({0, 0, 0, K, 0, 0}, true);          // [bS x K x N]
+  auto gradUF = (*gradU)({0, 0, K, 2 * K, 0, 0}, true);      // [bS x K x N]
+  auto gradUR = (*gradU)({0, 0, 2 * K, 3 * K, 0, 0}, true);  // [bS x K x N]
+
+  NDArray* ct_1 = nullptr;
+
+  std::vector<LongType> idx = {0, 0, 0, 0, 0, 0};
+
+  for (int t = N - 1; t >= 0; --t) {
+    // initialization
+    idx[4] = t;
+    idx[5] = t + 1;
+    auto xt = (*x)(idx);              // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto zt = (*wiZ)(idx);               // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto ft = (*wiF)(idx);               // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto rt = (*wiR)(idx);               // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto ct = (*c)(idx);              // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto inGradHt = (*inGradH)(idx);  // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradBRt = (*gradBR)(idx);       // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradBFt = (*gradBF)(idx);       // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradHXt = (*gradHX)(idx);    // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradUZt = (*gradUZ)(idx);       // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradUFt = (*gradUF)(idx);       // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+    auto gradURt = (*gradUR)(idx);       // [bS x K x N] -> [bS x K x 1] -> [bS x K]
+
+
+
+    if (t != 0) {
+      idx[4] = t - 1;
+      idx[5] = t;
+      ct_1 = new NDArray((*c)(idx));  // previous c_{t-1}
+    } else
+      ct_1 = c0;
+
+    ///////////////// forward
+    // ft = sigmoid(ft + bf), rt = sigmoid(rt + bR)
+    ft->addRowVector(bF, ft);
+    rt->addRowVector(bR, rt);
+    ft->applyTransform(transform::Sigmoid, ft);
+    rt->applyTransform(transform::Sigmoid, rt);
+
+    // TODO T val = (activation_type == 1) ? tanh(cur) : ((activation_type == 2) ? reluf(cur) : cur );
+    ct->applyTransform(transform::Tanh, gct);
+    // ftMinus = 1-ft,  rtMinus = 1-rt
+    ft->applyTransform(transform::OneMinus, ftMinus);
+    rt->applyTransform(transform::OneMinus, rtMinus);
+
+    ///////////////// backward
+    // bR, *grad_brt_ptr = inGradHt * (g_ct - xt) * (1.0f - rt) * rt;
+    gct->applyPairwiseTransform(pairwise::Subtract, xt, temp1);      // temp1 = (g_ct - xt)
+    rtMinus->applyPairwiseTransform(pairwise::Multiply, rt, temp2);  // temp2 = (1.0f - rt) * rt;
+    temp1->applyPairwiseTransform(pairwise::Multiply, temp2);        // temp1 = (g_ct - xt) * (1.0f - rt) * rt;
+    inGradHt->applyPairwiseTransform(pairwise::Multiply, temp1,
+                                    gradBRt);  // = inGradHt * (g_ct - xt) * (1.0f - rt) * rt;
+
+    // bF, TODO - tanh
+    // gradTanh = (1.0f - g_ct * g_ct);
+    gct->applyPairwiseTransform(pairwise::Multiply, gct, gradTanh);  // gradTanh = g_ct * g_ct
+    gradTanh->applyTransform(transform::OneMinus, gradTanh);          // gradTanh = (1.0f - g_ct * g_ct)
+    // gradCt  = inGradHt * rt * gradTanh
+    rt->applyPairwiseTransform(pairwise::Multiply, gradTanh, gradCt);      // gradCt = rt * gradTanh
+    inGradHt->applyPairwiseTransform(pairwise::Multiply, gradCt, gradCt);  // gradCt = inGradHt * rt * gradTanh
+    // gradBFt = (gradCt + inGradCt) * (ct_1 - zt) * (1 - ft) * ft;
+    gradCt->applyPairwiseTransform(pairwise::Add, inGradCt, temp1);     // temp1 = (gradCt + inGradCt)
+    ct_1->applyPairwiseTransform(pairwise::Subtract, zt, temp2);         // temp2 = (ct_1 - zt)
+    temp1->applyPairwiseTransform(pairwise::Multiply, ftMinus, temp1);  // temp1 = (gradCt + inGradCt)*(1-ft)
+    temp1->applyPairwiseTransform(pairwise::Multiply, ft, temp1);        // temp1 = (gradCt + inGradCt)*(1-ft)*ft
+    temp1->applyPairwiseTransform(pairwise::Multiply, temp2,
+                                  gradBFt);  // gradBFt = (gradCt + inGradCt) * (ct_1 - zt) * (1 - ft) * ft;
+
+    // x_t (highway connection), gradHXt = inGradHt * (1.0f - rt);
+    inGradHt->applyPairwiseTransform(pairwise::Multiply, rtMinus, gradHXt);
+
+    // U_t, gradUZt = (inGradHt * rt * grad_tanh + inGradCt) * (1.0f - ft);
+    rt->applyPairwiseTransform(pairwise::Multiply, gradTanh, temp1);     // temp1 = rt * grad_tanh
+    inGradHt->applyPairwiseTransform(pairwise::Multiply, temp1,temp1);  // temp1 = inGradHt * rt * grad_tanh
+    temp1->applyPairwiseTransform(pairwise::Add, inGradCt, temp1);  // temp1 = inGradHt * rt * grad_tanh + inGradCt
+    temp1->applyPairwiseTransform(pairwise::Multiply, ftMinus,
+                                  gradUZt);  // gradUZt = (inGradHt * rt * grad_tanh + inGradCt) * (1.0f - ft);
+    gradUFt->assign(gradBFt);
+    gradURt->assign(gradBRt);
+
+    // c_{t-1}, inGradCt = (gradCt + inGradCt) * ft;
+    gradCt->applyPairwiseTransform(pairwise::Add, inGradCt, temp1);  // temp1 = (gradCt + inGradCt)
+    temp1->applyPairwiseTransform(pairwise::Multiply, ft, inGradCt);  // inGradCt = (gradCt + inGradCt) * ft;
+
+    if (t != 0) delete ct_1;
+
+    delete xt;
+    delete zt;
+    delete ft;
+    delete rt;
+    delete ct;
+    delete inGradHt;
+    delete gradBRt;
+    delete gradHXt;
+    delete gradUZt;
+    delete gradUFt;
+    delete gradURt;
+  }
+
+  // gradInit
+  gradInit->assign(inGradCt);
+
+  // gradX
+  auto weightsT = w->transpose();                                                    // [K x 3K]
+  MmulHelper::mmul(weightsT, gradU, gradX, 1., 0.);                                 // [bS x K x N]
+  gradX->applyPairwiseTransform(pairwise::Add, gradHX, gradX);
+  std::vector<LongType> axes3 = {0, 1};
+  // + grad_highway_x
+  if (applyMask) gradX->applyBroadcast(broadcast::Multiply, &axes3, mask, gradX);  // apply mask
+
+  // gradB
+  std::vector<sd::LongType> gradBShape = { 2 * K};
+  auto gradB2 = gradB->reshape(gradB->ordering(), gradBShape);
+  std::vector<LongType> axes2;
+  axes.push_back(0);
+  axes.push_back(2);
+  gradBias->reduceAlongDimension(reduce::Sum, gradB2, &axes2);  // [1 x 2K]
+
+  // gradW [bS x 3K x K]
+  x->permutei({0, 2, 1}, false, false);                     // [bS x N x K]
+  MmulHelper::mmul(gradU, x, gradW, 1., 0.);  // [bS x 3K x K]
+
+  delete gct;
+  delete gradU;
+  delete gradHX;
+  delete temp1;
+  delete temp2;
+  delete gradCt;
+  delete wi;
+  delete gradTanh;
+  delete ftMinus;
+  delete rtMinus;
+  delete gradBias;
+  delete weightsT;
+  return Status::OK;
+}
+
+DECLARE_TYPES(sru_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(sru_bp) {
+  auto inShape = inputShape->at(0);  // [bS x inSize x time]
+  auto bS = inShape[1];
+  auto inSize = inShape[2];
+  auto time = inShape[3];
+  char order = (char)(inShape[9]);
+
+  auto ret = SHAPELIST(
+      ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
+                                                            std::vector<sd::LongType>{bS, inSize, time})->primary(),
+      ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
+                                                            std::vector<sd::LongType>{bS, 3 * inSize, inSize})->primary(),
+      ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
+                                                            std::vector<sd::LongType>{1, 2 * inSize})->primary(),
+      ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
+                                                            std::vector<sd::LongType>{bS, inSize})->primary()
+  );
+  return ret;
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sru_bi, 5, 2, true, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // X, input 3d tensor [time x bS x 2*inSize], time - number of time steps, bS - batch
+  // size, inSize - number of features
+  auto w = INPUT_VARIABLE(1);   // W, 2d tensor of weights [2*inSize x 6*inSize]
+  auto b = INPUT_VARIABLE(2);   // B, row of biases with twice length [1 x 4*inSize]
+  auto c0 = INPUT_VARIABLE(3);  // C_{0}, 2d tensor of initial state [bS x 2*inSize] at time t=0
+  NDArray* mask =
+      block.width() > 4 ? INPUT_VARIABLE(4) : nullptr;  // optional, 2d tensor of dropout mask [bS x 2*inSize]
+
+  auto ht = OUTPUT_VARIABLE(0);  // h_t, [time x bS x 2*inSize]
+  auto ct = OUTPUT_VARIABLE(1);  // c_t, [time x bS x 2*inSize]
+
+  // input shapes validation
+  const int rank = x->rankOf();
+  const LongType bS = x->sizeAt(1);
+  const LongType inSize = x->sizeAt(2) / 2;
+
+  REQUIRE_TRUE(x->rankOf() == rank, 0,
+               "SRU_BI operation: wrong rank of input array, expected is %i, but got %i instead !", rank, x->rankOf());
+  REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
+               "SRU_BI operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               w->rankOf());
+  REQUIRE_TRUE(b->rankOf() == 1, 0, "SRU_BI operation: wrong rank of biases array, expected is 1, but got %i instead !",
+               b->rankOf());
+  REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
+               "SRU_BI operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0->rankOf());
+  if (mask)
+  REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
+               "SRU_BI operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               mask->rankOf());
+
+  const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
+  const std::vector<LongType> bCorrectShape = {4 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
+
+  REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
+               "SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
+  if (mask)
+  REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
+
+  helpers::sruBI(block.launchContext(), x, w, b, c0, mask, ht, ct);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(sru_bi) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(sru_bi) {
+  auto xShapeInfo = inputShape->at(0);  // [time x bS x 2K ]
+  auto wShapeInfo = inputShape->at(1);
+  auto bShapeInfo = inputShape->at(2);
+  auto c0ShapeInfo = inputShape->at(3);
+  auto maskShapeInfo =
+      block.width() > 4 ? inputShape->at(4) : nullptr;  // optional,  2d tensor of dropout mask [bS x inSize]
+
+  const int rank = xShapeInfo[0];  // = 3
+  const LongType time = xShapeInfo[1];
+  const LongType bS = xShapeInfo[2];
+  const LongType inSize = xShapeInfo[3] / 2;
+
+  // input shapes validation
+  REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
+               "SRU_BI operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               wShapeInfo[0]);
+  REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
+               "SRU_BI operation: wrong rank of biases  array, expected is 1, but got %i instead !", bShapeInfo[0]);
+  REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
+               "SRU_BI operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0ShapeInfo[0]);
+  if (maskShapeInfo)
+  REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
+               "SRU_BI operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               maskShapeInfo[0]);
+
+  const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
+  const std::vector<LongType> bCorrectShape = {4 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
+               "SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
+               "SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
+  if (maskShapeInfo)
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
+
+  char order = shape::order(xShapeInfo);
+
+  ShapeDescriptor *descriptor = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, bS, 2 * inSize});
+  auto result = ConstantShapeHelper::getInstance().createShapeInfo(descriptor);
+  return SHAPELIST(result, result);
+}
+
+DECLARE_TYPES(sru_bi_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sru_bi_bp, 8, 4, true, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // X, input 3d tensor [time x bS x 2*inSize], time - number of time steps, bS - batch
+  // size, inSize - number of features
+  auto w = INPUT_VARIABLE(1);   // W, 2d tensor of weights [2*inSize x 6*inSize]
+  auto b = INPUT_VARIABLE(2);   // B, row of biases with twice length [4*inSize]
+  auto c0 = INPUT_VARIABLE(3);  // C_{0}, 2d tensor of initial state [bS x 2*inSize] at time t=0
+  auto ct = INPUT_VARIABLE(4);  // C, [time x bS x 2*inSize]
+  auto inGradC0 = INPUT_VARIABLE(5);  // [bS x 2*inSize]
+  auto inGradHt = INPUT_VARIABLE(6);  // [time x bS x 2*inSize]
+  NDArray* mask =
+      block.width() > 7 ? INPUT_VARIABLE(7) : nullptr;  // optional,  2d tensor of dropout mask [bS x 2*inSize]
+
+  // input shapes validation
+  const int rank = x->rankOf();
+  const LongType time = x->sizeAt(0);
+  const LongType bS = x->sizeAt(1);
+  const LongType inSize = x->sizeAt(2) / 2;
+
+  REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               w->rankOf());
+  REQUIRE_TRUE(b->rankOf() == 1, 0,
+               "SRU_BI_BP operation: wrong rank of biases array, expected is 1, but got %i instead !", b->rankOf());
+  REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0->rankOf());
+  REQUIRE_TRUE(ct->rankOf() == rank, 0,
+               "SRU_BI_BP operation: wrong rank of state array, expected is %i, but got %i instead !", rank,
+               ct->rankOf());
+  REQUIRE_TRUE(inGradC0->rankOf() == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of gradient c0, expected is %i, but got %i instead !", rank - 1,
+               inGradC0->rankOf());
+  REQUIRE_TRUE(inGradHt->rankOf() == rank, 0,
+               "SRU_BI_BP operation: wrong rank of gradient ht, expected is %i, but got %i instead !", rank,
+               inGradHt->rankOf());
+  if (mask)
+  REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               mask->rankOf());
+
+  const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
+  const std::vector<LongType> bCorrectShape = {4 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
+  const std::vector<LongType> ctCorrectShape = {time, bS, 2 * inSize};
+
+  REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
+               "SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
+  REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
+               "SRU_BI operation: wrong shape of biases  array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
+  REQUIRE_TRUE(ct->isSameShape(ctCorrectShape), 0,
+               "SRU_BI operation: wrong shape of state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(ctCorrectShape).c_str(), ShapeUtils::shapeAsString(ct).c_str());
+  if (mask)
+  REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
+
+  auto gradI = OUTPUT_VARIABLE(0);   // [time x bS x 2*inSize]
+  auto gradW = OUTPUT_VARIABLE(1);   // [time x 2*inSize x 6*inSize]
+  auto gradB = OUTPUT_VARIABLE(2);   // [1 x 4*inSize]
+  auto gradC0 = OUTPUT_VARIABLE(3);  // [bS x 2*inSize]
+
+  helpers::sruBIBP(block.launchContext(), x, w, b, c0, ct, inGradC0, inGradHt, mask, gradI, gradW, gradB, gradC0);
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(sru_bi_bp) {
+  auto xShapeInfo = inputShape->at(0);  // [time x bS x 2K ]
+  auto wShapeInfo = inputShape->at(1);
+  auto bShapeInfo = inputShape->at(2);
+  auto c0ShapeInfo = inputShape->at(3);
+  auto ctShapeInfo = inputShape->at(4);
+  auto inGradC0ShapeInfo = inputShape->at(5);
+  auto inGradHtShapeInfo = inputShape->at(6);
+  auto maskShapeInfo =
+      block.width() > 7 ? inputShape->at(7) : nullptr;  // optional,  2d tensor of dropout mask [bS x inSize]
+
+  // input shapes validation
+  const int rank = xShapeInfo[0];
+  const LongType time = xShapeInfo[1];
+  const LongType bS = xShapeInfo[2];
+  const LongType inSize = xShapeInfo[3] / 2;
+
+  REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
+               wShapeInfo[0]);
+  REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
+               "SRU_BI_BP operation: wrong rank of biases  array, expected is 1, but got %i instead !", bShapeInfo[0]);
+  REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
+               c0ShapeInfo);
+  REQUIRE_TRUE(ctShapeInfo[0] == rank, 0,
+               "SRU_BI_BP operation: wrong rank of state array, expected is %i, but got %i instead !", rank,
+               ctShapeInfo);
+  REQUIRE_TRUE(inGradC0ShapeInfo[0] == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of gradient c0, expected is %i, but got %i instead !", rank - 1,
+               inGradC0ShapeInfo[0]);
+  REQUIRE_TRUE(inGradHtShapeInfo[0] == rank, 0,
+               "SRU_BI_BP operation: wrong rank of gradient ht, expected is %i, but got %i instead !", rank,
+               inGradHtShapeInfo[0]);
+  if (maskShapeInfo)
+  REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
+               "SRU_BI_BP operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
+               maskShapeInfo[0]);
+
+  const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
+  const std::vector<LongType> bCorrectShape = {4 * inSize};
+  const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
+  const std::vector<LongType> ctCorrectShape = {time, bS, 2 * inSize};
+  const std::vector<LongType> inGradC0CorrectShape = {bS, 2 * inSize};
+  const std::vector<LongType> inGradHtCorrectShape = {time, bS, 2 * inSize};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
+               "SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
+               "SRU_BI operation: wrong shape of biases  array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(ctShapeInfo, ctCorrectShape), 0,
+               "SRU_BI operation: wrong shape of state array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(ctCorrectShape).c_str(), ShapeUtils::shapeAsString(ctShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(inGradC0ShapeInfo, inGradC0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of gradient c0 array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(inGradC0CorrectShape).c_str(),
+               ShapeUtils::shapeAsString(inGradC0ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(inGradHtShapeInfo, inGradHtCorrectShape), 0,
+               "SRU_BI operation: wrong shape of gradient ht array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(inGradHtCorrectShape).c_str(),
+               ShapeUtils::shapeAsString(inGradHtShapeInfo).c_str());
+  if (maskShapeInfo)
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
+               "SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
+
+  const char order = shape::order(xShapeInfo);
+
+  ShapeDescriptor *descriptor1 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, bS, 2 * inSize});
+  ShapeDescriptor *descriptor2 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, 2 * inSize, 6 * inSize});
+  ShapeDescriptor *descriptor3 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {4 * inSize});
+  ShapeDescriptor *descriptor4 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {bS, 2 * inSize});
+
+  return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(descriptor1),
+                   ConstantShapeHelper::getInstance().createShapeInfo(descriptor2),
+                   ConstantShapeHelper::getInstance().createShapeInfo(descriptor3),
+                   ConstantShapeHelper::getInstance().createShapeInfo(descriptor4));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/sruCell.cpp

@@ -0,0 +1,115 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author Yurii Shyrma, created on 05.12.2017
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_sruCell)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/sru.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(sruCell, 4, 2, false, 0, 0) {
+  auto xt = INPUT_VARIABLE(0);    // input [bS x inSize], bS - batch size, inSize - number of features
+  auto ct_1 = INPUT_VARIABLE(1);  // previous cell state ct  [bS x inSize], that is at previous time step t-1
+  auto w = INPUT_VARIABLE(2);     // weights [inSize x 3*inSize]
+  auto b = INPUT_VARIABLE(3);     // biases [2*inSize]
+
+  auto ht = OUTPUT_VARIABLE(0);  // current cell output [bS x inSize], that is at current time step t
+  auto ct = OUTPUT_VARIABLE(1);  // current cell state  [bS x inSize], that is at current time step t
+
+  const int rank = xt->rankOf();
+  const int bS = xt->sizeAt(0);
+  const int inSize = xt->sizeAt(1);  // inSize - number of features
+
+  // input shapes validation
+  const std::vector<LongType> correctCt_1Shape = {bS, inSize};
+  const std::vector<LongType> correctWShape = {inSize, 3 * inSize};
+  const std::vector<LongType> correctBShape = {2 * inSize};
+
+  REQUIRE_TRUE(ct_1->isSameShape(correctCt_1Shape), 0,
+               "SRUCELL operation: wrong shape of previous cell state, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctCt_1Shape).c_str(), ShapeUtils::shapeAsString(ct_1).c_str());
+  REQUIRE_TRUE(w->isSameShape(correctWShape), 0,
+               "SRUCELL operation: wrong shape of weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
+  REQUIRE_TRUE(b->isSameShape(correctBShape), 0,
+               "SRUCELL operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+
+  // fixme: shitty initializer lists
+  helpers::sruCell(block.launchContext(), xt, ct_1, w, b, ht, ct);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(sruCell) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(sruCell) {
+  auto xtShapeInfo = inputShape->at(0);    // input [bS x inSize], bS - batch size, inSize - number of features
+  auto ct_1ShapeInfo = inputShape->at(1);  // previous cell state ct  [bS x inSize], that is at previous time step t-1
+  auto wShapeInfo = inputShape->at(2);     // weights [inSize x 3*inSize]
+  auto bShapeInfo = inputShape->at(3);     // biases [2*inSize]
+
+  const int rank = xtShapeInfo[0];
+  const int bS = xtShapeInfo[1];
+  const int inSize = xtShapeInfo[2];  // inSize - number of features
+
+  // input shapes validation
+  const std::vector<LongType> correctCt_1Shape = {bS, inSize};
+  const std::vector<LongType> correctWShape = {inSize, 3 * inSize};
+  const std::vector<LongType> correctBShape = {2 * inSize};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(ct_1ShapeInfo, correctCt_1Shape), 0,
+               "SRUCELL operation: wrong shape of previous cell state, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctCt_1Shape).c_str(), ShapeUtils::shapeAsString(ct_1ShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, correctWShape), 0,
+               "SRUCELL operation: wrong shape of weights, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctWShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, correctBShape), 0,
+               "SRUCELL operation: wrong shape of biases, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(correctBShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *cShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x numProj]
+  ALLOCATE(cShapeInfo, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType);  // [bS x numUnits]
+
+  hShapeInfo[0] = cShapeInfo[0] = rank;
+  hShapeInfo[1] = cShapeInfo[1] = bS;
+  hShapeInfo[2] = cShapeInfo[2] = inSize;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, ct_1ShapeInfo, shape::order(ct_1ShapeInfo));
+  ShapeUtils::updateStridesAndType(cShapeInfo, ct_1ShapeInfo, shape::order(ct_1ShapeInfo));
+
+  return SHAPELIST(ConstantShapeHelper::getInstance().createFromExisting(hShapeInfo),
+                   ConstantShapeHelper::getInstance().createFromExisting(cShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/staticBidirectionalRNN.cpp

@@ -0,0 +1,286 @@
+/* ******************************************************************************
+ *
+ *
+ * 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, created on 03.04.2018
+//
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/reverse.h>
+#include <ops/declarable/helpers/rnn.h>
+#include <ops/declarable/helpers/transforms.h>
+
+#if NOT_EXCLUDED(OP_static_bi_directional_rnn)
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(static_bidirectional_rnn, 7, 3, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);     // input [time x bS x inSize]
+  auto WxFW = INPUT_VARIABLE(1);  // input-to-hidden  weights for forward  RNN, [inSize  x numUnitsFW]
+  auto WhFW = INPUT_VARIABLE(2);  // hidden-to-hidden weights for forward  RNN, [numUnitsFW x numUnitsFW]
+  auto bFW = INPUT_VARIABLE(3);   // biases for forward  RNN, [2*numUnitsFW]
+  auto WxBW = INPUT_VARIABLE(4);  // input-to-hidden  weights for backward RNN, [inSize  x numUnitsBW]
+  auto WhBW = INPUT_VARIABLE(5);  // hidden-to-hidden weights for backward RNN, [numUnitsBW x numUnitsBW]
+  auto bBW = INPUT_VARIABLE(6);   // biases for backward RNN, [2*v]
+
+  NDArray* h0FW = nullptr;  // initial cell output for forward  RNN (at time step = 0) [bS x numUnitsFW]
+  NDArray* h0BW = nullptr;  // initial cell output for backward RNN (at time step = 0) [bS x numUnitsBW]
+  NDArray* maxTimeStep =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+  // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  switch (block.width()) {
+    case 8:
+      maxTimeStep = INPUT_VARIABLE(7);
+      break;
+    case 9:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      break;
+    case 10:
+      h0FW = INPUT_VARIABLE(7);
+      h0BW = INPUT_VARIABLE(8);
+      maxTimeStep = INPUT_VARIABLE(9);
+      break;
+  }
+
+  auto h = OUTPUT_VARIABLE(0);  // cell outputs [time x bS x (numUnitsFW + numUnitsBW)], that is per each time step
+  auto hFWFinal = OUTPUT_VARIABLE(1);  // final cell out for forward  RNN [bS x numUnitsFW]
+  auto hBWFinal = OUTPUT_VARIABLE(2);  // final cell out for backward RNN [bS x numUnitsBF]
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input array must have rank = 3, but got %i instead !",
+               x->rankOf());
+  REQUIRE_TRUE(WxFW->rankOf() == 2, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for forward  RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxFW->rankOf());
+  REQUIRE_TRUE(WxBW->rankOf() == 2, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for backward RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxBW->rankOf());
+
+  const LongType inRank = x->rankOf();
+  const LongType time = x->sizeAt(0);
+  const LongType bS = x->sizeAt(1);
+  const LongType numUnitsFW = WxFW->sizeAt(1);
+  const LongType numUnitsBW = WxBW->sizeAt(1);
+
+  const std::vector<LongType> expectedWhFWshape = {numUnitsFW, numUnitsFW};
+  const std::vector<LongType> expectedWhBWshape = {numUnitsBW, numUnitsBW};
+  const std::vector<LongType> expectedbFWshape = {2 * numUnitsFW};
+  const std::vector<LongType> expectedbBWshape = {2 * numUnitsBW};
+
+  REQUIRE_TRUE(WhFW->isSameShape(expectedWhFWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for forward  "
+               "RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhFWshape).c_str(), ShapeUtils::shapeAsString(WhFW).c_str());
+  REQUIRE_TRUE(WhBW->isSameShape(expectedWhBWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for backward "
+               "RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhBWshape).c_str(), ShapeUtils::shapeAsString(WhBW).c_str());
+  REQUIRE_TRUE(bFW->isSameShape(expectedbFWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for forward  RNN), expected is "
+               "%s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbFWshape).c_str(), ShapeUtils::shapeAsString(bFW).c_str());
+  REQUIRE_TRUE(bBW->isSameShape(expectedbBWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for backward RNN), expected is "
+               "%s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbBWshape).c_str(), ShapeUtils::shapeAsString(bBW).c_str());
+  if (h0FW) {
+    const std::vector<LongType> expectedh0FWshape = {bS, numUnitsFW};
+    REQUIRE_TRUE(h0FW->isSameShape(expectedh0FWshape), 0,
+                 "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for forward  "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0FWshape).c_str(), ShapeUtils::shapeAsString(h0FW).c_str());
+  }
+  if (h0BW) {
+    const std::vector<LongType> expectedh0BWshape = {bS, numUnitsBW};
+    REQUIRE_TRUE(h0BW->isSameShape(expectedh0BWshape), 0,
+                 "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for backward "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0BWshape).c_str(), ShapeUtils::shapeAsString(h0BW).c_str());
+  }
+  if (maxTimeStep)
+  REQUIRE_TRUE(maxTimeStep->isSameShape({bS}), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of maxTimeStep array, expected is [%i], but "
+               "got %s instead !",
+               bS, ShapeUtils::shapeAsString(maxTimeStep).c_str());
+
+  // forward steps
+  std::vector<sd::LongType> expectedHshape = {time, bS, numUnitsFW + numUnitsBW};
+  auto hFW = new NDArray(x->ordering(),expectedHshape, x->dataType(), block.launchContext());
+  helpers::rnnTimeLoop(block.launchContext(), x, WxFW, WhFW, bFW, h0FW, maxTimeStep, hFW, hFWFinal);
+
+  auto seqLen = maxTimeStep;
+  if (seqLen == nullptr) {
+    std::vector<sd::LongType> seqShape = {x->sizeAt(1)};
+    seqLen = new NDArray(x->ordering(),seqShape, INT64, block.launchContext());  // [bS]
+    *seqLen = x->sizeAt(0);  // set each element of seqLen to be equal to time
+  }
+
+  // reverse x
+  auto revOut = new NDArray(x, false, block.launchContext());
+  helpers::reverseSequence(block.launchContext(), x, seqLen, revOut, 0, 1);
+
+  std::vector<sd::LongType> shape = {time, bS, numUnitsBW};
+  // backward steps
+  auto hBW = new NDArray(x->ordering(),shape, x->dataType(), block.launchContext());
+
+  helpers::rnnTimeLoop(block.launchContext(), revOut, WxBW, WhBW, bBW, h0BW, maxTimeStep, hBW, hBWFinal);
+
+  // reverse hBW
+  auto hBWcopy = new NDArray(*hBW);
+  helpers::reverseSequence(block.launchContext(), hBWcopy, seqLen, hBW, 0, 1);
+
+  // concatenate hFW and hBW along last third dimension
+  // NDArrayFactory<T>::concat({hFW, hBW}, 2, h);
+  helpers::concat(block.launchContext(), {hFW, hBW}, *h, 2);
+
+  delete hBW;
+  delete hFW;
+  delete hBWcopy;
+  delete revOut;
+
+  if (seqLen != maxTimeStep) delete seqLen;
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(static_bidirectional_rnn) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(static_bidirectional_rnn) {
+  auto xShapeInfo = inputShape->at(0);     // input [time x bS x inSize]
+  auto WxFWShapeInfo = inputShape->at(1);  // input-to-hidden  weights for forward  RNN, [inSize  x numUnitsFW]
+  auto WhFWShapeInfo = inputShape->at(2);  // hidden-to-hidden weights for forward  RNN, [numUnitsFW x numUnitsFW]
+  auto bFWShapeInfo = inputShape->at(3);   // biases for forward  RNN, [2*numUnitsFW]
+  auto WxBWShapeInfo = inputShape->at(4);  // input-to-hidden  weights for backward RNN, [inSize  x numUnitsBW]
+  auto WhBWShapeInfo = inputShape->at(5);  // hidden-to-hidden weights for backward RNN, [numUnitsBW x numUnitsBW]
+  auto bBWShapeInfo = inputShape->at(6);   // biases for backward RNN, [2*numUnitsBW]
+
+  LongType const* h0FWShapeInfo =
+      nullptr;  // initial cell output for forward  RNN (at time step = 0) [bS x numUnitsFW]
+  LongType const* h0BWShapeInfo =
+      nullptr;  // initial cell output for backward RNN (at time step = 0) [bS x numUnitsBW]
+  LongType const* maxTimeStepShapeInfo =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+  // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  switch (block.width()) {
+    case 8:
+      maxTimeStepShapeInfo = inputShape->at(7);
+      break;
+    case 9:
+      h0FWShapeInfo = inputShape->at(7);
+      h0BWShapeInfo = inputShape->at(8);
+      break;
+    case 10:
+      h0FWShapeInfo = inputShape->at(7);
+      h0BWShapeInfo = inputShape->at(8);
+      maxTimeStepShapeInfo = inputShape->at(9);
+      break;
+  }
+
+  REQUIRE_TRUE(xShapeInfo[0] == 3, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input array must have rank = 3, but got %i instead !",
+               xShapeInfo[0]);
+  REQUIRE_TRUE(WxFWShapeInfo[0] == 2, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for forward  RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxFWShapeInfo[0]);
+  REQUIRE_TRUE(WxBWShapeInfo[0] == 2, 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: input-to-hidden weights array (for backward RNN) must have "
+               "rank = 2, but got %i instead !",
+               WxBWShapeInfo[0]);
+
+  const int inRank = xShapeInfo[0];
+  const int time = xShapeInfo[1];
+  const int bS = xShapeInfo[2];
+  const int numUnitsFW = WxFWShapeInfo[2];
+  const int numUnitsBW = WxBWShapeInfo[2];
+
+  const std::vector<LongType> expectedWhFWshape = {numUnitsFW, numUnitsFW};
+  const std::vector<LongType> expectedWhBWshape = {numUnitsBW, numUnitsBW};
+  const std::vector<LongType> expectedbFWshape = {2 * numUnitsFW};
+  const std::vector<LongType> expectedbBWshape = {2 * numUnitsBW};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhFWShapeInfo, expectedWhFWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for forward  "
+               "RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhFWshape).c_str(), ShapeUtils::shapeAsString(WhFWShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhBWShapeInfo, expectedWhBWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of hidden-to-hidden weights array (for backward "
+               "RNN), expected is %s but got %s instead !",
+               ShapeUtils::shapeAsString(expectedWhBWshape).c_str(), ShapeUtils::shapeAsString(WhBWShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bFWShapeInfo, expectedbFWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for forward  RNN), expected is "
+               "%s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbFWshape).c_str(), ShapeUtils::shapeAsString(bFWShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bBWShapeInfo, expectedbBWshape), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of biases array (for backward RNN), expected is "
+               "%s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbBWshape).c_str(), ShapeUtils::shapeAsString(bBWShapeInfo).c_str());
+  if (h0FWShapeInfo) {
+    const std::vector<LongType> expectedh0FWshape = {bS, numUnitsFW};
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(h0FWShapeInfo, expectedh0FWshape), 0,
+                 "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for forward  "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0FWshape).c_str(),
+                 ShapeUtils::shapeAsString(h0FWShapeInfo).c_str());
+  }
+  if (h0BWShapeInfo) {
+    const std::vector<LongType> expectedh0BWshape = {bS, numUnitsBW};
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(h0BWShapeInfo, expectedh0BWshape), 0,
+                 "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of initial cell output array (for backward "
+                 "RNN), expected is %s but got %s instead !",
+                 ShapeUtils::shapeAsString(expectedh0BWshape).c_str(),
+                 ShapeUtils::shapeAsString(h0BWShapeInfo).c_str());
+  }
+  if (maxTimeStepShapeInfo)
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(maxTimeStepShapeInfo, {bS}), 0,
+               "STATIC_BIDIRECTIONAL_RNN custom operation: wrong shape of maxTimeStep array, expected is [%i], but "
+               "got %s instead !",
+               bS, ShapeUtils::shapeAsString(maxTimeStepShapeInfo).c_str());
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *hFWFinalPrevShapeInfo(nullptr), *hBWFinalPrevShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank), sd::LongType);
+  ALLOCATE(hFWFinalPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+  ALLOCATE(hBWFinalPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+
+  hShapeInfo[0] = inRank;
+  hFWFinalPrevShapeInfo[0] = hBWFinalPrevShapeInfo[0] = inRank - 1;
+  hShapeInfo[1] = time;
+  hShapeInfo[2] = hFWFinalPrevShapeInfo[1] = hBWFinalPrevShapeInfo[1] = bS;
+  hShapeInfo[3] = numUnitsFW + numUnitsBW;
+  hFWFinalPrevShapeInfo[2] = numUnitsFW;
+  hBWFinalPrevShapeInfo[2] = numUnitsBW;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, xShapeInfo, shape::order(xShapeInfo));
+  ShapeUtils::updateStridesAndType(hFWFinalPrevShapeInfo, xShapeInfo, shape::order(xShapeInfo));
+  ShapeUtils::updateStridesAndType(hBWFinalPrevShapeInfo, xShapeInfo, shape::order(xShapeInfo));
+
+  return SHAPELIST(CONSTANT(hShapeInfo), CONSTANT(hFWFinalPrevShapeInfo), CONSTANT(hBWFinalPrevShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+#endif

libnd4j/include/ops/declarable/generic/nn/recurrent/staticRNN.cpp

@@ -0,0 +1,169 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Yurii Shyrma, created on 02.04.2018
+//
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/rnn.h>
+#if NOT_EXCLUDED(OP_static_rnn)
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(static_rnn, 4, 2, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);   // input [time x bS x inSize]
+  auto Wx = INPUT_VARIABLE(1);  // input-to-hidden  weights, [inSize  x numUnits]
+  auto Wh = INPUT_VARIABLE(2);  // hidden-to-hidden weights, [numUnits x numUnits]
+  auto b = INPUT_VARIABLE(3);   // biases for, [2*numUnits]
+
+  NDArray* h0 = nullptr;  // initial cell output (at time step = 0) [bS x numUnits]
+  NDArray* maxTimeStep =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  if (block.width() == 5) {
+    if ((*INPUT_VARIABLE(4)).rankOf() == 2)
+      h0 = INPUT_VARIABLE(4);
+    else
+      maxTimeStep = INPUT_VARIABLE(4);
+  } else if (block.width() == 6) {
+    h0 = INPUT_VARIABLE(4);
+    maxTimeStep = INPUT_VARIABLE(5);
+  }
+
+  auto h = OUTPUT_VARIABLE(0);       // cell outputs [time x bS x numUnits]
+  auto hFinal = OUTPUT_VARIABLE(1);  // at the end it will store cell final non-zero output [bS x numUnits]
+
+  REQUIRE_TRUE(x->rankOf() == 3, 0,
+               "STATIC_RNN custom operation: input array x must have rank = 3, but got %i instead !", x->rankOf());
+  REQUIRE_TRUE(Wx->rankOf() == 2, 0,
+               "STATIC_RNN custom operation: input-to-hidden weights array must have rank = 2, but got %i instead !",
+               Wx->rankOf());
+
+  const int time = x->sizeAt(0);
+  const int bS = x->sizeAt(1);
+  const int inSize = x->sizeAt(2);
+  const int numUnits = Wx->sizeAt(1);
+
+  const std::vector<LongType> expectedWhShape = {numUnits, numUnits};
+  const std::vector<LongType> expectedbShape = {2 * numUnits};
+
+  REQUIRE_TRUE(Wh->isSameShape(expectedWhShape), 0,
+               "STATIC_RNN custom operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s "
+               "instead !",
+               ShapeUtils::shapeAsString(expectedWhShape).c_str(), ShapeUtils::shapeAsString(Wh).c_str());
+  REQUIRE_TRUE(b->isSameShape(expectedbShape), 0,
+               "STATIC_RNN custom operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
+  if (h0) {
+    const std::vector<LongType> expectedh0Shape = {bS, numUnits};
+    REQUIRE_TRUE(
+        h0->isSameShape(expectedh0Shape), 0,
+        "STATIC_RNN custom operation: wrong shape of initial cell output array, expected is %s but got %s instead !",
+        ShapeUtils::shapeAsString(expectedh0Shape).c_str(), ShapeUtils::shapeAsString(h0).c_str());
+  }
+  if (maxTimeStep)
+    REQUIRE_TRUE(maxTimeStep->isSameShape({bS}), 0,
+                 "STATIC_RNN custom operation: wrong shape of maxTimeStep array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS}).c_str(), ShapeUtils::shapeAsString(maxTimeStep).c_str());
+
+  helpers::rnnTimeLoop(block.launchContext(), x, Wx, Wh, b, h0, maxTimeStep, h, hFinal);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(static_rnn) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(static_rnn) {
+  auto xShapeInfo = inputShape->at(0);   // input [time x bS x inSize]
+  auto WxShapeInfo = inputShape->at(1);  // input-to-hidden  weights, [inSize  x numUnits]
+  auto WhShapeInfo = inputShape->at(2);  // hidden-to-hidden weights, [numUnits x numUnits]
+  auto bShapeInfo = inputShape->at(3);   // biases for, [2*numUnits]
+
+  const LongType* h0ShapeInfo = nullptr;  // initial cell output (at time step = 0) [bS x numUnits]
+  const LongType* maxTimeStepShapeInfo =
+      nullptr;  // vector [bS] containing integer values within [0,time), each element of this vector set max time step
+                // per each input in batch, this means there are no calculations for time >= maxTimeStep
+
+  if (block.width() == 5) {
+    if (inputShape->at(4)[0] == 2)
+      h0ShapeInfo = inputShape->at(4);
+    else
+      maxTimeStepShapeInfo = inputShape->at(4);
+  } else if (block.width() == 6) {
+    h0ShapeInfo = inputShape->at(4);
+    maxTimeStepShapeInfo = inputShape->at(5);
+  }
+
+  REQUIRE_TRUE(xShapeInfo[0] == 3, 0,
+               "STATIC_RNN custom operation: input array x must have rank = 3, but got %i instead !", xShapeInfo[0]);
+  REQUIRE_TRUE(WxShapeInfo[0] == 2, 0,
+               "STATIC_RNN custom operation: input-to-hidden weights array must have rank = 2, but got %i instead !",
+               WxShapeInfo[0]);
+
+  const int inRank = xShapeInfo[0];
+  const int time = xShapeInfo[1];
+  const int bS = xShapeInfo[2];
+  const int numUnits = WxShapeInfo[2];
+
+  const std::vector<LongType> expectedWhShape = {numUnits, numUnits};
+  const std::vector<LongType> expectedbShape = {2 * numUnits};
+
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(WhShapeInfo, expectedWhShape), 0,
+               "STATIC_RNN custom operation: wrong shape of hidden-to-hidden weights array, expected is %s, but got %s "
+               "instead !",
+               ShapeUtils::shapeAsString(expectedWhShape).c_str(), ShapeUtils::shapeAsString(WhShapeInfo).c_str());
+  REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, expectedbShape), 0,
+               "STATIC_RNN custom operation: wrong shape of biases array, expected is %s, but got %s instead !",
+               ShapeUtils::shapeAsString(expectedbShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
+  if (h0ShapeInfo) {
+    const std::vector<LongType> expectedh0Shape = {bS, numUnits};
+    REQUIRE_TRUE(
+        ShapeUtils::areShapesEqual(h0ShapeInfo, expectedh0Shape), 0,
+        "STATIC_RNN custom operation: wrong shape of initial cell output array, expected is %s but got %s instead !",
+        ShapeUtils::shapeAsString(expectedh0Shape).c_str(), ShapeUtils::shapeAsString(h0ShapeInfo).c_str());
+  }
+  if (maxTimeStepShapeInfo)
+    REQUIRE_TRUE(ShapeUtils::areShapesEqual(maxTimeStepShapeInfo, {bS}), 0,
+                 "STATIC_RNN custom operation: wrong shape of maxTimeStep array, expected is %s, but got %s instead !",
+                 ShapeUtils::shapeAsString({bS}).c_str(), ShapeUtils::shapeAsString(maxTimeStepShapeInfo).c_str());
+
+  // evaluate output shapeInfos
+  LongType *hShapeInfo(nullptr), *hPrevShapeInfo(nullptr);
+  ALLOCATE(hShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank), sd::LongType);
+  ALLOCATE(hPrevShapeInfo, block.getWorkspace(), shape::shapeInfoLength(inRank - 1), sd::LongType);
+
+  hShapeInfo[0] = inRank;
+  hPrevShapeInfo[0] = inRank - 1;
+  hShapeInfo[1] = time;
+  hShapeInfo[2] = hPrevShapeInfo[1] = bS;
+  hShapeInfo[3] = hPrevShapeInfo[2] = numUnits;
+
+  ShapeUtils::updateStridesAndType(hShapeInfo, xShapeInfo, shape::order(xShapeInfo));
+  ShapeUtils::updateStridesAndType(hPrevShapeInfo, xShapeInfo, shape::order(xShapeInfo));
+
+  return SHAPELIST(CONSTANT(hShapeInfo), CONSTANT(hPrevShapeInfo));
+}
+
+}  // namespace ops
+}  // namespace sd
+#endif

libnd4j/include/ops/declarable/generic/nn/relu_layer.cpp

@@ -0,0 +1,76 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author GS <sgazeos@gmail.com>
+//
+
+#include <ops/declarable/CustomOperations.h>
+#if NOT_EXCLUDED(OP_relu_layer)
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(relu_layer, 3, 1, false, 0, 0) {
+  auto x = INPUT_VARIABLE(0);
+  auto w = INPUT_VARIABLE(1);
+  auto b = INPUT_VARIABLE(2);
+
+  REQUIRE_TRUE(x->isMatrix(), 0, "relu_layer: x argument should be a 2D tensor, but got rank %i instead!", x->rankOf());
+  REQUIRE_TRUE(w->isMatrix(), 0, "relu_layer: weights argument should be a 2D tensor, but got rank %i instead!",
+               w->rankOf());
+  REQUIRE_TRUE(b->isVector(), 0, "relu_layer: biases argument should be a 1D tensor, but got rank %i instead!",
+               b->rankOf());
+  REQUIRE_TRUE(b->lengthOf() == w->sizeAt(1), 0,
+               "relu_layer: biases array length should match to columns of weights matrix, however got length = %i and "
+               "columns = %i!",
+               b->lengthOf(), w->sizeAt(1));
+  REQUIRE_TRUE(x->sizeAt(1) == w->sizeAt(0), 0,
+               "relu_layer: number of x columns should match to row number of weights matrix, but got x_columns = %i "
+               "and weights_rows = %i!",
+               x->sizeAt(1), w->sizeAt(0));
+
+  auto output = OUTPUT_VARIABLE(0);
+
+  sd::ops::xw_plus_b op;
+  auto status = op.execute({x, w, b}, {output});
+  REQUIRE_TRUE(sd::Status::OK == status, 0, "relu_layer: xw_plus_b op failed on input data.");
+
+  auto scalar = block.numT() > 0 ? block.getTArguments()->at(0) : 0.0;
+
+  output->applyScalar(sd::scalar::RELU, scalar, output);
+
+  return sd::Status::OK;
+}
+
+DECLARE_SHAPE_FN(relu_layer) {
+  auto inShape = inputShape->at(0);
+  auto weightsShape = inputShape->at(1);
+  auto outputShape = ShapeUtils::matrixProductShape(inShape, weightsShape, false, false,
+                                                    ArrayOptions::dataType(inShape), block.getWorkspace());
+
+  return SHAPELIST(outputShape);
+}
+
+DECLARE_TYPES(relu_layer) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(sd::DataType::ANY)
+      //                  ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedOutputTypes({ALL_FLOATS});
+}
+}  // namespace ops
+}  // namespace sd
+#endif

libnd4j/include/ops/declarable/generic/nn/softmax.cpp

@@ -0,0 +1,106 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 29/10/17
+// @author Yurii Shyrma (iuriish@yahoo.com)
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_softmax)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/activations.h>
+
+namespace sd {
+namespace ops {
+
+DECLARE_TYPES(softmax) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS})->setAllowedOutputTypes({ALL_FLOATS})->setSameMode(true);
+}
+
+CONFIGURABLE_OP_IMPL(softmax, 1, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto output = OUTPUT_VARIABLE(0);
+
+  const int rank = input->rankOf();
+  const int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
+
+  REQUIRE_TRUE(dim < rank, 0,
+               "SOFTMAX OP: the value of input integer parameter (dimension) must be less than input array rank %i, "
+               "but got dimension = %i instead !",
+               rank, dim);
+
+  helpers::softmax(block.launchContext(), input, output, dim);
+
+  return sd::Status::OK;
+}
+
+CONFIGURABLE_OP_IMPL(softmax_bp, 3, 1, true, 0, 0) {
+  auto input = INPUT_VARIABLE(0);
+  auto gradO = INPUT_VARIABLE(1);
+  auto softmaxedOut = INPUT_VARIABLE(2);
+  auto gradI = OUTPUT_VARIABLE(0);
+  
+  const int rank = input->rankOf();
+  const int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
+
+  REQUIRE_TRUE(dim < rank, 0,
+               "SOFTMAX_BP OP: the value of input integer parameter (dimension) must be less than input array rank %i, "
+               "but got dimension = %i instead !",
+               rank, dim);
+
+  // Refactored to minimize temporary allocations and ensure proper cleanup
+  gradI->assign(softmaxedOut);
+  
+  // Perform operations with proper pointer management
+  std::vector<sd::LongType> dimVector = {dim};
+  
+  // Compute gradI * gradO - returns pointer
+  auto* temp = (*gradI) * (*gradO);
+  
+  // Reduce along dimension - returns pointer
+  auto* sumAlongDim = temp->reduceAlongDimension(reduce::Sum, &dimVector, true);
+  delete temp;
+
+  // Compute gradO - sumAlongDim - returns pointer
+  auto* diff = (*gradO) - (*sumAlongDim);
+  // FIXED: Add defensive check - reduceAlongDimension may return view
+  if (sumAlongDim != nullptr && !sumAlongDim->isView()) {
+    delete sumAlongDim;
+  }
+  
+  // Compute final result: gradI * diff - returns pointer
+  auto* result = (*gradI) * (*diff);
+  delete diff;
+  
+  // Assign result to output
+  gradI->assign(result);
+  delete result;
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(softmax_bp) {
+  getOpDescriptor()->setAllowedInputTypes({ALL_FLOATS})->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/nn/xw_plus_b.cpp

@@ -0,0 +1,252 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  xw_plus_b op. Created by GS <george@skymind.io> 31.01.2018
+//  @author Oleg Semeniv <oleg.semeniv@gmail.com>
+//
+//  Fixed to handle higher-rank inputs (e.g., from ONNX Gemm with batched input)
+//  and corrected bias addition to always use row vector broadcasting.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_xw_plus_b)
+
+#include <helpers/MmulHelper.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/matmul.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(xw_plus_b, 3, 1, false, 0, 0) {
+  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
+  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
+  bool cTranspose = (block.getIArguments()->size() > 2 ? INT_ARG(2) == 1 : false);
+
+  auto x = INPUT_VARIABLE(0);
+  auto w = INPUT_VARIABLE(1);
+  auto b = INPUT_VARIABLE(2);
+  auto z = OUTPUT_VARIABLE(0);
+
+  if (x->isEmpty() || w->isEmpty() || b->isEmpty()) return Status::OK;
+
+  // Handle higher rank inputs by reshaping to 2D for matmul
+  // This supports inputs like [batch, 1, 1, hidden] from ONNX pooler operations
+  NDArray* xEffective = nullptr;
+  NDArray* wEffective = nullptr;
+  NDArray* zEffective = nullptr;
+  NDArray* bEffective = nullptr;
+
+  bool deleteX = false;
+  bool deleteW = false;
+  bool deleteZ = false;
+  bool deleteB = false;
+
+  sd::LongType batchSize = 1;
+  bool needsReshape = x->rankOf() > 2;
+
+  if (needsReshape) {
+    // Calculate the 2D shape: flatten all but last dimension
+    for (int i = 0; i < x->rankOf() - 1; i++) {
+      batchSize *= x->sizeAt(i);
+    }
+    sd::LongType inputLastDim = x->sizeAt(-1);
+    sd::LongType outputLastDim = z->sizeAt(-1);
+
+    // Reshape to 2D - these create new NDArray objects
+    std::vector<sd::LongType> xShape2D = {batchSize, inputLastDim};
+    std::vector<sd::LongType> zShape2D = {batchSize, outputLastDim};
+    xEffective = x->reshape('c', xShape2D);
+    zEffective = z->reshape('c', zShape2D);
+    deleteX = true;
+    deleteZ = true;
+  } else {
+    xEffective = x;
+    zEffective = z;
+  }
+
+  // Apply transposes if needed
+  if (aTranspose) {
+    auto xTransposed = new NDArray(xEffective->transpose());
+    if (deleteX) delete xEffective;
+    xEffective = xTransposed;
+    deleteX = true;
+  }
+  if (cTranspose) {
+    auto zTransposed = new NDArray(zEffective->transpose());
+    if (deleteZ) delete zEffective;
+    zEffective = zTransposed;
+    deleteZ = true;
+  }
+
+  // Handle weight transpose
+  if (bTranspose) {
+    wEffective = new NDArray(w->transpose());
+    deleteW = true;
+  } else {
+    wEffective = w;
+  }
+
+  REQUIRE_TRUE(xEffective->rankOf() == 2, 0,
+               "xw_plus_b: After reshaping, input x array should have rank equal 2, but got instead %i!",
+               xEffective->rankOf());
+  REQUIRE_TRUE(wEffective->rankOf() == 2, 0,
+               "xw_plus_b: Input weights array should have rank equal 2, but got instead %i!",
+               wEffective->rankOf());
+  REQUIRE_TRUE(zEffective->rankOf() == 2, 0,
+               "xw_plus_b: After reshaping, output array should have rank equal 2, but got instead %i!",
+               zEffective->rankOf());
+
+  // Perform matrix multiplication: z = x @ w
+  MmulHelper::mmul(xEffective, wEffective, zEffective, 1.0, 0.0);
+
+  // Add bias - ALWAYS as a row vector since output is [batch, features]
+  // The bias vector has shape [features] and should broadcast across the batch dimension
+  // This is the standard behavior for ONNX Gemm: Y = X @ W + B where B broadcasts
+  if (b->rankOf() == 1) {
+    std::vector<sd::LongType> bShape2D = {1, b->lengthOf()};
+    bEffective = b->reshape('c', bShape2D);
+    deleteB = true;
+  } else {
+    bEffective = b;
+  }
+
+  if (zEffective->isMatrix()) {
+    zEffective->addiRowVector(bEffective);
+  } else {
+    *zEffective += *bEffective;
+  }
+
+  // Cleanup heap-allocated arrays
+  if (deleteB) delete bEffective;
+  if (deleteW) delete wEffective;
+  if (deleteX) delete xEffective;
+  if (deleteZ) delete zEffective;
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(xw_plus_b) {
+  auto xShape = inputShape->at(0);
+  auto weights = INPUT_VARIABLE(1);
+  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
+  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
+  bool cTranspose = (block.getIArguments()->size() > 2 ? INT_ARG(2) == 1 : false);
+
+  int nWeightsFormat = block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;
+
+  auto weightsShape =
+      (1 == nWeightsFormat) ? ShapeUtils::evalTransposeShapeInfo(*weights, block.getWorkspace()) : inputShape->at(1);
+
+  // Handle higher rank inputs
+  if (shape::rank(xShape) > 2) {
+    // Calculate 2D shapes for matmul
+    sd::LongType batchSize = 1;
+    for (int i = 0; i < shape::rank(xShape) - 1; i++) {
+      batchSize *= shape::sizeAt(xShape, i);
+    }
+    sd::LongType lastDim = shape::sizeAt(xShape, shape::rank(xShape) - 1);
+    
+    // Create temporary 2D shape for x
+    std::vector<sd::LongType> x2dShape = {batchSize, lastDim};
+    auto x2dShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(xShape), 
+                                                                         'c', x2dShape);
+    
+    // Get the output shape from matmul
+    auto matmulOutput = ShapeUtils::matrixProductShape(x2dShapeInfo, const_cast<sd::LongType *>(weightsShape), 
+                                                       aTranspose, bTranspose,
+                                                       ArrayOptions::dataType(xShape), block.getWorkspace());
+    
+    // Calculate final output shape
+    std::vector<sd::LongType> outputShape;
+    for (int i = 0; i < shape::rank(xShape) - 1; i++) {
+      outputShape.push_back(shape::sizeAt(xShape, i));
+    }
+    // Add the output dimension from the weights
+    outputShape.push_back(shape::sizeAt(matmulOutput, 1));
+    
+    auto finalShape = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(xShape), 
+                                                                         'c', outputShape);
+    return SHAPELIST(finalShape);
+  } else {
+    // Original behavior for rank 2 inputs
+    auto outputShape = ShapeUtils::matrixProductShape(xShape, const_cast<sd::LongType *>(weightsShape), aTranspose,
+                                                      bTranspose,
+                                                      ArrayOptions::dataType(xShape), block.getWorkspace());
+    return SHAPELIST(outputShape);
+  }
+}
+
+DECLARE_TYPES(xw_plus_b) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+
+CUSTOM_OP_IMPL(xw_plus_b_bp, 4, 3, false, 0, 0) {
+
+  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
+  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
+  auto x = aTranspose ? INPUT_VARIABLE(0)->transpose() : INPUT_VARIABLE(0);  // transpose() already returns NDArray*
+  auto b = INPUT_VARIABLE(2);
+  auto dLdz = INPUT_VARIABLE(3);
+
+  if (x->isEmpty() || INPUT_VARIABLE(1)->isEmpty() || b->isEmpty() || dLdz->isEmpty()) return Status::OK;
+
+  auto w = bTranspose ? INPUT_VARIABLE(1)->transpose() : INPUT_VARIABLE(1);  // transpose() already returns NDArray*
+
+  REQUIRE_TRUE(x->rankOf() == 2, 0, "xw_plus_b BP: Input x array should have rank equal 2, but got instead %i!",
+               x->rankOf());
+  REQUIRE_TRUE(w->rankOf() == 2, 0, "xw_plus_b BP: Input weights array should have rank equal 2, but got instead %i!",
+               w->rankOf());
+  REQUIRE_TRUE(dLdz->rankOf() == 2, 0, "xw_plus_b BP: Output array should have rank equal 2, but got instead %i!",
+               dLdz->rankOf());
+
+  auto dLdx = aTranspose ? OUTPUT_VARIABLE(0)->transpose() : OUTPUT_VARIABLE(0);  // transpose() already returns NDArray*
+  auto dLdb = OUTPUT_VARIABLE(2);
+
+  auto dLdw = (bTranspose) ? OUTPUT_VARIABLE(1)->transpose() : OUTPUT_VARIABLE(1);  // transpose() already returns NDArray*
+
+  // dLdb - reduceAlongDimension returns pointer
+  std::vector<LongType> dims({0});
+  auto* assign = dLdz->reduceAlongDimension(reduce::Sum, &dims);
+  dLdb->assign(assign);
+  delete assign;
+
+  matmul_bp mmul_bp;
+  mmul_bp.execute({x, w, dLdz}, std::vector<NDArray*>{dLdx, dLdw}, {}, {}, {});
+
+  // Transpose views are managed by parent arrays - no deletion needed
+  // x is from INPUT_VARIABLE(0)->transpose() if aTranspose
+  // w is from INPUT_VARIABLE(1)->transpose() if bTranspose
+  // dLdx is from OUTPUT_VARIABLE(0)->transpose() if aTranspose
+  // dLdw is from OUTPUT_VARIABLE(1)->transpose() if bTranspose
+  // All are views managed by their parent arrays
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(xw_plus_b_bp) {
+  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)));
+}
+
+DECLARE_TYPES(xw_plus_b_bp) {
+  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif