[RELAY] Fix ops in packed layout (apache#2472)

* [RELAY] Fix ops in packed layout * Fix style
merrymercy · Feb 17, 2019 · a82117e · a82117e
1 parent 3440f21
commit a82117e
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/src/relay/op/nn/pooling.cc b/src/relay/op/nn/pooling.cc
@@ -83,7 +83,11 @@ bool Pool2DRel(const Array<Type>& types,
     return false;
   }
 
-  std::vector<IndexExpr> oshape({dshape[0], dshape[1], dshape[2], dshape[3]});
+  std::vector<IndexExpr> oshape;
+  for (const auto& e : dshape) {
+    oshape.push_back(e);
+  }
+
   if (param->ceil_mode) {
     oshape[hidx] = ((dshape[hidx] + pad_h - param->pool_size[0] +
                     param->strides[0] - 1) / param->strides[0]) + 1;

diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
@@ -76,7 +76,8 @@ RELAY_REGISTER_OP("cast")
 .set_support_level(3)
 .add_type_rel("Cast", CastRel)
 .set_attr<FTVMCompute>("FTVMCompute", CastCompute)
-.set_attr<TOpPattern>("TOpPattern", kElemWise);
+.set_attr<TOpPattern>("TOpPattern", kElemWise)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout);
 
 // relay.expand_dims
 TVM_REGISTER_NODE_TYPE(ExpandDimsAttrs);

diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py
@@ -82,6 +82,8 @@ def before():
         # a useless tuple, which will be eliminated
         y = relay.Tuple([y])[0]
         y = relay.nn.relu(y)
+        y = relay.nn.max_pool2d(y, pool_size=(2, 2))
+        y = relay.cast(y, 'int32')
         y = relay.nn.batch_flatten(y)
         y = relay.Function(free_vars(y), y)
         return y
@@ -112,6 +114,8 @@ def expected():
         y = relay.add(y, b)
 
         y = relay.nn.relu(y)
+        y = relay.nn.max_pool2d(y, pool_size=(2, 2), layout="NCHW16c")
+        y = relay.cast(y, 'int32')
         y = relay.layout_transform(y, "NCHW16c", "NCHW")
         y = relay.nn.batch_flatten(y)
         y = relay.Function(free_vars(y), y)