align torch new tensor (#7973)

* align torch new tensor

* refine

* refine

* change test case name

* change device type

* fix test case

* fix test case

* fix import error

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

python/oneflow/framework/tensor.py

@@ -1052,6 +1052,31 @@ def _isinf(self):
     return flow.isinf(self)
 
 
+def _new_tensor(
+    self, data, dtype=None, device=None, requires_grad=False, placement=None, sbp=None
+):
+    if dtype is None:
+        dtype = self.dtype
+    if self.is_local:
+        assert (
+            placement is None and sbp is None
+        ), "self is local tensor, placement and sbp are expected to be None."
+        if device is None:
+            device = self.device
+        return flow.tensor(
+            data, dtype=dtype, device=device, requires_grad=requires_grad
+        )
+    else:
+        assert device is None, "self is global tensor, device is expected to be None."
+        if placement is None:
+            placement = self.placement
+        if sbp is None:
+            sbp = self.sbp
+        return flow.tensor(
+            data, dtype=dtype, placement=placement, sbp=sbp, requires_grad=requires_grad
+        )
+
+
 def RegisterMethods():
     Tensor.__mul__ = lambda self, other: self.mul(other)
     Tensor.__rmul__ = lambda self, other: self.mul(other)
@@ -1266,6 +1291,7 @@ def RegisterMethods():
     Tensor.to_consistent = _to_consistent
     Tensor.isnan = _isnan
     Tensor.isinf = _isinf
+    Tensor.new_tensor = _new_tensor
 
 
 def register_tensor_op(op_name):

python/oneflow/test/tensor/test_new_tensor.py

@@ -0,0 +1,68 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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 os
+import unittest
+import numpy as np
+import oneflow as flow
+import oneflow.unittest
+
+
+class TestNewTensor(flow.unittest.TestCase):
+    @flow.unittest.skip_unless_1n1d()
+    def test_new_tensor_local_mode_with_default_args(test_case):
+        tensor = flow.randn(5)
+        data = [[1, 2], [3, 4]]
+        new_tensor = tensor.new_tensor(data)
+        test_case.assertEqual(new_tensor.dtype, tensor.dtype)
+        test_case.assertEqual(new_tensor.device, tensor.device)
+
+    @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")
+    @flow.unittest.skip_unless_1n1d()
+    def test_new_tensor_local_mode_with_spec_args(test_case):
+        tensor = flow.randn(5)
+        data = [[1, 2], [3, 4]]
+        new_tensor = tensor.new_tensor(data, flow.int64, "cuda")
+        test_case.assertEqual(new_tensor.dtype, flow.int64)
+        test_case.assertEqual(new_tensor.device, flow.device("cuda"))
+
+    @flow.unittest.skip_unless_1n2d()
+    def test_new_tensor_global_mode_with_default_args(test_case):
+        placement = flow.placement(type="cpu", ranks=[0, 1])
+        sbp = flow.sbp.split(0)
+        tensor = flow.randn(4, 4, placement=placement, sbp=sbp)
+        data = [[1, 2], [3, 4]]
+        new_tensor = tensor.new_tensor(data)
+        test_case.assertEqual(new_tensor.dtype, tensor.dtype)
+        test_case.assertEqual(new_tensor.placement, placement)
+        test_case.assertEqual(new_tensor.sbp, (sbp,))
+
+    @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")
+    @flow.unittest.skip_unless_1n2d()
+    def test_new_tensor_global_mode_with_spec_args(test_case):
+        placement = flow.placement(type="cuda", ranks=[0, 1])
+        sbp = flow.sbp.split(0)
+        tensor = flow.randn(4, 4, placement=placement, sbp=sbp)
+        data = [[1, 2], [3, 4]]
+        new_tensor = tensor.new_tensor(
+            data, placement=placement, sbp=flow.sbp.broadcast
+        )
+        test_case.assertEqual(new_tensor.dtype, tensor.dtype)
+        test_case.assertEqual(new_tensor.placement, placement)
+        test_case.assertEqual(new_tensor.sbp, (flow.sbp.broadcast,))
+
+
+if __name__ == "__main__":
+    unittest.main()