Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
e8d0f9db
MegEngine
项目概览
MegEngine 天元
/
MegEngine
接近 2 年 前同步成功
通知
414
Star
4708
Fork
583
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
MegEngine
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
提交
e8d0f9db
编写于
2月 28, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
docs(example): replace dist.help testcode with doctest format
GitOrigin-RevId: 79a7b72ca3b33d6566d7b169b040fe73b540223a
上级
f5f9249a
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
21 addition
and
44 deletion
+21
-44
imperative/python/megengine/distributed/helper.py
imperative/python/megengine/distributed/helper.py
+21
-44
未找到文件。
imperative/python/megengine/distributed/helper.py
浏览文件 @
e8d0f9db
...
...
@@ -28,14 +28,14 @@ from .group import WORLD, Group, group_barrier, is_distributed, override_backend
def
param_pack_split
(
inp
:
Tensor
,
offsets
:
list
,
shapes
:
list
):
r
"""Returns split tensor to
tensor list
as offsets and shapes described,
r
"""Returns split tensor to
list of tensors
as offsets and shapes described,
only used for ``parampack``.
Args:
inp: input tensor.
offsets: offsets of outputs, length of `
2 * n
`,
wh
ile n is tensor nums
you want to split,
format `
[begin0, end0, begin1, end1]
`.
offsets: offsets of outputs, length of `
`2 * n`
`,
wh
ere ``n`` is the number of tensor
you want to split,
format `
`[begin0, end0, begin1, end1]`
`.
shapes: tensor shapes of outputs.
Returns:
...
...
@@ -43,25 +43,14 @@ def param_pack_split(inp: Tensor, offsets: list, shapes: list):
Examples:
.. testcode::
import numpy as np
from megengine import tensor
from megengine.distributed.helper import param_pack_split
a = tensor(np.ones((10,), np.int32))
b, c = param_pack_split(a, [0, 1, 1, 10], [(1,), (3, 3)])
print(b.numpy())
print(c.numpy())
Outputs:
.. testoutput::
[1]
[[1 1 1]
[1 1 1]
[1 1 1]]
>>> a = F.ones(10)
>>> b, c = dist.helper.param_pack_split(a, [0, 1, 1, 10], [(1,), (3, 3)])
>>> b
Tensor([1.], device=xpux:0)
>>> c
Tensor([[1. 1. 1.]
[1. 1. 1.]
[1. 1. 1.]], device=xpux:0)
"""
op
=
ParamPackSplit
()
op
.
offsets
=
offsets
...
...
@@ -74,34 +63,22 @@ def param_pack_concat(inps: list, offsets: Tensor, offsets_val: list):
r
"""Returns concated tensor, only used for ``parampack``.
Args:
inps: input tensors.
inps:
list of
input tensors.
offsets: device value of offsets.
offsets_val: offsets of inputs, length of `
2 * n
`,
format `
[begin0, end0, begin1, end1]
`.
offsets_val: offsets of inputs, length of `
`2 * n`
`,
format `
`[begin0, end0, begin1, end1]`
`.
Returns:
concated tensor.
Examples:
.. testcode::
import numpy as np
from megengine import tensor
from megengine.distributed.helper import param_pack_concat
a = tensor(np.ones((1,), np.int32))
b = tensor(np.ones((3, 3), np.int32))
offsets_val = [0, 1, 1, 10]
offsets = tensor(offsets_val, np.int32)
c = param_pack_concat([a, b], offsets, offsets_val)
print(c.numpy())
Outputs:
.. testoutput::
[1 1 1 1 1 1 1 1 1 1]
>>> a = F.ones(1)
>>> b = F.ones((3, 3))
>>> offsets_val = [0, 1, 1, 10]
>>> offsets = Tensor(offsets_val)
>>> c = dist.helper.param_pack_concat([a, b], offsets, offsets_val) # doctest: +SKIP
Tensor([1. 1. 1. 1. 1. 1. 1. 1. 1. 1.], device=xpux:0)
"""
op
=
ParamPackConcat
()
op
.
offsets
=
offsets_val
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录