chore: import upstream snapshot with attribution

libnd4j/include/ops/declarable/generic/boolean/boolean_not.cpp

@@ -0,0 +1,45 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver on 6/6/2018.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_boolean_not)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+OP_IMPL(boolean_not, 1, 1, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto z = OUTPUT_VARIABLE(0);
+
+  x->applyTransform(transform::Not, z);
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(boolean_not) {
+  getOpDescriptor()->setAllowedInputTypes(0, BOOL)->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/choose.cpp

@@ -0,0 +1,87 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author Adam Gibson
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_choose)
+
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/choose.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(choose, -1, 2, false, -2, -1) {
+  int mode = INT_ARG(0);
+  auto result = OUTPUT_VARIABLE(0);
+  auto numResults = OUTPUT_VARIABLE(1);
+
+  if (block.width() > 1) {
+    auto arg = INPUT_VARIABLE(0);
+    auto comp = INPUT_VARIABLE(1);
+
+    helpers::chooseFunctorArray(block.launchContext(), arg, comp, mode, result, numResults);
+
+  }  // scalar case
+  else {
+    double scalar = T_ARG(0);
+    auto arg = INPUT_VARIABLE(0);
+    helpers::chooseFunctorScalar(block.launchContext(), arg, scalar, mode, result, numResults);
+  }
+
+  return Status::OK;
+}
+
+DECLARE_TYPES(choose) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, {ALL_FLOATS})
+      ->setAllowedInputTypes(1, {ALL_FLOATS})
+      ->setAllowedOutputTypes(0, {ALL_FLOATS})
+      ->setAllowedOutputTypes(1, {ALL_INTS});
+}
+
+DECLARE_SHAPE_FN(choose) {
+  int mode = INT_ARG(0);
+  auto numResults = NDArrayFactory::create<LongType>(0L);
+  if (block.width() > 1) {
+    auto first = INPUT_VARIABLE(0);
+    auto second = INPUT_VARIABLE(1);
+
+
+    helpers::chooseFunctorArray(block.launchContext(), first, second, mode, nullptr, numResults);
+  } else {
+    auto first = INPUT_VARIABLE(0);
+    double scalar = T_ARG(0);
+
+    helpers::chooseFunctorScalar(block.launchContext(), first, scalar, mode, nullptr, numResults);
+  }
+
+  auto newShape = ConstantShapeHelper::getInstance().vectorShapeInfo(numResults->e<LongType>(0),
+                                                                     ArrayOptions::dataType(inputShape->at(0)));
+
+  auto shapeScalar = ConstantShapeHelper::getInstance().scalarShapeInfo(INT64);
+  return SHAPELIST(newShape, shapeScalar);
+}
+
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/eq_scalar.cpp

@@ -0,0 +1,51 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_eq_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(eq_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<double>(0) == y->e<double>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+DECLARE_SYN(Equals, eq_scalar);
+// DECLARE_SYN(equals, eq_scalar);
+
+DECLARE_TYPES(eq_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/gt_scalar.cpp

@@ -0,0 +1,51 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_gt_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(gt_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<float>(0) > y->e<float>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+// DECLARE_SYN(Greater, gt_scalar);
+// DECLARE_SYN(greater, gt_scalar);
+
+DECLARE_TYPES(gt_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/gte_scalar.cpp

@@ -0,0 +1,51 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_gte_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(gte_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<float>(0) >= y->e<float>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+DECLARE_SYN(GreaterOrEquals, gte_scalar);
+DECLARE_SYN(greaterOrEquals, gte_scalar);
+
+DECLARE_TYPES(gte_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/is_non_decreasing.cpp

@@ -0,0 +1,53 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author @cpuheater
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_is_non_decreasing)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/compare_elem.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(is_non_decreasing, 1, true) {
+  auto input = INPUT_VARIABLE(0);
+
+  // in case of empty input there's nothing to do
+  if (input->isEmpty()) return Status::EQ_TRUE;
+
+  bool isNonDecreasing = true;
+
+  helpers::compare_elem(block.launchContext(), input, false, isNonDecreasing);
+
+  if (isNonDecreasing)
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+
+DECLARE_TYPES(is_non_decreasing) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/is_numeric_tensor.cpp

@@ -0,0 +1,43 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, 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 @cpuheater
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_is_numeric_tensor)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/compare_elem.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(is_numeric_tensor, 1, true) {
+  auto input = INPUT_VARIABLE(0);
+
+  return input->isR() || input->isZ() ? Status::EQ_TRUE : Status::EQ_FALSE;
+}
+
+DECLARE_TYPES(is_numeric_tensor) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/is_strictly_increasing.cpp

@@ -0,0 +1,53 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author @cpuheater
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_is_strictly_increasing)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/compare_elem.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(is_strictly_increasing, 1, true) {
+  auto input = INPUT_VARIABLE(0);
+
+  // in case of empty input there's nothing to do
+  if (input->isEmpty()) return Status::EQ_TRUE;
+
+  bool isStrictlyIncreasing = true;
+
+  helpers::compare_elem(block.launchContext(), input, true, isStrictlyIncreasing);
+
+  if (isStrictlyIncreasing)
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+
+DECLARE_TYPES(is_strictly_increasing) {
+  getOpDescriptor()->setAllowedInputTypes(0, ANY)->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/lt_scalar.cpp

@@ -0,0 +1,50 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lt_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(lt_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<float>(0) < y->e<float>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+
+
+DECLARE_TYPES(lt_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/lte_scalar.cpp

@@ -0,0 +1,51 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_lte_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(lte_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<float>(0) <= y->e<float>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+DECLARE_SYN(LessOrEquals, lte_scalar);
+DECLARE_SYN(lessorequals, lte_scalar);
+
+DECLARE_TYPES(lte_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/neq_scalar.cpp

@@ -0,0 +1,51 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// Created by raver119 on 13.10.2017.
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_neq_scalar)
+
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+BOOLEAN_OP_IMPL(neq_scalar, 2, true) {
+  auto x = INPUT_VARIABLE(0);
+  auto y = INPUT_VARIABLE(1);
+
+  if (x->e<float>(0) != y->e<float>(0))
+    return Status::EQ_TRUE;
+  else
+    return Status::EQ_FALSE;
+}
+DECLARE_SYN(NotEquals, neq_scalar);
+DECLARE_SYN(notequals, neq_scalar);
+
+DECLARE_TYPES(neq_scalar) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, ANY)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedOutputTypes(0, BOOL);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/select.cpp

@@ -0,0 +1,127 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author raver119@gmail.com
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_select)
+
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(select, 3, 1, false, 0, 0) {
+  auto cond = INPUT_VARIABLE(0);
+  auto x = INPUT_VARIABLE(1);
+  auto y = INPUT_VARIABLE(2);
+  //TODO: for some reason y being empty
+  //should not necessarily yield an empty result
+  //the loss test I'm currently dealing with seems
+  //to need to output a value yet it ends up being empty.
+  //we need to figure out why
+  if(x->isEmpty() || y->isEmpty() || cond->isEmpty()) {
+    return Status::OK;
+  }
+  if (x->isScalar()) {
+    REQUIRE_TRUE(cond->isScalar(), 0,
+                 "Select: Condition should gave either equal shape to X/Y first dimension or to be scalar");
+
+    auto z = OUTPUT_VARIABLE(0);
+
+    if (y->isR()) {
+#ifdef HAS_DOUBLE
+      auto v = !cond->e<bool>(0) ? y->e<double>(0) : x->e<double>(0);
+#elif defined(HAS_FLOAT32)
+      auto v = !cond->e<bool>(0) ? y->e<float>(0) : x->e<float>(0);
+#else
+#error "No floating-point type available for select operation"
+#endif
+      z->p(0, v);
+    } else {
+      auto v = !cond->e<bool>(0) ? y->e<LongType>(0) : x->e<LongType>(0);
+      z->p(0, v);
+    }
+  } else {
+    bool same = cond->isSameShape(x);
+    REQUIRE_TRUE(cond->isScalar() || cond->lengthOf() == x->sizeAt(0) || same, 0,
+                 "Select: Condition should gave either equal shape to X/Y first dimension or to be scalar");
+    if (same) {
+      auto z = OUTPUT_VARIABLE(0);
+
+      for (int e = 0; e < cond->lengthOf(); e++) {
+        if (y->isR()) {
+#ifdef HAS_DOUBLE
+          auto r = !cond->e<bool>(e) ? y->e<double>(e) : x->e<double>(e);
+#elif defined(HAS_FLOAT32)
+          auto r = !cond->e<bool>(e) ? y->e<float>(e) : x->e<float>(e);
+#else
+#error "No floating-point type available for select operation"
+#endif
+          z->p(e, r);
+        } else {
+          auto r = !cond->e<bool>(e) ? y->e<LongType>(e) : x->e<LongType>(e);
+          z->p(e, r);
+        }
+      }
+    } else {
+      REQUIRE_TRUE(cond->lengthOf() == x->sizeAt(0), 0,
+                   "Condition length should be equal to the dim0 of x/y to act as TAD-mask, but got %d instead",
+                   cond->lengthOf());
+
+      auto z = OUTPUT_VARIABLE(0);
+      std::vector<LongType> idxs;
+      idxs.push_back(0);
+      auto dims = ShapeUtils::evalDimsToExclude(x->rankOf() ,1,idxs.data());
+      auto tadsX = x->allTensorsAlongDimension(*dims);
+      auto tadsY = y->allTensorsAlongDimension(*dims);
+      auto tadsZ = z->allTensorsAlongDimension(*dims);
+
+      for (int e = 0; e < tadsX.size(); e++) {
+        if (!cond->e<bool>(e)) {
+          tadsZ.at(e)->assign(tadsY.at(e));
+        } else {
+          tadsZ.at(e)->assign(tadsX.at(e));
+        }
+      }
+
+      delete dims;
+    }
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(select) {
+  auto inShape = inputShape->at(1);
+  return SHAPELIST(CONSTANT(inShape));
+}
+
+DECLARE_TYPES(select) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, BOOL)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedInputTypes(2, ANY)
+      ->setAllowedOutputTypes(1, INHERIT);
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/where.cpp

@@ -0,0 +1,339 @@
+/* ******************************************************************************
+*
+*
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://www.apache.org/licenses/LICENSE-2.0.
+*
+*  See the NOTICE file distributed with this work for additional
+*  information regarding copyright ownership.
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+* WARRANTIES OR CONDITIONS OF ANY KIND, 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_Where)
+
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/where.h>
+
+namespace sd {
+namespace ops {
+
+// Helper function to evaluate condition regardless of underlying data type
+inline bool evaluateCondition(NDArray* condition, int index) {
+ switch(condition->dataType()) {
+#if defined(HAS_BOOL)
+   case DataType::BOOL:
+     return condition->e<bool>(index);
+#endif
+#if defined(HAS_INT8)
+   case DataType::INT8:
+     return condition->e<int8_t>(index) != 0;
+#endif
+#if defined(HAS_INT16)
+   case DataType::INT16:
+     return condition->e<int16_t>(index) != 0;
+#endif
+#if defined(HAS_INT32)
+   case DataType::INT32:
+     return condition->e<int32_t>(index) != 0;
+#endif
+#if defined(HAS_LONG)
+   case DataType::INT64:
+     return condition->e<sd::LongType>(index) != 0;
+#endif
+#if defined(HAS_UINT8)
+   case DataType::UINT8:
+     return condition->e<uint8_t>(index) != 0;
+#endif
+#if defined(HAS_UINT16)
+   case DataType::UINT16:
+     return condition->e<uint16_t>(index) != 0;
+#endif
+#if defined(HAS_UINT32)
+   case DataType::UINT32:
+     return condition->e<uint32_t>(index) != 0;
+#endif
+#if defined(HAS_UNSIGNEDLONG)
+   case DataType::UINT64:
+     return condition->e<sd::UnsignedLong>(index) != 0;
+#endif
+#if defined(HAS_FLOAT16)
+   case DataType::HALF:
+     return condition->e<float16>(index) != static_cast<float16>(0.0f);
+#endif
+#if defined(HAS_BFLOAT16)
+   case DataType::BFLOAT16:
+     return condition->e<bfloat16>(index) != static_cast<bfloat16>(0.0f);
+#endif
+#if defined(HAS_FLOAT32)
+   case DataType::FLOAT32:
+     return condition->e<float>(index) != 0.0f;
+#endif
+#if defined(HAS_DOUBLE)
+   case DataType::DOUBLE:
+     return condition->e<double>(index) != 0.0;
+#endif
+   default:
+     // Fallback: try to interpret as int32 and check if non-zero
+#if defined(HAS_INT32)
+#ifdef __cpp_exceptions
+     try {
+       return condition->e<int32_t>(index) != 0;
+     } catch (...) {
+       // Last resort: assume false to maintain safe behavior
+       return false;
+     }
+#else
+     return condition->e<int32_t>(index) != 0;
+#endif
+#else
+     // If INT32 is not available, return false as safe default
+     return false;
+#endif
+ }
+}
+
+// Helper function to perform element-wise where with proper broadcasting
+void performBroadcastedWhere(NDArray* condition, NDArray* x, NDArray* y, NDArray* z) {
+ // We'll process each element of the output array z
+ // and determine the appropriate indices for condition, x, and y based on broadcasting rules
+
+ auto* zShape = z->getShapeAsVector();
+ auto* condShape = condition->getShapeAsVector();
+ auto* xShape = x->getShapeAsVector();
+ auto* yShape = y->getShapeAsVector();
+
+ // For each element in the output array
+ for (LongType i = 0; i < z->lengthOf(); i++) {
+   // Convert linear index to multi-dimensional indices for output array
+   std::vector<LongType> zIndices(z->rankOf());
+   LongType remainder = i;
+   for (int dim = z->rankOf() - 1; dim >= 0; dim--) {
+     zIndices[dim] = remainder % z->sizeAt(dim);
+     remainder /= z->sizeAt(dim);
+   }
+
+   // Calculate corresponding indices in condition, x, and y arrays using broadcasting rules
+   auto getLinearIndex = [](const std::vector<LongType>& multiIndices, const std::vector<LongType>& shape, NDArray* array) -> LongType {
+     LongType linearIndex = 0;
+     LongType stride = 1;
+     int srcDim = shape.size() - 1;
+
+     for (int dim = multiIndices.size() - 1; dim >= 0; dim--) {
+       LongType srcIndex = 0;
+       if (srcDim >= 0) {
+         if (shape[srcDim] == 1) {
+           srcIndex = 0; // Broadcast dimension
+         } else {
+           srcIndex = multiIndices[dim];
+         }
+         srcDim--;
+       }
+       linearIndex += srcIndex * stride;
+       if (srcDim >= 0) {
+         stride *= shape[srcDim + 1];
+       }
+     }
+     return linearIndex;
+   };
+
+   LongType condIndex = condition->lengthOf() == 1 ? 0 : getLinearIndex(zIndices, *condShape, condition);
+   LongType xIndex = x->lengthOf() == 1 ? 0 : getLinearIndex(zIndices, *xShape, x);
+   LongType yIndex = y->lengthOf() == 1 ? 0 : getLinearIndex(zIndices, *yShape, y);
+
+   // Apply the where logic
+   if (z->isR()) {
+#ifdef HAS_DOUBLE
+     auto result = evaluateCondition(condition, condIndex) ?
+                                                           x->e<double>(xIndex) : y->e<double>(yIndex);
+#elif defined(HAS_FLOAT32)
+     auto result = evaluateCondition(condition, condIndex) ?
+                                                           x->e<float>(xIndex) : y->e<float>(yIndex);
+#else
+#error "No floating-point type available for where operation"
+#endif
+     z->p(i, result);
+   } else{
+     auto result = evaluateCondition(condition, condIndex) ?
+                                                           x->e<LongType>(xIndex) : y->e<LongType>(yIndex);
+     z->p(i, result);
+   }
+ }
+
+ delete zShape;
+ delete condShape;
+ delete xShape;
+ delete yShape;
+}
+
+CUSTOM_OP_IMPL(Where, 1, 1, false, 0, 0) {
+ auto condition = INPUT_VARIABLE(0);
+ auto z = OUTPUT_VARIABLE(0);
+ if (z->isEmpty()) return Status::OK;
+
+ if (block.width() == 3) {
+   auto x = INPUT_VARIABLE(1);
+   auto y = INPUT_VARIABLE(2);
+
+   // Check if x and y can be broadcast together (instead of requiring exact same shape)
+   REQUIRE_TRUE(x->isSameShape(y) || ShapeUtils::areShapesBroadcastable(*x, *y), 0,
+                "X and Y must have equal shapes or be broadcastable. X shape: %s, Y shape: %s",
+                ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str());
+
+   // Case 1: All arrays have exact shape matching (element-wise operation)
+   if (condition->isSameShape(x) && x->isSameShape(y)) {
+     // FIXME: for perf it might be better to issue memcpy here, and fill only mismatched values from either X or Y
+     for (int e = 0; e < condition->lengthOf(); e++) {
+       if (z->isR()) {
+#ifdef HAS_DOUBLE
+         auto r = !evaluateCondition(condition, e) ? y->e<double>(e) : x->e<double>(e);
+#elif defined(HAS_FLOAT32)
+         auto r = !evaluateCondition(condition, e) ? y->e<float>(e) : x->e<float>(e);
+#else
+#error "No floating-point type available for where operation"
+#endif
+         z->p(e, r);
+       } else {
+         auto r = !evaluateCondition(condition, e) ? y->e<LongType>(e) : x->e<LongType>(e);
+         z->p(e, r);
+       }
+     }
+   }
+   // Case 2: Broadcasting is possible (most flexible case)
+   else if (ShapeUtils::areShapesBroadcastable(*condition, *x) &&
+            ShapeUtils::areShapesBroadcastable(*condition, *y) &&
+            ShapeUtils::areShapesBroadcastable(*x, *y)) {
+     performBroadcastedWhere(condition, x, y, z);
+   }
+   // Case 3: TAD-mask operation (legacy behavior for specific cases)
+   else if (condition->rankOf() == 1 && condition->lengthOf() == x->sizeAt(0)) {
+     std::vector<LongType> zero({0});
+     auto dims = ShapeUtils::evalDimsToExclude(x->rankOf(), 1, zero.data());
+     auto tadsX = x->allTensorsAlongDimension(*dims);
+     auto tadsY = y->allTensorsAlongDimension(*dims);
+     auto tadsZ = z->allTensorsAlongDimension(*dims);
+
+     for (int e = 0; e < tadsX.size(); e++) {
+       if (!evaluateCondition(condition, e)) {
+         tadsZ.at(e)->assign(tadsY.at(e));
+       } else {
+         tadsZ.at(e)->assign(tadsX.at(e));
+       }
+     }
+
+     delete dims;
+   }
+   // Case 4: Invalid shapes - provide detailed error message
+   else {
+     std::string condShape = ShapeUtils::shapeAsString(condition);
+     std::string xShape = ShapeUtils::shapeAsString(x);
+     std::string yShape = ShapeUtils::shapeAsString(y);
+
+     REQUIRE_TRUE(false, 0,
+                  "Where operation: Invalid shapes for broadcasting. "
+                  "Condition shape: %s, X shape: %s, Y shape: %s. "
+                  "Condition must either: (1) match X/Y shapes exactly, "
+                  "(2) be broadcastable with X/Y shapes, or "
+                  "(3) be 1D with length equal to first dimension of X/Y for TAD-mask operation.",
+                  condShape.c_str(), xShape.c_str(), yShape.c_str());
+   }
+ } else {
+   // in this case we return 2D matrix, which basically contains coordinates fo true
+   REQUIRE_TRUE(block.width() == 1, 0, "Where op takes either 1 or 3 operands, But got %d operands instead",
+                block.width());
+   auto output = OUTPUT_VARIABLE(0);
+   std::vector<LongType> zero({0});
+
+   int width = condition->rankOf();
+   if (z->isEmpty()) return Status::OK;
+
+   std::vector<LongType> *dims = ShapeUtils::evalDimsToExclude(width,1,zero.data());
+
+   helpers::_where(block.launchContext(), *condition, *output, block.workspace());
+   delete dims;
+ }
+ return Status::OK;
+}
+
+DECLARE_SHAPE_FN(Where) {
+ if (block.width() == 3) {
+   auto x = INPUT_VARIABLE(1);
+   auto y = INPUT_VARIABLE(2);
+
+   // Calculate the broadcast result shape for x and y
+   LongType* resultShapeInfo = nullptr;
+   bool canBroadcast = ShapeUtils::evalBroadcastShapeInfo(*x, *y, true, resultShapeInfo, block.getWorkspace());
+
+   if (canBroadcast && resultShapeInfo != nullptr) {
+     return SHAPELIST(CONSTANT(resultShapeInfo));
+   } else {
+     // Fallback to x's shape if broadcasting fails (should have been caught in validation)
+     auto inShape = inputShape->at(1);
+     return SHAPELIST(CONSTANT(inShape));
+   }
+ } else {
+   // FIXME: we can't estimate result here in this case
+   // output shape is the 2D tensor num_true x rankOf (inShape)
+   auto condition = INPUT_VARIABLE(0);
+   auto inShape = inputShape->at(0);
+   LongType numOfTrue = 0;  // condition->reduceNumber(reduce::CountNonZero, nullptr).e<sd::LongType>(0);
+   for (LongType i = 0; i < condition->lengthOf(); i++)
+     if (evaluateCondition(condition, i)) numOfTrue++;
+
+   LongType * theNewShape;
+   if (numOfTrue > 0) {
+     LongType* newShape;
+     ALLOCATE(newShape, block.getWorkspace(), shape::shapeInfoLength(2), sd::LongType);
+     newShape[0] = 2;
+     newShape[1] = numOfTrue;
+     newShape[2] = shape::rank(inShape);
+     newShape[3] = 1;
+     newShape[4] = 1;
+     newShape[5] = 0;
+     newShape[6] = 1;
+     newShape[7] = 99;
+#if defined(HAS_LONG)
+     ShapeUtils::updateStridesAndType(newShape, INT64, 'c');
+#else
+     // Fallback to INT32 if INT64 is not available
+     ShapeUtils::updateStridesAndType(newShape, INT32, 'c');
+#endif
+
+     theNewShape = CONSTANT(newShape);
+     RELEASE(newShape, block.getWorkspace());
+   } else {
+#if defined(HAS_LONG)
+     theNewShape = ConstantShapeHelper::getInstance().emptyShapeInfo(INT64);
+#else
+     // Fallback to INT32 if INT64 is not available
+     theNewShape = ConstantShapeHelper::getInstance().emptyShapeInfo(INT32);
+#endif
+   }
+
+   return SHAPELIST(theNewShape);
+ }
+}
+
+DECLARE_TYPES(Where) {
+ getOpDescriptor()
+     ->setAllowedInputTypes(0, ANY)  // bool
+     ->setAllowedInputTypes(1, ANY)
+     ->setAllowedInputTypes(2, ANY)
+     ->setAllowedOutputTypes(0, {ALL_INTS, ALL_FLOATS,BOOL});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif

libnd4j/include/ops/declarable/generic/boolean/where_np.cpp

@@ -0,0 +1,179 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+//  @author Adam Gibson
+//
+
+#include <ops/declarable/headers/boolean.h>
+#include <system/op_boilerplate.h>
+
+#if NOT_EXCLUDED(OP_where_np)
+
+#include <helpers/ShapeUtils.h>
+#include <ops/declarable/CustomOperations.h>
+
+namespace sd {
+namespace ops {
+CUSTOM_OP_IMPL(where_np, -1, 1, false, 0, 0) {
+  auto condition = INPUT_VARIABLE(0);
+
+  if (block.width() == 3) {
+    auto x = INPUT_VARIABLE(1);
+    auto y = INPUT_VARIABLE(2);
+
+    auto z = OUTPUT_VARIABLE(0);
+    int numMatches = 0;
+    // if cond matches x/y shape - we have per-element mask
+    if (condition->isSameShape(x)) {
+      // FIXME: for perf it might be better to issue memcpy here, and fill only mismatched values from either X or Y
+      if (y->isScalar()) {
+        if (y->isR()) {
+          for (int e = 0; e < condition->lengthOf(); e++) {
+#ifdef HAS_DOUBLE
+            auto r = condition->e<bool>(e) ? y->e<double>(0) : x->e<double>(e);
+#elif defined(HAS_FLOAT32)
+            auto r = condition->e<bool>(e) ? y->e<float>(0) : x->e<float>(e);
+#else
+#error "No floating-point type available for where_np operation"
+#endif
+            z->p(e, r);
+          }
+        } else{
+          for (int e = 0; e < condition->lengthOf(); e++) {
+            auto r = condition->e<bool>(e) ? y->e<LongType>(0) : x->e<LongType>(e);
+            z->p(e, r);
+          }
+        }
+      } else {
+        if (y->isR()) {
+          for (int e = 0; e < condition->lengthOf(); e++) {
+            if (condition->e<bool>(e)) {
+#ifdef HAS_DOUBLE
+              auto r = y->e<double>(numMatches);
+#elif defined(HAS_FLOAT32)
+              auto r = y->e<float>(numMatches);
+#else
+#error "No floating-point type available for where_np operation"
+#endif
+              z->p(e, r);
+              numMatches++;
+            } else {
+#ifdef HAS_DOUBLE
+              auto r = x->e<double>(e);
+#elif defined(HAS_FLOAT32)
+              auto r = x->e<float>(e);
+#else
+#error "No floating-point type available for where_np operation"
+#endif
+              z->p(e, r);
+            }
+          }
+        } else {
+          for (int e = 0; e < condition->lengthOf(); e++) {
+            if (condition->e<bool>(e)) {
+              auto r = y->e<LongType>(numMatches);
+              z->p(e, r);
+              numMatches++;
+            } else {
+              auto r = x->e<LongType>(e);
+              z->p(e, r);
+            }
+          }
+        }
+      }
+    } else {
+      REQUIRE_TRUE(condition->lengthOf() == x->sizeAt(0), 0,
+                   "Condition length should be equal to the dim0 of x/y to act as TAD-mask, but got %d instead",
+                   condition->lengthOf());
+
+      std::vector<LongType> idxs;
+      idxs.push_back(0);
+      auto dims = ShapeUtils::evalDimsToExclude(x->rankOf(), 1,idxs.data());
+      auto tadsX = x->allTensorsAlongDimension(*dims);
+      auto tadsY = y->allTensorsAlongDimension(*dims);
+      auto tadsZ = z->allTensorsAlongDimension(*dims);
+
+      for (int e = 0; e < tadsX.size(); e++) {
+        if (!condition->e<bool>(e))
+          tadsZ.at(e)->assign(tadsY.at(e));
+        else
+          tadsZ.at(e)->assign(tadsX.at(e));
+      }
+
+      delete dims;
+    }
+  } else {
+    // in this case we return 2D matrix, which basically contains coordinates fo true
+
+    REQUIRE_TRUE(block.width() == 1, 0, "Where op takes either 1 or 3 operands, But got %d operands instead",
+                 block.width());
+    LongType width = condition->rankOf();
+
+    Where op;
+    auto res(op.evaluate({condition}));
+    REQUIRE_OK(res.status());
+    NDArray* whereTrue = res.at(0);
+
+    if (whereTrue->isEmpty()) return Status::OK;
+    for (LongType outNext = 0; outNext < width; ++outNext) {
+      auto output = OUTPUT_VARIABLE(outNext);
+      for (LongType e = 0; e < output->lengthOf(); ++e) {
+        output->p<LongType>(e, whereTrue->e<LongType>(e, outNext));
+      }
+    }
+  }
+
+  return Status::OK;
+}
+
+DECLARE_SHAPE_FN(where_np) {
+  auto shapes = SHAPELIST();
+  if (block.width() == 3) {
+    auto inShape = inputShape->at(1);
+    shapes->push_back(CONSTANT(inShape));
+  } else {
+    auto condition = INPUT_VARIABLE(0);
+
+    LongType numOfTrue = 0LL;  // condition->reduceNumber(reduce::CountNonZero).e<sd::LongType>(0);
+    for (LongType i = 0; i < condition->lengthOf(); ++i)
+      if (condition->e<bool>(i)) numOfTrue++;
+
+    // output shape - a tuple of rank(inShape) 1D tensors with numOfTrue len
+    if (numOfTrue) {
+      for (LongType e = 0; e < condition->rankOf(); ++e) {
+        shapes->push_back(ConstantShapeHelper::getInstance().vectorShapeInfo(numOfTrue, INT64));
+      }
+    } else {
+      shapes->push_back(ConstantShapeHelper::getInstance().emptyShapeInfo(INT64));
+    }
+  }
+  return shapes;
+}
+
+DECLARE_TYPES(where_np) {
+  getOpDescriptor()
+      ->setAllowedInputTypes(0, BOOL)
+      ->setAllowedInputTypes(1, ANY)
+      ->setAllowedInputTypes(2, ANY)
+      ->setAllowedOutputTypes({ALL_FLOATS, ALL_INTS});
+}
+}  // namespace ops
+}  // namespace sd
+
+#endif