Fix paddle.diff when n!=1 and has prepend or append 易用性提升 (#67800)

NKNaN · web-flow · commit c4319b4ed170 · 2024-09-03T16:49:53.000+08:00
* fix diff

* update test

* fix docs
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -6311,6 +6311,16 @@ def diff(
             Tensor(shape=[5], dtype=int64, place=Place(cpu), stop_gradient=True,
             [ 3,  1, -3,  5,  2])
 
+            >>> out = paddle.diff(x, n=2, append=y)
+            >>> out
+            Tensor(shape=[4], dtype=int64, place=Place(cpu), stop_gradient=True,
+            [-2, -4,  8, -3])
+
+            >>> out = paddle.diff(x, n=3, append=y)
+            >>> out
+            Tensor(shape=[3], dtype=int64, place=Place(cpu), stop_gradient=True,
+            [-2 ,  12, -11])
+
             >>> z = paddle.to_tensor([[1, 2, 3], [4, 5, 6]])
             >>> out = paddle.diff(z, axis=0)
             >>> out
@@ -6455,7 +6465,7 @@ def _diff_handler(x, n=1, axis=-1, prepend=None, append=None, name=None):
     if n > 1:
         for _ in range(n - 1):
             out = _diff_handler(
-                out, n=1, axis=axis, prepend=prepend, append=append, name=name
+                out, n=1, axis=axis, prepend=None, append=None, name=name
             )
     return out
 
diff --git a/test/legacy_test/test_diff_op.py b/test/legacy_test/test_diff_op.py
@@ -171,6 +171,77 @@ def set_args(self):
         self.append = None
 
 
+class TestDiffOpNAxis(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = 1
+        self.prepend = None
+        self.append = None
+
+
+class TestDiffOpNPrepend(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = -1
+        self.prepend = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype(
+            'float32'
+        )
+        self.append = None
+
+
+class TestDiffOpNAppend(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = -1
+        self.prepend = None
+        self.append = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype('float32')
+
+
+class TestDiffOpNPreAppend(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = -1
+        self.prepend = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype(
+            'float32'
+        )
+        self.append = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype('float32')
+
+
+class TestDiffOpNPrependAxis(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = 0
+        self.prepend = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype(
+            'float32'
+        )
+        self.append = None
+
+
+class TestDiffOpNAppendAxis(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = 0
+        self.prepend = None
+        self.append = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype('float32')
+
+
+class TestDiffOpNPreAppendAxis(TestDiffOp):
+    def set_args(self):
+        self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
+        self.n = 2
+        self.axis = 0
+        self.prepend = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype(
+            'float32'
+        )
+        self.append = np.array([[2, 3, 4, 11], [1, 3, 5, 10]]).astype('float32')
+
+
 class TestDiffOpAxis(TestDiffOp):
     def set_args(self):
         self.input = np.array([[1, 4, 5, 2], [1, 5, 4, 2]]).astype('float32')
