Add backward_decomp and support dynamic shape for kthvalue_grad (#68827)

eggman-1024 · web-flow · commit 3e4bafd31332 · 2024-10-25T10:22:05.000+08:00
* Add backward_decomp and support dynamic shape for kthvalue_grad

* modify codes

* modified some files

* fixed some bugs

* modified details.h
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -88,6 +88,7 @@
     'CropInferMeta',
     'EigvalsInferMeta',
     'FractionalMaxPoolInferMeta',
+    'KthvalueInferMeta',
     'MaxPoolWithIndexInferMeta',
     'MaxPoolV2InferMeta',
     'MultinomialInferMeta',
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -128,6 +128,7 @@
     'where_grad',
     'logcumsumexp_grad',
     'logsumexp_grad',
+    'kthvalue_grad',
 ]
 
 # whole vjp list of primitive op vjp
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2923,6 +2923,66 @@ void trunc_grad(const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void kthvalue_grad(const Tensor& x,
+                   const Tensor& indices,
+                   const Tensor& out_grad,
+                   int k,
+                   int axis,
+                   bool keepdim,
+                   Tensor* x_grad) {
+  if (x_grad) {
+    auto x_cast = ConverToMT<T>(x);
+    auto out_grad_cast = ConverToMT<T>(out_grad);
+    // put_along_axis doesn't support zero dim
+    if (x.dims().size() == 0) {
+      by_pass<T>(out_grad, x_grad);
+      return;
+    }
+
+    // function `put_along_axis` requires a non-negative axis
+    if (axis < 0) {
+      axis += x.dims().size();
+    }
+
+    Tensor zero_tensor;
+    Tensor x_grad_tmp;
+    if (has_dynamic_shape(x_cast.shape())) {
+      zero_tensor =
+          backend::full_with_tensor<T>(shape<T>(x_cast), 0, x_cast.dtype());
+
+      if (keepdim) {
+        x_grad_tmp = backend::put_along_axis<T>(
+            zero_tensor, indices, out_grad_cast, axis);
+      } else {
+        auto axis_ = std::vector<int64_t>(1, axis);
+        auto out_grad_shape =
+            get_unsqueeze_dims<T>(shape<T>(out_grad_cast), axis_);
+        auto out_grad_ = backend::reshape<T>(out_grad_cast, out_grad_shape);
+        auto indices_shape = get_unsqueeze_dims<T>(shape<T>(indices), axis_);
+        auto indices_ = backend::reshape<T>(indices, indices_shape);
+        x_grad_tmp =
+            backend::put_along_axis<T>(zero_tensor, indices_, out_grad_, axis);
+      }
+    } else {
+      zero_tensor =
+          full<T>(common::vectorize(x_cast.dims()), 0, x_cast.dtype());
+      if (keepdim) {
+        x_grad_tmp =
+            put_along_axis<T>(zero_tensor, indices, out_grad_cast, axis);
+      } else {
+        auto axis_ = std::vector<int64_t>(1, axis);
+        auto out_grad_shape = get_unsqueeze_dims(out_grad_cast, axis_);
+        auto out_grad_ = reshape<T>(out_grad_cast, out_grad_shape);
+        auto indices_shape = get_unsqueeze_dims(indices, axis_);
+        auto indices_ = reshape<T>(indices, indices_shape);
+        x_grad_tmp = put_along_axis<T>(zero_tensor, indices_, out_grad_, axis);
+      }
+    }
+    set_output<T>(ConverToOrig<T>(x_grad_tmp, x.dtype()), x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -52,6 +52,7 @@
     "pd_op.gather_nd",
     "pd_op.gelu",
     "pd_op.hardswish",
+    "pd_op.kthvalue",
     "pd_op.leaky_relu",
     "pd_op.log",
     "pd_op.logcumsumexp",
diff --git a/test/legacy_test/test_kthvalue_op.py b/test/legacy_test/test_kthvalue_op.py
@@ -45,9 +45,11 @@ def init_dtype(self):
 
     def setUp(self):
         self.op_type = "kthvalue"
+        self.prim_op_type = "prim"
         self.python_api = paddle.kthvalue
+        self.public_python_api = paddle.kthvalue
         self.init_dtype()
-        self.input_data = np.random.random((2, 1, 2, 4, 10)).astype(self.dtype)
+        self.input_data = np.random.random([2, 1, 2, 4, 10]).astype(self.dtype)
         self.init_args()
         self.inputs = {'X': self.input_data}
         self.attrs = {'k': self.k, 'axis': self.axis}
@@ -62,7 +64,12 @@ def test_check_output(self):
 
     def test_check_grad(self):
         paddle.enable_static()
-        self.check_grad({'X'}, 'Out', check_pir=True)
+        self.check_grad(
+            ['X'],
+            'Out',
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestKthvalueOpFp16(TestKthvalueOp):
@@ -82,8 +89,10 @@ def setUp(self):
         self.init_args()
         self.init_dtype()
         self.op_type = "kthvalue"
+        self.prim_op_type = "prim"
         self.python_api = paddle.kthvalue
-        self.input_data = np.random.random((1, 3, 2, 4, 10)).astype(self.dtype)
+        self.public_python_api = paddle.kthvalue
+        self.input_data = np.random.random([1, 3, 2, 4, 10]).astype(self.dtype)
         self.inputs = {'X': self.input_data}
         self.attrs = {'k': self.k, 'axis': self.axis, 'keepdim': True}
         output, indices = cal_kthvalue(
@@ -97,7 +106,12 @@ def test_check_output(self):
 
     def test_check_grad(self):
         paddle.enable_static()
-        self.check_grad({'X'}, 'Out', check_pir=True)
+        self.check_grad(
+            ['X'],
+            'Out',
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestKthvalueOpWithKeepdimFp16(TestKthvalueOpWithKeepdim):
diff --git a/test/prim/pir_prim/test_prim_sub_graph_klmno_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_klmno_backward_dynamic_shape.py
@@ -23,6 +23,18 @@
 import paddle
 
 
+def kthvalue_net1(x):
+    return paddle.kthvalue(x, k=2)[0]
+
+
+def kthvalue_net2(x):
+    return paddle.kthvalue(x, k=1, axis=1)[0]
+
+
+def kthvalue_net3(x):
+    return paddle.kthvalue(x, k=1, axis=1, keepdim=True)[0]
+
+
 def leaky_relu_net(x):
     return paddle.nn.functional.leaky_relu(x)
 
@@ -107,6 +119,58 @@ def multiply_net(x, y):
     return x * y
 
 
+class TestPrimKthvalueWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.kthvalue_grad"
+        self.dtype = "float32"
+        self.x_shape = [30]
+        self.init_x_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = kthvalue_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimKthvalueWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.kthvalue_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = kthvalue_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimKthvalueWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.kthvalue_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = kthvalue_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimKthvalueWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.kthvalue_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = kthvalue_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimLeakyReluWithGrad(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2024)

Original file line number	Diff line number	Diff line change
`@@ -128,6 +128,7 @@`
`128`	`128`	`'where_grad',`
`129`	`129`	`'logcumsumexp_grad',`
`130`	`130`	`'logsumexp_grad',`
	`131`	`+ 'kthvalue_grad',`
`131`	`132`	`]`
`132`	`133`
`133`	`134`	`# whole vjp list of primitive op vjp`