megengine.functional.stack

stack(inps, axis=0, device=None)[source]

Concats a sequence of tensors along a new axis. The input tensors must have the same shape.

Parameters
  • inps – input tensors.

  • axis – which axis will be concatenated.

  • device – the device output will be. Default: None

Returns

output concatenated tensor.

Examples

>>> import numpy as np
>>> 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)
>>> out.numpy()
array([[0., 1., 2.],
       [6., 7., 8.]], dtype=float32)