megengine.distributed.helper.param_pack_split

param_pack_split(inp, offsets, shapes)[source]

Returns split tensor to list of tensors as offsets and shapes described, only used for parampack.

Parameters

  • inp (Tensor) – input tensor.

  • offsets (list) – offsets of outputs, length of 2 * n, where n is the number of tensor you want to split, format [begin0, end0, begin1, end1].

  • shapes (list) – tensor shapes of outputs.

Returns

splitted tensors.

Examples

>>> 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)

previous

AllreduceCallback

next

megengine.distributed.helper.param_pack_concat