paddle.put_along_axis_ add attr broadcast 易用性提升 (#66687)

NKNaN · web-flow · commit c9389764bab1 · 2024-07-30T10:05:18.000+08:00
* update put_along_axis_

* fix test
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6276,7 +6276,7 @@ def take_along_axis(
 def put_along_axis(
     arr: Tensor,
     indices: Tensor,
-    values: int | Tensor,
+    values: float | Tensor,
     axis: int,
     reduce: Literal[
         'assign', 'add', 'mul', 'multiply', 'mean', 'amin', 'amax'
@@ -6291,7 +6291,7 @@ def put_along_axis(
         arr (Tensor) : The Destination Tensor. Supported data types are float32 and float64.
         indices (Tensor) : Indices to put along each 1d slice of arr. This must match the dimension of arr,
             and need to broadcast against arr if broadcast is 'True'. Supported data type are int and int64.
-        values (int|Tensor) : The value element(s) to put. The data types should be same as arr.
+        values (scalar|Tensor) : The value element(s) to put. The data types should be same as arr.
         axis (int) : The axis to put 1d slices along.
         reduce (str, optional): The reduce operation, default is 'assign', support 'add', 'assign', 'mul', 'multiply', 'mean', 'amin' and 'amax'.
         include_self (bool, optional): whether to reduce with the elements of arr, default is 'True'.
@@ -6442,12 +6442,13 @@ def put_along_axis(
 def put_along_axis_(
     arr: Tensor,
     indices: Tensor,
-    values: Tensor,
+    values: float | Tensor,
     axis: int,
     reduce: Literal[
         'assign', 'add', 'mul', 'multiply', 'mean', 'amin', 'amax'
     ] = "assign",
     include_self: bool = True,
+    broadcast: bool = True,
 ):
     r"""
     Inplace version of ``put_along_axis`` API, the output Tensor will be inplaced with input ``arr``.
@@ -6458,15 +6459,41 @@ def put_along_axis_(
             "`indices` and `arr` must have the same number of dimensions!"
         )
     axis = non_negative_axis(arr, axis)
-    broadcast_shape = infer_broadcast_shape(arr, indices, axis)
-    values = (
-        paddle.to_tensor(values)
-        if not isinstance(values, paddle.Tensor)
-        else values
-    )
-    if broadcast_shape:
-        indices = paddle.broadcast_to(indices, broadcast_shape)
-    values = paddle.broadcast_to(values, indices.shape)
+    if broadcast:
+        broadcast_shape = infer_broadcast_shape(arr, indices, axis)
+        values = (
+            paddle.to_tensor(values)
+            if not isinstance(values, paddle.Tensor)
+            else values
+        )
+        if broadcast_shape:
+            indices = paddle.broadcast_to(indices, broadcast_shape)
+        values = paddle.broadcast_to(values, indices.shape)
+    else:
+        if isinstance(values, (paddle.Tensor, paddle.pir.Value)):
+            if len(indices.shape) != len(values.shape):
+                raise ValueError(
+                    "`indices` and `values` must have the same number of dimensions!"
+                )
+            for i in range(len(arr.shape)):
+                if (
+                    i != axis and arr.shape[i] < indices.shape[i]
+                ) or indices.shape[i] > values.shape[i]:
+                    raise RuntimeError(
+                        f"Size does not match at dimension {i} expected index {indices.shape} to be smaller than self {arr.shape} apart from dimension {axis} and to be smaller size than values {values.shape}"
+                    )
+        else:
+            values = paddle.to_tensor(values).astype(arr.dtype)
+            elements = 1
+            for num in values.shape:
+                elements *= num
+            if elements == 1:  # paddle.pir.Value has no attribute 'size'
+                values = paddle.broadcast_to(values, indices.shape)
+        axis_max_size = arr.shape[axis]
+        if not (indices < axis_max_size).all():
+            raise RuntimeError(
+                f"one of element of indices is out of bounds for dimension {axis} with size {axis_max_size}"
+            )
     return _C_ops.put_along_axis_(
         arr, indices, values, axis, reduce, include_self
     )
diff --git a/test/legacy_test/test_put_along_axis_op.py b/test/legacy_test/test_put_along_axis_op.py
@@ -899,8 +899,36 @@ def run(place):
 
             paddle.enable_static()
 
+        def run_inplace(place):
+            paddle.disable_static(place)
+            x_tensor = paddle.to_tensor(self.x_np)
+            index_tensor1 = paddle.to_tensor(self.index_np1)
+            value_tensor = paddle.to_tensor(self.value)
+            x_tensor.put_along_axis_(
+                index_tensor1, value_tensor, 0, 'assign', True, False
+            )
+            out_ref = copy.deepcopy(self.x_np)
+            for i in range(self.index1_shape[0]):
+                for j in range(self.index1_shape[1]):
+                    out_ref[self.index_np1[i, j], j] = self.value[i, j]
+            np.testing.assert_allclose(x_tensor.numpy(), out_ref, rtol=0.001)
+
+            x_tensor = paddle.to_tensor(self.x_np)
+            index_tensor2 = paddle.to_tensor(self.index_np2)
+            x_tensor.put_along_axis_(
+                index_tensor2, 10, 1, 'assign', True, False
+            )
+            out_ref = copy.deepcopy(self.x_np)
+            for i in range(self.index2_shape[0]):
+                for j in range(self.index2_shape[1]):
+                    out_ref[i, self.index_np2[i, j]] = 10
+            np.testing.assert_allclose(x_tensor.numpy(), out_ref, rtol=0.001)
+
+            paddle.enable_static()
+
         for place in self.place:
             run(place)
+            run_inplace(place)
 
     @test_with_pir_api
     def test_api_static(self):
@@ -977,6 +1005,11 @@ def test_error(self):
             )
         except Exception as error:
             self.assertIsInstance(error, ValueError)
+        # len(values.shape) != len(indices.shape)
+        try:
+            tensorx.put_along_axis_(indices, values, 0, 'assign', True, False)
+        except Exception as error:
+            self.assertIsInstance(error, ValueError)
         indices = paddle.to_tensor(
             [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]
         ).astype("int32")
@@ -987,6 +1020,11 @@ def test_error(self):
             )
         except Exception as error:
             self.assertIsInstance(error, RuntimeError)
+        # indices too large
+        try:
+            tensorx.put_along_axis_(indices, 1.0, 0, 'assign', True, False)
+        except Exception as error:
+            self.assertIsInstance(error, RuntimeError)
         indices = paddle.to_tensor([[10]]).astype("int32")
         # the element of indices out of range
         try:
@@ -995,6 +1033,11 @@ def test_error(self):
             )
         except Exception as error:
             self.assertIsInstance(error, RuntimeError)
+        # the element of indices out of range
+        try:
+            tensorx.put_along_axis_(indices, 1.0, 0, 'assign', True, False)
+        except Exception as error:
+            self.assertIsInstance(error, RuntimeError)
 
     def test_index_type_error(self):
         tensorx = paddle.to_tensor([[1, 2, 3], [4, 5, 6]]).astype("float32")
