[Prim&Comp] Support repeat_interleave double grad (#73056)

HydrogenSulfate · web-flow · commit b0ce2a5d893f · 2025-06-04T15:05:10.000+08:00
* support repeat_interleave double grad

* remove redundant config
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2704,6 +2704,12 @@
   kernel :
     func : renorm_grad
 
+- backward_op : repeat_interleave_double_grad
+  forward : repeat_interleave_grad(Tensor x, Tensor grad_out, int repeats, int axis) -> Tensor(grad_x)
+  args : (Tensor grad_x_grad, int repeats, int axis)
+  output : Tensor(grad_out_grad)
+  invoke: repeat_interleave(grad_x_grad, repeats, axis)
+
 - backward_op : repeat_interleave_grad
   forward : repeat_interleave(Tensor x, int repeats, int axis) -> Tensor(out)
   args : (Tensor x, Tensor out_grad, int repeats, int axis)
@@ -2713,6 +2719,13 @@
     param : [x]
   kernel :
     func : repeat_interleave_grad
+  backward: repeat_interleave_double_grad
+
+- backward_op : repeat_interleave_with_tensor_index_double_grad
+  forward : repeat_interleave_with_tensor_index_grad(Tensor x, Tensor repeats, Tensor grad_out, int axis) -> Tensor(grad_x)
+  args : (Tensor repeats, Tensor grad_x_grad, int axis)
+  output : Tensor(grad_out_grad)
+  invoke: repeat_interleave_with_tensor_index(grad_x_grad, repeats, axis)
 
 - backward_op : repeat_interleave_with_tensor_index_grad
   forward : repeat_interleave_with_tensor_index(Tensor x, Tensor repeats, int axis) -> Tensor(out)
@@ -2724,6 +2737,7 @@
   kernel :
     func : repeat_interleave_with_tensor_index_grad
     data_type : x
+  backward: repeat_interleave_with_tensor_index_double_grad
 
 - backward_op : reshape_double_grad
   forward : reshape_grad (Tensor x, Tensor grad_out) -> Tensor(grad_x)
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -3133,7 +3133,7 @@
     repeats : Repeats
 
 - op : repeat_interleave
-  backward : repeat_interleave_grad
+  backward : repeat_interleave_grad, repeat_interleave_double_grad
   inputs :
     x : X
   outputs :
@@ -3142,7 +3142,7 @@
     {repeats : Repeats, axis : dim}
 
 - op : repeat_interleave_with_tensor_index
-  backward : repeat_interleave_with_tensor_index_grad
+  backward : repeat_interleave_with_tensor_index_grad, repeat_interleave_with_tensor_index_double_grad
   inputs :
     {x : X, repeats: RepeatTensor}
   outputs:
diff --git a/test/legacy_test/test_repeat_interleave_op.py b/test/legacy_test/test_repeat_interleave_op.py
@@ -316,6 +316,50 @@ def test_dygraph_api(self):
         )
         np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
 
+        # case 5 repeat_interleave_with_tensor_double_grad
+        with base.dygraph.guard():
+            x_pd = paddle.randn([40, 50])
+            x_pd.stop_gradient = False
+            axis = 1
+            repeats_pd = paddle.randint(1, 50, [x_pd.shape[axis]])
+
+            y_pd = paddle.repeat_interleave(x_pd, repeats_pd, axis)
+            dy_pd = paddle.randn_like(y_pd)
+            dy_pd.stop_gradient = False
+            g_pd = paddle.grad(y_pd, x_pd, dy_pd, create_graph=True)[0]
+
+            ddx_pd = paddle.randn_like(x_pd)
+            gg_pd = paddle.grad(g_pd, dy_pd, ddx_pd)[0]
+
+            np.testing.assert_allclose(
+                gg_pd.numpy(),
+                paddle.repeat_interleave(ddx_pd, repeats_pd, axis).numpy(),
+                1e-5,
+                1e-5,
+            )
+
+        # case 6 repeat_interleave_double_grad
+        with base.dygraph.guard():
+            x_pd = paddle.randn([40, 50])
+            x_pd.stop_gradient = False
+            axis = 1
+            repeats_pd = 4
+
+            y_pd = paddle.repeat_interleave(x_pd, repeats_pd, axis)
+            dy_pd = paddle.randn_like(y_pd)
+            dy_pd.stop_gradient = False
+            g_pd = paddle.grad(y_pd, x_pd, dy_pd, create_graph=True)[0]
+
+            ddx_pd = paddle.randn_like(x_pd)
+            gg_pd = paddle.grad(g_pd, dy_pd, ddx_pd)[0]
+
+            np.testing.assert_allclose(
+                gg_pd.numpy(),
+                paddle.repeat_interleave(ddx_pd, repeats_pd, axis).numpy(),
+                1e-5,
+                1e-5,
+            )
+
 
 if __name__ == '__main__':
     unittest.main()
