You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by an...@apache.org on 2022/05/19 16:15:26 UTC
[tvm] branch main updated: nn.batch_flatten is a reshape op (#11367)
This is an automated email from the ASF dual-hosted git repository.
andrewzhaoluo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 16c4faf86c nn.batch_flatten is a reshape op (#11367)
16c4faf86c is described below
commit 16c4faf86c584b22dbeaf304108cee5103ac23c2
Author: Altan Haan <31...@users.noreply.github.com>
AuthorDate: Thu May 19 09:15:18 2022 -0700
nn.batch_flatten is a reshape op (#11367)
---
src/relay/op/nn/nn.cc | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc
index 27f295b8b3..234cafdca1 100644
--- a/src/relay/op/nn/nn.cc
+++ b/src/relay/op/nn/nn.cc
@@ -540,10 +540,12 @@ Example::
.add_argument("data", "Tensor", "The input tensor.")
.set_support_level(2)
.add_type_rel("BatchFlatten", BatchFlattenRel)
- .set_attr<FTVMCompute>("FTVMCompute", [](const Attrs& attrs, const Array<te::Tensor>& inputs,
- const Type& out_type) {
- return Array<te::Tensor>{topi::nn::flatten(inputs[0])};
- });
+ .set_attr<FTVMCompute>("FTVMCompute",
+ [](const Attrs& attrs, const Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ return Array<te::Tensor>{topi::nn::flatten(inputs[0])};
+ })
+ .set_attr<TReshapeOp>("TReshapeOp", true);
// relu
TVM_REGISTER_GLOBAL("relay.op.nn._make.relu").set_body_typed([](Expr data) {