cavis/libnd4j/include/ops/declarable/generic/parity_ops/embedding_lookup.cpp

119 lines
4.1 KiB
C++

/*******************************************************************************
* Copyright (c) 2015-2018 Skymind, Inc.
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// Created by GS <sgazeos@gmail.com>
//
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_embedding_lookup)
#include <ops/declarable/CustomOperations.h>
#include <helpers/ShapeUtils.h>
#include <vector>
#include <numeric>
namespace sd {
namespace ops {
//////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(embedding_lookup, 2, 1, false, 0, 1) {
auto input = INPUT_VARIABLE(0); // lookup param
auto indeces = INPUT_VARIABLE(1); // indeces, as is
auto output = OUTPUT_VARIABLE(0); //
if (block.width() > 2) { // multiple input
indeces = INPUT_VARIABLE(block.width() - 1);
std::vector<int> dims(input->rankOf());
int i = output->rankOf() - input->rankOf();
for (auto& v: dims){
v = i++;
}
ResultSet outputView = output->allTensorsAlongDimension(dims);
REQUIRE_TRUE(block.width() > output->sizeAt(0), 0, "embedding_lookup: input list should be greater then %i, but %i given.",
output->sizeAt(0), block.width()
);
for (Nd4jLong e = 0; e < indeces->lengthOf(); ++e) {
Nd4jLong thisIndex = (*indeces).e<Nd4jLong>(e);
input = INPUT_VARIABLE(thisIndex); // lookup param
outputView.at(e)->assign(input);
}
}
else {
int indexRank = indeces->rankOf();
REQUIRE_TRUE(indexRank > 0, 0, "embeded_lookup: input array of indexes can't be single scalar, the requirement is: rank > 0 !");
int inputRank = input->rankOf();
int lastIndDim = indeces->lengthOf();
int partition_mode = INT_ARG(0); // partition_mode == 0 - i.e. 'mod' , 1 - 'div'
sd::ops::gather op;
std::unique_ptr<ResultSet> result(op.evaluate({input, indeces}, {0}));
REQUIRE_TRUE(result->status() == 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 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 indecesShapeInfo = inputShape->at(1);
int inRank = shape::rank(inShapeInfo);
if (inputShape->size() == 2u) {
int outRank = inRank;
std::vector<Nd4jLong> shapeInfo(outRank);
shapeInfo[0] = indecesShapeInfo[1]; // vector - how many elements
for (int 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<Nd4jLong> shapeInfo(outRank);
auto indeces = INPUT_VARIABLE(block.width() - 1);
shapeInfo[0] = indeces->lengthOf(); // vector - how many elements
for (int 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);
}
}
}
#endif