strided_slice without view (#288)

Signed-off-by: raver119 <raver119@gmail.com>
master
raver119 2020-03-05 09:56:52 +03:00 committed by GitHub
parent ca96a13ed0
commit 3bb22a6ff8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 2 deletions

View File

@ -411,8 +411,24 @@ namespace sd {
// }
// else {
if (indices.size()) {
auto sub = (*x)(indices, true, true);
z->assign(sub);
Nd4jLong* subArrShapeInfo = nullptr;
ALLOCATE(subArrShapeInfo, block.getWorkspace(), shape::shapeInfoLength(x->rankOf()), Nd4jLong);
Nd4jLong offset;
shape::calcSubArrShapeInfoAndOffset(indices.data(), x->getShapeInfo(), subArrShapeInfo, offset, true, true);
auto subArrShapeInfoPack = ConstantShapeHelper::getInstance()->bufferForShapeInfo(subArrShapeInfo);
NDArray::prepareSpecialUse({z}, {x});
NativeOpExecutioner::execTransformAny(block.launchContext(), sd::transform::Assign,
x->bufferWithOffset(offset), reinterpret_cast<Nd4jLong *>(subArrShapeInfoPack.primary()),
x->specialBufferWithOffset(offset), reinterpret_cast<Nd4jLong *>(subArrShapeInfoPack.special()),
z->buffer(), z->shapeInfo(), z->specialBuffer(), z->specialShapeInfo(),
nullptr, nullptr, nullptr, true);
NDArray::registerSpecialUse({z}, {x});
RELEASE(subArrShapeInfo, block.getWorkspace());
}
else if (!z->isEmpty()){
z->assign(x->e(0));