megengine.functional.stack

stack(inps, axis=0, device=None)[源代码]

在一个新的轴上连接一个序列中的张量。输入张量必须有相同的形状。

参数
  • inps – 输入张量。

  • axis – 用来连接的轴。

  • device – 输出所在的计算节点。 默认:None

返回

输入拼接好的张量。

实际案例

import numpy as np
from megengine import tensor
import megengine.functional as F

x1 = tensor(np.arange(0, 3, dtype=np.float32).reshape((3)))
x2 = tensor(np.arange(6, 9, dtype=np.float32).reshape((3)))
out = F.stack([x1, x2], axis=0)
print(out.numpy())

输出:

[[0. 1. 2.]
 [6. 7. 8.]]