Fix bugs in topk_grad that occur when using negative axis values (#68741)

cubehan3 · web-flow · commit f635fef63ced · 2024-10-17T14:10:49.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1685,7 +1685,7 @@ void topk_grad(const Tensor& x,
                const Tensor& indices,
                const Tensor& out_grad,
                const Scalar& k,
-               const int& axis,
+               int axis,
                const bool& largest,
                const bool& sorted,
                Tensor* x_grad) {
@@ -1695,6 +1695,12 @@ void topk_grad(const Tensor& x,
       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;
     if (has_dynamic_shape(x.shape())) {
       zero_tensor = backend::full_with_tensor<T>(shape<T>(x), 0, x.dtype());
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
@@ -171,7 +171,7 @@ def tanh_net(x):
 
 
 def topk_net(x):
-    return paddle.topk(x, k=3, axis=1)[0]
+    return paddle.topk(x, k=3, axis=-1)[0]
 
 
 def transpose_net(x):
