From 6727ab5c07e827b259ea6a4449ab12fccc5a825f Mon Sep 17 00:00:00 2001
From: wuhuachaocoding <77733235+wuhuachaocoding@users.noreply.github.com>
Date: Thu, 22 Dec 2022 10:25:06 +0800
Subject: [PATCH] update 0-d tensor test for send_recv. (#49218)

---
 .../unittests/collective/process_group_nccl.py   | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py b/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py
index 6fcd90158f5..6debf636958 100644
--- a/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py
+++ b/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py
@@ -568,6 +568,22 @@ class TestProcessGroupFp32(unittest.TestCase):
 
         print("test send api ok")
 
+        # test send 0-d tensor
+        # rank 0
+        x = np.random.uniform(-1, 1, []).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        # rank 1
+        y = np.array(0.2022).astype(self.dtype)
+        tensor_y = paddle.to_tensor(y)
+
+        if pg.rank() == 0:
+            task = dist.send(tensor_x, 1, sync_op=True)
+        else:
+            task = dist.recv(tensor_y, 0, sync_op=True)
+            assert np.array_equal(tensor_y, tensor_x) and tensor_y.shape == []
+
+        print("test send & recv 0-d tensor ok")
+
 
 class TestProcessGroupFp16(TestProcessGroupFp32):
     def setUp(self):
-- 
GitLab