zheng-da commented on a change in pull request #13687: split_v2 operator
URL: https://github.com/apache/incubator-mxnet/pull/13687#discussion_r243618807
##########
File path: src/operator/tensor/matrix_op-inl.h
##########
@@ -2520,6 +2520,323 @@ void SpaceToDepthOpForward(const nnvm::NodeAttrs&
attrs,
});
}
+namespace split_enum {
+enum SplitOpInputs {kData};
+} // namespace split_enum
+
+struct SplitParam : public dmlc::Parameter<SplitParam> {
+ TShape indices;
+ int axis;
+ bool squeeze_axis;
+ int sections;
+ DMLC_DECLARE_PARAMETER(SplitParam) {
+ DMLC_DECLARE_FIELD(indices)
+ .describe("Indices of splits. The elements should denote the boundaries of
at which split"
+ " is performed along the `axis`.");
+ DMLC_DECLARE_FIELD(axis).set_default(1)
+ .describe("Axis along which to split.");
+ DMLC_DECLARE_FIELD(squeeze_axis).set_default(0)
+ .describe("If true, Removes the axis with length 1 from the shapes of the
output arrays."
+ " **Note** that setting `squeeze_axis` to ``true`` removes axis
with length 1"
+ " only along the `axis` which it is split."
+ " Also `squeeze_axis` can be set to ``true``"
+ " only if ``input.shape[axis] == num_outputs``.");
+ DMLC_DECLARE_FIELD(sections).set_default(0)
+ .describe("Number of sections if equally splitted. Default to 0 which
means split by indices.");
+ }
+}; // struct SplitParam
+
+inline TShape GetSplitIndices(const TShape& ishape, int axis, int sections) {
Review comment:
i know TShape uses Tuple. I think it's better to use Tuple for the return
result. TShape makes the return result look like a shape.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
