[Pir] Fix bugs in put_along_axis for dynamic shape handling (#68722)

cubehan3 · web-flow · commit 4382729b0d9b · 2024-10-17T19:44:55.000+08:00
* Fix bugs in put_along_axis for dynamic shape handling

* polish manipulation.py

* Add unit test for dynamic shape of put_along_axis

* fixed test case bugs
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2439,12 +2439,14 @@ void put_along_axis_grad(const Tensor& x,
     if (include_self == false || reduce == "assign") {
       Tensor zero_tensor = full<T>(index.shape(), 0, out_grad.dtype());
       x_grad_tmp = put_along_axis<T>(out_grad, index, zero_tensor, axis);
+      set_output<T>(x_grad_tmp, x_grad);
     } else if (reduce == "multiply" || reduce == "mul") {
       Tensor zero_tensor_x = full<T>(x.shape(), 0, x.dtype());
       Tensor one_tensor_idx = full<T>(index.shape(), 1, x.dtype());
       Tensor mask =
           put_along_axis<T>(zero_tensor_x, index, one_tensor_idx, axis);
       x_grad_tmp = where<T>(mask > zero_tensor_x, out_grad * out / x, out_grad);
+      set_output<T>(x_grad_tmp, x_grad);
     } else if (reduce == "amin" || reduce == "amax") {
       Tensor zero_tensor = full<T>(x.shape(), 0, x.dtype());
       Tensor one_tensor = full<T>(x.shape(), 1, x.dtype());
@@ -2461,6 +2463,7 @@ void put_along_axis_grad(const Tensor& x,
         num = num + put_along_axis<T>(zero_tensor, sub_index, sub_count, axis);
       }
       x_grad_tmp = zero_result * out_grad / (num + 1);
+      set_output<T>(x_grad_tmp, x_grad);
     } else if (reduce == "mean") {
       Tensor zero_tensor_x = full<T>(x.shape(), 0, x.dtype());
 
@@ -2474,9 +2477,12 @@ void put_along_axis_grad(const Tensor& x,
       }
       x_grad_tmp =
           where<T>(num > zero_tensor_x, out_grad / (num + 1), out_grad);
+      set_output<T>(x_grad_tmp, x_grad);
+    } else if (reduce == "add") {
+      by_pass<T>(out_grad, x_grad);
     }
-    set_output<T>(x_grad_tmp, x_grad);
   }
+
   if (value_grad) {
     Tensor value_grad_tmp = full<T>(index.shape(), 0, x.dtype());
     if (reduce == "assign") {
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6342,6 +6342,24 @@ def non_negative_axis(arr, axis):
     return axis
 
 
+def infer_dynamic_broadcast_shape(
+    arr_shape: Tensor, indices_shape: Tensor, axis: int
+) -> Tensor:
+    """
+    Find the broadcast shape for indices when `arr` has a dynamic shape.
+
+    Args:
+        arr_shape (Tensor): Shape tensor of arr.
+        indices_shape (Tensor): Shape tensor of indices.
+        axis (int): The axis to put 1d slices along.
+
+    Returns:
+        Tensor: The shape tensor for later broadcasting
+    """
+    arr_shape[axis] = indices_shape[axis]
+    return arr_shape
+
+
 def infer_broadcast_shape(
     arr: Tensor, indices: Tensor, axis: int | Tensor
 ) -> tuple[int] | None:
@@ -6525,23 +6543,45 @@ def put_along_axis(
              [60, 40, 50]])
 
     """
+
+    def has_dynamic_shape(shapes):
+        return any(shape < 0 for shape in shapes)
+
     if len(arr.shape) != len(indices.shape):
         raise ValueError(
             "`indices` and `arr` must have the same number of dimensions!"
         )
     axis = non_negative_axis(arr, axis)
     if broadcast:
-        broadcast_shape = infer_broadcast_shape(arr, indices, axis)
-        values = (
-            paddle.to_tensor(values)
-            if not isinstance(
-                values, (paddle.Tensor, paddle.pir.Value, Variable)
+        if has_dynamic_shape(arr.shape) or has_dynamic_shape(indices.shape):
+            arr_shape = paddle.shape(arr)
+            indices_shape = paddle.shape(indices)
+            broadcast_shape = infer_dynamic_broadcast_shape(
+                arr_shape, indices_shape, axis
+            )
+            values = (
+                paddle.to_tensor(values)
+                if not isinstance(
+                    values, (paddle.Tensor, paddle.pir.Value, Variable)
+                )
+                else values
             )
-            else values
-        )
-        if broadcast_shape:
             indices = paddle.broadcast_to(indices, broadcast_shape)
-        values = paddle.broadcast_to(values, indices.shape)
+            values = paddle.broadcast_to(values, broadcast_shape)
+        else:
+            broadcast_shape = infer_broadcast_shape(arr, indices, axis)
+            values = (
+                paddle.to_tensor(values)
+                if not isinstance(
+                    values, (paddle.Tensor, paddle.pir.Value, Variable)
+                )
+                else values
+            )
+            if broadcast_shape:
+                indices = paddle.broadcast_to(indices, broadcast_shape)
+                values = paddle.broadcast_to(values, broadcast_shape)
+            else:
+                values = paddle.broadcast_to(values, indices.shape)
     else:
         if isinstance(values, (paddle.Tensor, paddle.pir.Value)):
             if len(indices.shape) != len(values.shape):
diff --git a/test/legacy_test/test_put_along_axis_op.py b/test/legacy_test/test_put_along_axis_op.py
@@ -18,13 +18,22 @@
 
 import numpy as np
 from op_test import OpTest, convert_float_to_uint16
+from utils import dygraph_guard
 
 import paddle
 from paddle.framework import core
+from paddle.static import InputSpec
 
 paddle.enable_static()
 
 
+def put_along_axis_net(arr):
+    indices = paddle.to_tensor([[[[2]]]], dtype='int32', stop_gradient=False)
+    return paddle.tensor.put_along_axis(
+        arr, indices=indices, values=-4.0, axis=-2, reduce='add'
+    )
+
+
 class TestPutAlongAxisOp(OpTest):
     def setUp(self):
         self.init_data()
@@ -1328,6 +1337,56 @@ def run(place):
         run(paddle.CUDAPlace(0))
 
 
+class TestPutAlongAxisDynamicShape(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2024)
+        self.net = put_along_axis_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+        self.dtype = "float32"
+        self.input_specs = [
+            InputSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=self.dtype,
+                stop_gradient=False,
+            )
+        ]
+        self.arr = np.random.random([10, 10, 10, 10]).astype(self.dtype)
+
+    def train(self, to_static):
+        arr = paddle.to_tensor(self.arr, stop_gradient=False)
+        if to_static:
+            build_strategy = paddle.static.BuildStrategy()
+            build_strategy.build_cinn_pass = self.enable_cinn
+            net = paddle.jit.to_static(
+                self.net,
+                input_spec=self.input_specs,
+                build_strategy=build_strategy,
+                full_graph=True,
+            )
+            net.train()
+        else:
+            net = self.net
+
+        res = net(arr)
+        res.backward()
+        arr_grad = arr.gradient()
+        return res, arr_grad
+
+    def test_dynamic_static(self):
+        with dygraph_guard():
+            st_out, st_grads = self.train(to_static=True)
+            dy_out, dy_grads = self.train(to_static=False)
+
+            for ref, actual in zip(dy_out, st_out):
+                np.testing.assert_allclose(
+                    ref, actual, rtol=self.tol, atol=self.tol
+                )
+
+            for dr, d in zip(dy_grads, st_grads):
+                np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
