stack

Documentation

treetensor.torch.stack(tensors, *args, **kwargs)

Concatenates a sequence of tensors along a new dimension.

Examples:

>>> import torch
>>> import treetensor.torch as ttorch
>>> t1 = torch.randint(10, 30, (2, 3))
>>> t1
tensor([[17, 15, 27],
        [12, 17, 29]])
>>> t2 = torch.randint(30, 50, (2, 3))
>>> t2
tensor([[45, 41, 47],
        [37, 37, 36]])
>>> t3 = torch.randint(50, 70, (2, 3))
>>> t3
tensor([[60, 50, 55],
        [69, 54, 58]])
>>> ttorch.stack((t1, t2, t3))
tensor([[[17, 15, 27],
         [12, 17, 29]],

        [[45, 41, 47],
         [37, 37, 36]],

        [[60, 50, 55],
         [69, 54, 58]]])

>>> tt1 = ttorch.randint(10, 30, {
...     'a': (2,  3),
...     'b': {'x': (3, 4)},
... })
>>> tt1
<Tensor 0x7f4c8eba9630>
├── a --> tensor([[25, 22, 29],
│                 [19, 21, 27]])
└── b --> <Tensor 0x7f4c8eba9550>
    └── x --> tensor([[20, 17, 28, 10],
                      [28, 16, 27, 27],
                      [18, 21, 17, 12]])
>>> tt2 = ttorch.randint(30, 50, {
...     'a': (2,  3),
...     'b': {'x': (3, 4)},
... })
>>> tt2
<Tensor 0x7f4c8eba97b8>
├── a --> tensor([[40, 44, 41],
│                 [39, 44, 40]])
└── b --> <Tensor 0x7f4c8eba9710>
    └── x --> tensor([[44, 42, 38, 44],
                      [30, 44, 42, 31],
                      [36, 30, 33, 31]])
>>> ttorch.stack((tt1, tt2))
<Tensor 0x7f4c8eb411d0>
├── a --> tensor([[[25, 22, 29],
│                  [19, 21, 27]],
│
│                 [[40, 44, 41],
│                  [39, 44, 40]]])
└── b --> <Tensor 0x7f4c8eb410b8>
    └── x --> tensor([[[20, 17, 28, 10],
                       [28, 16, 27, 27],
                       [18, 21, 17, 12]],

                      [[44, 42, 38, 44],
                       [30, 44, 42, 31],
                       [36, 30, 33, 31]]])
>>> ttorch.stack((tt1, tt2), dim=1)
<Tensor 0x7f4c8eba9da0>
├── a --> tensor([[[25, 22, 29],
│                  [40, 44, 41]],
│
│                 [[19, 21, 27],
│                  [39, 44, 40]]])
└── b --> <Tensor 0x7f4d01fb4898>
    └── x --> tensor([[[20, 17, 28, 10],
                       [44, 42, 38, 44]],

                      [[28, 16, 27, 27],
                       [30, 44, 42, 31]],

                      [[18, 21, 17, 12],
                       [36, 30, 33, 31]]])

Torch Version Related

This documentation is based on torch.stack in torch v1.9.0+cu102. Its arguments’ arrangements depend on the version of pytorch you installed.

If some arguments listed here are not working properly, please check your pytorch’s version with the following command and find its documentation.

python -c 'import torch;print(torch.__version__)'

The arguments and keyword arguments supported in torch v1.9.0+cu102 is listed below.

Description From Torch v1.9.0+cu102

torch.stack(tensors, dim=0, *, out=None) → Tensor

Concatenates a sequence of tensors along a new dimension.

All tensors need to be of the same size.

Arguments:

tensors (sequence of Tensors): sequence of tensors to concatenate dim (int): dimension to insert. Has to be between 0 and the number

of dimensions of concatenated tensors (inclusive)

Keyword args:

out (Tensor, optional): the output tensor.