[Prim] Support take along axis double grad (#69373)

HydrogenSulfate · web-flow · commit e93033ed071b · 2024-11-14T14:18:37.000+08:00
* support take_along_axis double_grad

* add unitest
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -83,6 +83,7 @@
     "index_put_double_grad",
     "gather_nd_double_grad",
     "reshape_double_grad",
+    "take_along_axis_double_grad",
 ]
 
 # white ops list whose kernel can automatically do type promotion.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py
@@ -35,4 +35,5 @@
     'bmm_grad',
     'index_put_grad',
     'gather_nd_grad',
+    'take_along_axis_grad',
 ]
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -53,3 +53,4 @@
 - index_put
 - isnan
 - isfinite
+- take_along_axis
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -1072,5 +1072,20 @@ void reshape_double_grad(const Tensor& grad_out,
   }
 }
 
+template <typename T>
+void take_along_axis_double_grad(const Tensor& indices,
+                                 const Tensor& grad_arr_grad,
+                                 int axis,
+                                 Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    if (axis < 0) {
+      axis += grad_arr_grad.dims().size();
+    }
+    // ddout = take_along_axis(ddx, index)
+    auto grad_out_grad_tmp = take_along_axis<T>(grad_arr_grad, indices, axis);
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 }  // namespace prim
 }  // namespace paddle
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -3277,6 +3277,15 @@
     data_type : out_grad
   optional : reserve_space
 
+- backward_op : take_along_axis_double_grad
+  forward : take_along_axis_grad (Tensor arr, Tensor indices, Tensor grad_out, int axis) -> Tensor(grad_arr)
+  args : (Tensor indices, Tensor grad_arr_grad, int axis)
+  output : Tensor(grad_out_grad)
+  infer_meta :
+    func : TakeAlongAxisInferMeta
+    param : [grad_arr_grad, indices, axis]
+  composite : take_along_axis_double_grad(indices, grad_arr_grad, axis, grad_out_grad)
+
 - backward_op : take_along_axis_grad
   forward : take_along_axis (Tensor arr, Tensor indices, int axis) -> Tensor(out)
   args : (Tensor arr, Tensor indices, Tensor out_grad, int axis)
@@ -3286,6 +3295,7 @@
     param : [arr]
   kernel :
     func : take_along_axis_grad
+  backward : take_along_axis_double_grad
 
 - backward_op : tan_grad
   forward : tan (Tensor x) -> Tensor(out)
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_take_along_axis_double_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_take_along_axis_double_grad.py
@@ -0,0 +1,54 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+
+
+@param.parameterized_class(
+    ('arr', 'indices', 'axis', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.randn(4, 3, 2),  # arr
+            np.random.randint(0, 3, size=(1, 3, 2)),  # indices
+            1,  # axis
+            np.random.rand(1, 3, 2),  # cotangent
+            np.float32,  # dtype
+        ),
+    ],
+)
+class TestTakeAlongAxisTanhDoubleGrad(unittest.TestCase):
+    def test_take_along_axis_tanh_double_grad(self):
+        arr_tensor = paddle.to_tensor(
+            self.arr, dtype=self.dtype, stop_gradient=False
+        )
+        indices_tensor = paddle.to_tensor(self.indices, dtype="int64")
+        dout_tensor = paddle.to_tensor(
+            self.cotangent, dtype=self.dtype, stop_gradient=False
+        )
+        out = paddle.take_along_axis(arr_tensor, indices_tensor, axis=self.axis)
+
+        out = paddle.tanh(out)
+
+        dx = paddle.grad(out, arr_tensor, dout_tensor, create_graph=True)[0]
+
+        ddx = paddle.grad(dx, dout_tensor, create_graph=True)[0]
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,7 @@`
`83`	`83`	`"index_put_double_grad",`
`84`	`84`	`"gather_nd_double_grad",`
`85`	`85`	`"reshape_double_grad",`
	`86`	`+ "take_along_axis_double_grad",`
`86`	`87`	`]`
`87`	`88`
`88`	`89`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -35,4 +35,5 @@`
`35`	`35`	`'bmm_grad',`
`36`	`36`	`'index_put_grad',`
`37`	`37`	`'gather_nd_grad',`
	`38`	`+ 'take_along_axis_grad',`
`38`	`39`	`]`