150 lines
4.6 KiB
C++
150 lines
4.6 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 raver119@gmail.com
|
|
//
|
|
|
|
#include <system/op_boilerplate.h>
|
|
#if NOT_EXCLUDED(OP_split)
|
|
|
|
#include <ops/declarable/headers/parity_ops.h>
|
|
#include<ops/declarable/helpers/transforms.h>
|
|
#include <array>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
CUSTOM_OP_IMPL(split, 1, -1, false, 0, 1) {
|
|
NDArray *input = nullptr;
|
|
int num_splits = INT_ARG(0);
|
|
|
|
// axis is 0 by default
|
|
int axis = 0;
|
|
|
|
if (block.width() == 1) {
|
|
input = INPUT_VARIABLE(0);
|
|
} else {
|
|
auto a = INPUT_VARIABLE(0);
|
|
auto b = INPUT_VARIABLE(1);
|
|
|
|
if (a->isScalar()) {
|
|
// axis goes first
|
|
axis = a->e<int>(0);
|
|
input = b;
|
|
} else if (b->isScalar()) {
|
|
axis = b->e<int>(0);
|
|
input = a;
|
|
}
|
|
}
|
|
|
|
//Edge case: splitting empty array (mainly for TF import compatibility) -> return N empty arrays
|
|
if(input->isEmpty()){
|
|
for( int i=0; i< num_splits; i++ ){
|
|
REQUIRE_TRUE(OUTPUT_VARIABLE(i)->isEmpty(), 0, "Split: When input array is empty, all output arrays must be empty");
|
|
}
|
|
//No op
|
|
return Status::OK();
|
|
}
|
|
|
|
if (block.numI() == 2)
|
|
axis = INT_ARG(1);
|
|
|
|
if(axis < 0) axis += input->rankOf();
|
|
|
|
REQUIRE_TRUE(input->sizeAt(axis) % num_splits == 0, 0, "Split: num_splits has wrong value, remainder of division should be 0, but it's %i", input->sizeAt(axis) % num_splits);
|
|
|
|
std::vector<NDArray*> outArrs(num_splits);
|
|
for (int e = 0; e < num_splits; e++) {
|
|
outArrs[e] = OUTPUT_VARIABLE(e);
|
|
}
|
|
|
|
helpers::split(block.launchContext(), *input, outArrs, axis);
|
|
|
|
return Status::OK();
|
|
}
|
|
|
|
DECLARE_TYPES(split) {
|
|
getOpDescriptor()
|
|
->setAllowedInputTypes({ALL_INTS, ALL_FLOATS})
|
|
->setAllowedOutputTypes({ALL_INTS, ALL_FLOATS});
|
|
}
|
|
|
|
DECLARE_SHAPE_FN(split) {
|
|
int num_splits = INT_ARG(0);
|
|
auto input = inputShape->at(0);
|
|
sd::DataType dataType = ArrayOptions::dataType(input);
|
|
|
|
// axis is 0 by default
|
|
int axis = 0;
|
|
|
|
int inputVar = 0;
|
|
if (inputShape->size() != 1) {
|
|
auto shape0 = inputShape->at(0);
|
|
auto shape1 = inputShape->at(1);
|
|
|
|
if (shape::isScalar(shape0)) {
|
|
input = shape1;
|
|
auto _a = INPUT_VARIABLE(0);
|
|
axis = _a->e<int>(0);
|
|
dataType = ArrayOptions::dataType(shape1);
|
|
inputVar = 1;
|
|
} else if (shape::isScalar(shape1)) {
|
|
input = shape0;
|
|
auto _a = INPUT_VARIABLE(1);
|
|
axis = _a->e<int>(0);
|
|
dataType = ArrayOptions::dataType(shape0);
|
|
inputVar = 0;
|
|
}
|
|
}
|
|
|
|
auto shapes = SHAPELIST();
|
|
|
|
//Edge case: splitting empty array (mainly for TF import compatibility) -> return N empty arrays
|
|
// if(INPUT_VARIABLE(inputVar)->isEmpty()){
|
|
// for (int e = 0; e < num_splits; e++) {
|
|
// auto empty = ConstantShapeHelper::getInstance().emptyShapeInfo(dataType);
|
|
// shapes->push_back(empty);
|
|
// }
|
|
// return shapes;
|
|
// }
|
|
|
|
if (block.numI() == 2)
|
|
axis = INT_ARG(1);
|
|
|
|
if (axis < 0)
|
|
axis += shape::rank(input);
|
|
|
|
std::vector<Nd4jLong> shape(shape::rank(input));
|
|
|
|
for (int e = 0; e < shape::rank(input); e++)
|
|
if (e == axis)
|
|
shape[e] = shape::sizeAt(input, e) / num_splits;
|
|
else
|
|
shape[e] = shape::sizeAt(input, e);
|
|
|
|
for (int e = 0; e < num_splits; e++) {
|
|
auto newShape = ConstantShapeHelper::getInstance().createShapeInfo(dataType, shape::order(input), shape);
|
|
shapes->push_back(newShape);
|
|
}
|
|
|
|
return shapes;
|
|
}
|
|
}
|
|
}
|
|
|
|
#endif |