update diagonal scatter dygraph only

PaddlePaddle · Oct 5, 2023 · bab7de6 · bab7de6
1 parent 1c74ed5
commit bab7de6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 40 deletions.
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -5155,6 +5155,7 @@ def unfold(x, axis, size, step, name=None):
     return _C_ops.tensor_unfold(x, axis, size, step)
 
 
+@dygraph_only
 def diagonal_scatter(x, y, offset=0, axis1=0, axis2=1, name=None):
     """
     Embed the values of Tensor ``y`` into Tensor ``x`` along the diagonal elements
@@ -5200,14 +5201,7 @@ def diagonal_scatter(x, y, offset=0, axis1=0, axis2=1, name=None):
                 diagonal_slice.shape, y.shape
             )
         )
-    if in_dynamic_mode():
-        diagonal_slice[:] = y
-    else:
-        from builtins import slice
-
-        diagonal_slice = paddle.static.setitem(
-            diagonal_slice, slice(None, None, None), y
-        )
+    diagonal_slice[:] = y
     return output
 
 

diff --git a/test/legacy_test/test_diagonal_scatter.py b/test/legacy_test/test_diagonal_scatter.py
@@ -17,8 +17,6 @@
 import numpy as np
 
 import paddle
-from paddle import base
-from paddle.base import core
 
 
 class TestDiagonalScatterAPI(unittest.TestCase):
@@ -46,35 +44,6 @@ def setUp(self):
             ]
         )
 
-    def test_static_graph(self):
-        paddle.enable_static()
-        startup_program = base.Program()
-        train_program = base.Program()
-        with base.program_guard(startup_program, train_program):
-            x = paddle.static.data(
-                name='x', shape=self.x_shape, dtype=self.dtype
-            )
-            y = paddle.static.data(
-                name='y', shape=self.y_shape, dtype=self.dtype
-            )
-            out = paddle.diagonal_scatter(
-                x, y, offset=self.offset, axis1=self.axis1, axis2=self.axis2
-            )
-
-            place = (
-                base.CUDAPlace(0)
-                if core.is_compiled_with_cuda()
-                else base.CPUPlace()
-            )
-            exe = base.Executor(place)
-            res = exe.run(
-                base.default_main_program(),
-                feed={'x': self.x, 'y': self.y},
-                fetch_list=[out],
-            )
-            np.testing.assert_allclose(res[0], self.out, atol=1e-5, rtol=1e-5)
-            paddle.disable_static()
-
     def test_dygraph(self):
         paddle.disable_static()
         x = paddle.to_tensor(self.x)
@@ -83,14 +52,15 @@ def test_dygraph(self):
             x, y, offset=self.offset, axis1=self.axis1, axis2=self.axis2
         )
         np.testing.assert_allclose(self.out, result.numpy(), rtol=1e-5)
-
         paddle.enable_static()
 
     def test_error(self):
+        paddle.disable_static()
         x = paddle.to_tensor(self.x)
         y = paddle.ones((1, 2))
         self.assertRaises(ValueError, paddle.diagonal_scatter, x, y)
         self.assertRaises(ValueError, paddle.diagonal_scatter, y, x)
+        paddle.enable_static()
 
 
 class TestDiagonalScatterAPI1(TestDiagonalScatterAPI):