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 v2.4.1+cu121. 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 v2.4.1+cu121 is listed below.

Description From Torch v2.4.1+cu121

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.