chore: import upstream snapshot with attribution

2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
@@ -0,0 +1,175 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author Shyrma Yurii (iuriish@yahoo.com), created on 16.11.2017
+//
+
+#include <system/op_boilerplate.h>
+#if NOT_EXCLUDED(OP_gather)
+
+#include <ops/declarable/CustomOperations.h>
+#include <ops/declarable/helpers/gather.h>
+#include <ops/declarable/helpers/scatter.h>
+
+namespace sd {
+namespace ops {
+
+//////////////////////////////////////////////////////////////////////////
+CUSTOM_OP_IMPL(gather, 1, 1, false, 0, -2) {
+  auto input = INPUT_VARIABLE(0);
+  auto indices = block.width() > 1 ? INPUT_VARIABLE(1) : nullptr;
+  auto output = OUTPUT_VARIABLE(0);
+
+  const bool checkIndices = block.getBArguments()->empty() ? true : B_ARG(0);
+
+  // Edge case: empty indices -> empty output
+  if (indices != nullptr && indices->isEmpty()) {
+    REQUIRE_TRUE(output->isEmpty(), 0, "Gather op: If indices are empty, output must also be empty");
+    return sd::Status::OK;  // No op
+  }
+
+  const sd::LongType numOfIntArgs = block.numI();
+
+  std::vector<sd::LongType> intArgs;
+  if (block.width() > 2) {
+    intArgs = INPUT_VARIABLE(2)->template asVectorT<sd::LongType>();
+  } else {
+    if (numOfIntArgs == 0)
+      intArgs.emplace_back(0);
+    else
+      for (sd::LongType i = 0; i < numOfIntArgs; i++) intArgs.emplace_back(block.getIArguments()->at(i));
+  }
+
+  const sd::LongType inputRank = input->rankOf();
+  if (intArgs[0] < 0) intArgs[0] += inputRank;
+
+  // input validation
+  REQUIRE_TRUE(intArgs[0] < inputRank, 0,
+               "GATHER op: input axis must be smaller than input array rank, but got %i and %i correspondingly!",
+               intArgs[0], inputRank);
+  REQUIRE_TRUE(indices != nullptr || numOfIntArgs > 1, 0,
+               "GATHER op: indices should be provided either as additional input array or as IntArguments !");
+
+  if (checkIndices) {
+    NDArray* pIndices = indices;
+    bool ownsIndices = false;
+    
+    if (indices == nullptr) {
+      std::vector<sd::LongType> shape = {static_cast<sd::LongType>(intArgs.size()) - 1};
+      std::vector<double> inputVec = std::vector<double>(intArgs.begin() + 1, intArgs.end());
+      pIndices = new NDArray(input->ordering(), shape, inputVec, DataType::INT64, block.launchContext());
+      ownsIndices = true;
+    }
+    
+    const sd::LongType numOfBadIndx = helpers::checkIndices(block.launchContext(), *pIndices, *input, intArgs[0]);
+
+    // FIXED: Cleanup BEFORE checking condition (REQUIRE_TRUE can throw)
+    if (ownsIndices) {
+      delete pIndices;
+      pIndices = nullptr;
+    }
+
+    // Check condition after cleanup
+    REQUIRE_TRUE(numOfBadIndx == 0, 0,
+                 "GATHER OP: please check elements of indices-array, total number of wrong elements is %lld!",
+                 numOfBadIndx);
+  }
+
+  helpers::gather(block.launchContext(), input, indices, output, intArgs);
+
+  return sd::Status::OK;
+}
+
+DECLARE_TYPES(gather) {
+  getOpDescriptor()->setAllowedInputTypes(0, {ALL_INTS, ALL_FLOATS});
+  getOpDescriptor()->setAllowedInputTypes(1, {ALL_INTS,ALL_FLOATS});
+  getOpDescriptor()->setAllowedOutputTypes(0, {ALL_INTS, ALL_FLOATS});
+}
+
+DECLARE_SHAPE_FN(gather) {
+  // check shape of paddings
+  auto inputShapeInfo = inputShape->at(0);
+  sd::LongType* outputShapeInfo = nullptr;
+
+  sd::LongType axis = 0;
+
+  if (block.width() > 2) {
+    axis = INPUT_VARIABLE(2)->e<sd::LongType>(0);
+  } else
+    axis = block.numI() > 0 ? block.getIArguments()->at(0) : 0;
+
+  sd::LongType inputRank = shape::rank(inputShapeInfo);
+  if (axis < 0) axis += inputRank;
+
+  REQUIRE_TRUE(axis < inputRank, 0,
+               "GATHER op: input axis must be smaller than input array rank, but got %i and %i correspondingly!", axis,
+               inputRank);
+
+  bool isEmpty = false;
+
+  if (block.width() > 1) {
+    auto indicesShapeInfo = inputShape->at(1);
+
+    sd::LongType indicesRank = shape::rank(indicesShapeInfo);
+
+    sd::LongType outputRank = inputRank + indicesRank - 1;
+
+    ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(outputRank), sd::LongType);
+
+    // fill output shapeInfo
+    outputShapeInfo[0] = outputRank;
+    sd::LongType shapeIdx = 1;
+
+    for (sd::LongType i = 0; i < axis; ++i) outputShapeInfo[shapeIdx++] = inputShapeInfo[i + 1];
+
+    for (sd::LongType i = 0; i < indicesRank; ++i) outputShapeInfo[shapeIdx++] = indicesShapeInfo[i + 1];
+
+    for (sd::LongType i = axis + 1; i < inputRank; ++i) outputShapeInfo[shapeIdx++] = inputShapeInfo[i + 1];
+  } else if (block.numI() > 1) {
+    int indicesRank = block.numI() == 2 ? 0 : 1;
+
+    sd::LongType outputRank = inputRank + indicesRank - 1;
+    ALLOCATE(outputShapeInfo, block.getWorkspace(), shape::shapeInfoLength(outputRank), sd::LongType);
+
+    // building shape manually
+    outputShapeInfo[0] = outputRank;
+    int shapeIdx = 1;
+    for (sd::LongType i = 0; i < axis; ++i) outputShapeInfo[shapeIdx++] = inputShapeInfo[i + 1];
+
+    if (block.numI() > 2) outputShapeInfo[shapeIdx++] = block.numI() - 1;
+
+    for (sd::LongType i = axis + 1; i < inputRank; ++i) outputShapeInfo[shapeIdx++] = inputShapeInfo[i + 1];
+  } else
+    REQUIRE_TRUE(false, 0,
+                 "GATHER op: indices should be provided either as additional input array or as IntArguments !");
+
+  ShapeUtils::updateStridesAndType(outputShapeInfo, inputShapeInfo, shape::order(inputShapeInfo));
+
+  if (isEmpty) {
+    ArrayOptions::setPropertyBit(outputShapeInfo, ARRAY_EMPTY);
+  }
+
+  auto result = ConstantShapeHelper::getInstance().bufferForShapeInfo(outputShapeInfo)->primary();
+  RELEASE(outputShapeInfo, block.getWorkspace());
+  return SHAPELIST(result);
+}
+
+}  // namespace ops
+}  // namespace sd
+#endif