[Dy2St] Add support for x.to(device=y.place) in dy2st (#68030)

gouzil · web-flow · commit 6b5e3b87a913 · 2024-09-06T01:22:30.000+08:00
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -780,7 +780,8 @@ def is_dtype(arg):
             )
 
         def is_device(arg):
-            return isinstance(arg, (paddle.core.Place, str))
+            # in dy2static, arg can be None
+            return arg is None or isinstance(arg, (paddle.core.Place, str))
 
         def is_tensor(arg):
             return isinstance(arg, paddle.pir.Value)
diff --git a/test/dygraph_to_static/test_tensor_to.py b/test/dygraph_to_static/test_tensor_to.py
@@ -57,62 +57,77 @@ def res():
 get_place = place_res()
 
 
-def to_dtype(tensorx, dtype):
-    return tensorx.to(dtype)
+def to_dtype(tensor_x, dtype):
+    return tensor_x.to(dtype)
 
 
-def to_device(tensorx, device):
-    return tensorx.to(device)
+def to_device(tensor_x, device):
+    return tensor_x.to(device)
 
 
-def to__device(tensorx, device):
-    return tensorx._to(device)
+def to__device(tensor_x, device):
+    return tensor_x._to(device)
 
 
-def to_device_dtype(tensorx, device, dtype):
-    return tensorx.to(device, dtype)
+def to_device_dtype(tensor_x, device, dtype):
+    return tensor_x.to(device, dtype)
 
 
-def to_other(tensorx, other):
-    return tensorx.to(other)
+def to_other(tensor_x, other):
+    return tensor_x.to(other)
 
 
-def to_other_blocking(tensorx, other, blocking):
-    return tensorx.to(other, blocking)
+def to_other_blocking(tensor_x, other, blocking):
+    return tensor_x.to(other, blocking)
 
 
-def to_dtype_blocking(tensorx, dtype, blocking):
-    return tensorx.to(dtype, blocking)
+def to_dtype_blocking(tensor_x, dtype, blocking):
+    return tensor_x.to(dtype, blocking)
 
 
-def to_device_dtype_blocking(tensorx, device, dtype, blocking):
-    return tensorx.to(device, dtype, blocking)
+def to_device_dtype_blocking(tensor_x, device, dtype, blocking):
+    return tensor_x.to(device, dtype, blocking)
 
 
-def to_kwargs_device_dtype_blocking(tensorx, device, dtype, blocking):
-    return tensorx.to(device=device, dtype=dtype, blocking=blocking)
+def to_kwargs_tesnor_device(tensor_x, tensor_y):
+    return tensor_x.to(device=tensor_y.place)
 
 
-def to_kwargs_other(tensorx, other):
-    return tensorx.to(other=other)
+def to_kwargs_device_dtype_blocking(tensor_x, device, dtype, blocking):
+    return tensor_x.to(device=device, dtype=dtype, blocking=blocking)
 
 
-def to_invalid_key_error(tensorx, device, dtype, test_key):
-    return tensorx.to(device, dtype, test_key=test_key)
+def to_kwargs_other(tensor_x, other):
+    return tensor_x.to(other=other)
 
 
-def to_many_key_error(tensorx, device, dtype):
-    return tensorx.to(device, dtype, device, dtype)
+def to_invalid_key_error(tensor_x, device, dtype, test_key):
+    return tensor_x.to(device, dtype, test_key=test_key)
+
+
+def to_many_key_error(tensor_x, device, dtype):
+    return tensor_x.to(device, dtype, device, dtype)
 
 
 class TensorToTest(Dy2StTestBase):
     @test_pir_only
     def test_tensor_to_dtype(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.to_tensor([1, 2, 3])
         for dtype in _valid_dtypes:
-            t = paddle.jit.to_static(to_dtype)(tensorx, dtype)
-            typex_str = str(t.dtype)
-            self.assertEqual(typex_str, "paddle." + dtype)
+            t = paddle.jit.to_static(to_dtype)(tensor_x, dtype)
+            type_x_str = str(t.dtype)
+            self.assertEqual(type_x_str, "paddle." + dtype)
+
+    @test_pir_only
+    def test_tensor_to_device(self):
+        if paddle.is_compiled_with_cuda():
+            x = paddle.to_tensor([1, 2, 3], place="gpu")
+        else:
+            x = paddle.to_tensor([1, 2, 3])
+
+        y = paddle.to_tensor([1, 2, 3], place="cpu")
+        y = paddle.jit.to_static(to_kwargs_tesnor_device)(y, x)
+        self.assertEqual(str(x.place), str(y.place))
 
     @test_pir_only
     def test_tensor_to_device2(self):
@@ -128,34 +143,34 @@ def test_tensor_to_device2(self):
 
     @test_pir_only
     def test_tensor_to_device_dtype(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.to_tensor([1, 2, 3])
         places = ["cpu"]
         if paddle.is_compiled_with_cuda():
             places.append("gpu")
         for dtype in _valid_dtypes:
             for place in places:
-                tensorx = paddle.jit.to_static(to_device_dtype)(
-                    tensorx, place, dtype
+                tensor_x = paddle.jit.to_static(to_device_dtype)(
+                    tensor_x, place, dtype
                 )
-                placex_str = str(tensorx.place)
+                place_x_str = str(tensor_x.place)
                 if "gpu" == place:
-                    self.assertEqual(placex_str, _gpu_place)
+                    self.assertEqual(place_x_str, _gpu_place)
                 else:
-                    self.assertEqual(placex_str, _cpu_place)
-                typex_str = str(tensorx.dtype)
-                self.assertEqual(typex_str, "paddle." + dtype)
+                    self.assertEqual(place_x_str, _cpu_place)
+                type_x_str = str(tensor_x.dtype)
+                self.assertEqual(type_x_str, "paddle." + dtype)
 
     @test_pir_only
     def test_tensor_to_blocking(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
-        tensorx = paddle.jit.to_static(to_device_dtype_blocking)(
-            tensorx, "cpu", "int32", False
+        tensor_x = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.jit.to_static(to_device_dtype_blocking)(
+            tensor_x, "cpu", "int32", False
         )
-        self.assertEqual(str(tensorx.place), _cpu_place)
-        self.assertEqual(tensorx.dtype, paddle.int32)
+        self.assertEqual(str(tensor_x.place), _cpu_place)
+        self.assertEqual(tensor_x.dtype, paddle.int32)
         tensor2 = paddle.to_tensor([4, 5, 6])
         tensor2 = paddle.jit.to_static(to_other_blocking)(
-            tensor2, tensorx, False
+            tensor2, tensor_x, False
         )
         # Note: in static mode, the place of tensor2 is not changed
         self.assertEqual(str(tensor2.place), get_place())
@@ -177,90 +192,90 @@ def test_tensor_to_other(self):
 
     @test_pir_only
     def test_kwargs(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
-        tensorx = paddle.jit.to_static(to_kwargs_device_dtype_blocking)(
-            tensorx, device="cpu", dtype="int8", blocking=True
+        tensor_x = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.jit.to_static(to_kwargs_device_dtype_blocking)(
+            tensor_x, device="cpu", dtype="int8", blocking=True
         )
-        self.assertEqual(str(tensorx.place), _cpu_place)
-        self.assertEqual(tensorx.dtype, paddle.int8)
+        self.assertEqual(str(tensor_x.place), _cpu_place)
+        self.assertEqual(tensor_x.dtype, paddle.int8)
         tensor2 = paddle.to_tensor([4, 5, 6])
-        tensor2 = paddle.jit.to_static(to_kwargs_other)(tensor2, other=tensorx)
+        tensor2 = paddle.jit.to_static(to_kwargs_other)(tensor2, other=tensor_x)
         # Note: in static mode, the place of tensor2 is not changed
         self.assertEqual(str(tensor2.place), get_place())
         self.assertEqual(tensor2.dtype, paddle.int8)
 
     @test_ast_only
     @test_pir_only
     def test_ast_error(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.to_tensor([1, 2, 3])
         # device value error
         with self.assertRaises(ValueError) as context1:
-            paddle.jit.to_static(to_device)(tensorx, "error_device")
+            paddle.jit.to_static(to_device)(tensor_x, "error_device")
         self.assertTrue(
             "The device must be a string which is like"
             in str(context1.exception)
         )
         # no matching signature error
         with self.assertRaises(ValueError) as context2:
-            paddle.jit.to_static(to_device)(tensorx, int)
+            paddle.jit.to_static(to_device)(tensor_x, int)
         self.assertTrue(
             "No matching signature found" in str(context2.exception)
         )
         # invalid key error
         with self.assertRaises(TypeError) as context3:
             paddle.jit.to_static(to_invalid_key_error)(
-                tensorx, "cpu", "int32", test_key=False
+                tensor_x, "cpu", "int32", test_key=False
             )
         self.assertTrue(
             "to() got an unexpected keyword argument" in str(context3.exception)
         )
         # device value error
         with self.assertRaises(ValueError) as context4:
-            paddle.jit.to_static(to__device)(tensorx, int)
+            paddle.jit.to_static(to__device)(tensor_x, int)
         self.assertTrue(
             "device value error, must be str" in str(context4.exception)
         )
         # too many key error
         with self.assertRaises(TypeError) as context5:
-            paddle.jit.to_static(to_many_key_error)(tensorx, "cpu", "int32")
+            paddle.jit.to_static(to_many_key_error)(tensor_x, "cpu", "int32")
         self.assertTrue(
             "to() received too many arguments" in str(context5.exception)
         )
 
     @test_sot_only
     @test_pir_only
     def test_sot_error(self):
-        tensorx = paddle.to_tensor([1, 2, 3])
+        tensor_x = paddle.to_tensor([1, 2, 3])
         # device value error
         with self.assertRaises(Exception) as context1:
-            paddle.jit.to_static(to_device)(tensorx, "error_device")
+            paddle.jit.to_static(to_device)(tensor_x, "error_device")
         self.assertTrue(
             "The device must be a string which is like"
             in str(context1.exception)
         )
         # no matching signature error
         with self.assertRaises(Exception) as context2:
-            paddle.jit.to_static(to_device)(tensorx, int)
+            paddle.jit.to_static(to_device)(tensor_x, int)
         self.assertTrue(
             "No matching signature found" in str(context2.exception)
         )
         # invalid key error
         with self.assertRaises(Exception) as context3:
             paddle.jit.to_static(to_invalid_key_error)(
-                tensorx, "cpu", "int32", test_key=False
+                tensor_x, "cpu", "int32", test_key=False
             )
         self.assertTrue(
             "to() got an unexpected keyword argument" in str(context3.exception)
         )
         # device value error
         with self.assertRaises(Exception) as context4:
-            paddle.jit.to_static(to__device)(tensorx, int)
+            paddle.jit.to_static(to__device)(tensor_x, int)
         self.assertTrue(
             "device value error, must be str" in str(context4.exception)
         )
         # too many key error
         with self.assertRaises(Exception) as context5:
-            paddle.jit.to_static(to_many_key_error)(tensorx, "cpu", "int32")
+            paddle.jit.to_static(to_many_key_error)(tensor_x, "cpu", "int32")
         self.assertTrue(
             "to() received too many arguments" in str(context5.exception)
         )

Original file line number	Diff line number	Diff line change
`@@ -780,7 +780,8 @@ def is_dtype(arg):`
`780`	`780`	`)`
`781`	`781`
`782`	`782`	`def is_device(arg):`
`783`		`- return isinstance(arg, (paddle.core.Place, str))`
	`783`	`+ # in dy2static, arg can be None`
	`784`	`+ return arg is None or isinstance(arg, (paddle.core.Place, str))`
`784`	`785`
`785`	`786`	`def is_tensor(arg):`
`786`	`787`	`return isinstance(arg, paddle.pir.Value)`