From d1bbbcc6a405d7403a5b319441387a2bef2b15ae Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 28 Aug 2020 18:06:37 +0800
Subject: [PATCH] docs(mge/functional): add doc for copy, param pack split and
 concat

GitOrigin-RevId: c93bdae29015e20f4ac92d3c713771c512ac0ab6
---
 imperative/python/megengine/functional/tensor.py | 61 ++++++++++++++++++++++++
 imperative/python/megengine/functional/utils.py  | 24 ++++++++++
 2 files changed, 85 insertions(+)

diff --git a/imperative/python/megengine/functional/tensor.py b/imperative/python/megengine/functional/tensor.py
index 8624f73f..d0f53bf4 100644
--- a/imperative/python/megengine/functional/tensor.py
+++ b/imperative/python/megengine/functional/tensor.py
@@ -937,6 +937,38 @@ def arange(
 
 
 def param_pack_split(inp: Tensor, offsets: List, shapes: List) -> Tensor:
+    r"""
+    Returns split Tensor to Tensor list as offsets and shapes described,
+            only used for parampack.
+
+    :param inp: Input tensor
+    :param offsets: offsets of outputs, length of 2 * n,
+            while n is tensor nums you want to split,
+            format [begin0, end0, begin1, end1].
+    :param shapes: tensor shapes of outputs
+    :return: split tensors
+
+    Examples:
+
+    .. testcode::
+
+        import numpy as np
+        import megengine.functional as F
+        from megengine import tensor
+
+        a = tensor(np.ones((10,), np.int32))
+        b, c = F.param_pack_split(a, [0, 1, 1, 10], [(1,), (3, 3)])
+        print(b.numpy())
+        print(c.numpy())
+
+    .. testoutput::
+
+        [1]
+        [[1 1 1]
+         [1 1 1]
+         [1 1 1]]
+
+    """
     op = builtin.ParamPackSplit()
     op.offsets = offsets
     op.shapes = shapes
@@ -944,6 +976,35 @@ def param_pack_split(inp: Tensor, offsets: List, shapes: List) -> Tensor:
 
 
 def param_pack_concat(inps: List, offsets: Tensor, offsets_val: List) -> Tensor:
+    r"""
+    Returns concat Tensor, only used for parampack.
+
+    :param inps: Input tensors
+    :param offsets: offsets of inputs, length of 2 * n,
+            format [begin0, end0, begin1, end1].
+    :param offsets_val: device value of offsets
+    :return: split tensors
+
+    Examples:
+
+    .. testcode::
+
+        import numpy as np
+        import megengine.functional as F
+        from megengine import tensor
+
+        a = tensor(np.ones((1,), np.int32))
+        b = tensor(np.ones((3, 3), np.int32))
+        offsets = [0, 1, 1, 10]
+        offsets_val = tensor(offsets, np.int32)
+        c = F.param_pack_concat([a, b], offsets, offsets_val)
+        print(c.numpy())
+
+    .. testoutput::
+
+        [1 1 1 1 1 1 1 1 1 1]
+
+    """
     op = builtin.ParamPackConcat()
     op.offsets = offsets_val
     return apply(op, *inps, offsets)[0]
diff --git a/imperative/python/megengine/functional/utils.py b/imperative/python/megengine/functional/utils.py
index 85c85d43..196687be 100644
--- a/imperative/python/megengine/functional/utils.py
+++ b/imperative/python/megengine/functional/utils.py
@@ -77,4 +77,28 @@ def zero_grad(inp: Tensor) -> Tensor:
 
 
 def copy(inp, cn):
+    r"""
+    Copy tensor to another device.
+
+    :param inp: Input tensor.
+    :param cn: device that you copy to
+
+    Examples:
+
+    .. testcode::
+
+        import numpy as np
+        from megengine import tensor
+        import megengine.functional as F
+
+        x = tensor([1, 2, 3], np.int32)
+        y = F.copy(x, "xpu1")
+        print(y.numpy())
+
+    Outputs:
+
+    .. testoutput::
+
+        [1 2 3]
+    """
     return apply(Copy(comp_node=cn), inp)[0]