align view input style with torch

python/oneflow/nn/modules/reshape.py

@@ -18,6 +18,15 @@
 import oneflow as flow
 from oneflow.framework.tensor import register_tensor_op
 from oneflow.nn.module import Module
+from oneflow.nn.modules.utils import _single
+
+
+def _input_args_is_int(args):
+    return all((isinstance(x, int) for x in args))
+
+
+def _input_args_is_flow_size(args):
+    return all((isinstance(x, flow.Size) for x in args)) and len(args) == 1
 
 
 class Reshape(Module):
@@ -61,7 +70,7 @@ def reshape_op(x, shape: Sequence[int] = None):
 
 
 @register_tensor_op("view")
-def view_op(x, shape: Sequence[int] = None):
+def view_op(input, *shape):
     """
     The interface is consistent with PyTorch.
     The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
@@ -87,10 +96,10 @@ def view_op(x, shape: Sequence[int] = None):
     :meth:`contiguous`) otherwise.
 
     Args:
-        x: A Tensor.
-        shape: Shape of the output tensor.
+        input: A Tensor.
+        *shape: flow.Size or int...
     Returns:
-        A Tensor has the same type as `x`.
+        A Tensor has the same type as `input`.
 
     For example:
 
@@ -104,12 +113,18 @@ def view_op(x, shape: Sequence[int] = None):
         ... ).astype(np.float32)
         >>> input = flow.Tensor(x)
 
-        >>> y = flow.view(input, shape=[2, 2, 2, -1]).numpy().shape
+        >>> y = input.view(2, 2, 2, -1).numpy().shape
         >>> y
         (2, 2, 2, 2)
 
     """
-    return Reshape(shape=shape)(x)
+    if _input_args_is_int(shape):
+        new_shape = _single(shape)
+    elif _input_args_is_flow_size(shape):
+        new_shape = _single(*shape)
+    else:
+        raise ValueError("the input shape parameter of view is not illegal!")
+    return Reshape(shape=new_shape)(input)
 
 
 if __name__ == "__main__":

python/oneflow/test/modules/test_view.py

@@ -28,39 +28,34 @@ def _test_view(test_case, device):
     x = np.array(
         [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
     ).astype(np.float32)
-    input = flow.Tensor(x, device=flow.device(device))
-    of_shape = flow.view(input, shape=[2, 2, 2, -1]).numpy().shape
-    of_shape = flow.view(input, shape=[2, 2, 2, -1]).numpy().shape
-    np_shape = (2, 2, 2, 2)
-    test_case.assertTrue(np.array_equal(of_shape, np_shape))
-
-
-def _test_view_tuple(test_case, device):
-    x = np.array(
-        [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
-    ).astype(np.float32)
-    input = flow.Tensor(x, device=flow.device(device))
-    of_shape = flow.view(input, shape=(2, 2, 2, -1)).numpy().shape
-    np_shape = (2, 2, 2, 2)
-    test_case.assertTrue(np.array_equal(of_shape, np_shape))
-
-
-def _test_tensor_view(test_case, device):
-    x = np.array(
-        [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
-    ).astype(np.float32)
-    input = flow.Tensor(x, device=flow.device(device))
-    of_shape = input.view(shape=[2, 2, 2, -1]).numpy().shape
+    input = flow.Tensor(x, device=flow.device(device), requires_grad=True)
+    of_out = input.view(2, 2, 2, -1)
+    of_shape = of_out.numpy().shape
     np_shape = (2, 2, 2, 2)
     test_case.assertTrue(np.array_equal(of_shape, np_shape))
+    of_out = of_out.sum()
+    of_out.backward()
+    np_grad = np.array(
+        [
+            [1.0, 1.0, 1.0, 1.0],
+            [1.0, 1.0, 1.0, 1.0],
+            [1.0, 1.0, 1.0, 1.0],
+            [1.0, 1.0, 1.0, 1.0],
+        ]
+    )
+    test_case.assertTrue(np.allclose(np_grad, input.grad.numpy(), 0.0001, 0.0001))
 
 
-def _test_view_backward(test_case, device):
+def _test_view_flow_size(test_case, device):
     x = np.array(
         [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
     ).astype(np.float32)
     input = flow.Tensor(x, device=flow.device(device), requires_grad=True)
-    of_out = flow.view(input, shape=[2, 2, 2, -1]).sum()
+    shape = flow.Size([2, 2, 2, -1])
+    of_out = input.view(shape)
+    np_shape = (2, 2, 2, 2)
+    test_case.assertTrue(np.array_equal(of_out.numpy().shape, np_shape))
+    of_out = of_out.sum()
     of_out.backward()
     np_grad = np.array(
         [
@@ -74,14 +69,12 @@ def _test_view_backward(test_case, device):
 
 
 @flow.unittest.skip_unless_1n1d()
-class TestModule(flow.unittest.TestCase):
+class TestView(flow.unittest.TestCase):
     def test_view(test_case):
         arg_dict = OrderedDict()
         arg_dict["test_fun"] = [
             _test_view,
-            _test_view_tuple,
-            _test_tensor_view,
-            _test_view_backward,
+            _test_view_flow_size,
         ]
         arg_dict["device"] = ["cpu", "cuda"]
         for arg in GenArgList(arg_dict):