update dygraph auto_parallel en API docs. (#59557)

wuhuachaocoding · web-flow · commit 6577edbca5a7 · 2023-12-04T16:36:53.000+08:00
diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc
@@ -372,7 +372,30 @@ void BindAutoParallel(py::module *m) {
           py::arg("memo"))
       .def("__str__", &DeviceMesh::to_string);
 
-  py::enum_<phi::ReduceType>(*m, "ReduceType")
+  py::enum_<phi::ReduceType>(*m, "ReduceType", R"DOC(
+    Specify the type of operation used for paddle.distributed.Partial().
+    It should be one of the following values:
+
+        - ReduceType.kRedSum
+        - ReduceType.kRedMax
+        - ReduceType.kRedMin
+        - ReduceType.kRedProd
+        - ReduceType.kRedAvg
+        - ReduceType.kRedAny
+        - ReduceType.kRedAll
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+            >>> a = paddle.ones([10, 20])
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # distributed tensor
+            >>> d_tensor = dist.shard_tensor(a, mesh, [dist.Partial(dist.ReduceType.kRedSum)])
+
+      )DOC")
       .value("kRedSum", phi::ReduceType::kRedSum)
       .value("kRedMax", phi::ReduceType::kRedMax)
       .value("kRedMin", phi::ReduceType::kRedMin)
@@ -383,7 +406,25 @@ void BindAutoParallel(py::module *m) {
 
   auto Placement =
       py::class_<phi::distributed::Placement,
-                 std::shared_ptr<phi::distributed::Placement>>(*m, "Placement")
+                 std::shared_ptr<phi::distributed::Placement>>(
+          *m, "Placement", R"DOC(
+        The `Placement` is base class that describes how to place the tensor on ProcessMesh. it has three subclass: `Replicate`, `Shard` and `Partial`.
+
+        Examples:
+            .. code-block:: python
+
+                >>> import paddle.distributed as dist
+                >>> placements = [dist.Replicate(), dist.Shard(0), dist.Partial()]
+                >>> for p in placements:
+                >>>     if isinstance(p, dist.Placement):
+                >>>         if p.is_replicated():
+                >>>             print("replicate.")
+                >>>         elif p.is_shard():
+                >>>             print("shard.")
+                >>>         elif p.is_partial():
+                >>>             print("partial.")
+
+      )DOC")
           .def(py::init<>())
           .def("is_shard",
                &phi::distributed::Placement::is_shard,
@@ -397,7 +438,24 @@ void BindAutoParallel(py::module *m) {
 
   auto Shard = py::class_<phi::distributed::Shard,
                           std::shared_ptr<phi::distributed::Shard>>(
-                   *m, "Shard", Placement)
+                   *m, "Shard", Placement, R"DOC(
+               The `Shard` describes how `Tensor` splitted across multiple devices according to specified dimensions.
+
+               Parameters:
+                   dim (int): specify the slicing dimension of the tensor.
+
+               Examples:
+                   .. code-block:: python
+
+                       >>> import paddle
+                       >>> import paddle.distributed as dist
+                       >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=['x', 'y'])
+                       >>> a = paddle.to_tensor([[1,2,3],[5,6,7]])
+                       >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                       >>> # distributed tensor
+                       >>> d_tensor = dist.shard_tensor(a, mesh, [dist.Shard(0), dist.Shard(1)])
+
+               )DOC")
                    .def(py::init([](int64_t dim) {
                      return std::make_shared<phi::distributed::Shard>(dim);
                    }))
@@ -409,7 +467,21 @@ void BindAutoParallel(py::module *m) {
 
   auto Replicate = py::class_<phi::distributed::Replicate,
                               std::shared_ptr<phi::distributed::Replicate>>(
-                       *m, "Replicate", Placement)
+                       *m, "Replicate", Placement, R"DOC(
+                   The `Replicate` describes the tensor placed repeatedly on ProcessMesh.
+
+                   Examples:
+                       .. code-block:: python
+
+                           >>> import paddle
+                           >>> import paddle.distributed as dist
+                           >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+                           >>> a = paddle.ones([10, 20])
+                           >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                           >>> # distributed tensor
+                           >>> d_tensor = dist.shard_tensor(a, mesh, [dist.Replicate()])
+
+                   )DOC")
                        .def(py::init<>())
                        .def("__hash__", &phi::distributed::Replicate::hash)
                        .def("__str__", &phi::distributed::Replicate::to_string)
@@ -418,7 +490,24 @@ void BindAutoParallel(py::module *m) {
 
   auto Partial = py::class_<phi::distributed::Partial,
                             std::shared_ptr<phi::distributed::Partial>>(
-                     *m, "Partial", Placement)
+                     *m, "Partial", Placement, R"DOC(
+                 The `Partial` describes `Tensor` across multiple devices, this type of tensor has the same shape but only a fraction of the value, which can be further reduce (e.g. sum/min/max) to obtain dist_tensor, often used as an intermediate representation.
+
+                 Parameters:
+                   reduce_type (paddle.distributed.ReduceType): the reduce type of the Partial state, default `paddle.distributed.ReduceType.kRedSum`.
+
+                 Examples:
+                     .. code-block:: python
+
+                         >>> import paddle
+                         >>> import paddle.distributed as dist
+                         >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+                         >>> a = paddle.ones([10, 20])
+                         >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                         >>> # distributed tensor
+                         >>> d_tensor = dist.shard_tensor(a, mesh, [dist.Partial()])
+
+                 )DOC")
                      .def(py::init<phi::ReduceType>(),
                           py::arg("reduce_type") = phi::ReduceType::kRedSum)
                      .def("__hash__", &phi::distributed::Partial::hash)
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -68,9 +68,8 @@
 
 from .auto_parallel.process_mesh import ProcessMesh
 
+from paddle.base.core import ReduceType, Placement
 from .auto_parallel.placement_type import (
-    ReduceType,
-    Placement,
     Shard,
     Replicate,
     Partial,
diff --git a/python/paddle/distributed/auto_parallel/placement_type.py b/python/paddle/distributed/auto_parallel/placement_type.py
@@ -13,9 +13,7 @@
 # limitations under the License.
 from typing import cast
 
-from paddle.base.core import Partial, Placement, ReduceType, Replicate, Shard
-
-__all__ = ["ReduceType", "Placement", "Replicate", "Shard", "Partial"]
+from paddle.base.core import Partial, Replicate, Shard
 
 
 def to_placements(dim_map, mesh, partial_idx=[]):
