[Dy2St] mark pd_op.memcpy as forward only (#71193)

DrRyanHuang · web-flow · commit dca84620539a · 2025-02-19T22:49:49.000+08:00
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -532,6 +532,7 @@
     func : memcpy
     param: [x, dst_place_type]
   interfaces : paddle::dialect::InferSymbolicShapeInterface
+  traits : paddle::dialect::ForwardOnlyTrait
 
 - op : min
   args : (Tensor x, IntArray axis={}, bool keepdim=false)
diff --git a/test/dygraph_to_static/test_tensor_memcpy_on_gpu.py b/test/dygraph_to_static/test_tensor_memcpy_on_gpu.py
@@ -152,5 +152,34 @@ def test_tensor_cpu_with_compute_on_default_gpu(self):
         self.assertTrue(static_place.is_cpu_place())
 
 
+class TestMemcpyGrad(Dy2StTestBase):
+    def test_memcpy_grad(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+
+        def fn(x):
+            return x.cpu()
+
+        paddle.seed(2)
+        x_dy = paddle.rand([3, 2])
+        x_dy.stop_gradient = False
+        paddle.seed(2)
+        x_st = paddle.rand([3, 2])
+        x_st.stop_gradient = False
+
+        static_fn = paddle.jit.to_static(fn)
+
+        dy_out = fn(x_dy)
+        st_out = static_fn(x_st)
+        np.testing.assert_allclose(dy_out, st_out)
+
+        st_out.backward()
+        dy_out.backward()
+
+        self.assertIsNone(x_dy.grad)
+        # TODO: The static graph result needs to align with the dynamic graph result.
+        # self.assertIsNone(x_st.grad)
+
+
 if __name__ == '__main__':
     unittest.main()
