megengine.distributed.helper.param_pack_concat

param_pack_concat(inps, offsets, offsets_val)[source]

Returns concated tensor, only used for parampack.

Parameters

  • inps (list) – list of input tensors.

  • offsets (Tensor) – device value of offsets.

  • offsets_val (list) – offsets of inputs, length of 2 * n, format [begin0, end0, begin1, end1].

Returns

concated tensor.

Examples

>>> 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)  
Tensor([1. 1. 1. 1. 1. 1. 1. 1. 1. 1.], device=xpux:0)

previous

megengine.distributed.helper.param_pack_split

next

megengine.distributed.helper.pack_allreduce_split