chore: import upstream snapshot with attribution

test/auto_parallel/test_shard_tensor_api.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.dygraph.base import switch_to_static_graph
+from paddle.distributed import Replicate, Shard
+
+in_pir_mode = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+    "FLAGS_enable_pir_api"
+]
+
+
+class TestDistAttrBasic(unittest.TestCase):
+    def test_mesh_argument_error(self):
+        exception = None
+        try:
+            mesh = [[0, 1], [2, 3]]
+            dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=[None, None])
+        except ValueError as ex:
+            self.assertIn(
+                "The mesh must be an instance of paddle.distributed.ProcessMesh",
+                str(ex),
+            )
+            exception = ex
+
+        self.assertIsNotNone(exception)
+
+    def test_sharding_specs_argument_error(self):
+        exception = None
+        try:
+            mesh = dist.ProcessMesh(
+                [[2, 4, 5], [0, 1, 3]], dim_names=["x", "y"]
+            )
+            dist_attr = dist.DistAttr(
+                mesh=mesh, sharding_specs={"x": None, "y": None}
+            )
+        except ValueError as ex:
+            self.assertIn(
+                "The sharding_specs must be an instance of list", str(ex)
+            )
+            exception = ex
+
+        self.assertIsNotNone(exception)
+
+
+class TestShardTensorDynamic(unittest.TestCase):
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["x", "y"]
+        )
+
+    def test_dynamic_mode_basic(self):
+        input = paddle.rand([4, 1024, 512])
+        d_tensor = dist.shard_tensor(
+            input, self.mesh, [Replicate(), Replicate()]
+        )
+
+        self.assertEqual(d_tensor.process_mesh, self.mesh)
+
+    def test_dynamic_mode_property_change(self):
+        x = np.random.random([4, 1024, 512]).astype("float32")
+        input = paddle.to_tensor(
+            x, dtype="float32", place='cpu', stop_gradient=False
+        )
+        d_tensor = dist.shard_tensor(
+            input,
+            dtype="float64",
+            place='gpu:0',
+            stop_gradient=True,
+            mesh=self.mesh,
+            placements=[Replicate(), Replicate()],
+        )
+
+        self.assertEqual(d_tensor.dtype, paddle.float64)
+        self.assertTrue(d_tensor.place.is_gpu_place())
+        self.assertEqual(d_tensor.stop_gradient, True)
+
+        self.assertEqual(d_tensor.process_mesh, self.mesh)
+
+    def test_stop_gradient(self):
+        x = paddle.ones([4, 1024, 512])
+        x.stop_gradient = False
+        x = dist.shard_tensor(x, self.mesh, [Shard(0), Replicate()])
+        assert not x.stop_gradient
+
+
+class TestShardTensorStatic(unittest.TestCase):
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["x", "y"]
+        )
+
+    @switch_to_static_graph
+    def test_static_mode(self):
+        input = paddle.static.data(
+            name="input",
+            shape=[4, 1024, 512],
+            dtype='float32',
+        )
+        d_tensor = dist.shard_tensor(input, self.mesh, [Shard(0), Replicate()])
+        self.assertEqual(d_tensor.dist_attr().process_mesh, self.mesh)
+
+
+class TestShardTensorStaticDy2Static(unittest.TestCase):
+    def test_dy2static(self):
+        @paddle.jit.to_static(full_graph=True, input_spec=[])
+        def func():
+            mesh = dist.ProcessMesh(
+                [[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["x", "y"]
+            )
+            input = paddle.rand([4, 1024, 512])
+            d_tensor = dist.shard_tensor(
+                input, mesh, [Replicate(), Replicate()]
+            )
+            return d_tensor, mesh
+
+        # dy_tensor, mesh = func()
+        static_tensor = func.outputs[0]  # get the inputs of static program
+        mesh = dist.ProcessMesh(
+            [[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["x", "y"]
+        )
+        self.assertEqual(static_tensor.dist_attr().process_mesh, mesh)
+
+
+class DemoNet(paddle.nn.Layer):
+    def __init__(self, dist_attr):
+        super().__init__()
+        self.w0 = dist.shard_tensor(
+            self.create_parameter(shape=[784, 784]), *dist_attr
+        )
+
+    def forward(self, x):
+        return paddle.matmul(x, self.w0)
+
+
+class TestShardTensorParameter(unittest.TestCase):
+    def setUp(self):
+        self.mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        self.placements_and_mesh = (self.mesh, [Replicate()])
+
+    def test_shard_parameter(self):
+        x = np.random.random(size=[16, 784]).astype("float32")
+        dist_x = dist.shard_tensor(x, *self.placements_and_mesh)
+        net = DemoNet(self.placements_and_mesh)
+        out = net(dist_x)
+        self.assertEqual(out.shape, [16, 784])
+        self.assertEqual(out.is_dist(), True)
+
+
+if __name__ == "__main__":
+    unittest.main()