[Prim][PIR] Add back_decomp and support dynamic shape for take_along_axis_grad (#69372)

eggman-1024 · web-flow · commit c27d8ff55351 · 2024-11-20T19:14:40.000+08:00
* add back_decomp for argsort_grad(unfinished)

* modified details.h and add dynamic shape test file

* modified dynamic shape test file

* modified dynamic shape test file

* modified details.h

* modified test files

* add dynamic test file
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -124,6 +124,7 @@
     'scatter_nd_add_grad',
     'slice_grad',
     'squeeze_grad',
+    'take_along_axis_grad',
     'tile_grad',
     'topk_grad',
     'unsqueeze_grad',
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -3215,6 +3215,40 @@ void kron_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void take_along_axis_grad(const Tensor& arr,
+                          const Tensor& indices,
+                          const Tensor& out_grad,
+                          int axis,
+                          Tensor* arr_grad) {
+  if (arr_grad) {
+    auto arr_cast = ConverToMT<T>(arr);
+    auto out_grad_cast = ConverToMT<T>(out_grad);
+    // put_along_axis doesn't support zero dim
+    if (arr_cast.dims().size() == 0) {
+      by_pass<T>(ConverToOrig<T>(out_grad_cast, out_grad.dtype()), arr_grad);
+      return;
+    }
+
+    // function `put_along_axis` requires a non-negative axis
+    if (axis < 0) {
+      axis += arr_cast.dims().size();
+    }
+
+    Tensor zero_tensor;
+    if (has_dynamic_shape(arr_cast.shape())) {
+      zero_tensor =
+          backend::full_with_tensor<T>(shape<T>(arr_cast), 0, arr_cast.dtype());
+    } else {
+      zero_tensor =
+          full<T>(common::vectorize(arr_cast.dims()), 0, arr_cast.dtype());
+    }
+    auto arr_grad_tmp =
+        put_along_axis<T>(zero_tensor, indices, out_grad_cast, axis);
+    set_output<T>(ConverToOrig<T>(arr_grad_tmp, arr.dtype()), arr_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -91,6 +91,7 @@
     "pd_op.sum",
     "pd_op.swiglu",
     "pd_op.swish",
+    "pd_op.take_along_axis",
     "pd_op.tanh",
     "pd_op.topk",
     "pd_op.transpose",
diff --git a/test/legacy_test/test_take_along_axis_op.py b/test/legacy_test/test_take_along_axis_op.py
@@ -29,7 +29,9 @@ class TestTakeAlongAxisOp(OpTest):
     def setUp(self):
         self.init_data()
         self.op_type = "take_along_axis"
+        self.prim_op_type = "prim"
         self.python_api = paddle.tensor.take_along_axis
+        self.public_python_api = paddle.tensor.take_along_axis
         self.check_cinn = True
         self.xnp = np.random.random(self.x_shape).astype(self.x_type)
         self.target = np.take_along_axis(self.xnp, self.index, self.axis)
@@ -49,7 +51,11 @@ def test_check_output(self):
 
     def test_check_grad(self):
         self.check_grad(
-            ['Input'], 'Result', check_cinn=self.check_cinn, check_pir=True
+            ['Input'],
+            'Result',
+            check_cinn=self.check_cinn,
+            check_pir=True,
+            check_prim_pir=True,
         )
 
     def init_data(self):
@@ -78,7 +84,7 @@ def init_data(self):
         self.axis_type = "int64"
 
 
-class TestTakeAlongAxisOp(OpTest):
+class TestTakeAlongAxisOp2(OpTest):
     def setUp(self):
         self.init_data()
         self.op_type = "take_along_axis"
@@ -118,7 +124,9 @@ class TestTakeAlongAxisBF16Op(OpTest):
     def setUp(self):
         self.init_data()
         self.op_type = "take_along_axis"
+        self.prim_op_type = "prim"
         self.python_api = paddle.tensor.take_along_axis
+        self.public_python_api = paddle.tensor.take_along_axis
         self.check_cinn = True
         self.xnp = np.random.random(self.x_shape).astype(self.x_type)
         self.target = np.take_along_axis(self.xnp, self.index, self.axis)
@@ -149,6 +157,7 @@ def test_check_grad(self):
             'Result',
             check_cinn=self.check_cinn,
             check_pir=True,
+            check_prim_pir=True,
         )
 
     def init_data(self):
diff --git a/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py
@@ -166,6 +166,18 @@ def swish_net(x):
     return paddle.nn.functional.swish(x)
 
 
+def take_along_axis_net1(x, y):
+    return paddle.take_along_axis(x, y, 0, True)
+
+
+def take_along_axis_net2(x, y):
+    return paddle.take_along_axis(x, y, 1, True)
+
+
+def take_along_axis_net3(x, y):
+    return paddle.take_along_axis(x, y, 1, False)
+
+
 def tanh_net(x):
     return paddle.tanh(x)
 
@@ -1043,6 +1055,118 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimTakeAlongAxisWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.take_along_axis_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.y_shape = [1, 1, 1]
+        self.init_y_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([[[2]]], dtype="int32")
+        self.net = take_along_axis_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y)
+        y = paddle.broadcast_to(y, [1, 200, 40])
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='int32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        x_grad = x.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x, y)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_all_enabled(False)
+        return res, [x_grad]
+
+
+class TestPrimTakeAlongAxisWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.take_along_axis_grad"
+        self.dtype = "float32"
+        self.x_shape = [2, 40, 200]
+        self.init_x_shape = [None, None, None]
+        self.y_shape = [2, 1, 1]
+        self.init_y_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([[[1]], [[3]]], dtype="int32")
+        self.net = take_along_axis_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y)
+        y = paddle.broadcast_to(y, [2, 1, 200])
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='int32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        x_grad = x.gradient()
+        if flag == "prim":
+            ops = [
+                op.name()
+                for op in fn.get_concrete_program(x, y)[-1]
+                .program.backward_program.global_block()
+                .ops
+            ]
+            assert self.op_name not in ops
+            core._set_prim_all_enabled(False)
+        return res, [x_grad]
+
+
+class TestPrimTakeAlongAxisWithGrad3(TestPrimTakeAlongAxisWithGrad2):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.take_along_axis_grad"
+        self.dtype = "float32"
+        self.x_shape = [2, 40, 200]
+        self.init_x_shape = [None, None, None]
+        self.y_shape = [2, 1, 1]
+        self.init_y_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.array([[[1]], [[3]]], dtype="int32")
+        self.net = take_along_axis_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimTanhWithGrad(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2024)
