[Paddle Tensor No.23] 为TensorA[mask] = TensorB补充单测 (#69534)

LittleHeroZZZX · web-flow · commit 074be90d3ae5 · 2024-11-21T16:26:00.000+08:00
* add test_boolean_mask_tensor_broadcast_v for tenor.__setitem__

* fix test setitem error

* fix test setitem tolerance error
diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py
@@ -533,6 +533,17 @@ def test_basic_setitem_bwd_2(self):
         else:
             np.testing.assert_equal(v.grad.numpy(), expected_v_grad)
 
+    def test_boolean_mask_tensor_broadcast_v(self):
+        tensor_np = np.zeros((2, 2, 3)).astype(np.float32)
+        mask_np = np.array([[True, False], [False, True]])
+        value_np = np.array([100] * 3).astype(np.float32)
+        tensor = paddle.to_tensor(tensor_np)
+        mask = paddle.to_tensor(mask_np)
+        value = paddle.to_tensor(value_np)
+        tensor[mask] = value
+        tensor_np[mask_np] = value_np
+        np.testing.assert_allclose(tensor.numpy(), tensor_np)
+
     def test_boolean_mask_scalar(self):
         tensor_np = np.arange(2 * 3).reshape(2, 3)
         tensor = paddle.to_tensor(tensor_np)
@@ -995,6 +1006,22 @@ def test_boolean_mask_tensor(self):
         tensor_np[mask_np] = value_np
         np.testing.assert_equal(res[0], tensor_np)
 
+    def test_boolean_mask_tensor_broadcast_v(self):
+        tensor_np = np.zeros((2, 2, 3)).astype(np.float32)
+        mask_np = np.array([[True, False], [False, True]])
+        value_np = np.array([100] * 3).astype(np.float32)
+
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            tensor = paddle.to_tensor(tensor_np)
+            mask = paddle.to_tensor(mask_np)
+            value = paddle.to_tensor(value_np)
+            tensor[mask] = value
+            res = self.exe.run(fetch_list=[tensor])[0]
+        tensor_np[mask_np] = value_np
+        np.testing.assert_allclose(res, tensor_np)
+
 
 if __name__ == '__main__':
     unittest.main()
