Add new SPLIT op
Also add tests for it.
Bug: 113563597
Test: NeuralNetworksTest_static with new tests
Change-Id: I32cb7ccd3fae023e97d207bdfbad29c497dc4044
Merged-In: I32cb7ccd3fae023e97d207bdfbad29c497dc4044
(cherry picked from commit a85cac86b4b5d53f0173197ff0882eb70ce0b2e1)
diff --git a/common/OperationsUtils.cpp b/common/OperationsUtils.cpp
index 9601853..cfd1c92 100644
--- a/common/OperationsUtils.cpp
+++ b/common/OperationsUtils.cpp
@@ -921,5 +921,24 @@
return true;
}
+
+bool splitPrepare(const Shape& input, int32_t axis, int32_t numOutputs,
+ std::vector<Shape>* output) {
+ axis = getDimensionIndex(input, axis);
+
+ const int32_t sizeOfAxisToSplit = input.dimensions[axis];
+ NN_OPS_CHECK(sizeOfAxisToSplit % numOutputs == 0);
+ const int32_t sliceSize = sizeOfAxisToSplit / numOutputs;
+
+ for (int i = 0; i < numOutputs; ++i) {
+ output->at(i).type = input.type;
+ output->at(i).dimensions = input.dimensions;
+ output->at(i).dimensions[axis] = sliceSize;
+ output->at(i).offset = input.offset;
+ output->at(i).scale = input.scale;
+ }
+ return true;
+}
+
} // namespace nn
} // namespace android
