oneflow.stack

oneflow.stack()

Concatenates a sequence of tensors along a new dimension. The returned tensor shares the same underlying data with input tensors.

A dim value within the range [-input.ndimension() - 1, input.ndimension() + 1] can be used. Negative dim will correspond to stack() applied at dim = dim + input.ndimension() + 1.

Parameters
  • inputs (List[oneflow.Tensor]) – the list of input tensors. Each tensor should have the same shape.

  • dim (int) – the index at which to insert the concatenated dimension.

Returns

A Tensor

For example:

>>> import oneflow as flow
>>> import numpy as np

>>> x1 = flow.tensor(np.random.rand(1, 3, 5))
>>> x2 = flow.tensor(np.random.rand(1, 3, 5))
>>> y = flow.stack([x1, x2], dim = -1)
>>> y.shape
oneflow.Size([1, 3, 5, 2])