add reshape op

examples/oneflow2onnx/nodes/test_reshape.py

@@ -13,21 +13,35 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import tempfile
 import oneflow as flow
-import oneflow.typing as tp
 from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
 
+class Reshape(flow.nn.Module):
+    def __init__(self) -> None:
+        super(Reshape, self).__init__()
+    
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        return flow.reshape(x, (1, 3, -1))
 
-def test_reshape():
-    @flow.global_function()
-    def reshape(x: tp.Numpy.Placeholder((3, 4, 2, 5))):
-        return flow.reshape(x, (4, 30))
+reshape = Reshape()
+class reshapeOpGraph(flow.nn.Graph):
+    def __init__(self):
+        super().__init__()
+        self.m = reshape
+
+    def build(self, x):
+        out = self.m(x)
+        return out
 
-    convert_to_onnx_and_check(reshape)
 
-def test_reshape_negative_dim():
-    @flow.global_function()
-    def reshape(x: tp.Numpy.Placeholder((3, 4, 2, 5))):
-        return flow.reshape(x, (3, -1))
+def test_reshape():
+    
+    reshape_graph = reshapeOpGraph()
+    reshape_graph._compile(flow.randn(1, 3, 224, 224))
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        flow.save(reshape.state_dict(), tmpdirname)
+        convert_to_onnx_and_check(reshape_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp")
 
-    convert_to_onnx_and_check(reshape)
+test_reshape()