chore: import upstream snapshot with attribution

test/auto_parallel/semi_auto_parallel_util.py

@@ -0,0 +1,146 @@
+# 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 os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.placement_type import (
+    dims_mapping_to_placements,
+)
+
+
+class SemiAutoParallelTestBase:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def check_tensor_eq(self, a, b):
+        if a is None:
+            assert b is None
+            return
+        np1 = a.numpy()
+        np2 = b.numpy()
+        np.testing.assert_allclose(np1, np2, rtol=1e-05, verbose=True)
+
+    def flatten(self, inputs, terminal_cond):
+        """
+        inputs may be single tensor、tuple
+        """
+
+        if terminal_cond(inputs):
+            return [inputs], "i"
+
+        assert isinstance(inputs, (tuple, list))
+        flattened = []
+        structure = []
+        for i in range(len(inputs)):
+            tmp, tmp_structure = self.flatten(inputs[i], terminal_cond)
+            flattened.extend(tmp)
+            structure.append(tmp_structure)
+
+        if isinstance(inputs, tuple):
+            structure = tuple(structure)
+        return flattened, structure
+
+    def unflatten(self, inputs, structure, offset=0):
+        """
+        inputs may be single tensor
+        """
+        assert isinstance(inputs, list)
+        assert offset < len(inputs)
+        if structure == "i":
+            offset = offset + 1
+            # return a list
+            return inputs[offset - 1], offset
+        assert isinstance(structure, (tuple, list))
+        unflattened = []
+        for i in range(len(structure)):
+            tmp, offset = self.unflatten(inputs, structure[i], offset)
+            unflattened.append(tmp)
+        if isinstance(structure, tuple):
+            unflattened = tuple(unflattened)
+        return unflattened, offset
+
+    def runfunc_and_check(
+        self, inputs_shape, inputs_specs, op_func, with_backward, **kwargs
+    ):
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+
+        flat_inputs = []
+        flat_dist_inputs = []
+
+        def terminal_cond(x):
+            return isinstance(x, list) and all(
+                not isinstance(e, (list, tuple)) for e in x
+            )
+
+        flat_inputs_specs, inputs_structure = self.flatten(
+            inputs_specs, terminal_cond
+        )
+        flat_inputs_shape, _ = self.flatten(inputs_shape, terminal_cond)
+        assert len(flat_inputs_specs) == len(flat_inputs_shape)
+
+        for shape, spec in zip(flat_inputs_shape, flat_inputs_specs):
+            input_np = np.random.random(size=shape).astype(self._dtype)
+            input = paddle.to_tensor(input_np)
+            input.stop_gradient = not with_backward
+            # retain dist_attr here.
+            input_dist_attr = dist.DistAttr(
+                mesh=self._mesh, sharding_specs=spec
+            )
+            # for dygraph auto_parallel, get placements by using to_placements
+            placements = dims_mapping_to_placements(
+                input_dist_attr.multi_dims_mapping, self._mesh
+            )
+            dist_input = dist.shard_tensor(input, self._mesh, placements)
+            dist_input.stop_gradient = not with_backward
+            flat_inputs.append(input)
+            flat_dist_inputs.append(dist_input)
+
+        inputs, _ = self.unflatten(flat_inputs, inputs_structure)
+        dist_inputs, _ = self.unflatten(flat_dist_inputs, inputs_structure)
+
+        def wrap_tuple(e):
+            return e if isinstance(e, tuple) else (e,)
+
+        op_inputs = wrap_tuple(inputs)
+        op_dist_input = wrap_tuple(dist_inputs)
+
+        out = op_func(*op_inputs, **kwargs)
+        dist_out = op_func(*op_dist_input, **kwargs)
+
+        if with_backward:
+
+            def terminal_cond2(x):
+                return not isinstance(x, (list, tuple))
+
+            flat_out, _ = self.flatten(out, terminal_cond2)
+            flat_dist_out, _ = self.flatten(dist_out, terminal_cond2)
+            assert len(flat_out) == len(flat_dist_out)
+            for output, dist_output in zip(flat_out, flat_dist_out):
+                self.check_tensor_eq(output, dist_output)
+                if output is not None:
+                    output.backward()
+                    dist_output.backward()
+
+            for x, dist_x in zip(flat_inputs, flat_dist_inputs):
+                self.check_tensor_eq(x.grad, dist_x.grad)
+
+        return dist_inputs, dist_out