New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Relay] [Op] Squeeze #1858
[Relay] [Op] Squeeze #1858
Conversation
include/tvm/relay/attrs/transform.h
Outdated
struct SqueezeAttrs : public tvm::AttrsNode<SqueezeAttrs> { | ||
Array<IndexExpr> axis; | ||
|
||
TVM_DECLARE_ATTRS(ExpandDimsAttrs, "relay.attrs.ExpandDimsAttrs") { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
wrong name: ExpandDimAttrs
@joshpoll @siju-samuel @junrushao1994 please review |
Will do tmr morning |
python/tvm/relay/op/transform.py
Outdated
""" | ||
axes = axes or [] | ||
return _make.transpose(data, list(axes)) | ||
|
||
|
||
def squeeze(data, axes=None): | ||
"""Remove the axis that equals to 1 in an array |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can rephrase to "Squeeze axes in the array."
src/relay/op/tensor/transform.cc
Outdated
} | ||
|
||
RELAY_REGISTER_OP("squeeze") | ||
.describe(R"code(squeeze the input tensor at the dimensions given by `axis` |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Squeeze the input tensor at the dimensions given by axes
89ba54b
to
6c85204
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please improve the error message, other parts look good
dbc74e6
to
ce44854
Compare
ce44854
to
1c27c23
Compare
@tqchen can you give it another look? I had addressed all comment. |
Thanks for contributing to TVM! Please refer to guideline https://docs.tvm.ai/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from others in the community.