提交 d1bbbcc6 编写于 作者: M Megvii Engine Team

docs(mge/functional): add doc for copy, param pack split and concat

GitOrigin-RevId: c93bdae29015e20f4ac92d3c713771c512ac0ab6
上级 7e8f7209
......@@ -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]
......@@ -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]
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册