Contents

Pad a list of variable length tensors with zeros, or some other value, and pack them into a single tensor.

Parameters:

  • tensors (list[Tensor] | tuple[Tensor] | Tensor) – List of tensors to be packed, all with identical dtypes.

  • axis (int) – Axis along which tensors should be packed; 0 for first axis -1 for the last axis, etc. This will be a new dimension.

  • value (int | float) – The value with which the tensor is to be padded.

  • size (list[int] | tuple[int, …] | torch.Size | None) – Size of each dimension to which tensors should be padded. This to the largest size encountered along each dimension.

  • return_mask (bool, optional) – If True, a mask identifying the padding values is returned. Defaults to False.

returns:

Input tensors padded and packed into a single tensor. Optionally, the mask is also returned.

rtype:

Tensor | tuple[Tensor, Tensor]

Examples

Multiple tensors can be packed into a single tensor like so:

>>> from tbmalt.common.batch import pack
>>> import torch
>>> a, b, c = torch.rand(2,2), torch.rand(3,3), torch.rand(4,4)
>>> abc_packed_a = pack([a, b, c])
>>> print(abc_packed_a.shape)
torch.Size([3, 4, 4])
>>> abc_packed_b = pack([a, b, c], axis=1)
>>> print(abc_packed_b.shape)
torch.Size([4, 3, 4])
>>> abc_packed_c = pack([a, b, c], axis=-1)
>>> print(abc_packed_c.shape)
torch.Size([4, 4, 3])

An optional mask identifying the padding values can also be returned:

>>> packed, mask = pack(
...     [
...         torch.tensor([1.0]),
...         torch.tensor([2.0, 2.0]),
...         torch.tensor([3.0, 3.0, 3.0]),
...     ],
...     return_mask=True,
... )
>>> print(packed)
tensor([[1., 0., 0.],
        [2., 2., 0.],
        [3., 3., 3.]])
>>> print(mask)
tensor([[ True, False, False],
        [ True,  True, False],
        [ True,  True,  True]])
Contents