kevinthesun commented on a change in pull request #4312: URL: https://github.com/apache/incubator-tvm/pull/4312#discussion_r417026419
########## File path: src/relay/op/tensor/transform.cc ########## @@ -1891,81 +1952,163 @@ Array<Array<Layout> > StridedSliceInferCorrectLayout( } CHECK(old_in_layouts.defined()); - CHECK_EQ(old_in_layouts.size(), 1); + CHECK_GE(old_in_layouts.size(), 1); CHECK(old_in_shapes.defined()); - CHECK_EQ(old_in_shapes.size(), 1); + CHECK_GE(old_in_shapes.size(), 1); auto layout = old_in_layouts[0]; if (layout.defined() && new_in_layouts.defined()) { - CHECK_EQ(new_in_layouts.size(), 1); + CHECK_GE(new_in_layouts.size(), 1); auto new_layout = new_in_layouts[0]; auto shape = old_in_shapes[0]; // NOTE: Discard "const" qualifier here. auto *params = const_cast<StridedSliceAttrs*>(attrs.as<StridedSliceAttrs>()); + CHECK(params != nullptr); + Array<Integer> begin, end, strides; + const ConstantNode *cbegin, *cend, *cstrides; + if ((cbegin = params->begin.as<ConstantNode>()) && + (cend = params->end.as<ConstantNode>()) && + (cstrides = params->strides.as<ConstantNode>())) { + int64_t* strides_val = ToVector(cstrides->data); + for (int64_t i = 0; i < cstrides->data.Shape().front(); ++i) { + strides.push_back(strides_val[i]); + } + int64_t* begin_val = ToVector(cbegin->data); + for (int64_t i = 0; i < cbegin->data.Shape().front(); ++i) { + begin.push_back(begin_val[i]); + } + int64_t* end_val = ToVector(cend->data); + for (int64_t i = 0; i < cend->data.Shape().front(); ++i) { + end.push_back(end_val[i]); + } + } Review comment: For else case, I think we should directly return the original layout as new layout, since we can't compute for symbolic begin/end/strides. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org