Speedup reshape module by 5x.

oneflow/core/functional/impl/array_functor.cpp

@@ -270,8 +270,30 @@ class ReshapeFunctor {
     op_ = CHECK_JUST(one::OpBuilder("reshape").Input("in").Output("out").Build());
   }
   Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const Shape& shape) const {
+    int need_infer_axis = -1;
+    size_t count = 1;
+    for (int i = 0; i < shape.NumAxes(); ++i) {
+      if (shape.At(i) == -1) {
+        CHECK_EQ_OR_RETURN(need_infer_axis, -1)
+            << "Shape " << shape.ToString() << " has more than 1 axis that needs to be infered.";
+        need_infer_axis = i;
+      } else {
+        count *= shape.At(i);
+      }
+    }
+    size_t x_count = x->shape()->Count(0);
     MutableAttrMap attrs;
-    JUST(attrs.SetAttr<Shape>("shape", shape));
+    if (need_infer_axis == -1) {
+      CHECK_EQ_OR_RETURN(shape.Count(0), x_count);
+      JUST(attrs.SetAttr<Shape>("shape", shape));
+    } else {
+      Shape infered_shape = shape;
+      infered_shape.Set(need_infer_axis, x_count / count);
+      CHECK_EQ_OR_RETURN(infered_shape.Count(0), x_count)
+          << "Shape " << shape.ToString() << " is invalid for input of shape "
+          << x->shape()->ToString();
+      JUST(attrs.SetAttr<Shape>("shape", infered_shape));
+    }
     return OpInterpUtil::Dispatch<Tensor>(*op_, {x}, attrs);
   }
 

oneflow/python/nn/modules/reshape.py

@@ -18,35 +18,15 @@
 from oneflow.python.oneflow_export import oneflow_export, experimental_api
 from oneflow.python.framework.tensor import register_tensor_op
 from typing import Sequence
-from functools import reduce
-import operator
-
-
-def infer_shape(x, shape):
-    dim_index_need_infer = shape.index(-1) if shape.count(-1) == 1 else None
-    in_elem_cnt = reduce(operator.mul, x.shape, 1)
-    out_elem_cnt = reduce(operator.mul, shape, 1)
-    if dim_index_need_infer is not None:
-        assert (in_elem_cnt % out_elem_cnt) == 0
-        shape[dim_index_need_infer] = int(abs(in_elem_cnt / out_elem_cnt))
-    else:
-        assert in_elem_cnt == out_elem_cnt
-    return shape
 
 
 class Reshape(Module):
     def __init__(self, shape: Sequence[int]) -> None:
         super().__init__()
-
-        assert isinstance(shape, tuple) or isinstance(shape, list)
-        shape = list(shape)
-        assert all(dim == -1 or dim > 0 for dim in shape)
-        assert shape.count(-1) <= 1
         self.shape = shape
 
     def forward(self, x):
-        new_shape = infer_shape(x, self.shape)
-        return flow.F.reshape(x, shape=new_shape)
+        return flow.F.reshape(x, shape=self.shape)
 
 
 @oneflow_export("reshape")