wehub-resource-sync
98 lines
3.4 KiB
C++
98 lines
3.4 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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 30.05.2019
|
|
//
|
|
// CPU implementation of one_hot helper
|
|
//
|
|
|
|
#include <execution/Threads.h>
|
|
#include <ops/declarable/helpers/one_hot.h>
|
|
#include <system/op_boilerplate.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
template <typename X, typename Z>
|
|
static void onehotImpl_(NDArray* indices, NDArray* output,
|
|
const LongType axis, const LongType depth,
|
|
const double on, const double off) {
|
|
auto xBuffer = indices->bufferAsT<X>();
|
|
auto zBuffer = output->bufferAsT<Z>();
|
|
|
|
auto xShapeInfo = indices->shapeInfo();
|
|
auto zShapeInfo = output->shapeInfo();
|
|
|
|
const int xRank = shape::rank(xShapeInfo);
|
|
const int zRank = shape::rank(zShapeInfo);
|
|
const sd::LongType* xShape = shape::shapeOf(xShapeInfo);
|
|
const sd::LongType* zShape = shape::shapeOf(zShapeInfo);
|
|
const sd::LongType* xStride = shape::stride(xShapeInfo);
|
|
const sd::LongType* zStride = shape::stride(zShapeInfo);
|
|
const sd::LongType zLen = output->lengthOf();
|
|
|
|
const Z onVal = static_cast<Z>(on);
|
|
const Z offVal = static_cast<Z>(off);
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; i++) {
|
|
sd::LongType coord[SD_MAX_RANK];
|
|
|
|
// Compute output coordinate and offset
|
|
INDEX2COORDS(i, zRank, zShape, coord);
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, coord, zOffset);
|
|
|
|
// Extract depth coordinate and shift axis
|
|
const auto depthCoord = coord[axis];
|
|
for (int j = axis; j < zRank - 1; ++j) {
|
|
coord[j] = coord[j + 1];
|
|
}
|
|
|
|
// Compute input offset
|
|
sd::LongType xOffset;
|
|
COORDS2INDEX(xRank, xStride, coord, xOffset);
|
|
|
|
// Check if the depth matches the index
|
|
const LongType idx = static_cast<LongType>(xBuffer[xOffset]);
|
|
zBuffer[zOffset] = (depthCoord == idx) ? onVal : offVal;
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, zLen);
|
|
}
|
|
|
|
void onehot(const LaunchContext* context, NDArray* indices, NDArray* output,
|
|
const LongType axis, const LongType depth, const double on, const double off) {
|
|
const auto xType = indices->dataType();
|
|
const auto zType = output->dataType();
|
|
|
|
NDArray::prepareSpecialUse({output}, {indices});
|
|
|
|
BUILD_DOUBLE_SELECTOR(xType, zType, onehotImpl_, (indices, output, axis, depth, on, off),
|
|
SD_COMMON_TYPES, SD_COMMON_TYPES);
|
|
|
|
NDArray::registerSpecialUse({output}, {indices});
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|