oneflow.nn.utils.rnn.pad_sequence

oneflow.nn.utils.rnn.pad_sequence(sequences: Union[oneflow.Tensor, List[oneflow.Tensor]], batch_first: bool = False, padding_value: float = 0.0) → oneflow.Tensor

The interface is consistent with PyTorch. The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.nn.utils.rnn.pad_sequence.html.

Pad a list of variable length Tensors with padding_value

pad_sequence stacks a list of Tensors along a new dimension, and pads them to equal length. For example, if the input is list of sequences with size L x * and if batch_first is False, and T x B x * otherwise.

B is batch size. It is equal to the number of elements in sequences. T is length of the longest sequence. L is length of the sequence. * is any number of trailing dimensions, including none.

Note

This function returns a Tensor of size T x B x * or B x T x * where T is the length of the longest sequence. This function assumes trailing dimensions and type of all the Tensors in sequences are same.

Parameters

  • sequences (list[Tensor]) – list of variable length sequences.

  • batch_first (bool, optional) – output will be in B x T x * if True, or in T x B x * otherwise. Default: False.

  • padding_value (float, optional) – value for padded elements. Default: 0.

Returns

Tensor of size T x B x * if batch_first is False. Tensor of size B x T x * otherwise

For example:

>>> from oneflow.nn.utils.rnn import pad_sequence
>>> import oneflow as flow

>>> a = flow.ones(25, 300)
>>> b = flow.ones(22, 300)
>>> c = flow.ones(15, 300)
>>> out = pad_sequence([a, b, c])
>>> out.size()
oneflow.Size([25, 3, 300])