chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,549 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import Tensor
|
||||
|
||||
__all__ = [
|
||||
"PackedSequence",
|
||||
"invert_permutation",
|
||||
"pack_padded_sequence",
|
||||
"pad_packed_sequence",
|
||||
"pad_sequence",
|
||||
"unpad_sequence",
|
||||
"pack_sequence",
|
||||
"unpack_sequence",
|
||||
]
|
||||
|
||||
|
||||
def invert_permutation(permutation: Tensor | None) -> Tensor | None:
|
||||
"""Returns the inverse of ``permutation``.
|
||||
|
||||
This is useful for converting between sorted and unsorted indices in
|
||||
a :class:`~nn.utils.rnn.PackedSequence`.
|
||||
|
||||
Args:
|
||||
permutation (Tensor|None): a 1-D tensor of indices to invert.
|
||||
|
||||
Returns:
|
||||
Tensor|None: the inverse permutation tensor, or None if input is None.
|
||||
|
||||
Examples:
|
||||
>>> import paddle
|
||||
"""
|
||||
if permutation is None:
|
||||
return None
|
||||
# Use paddle.scatter instead of scatter_ for better static mode support
|
||||
output = paddle.scatter(
|
||||
paddle.zeros_like(permutation),
|
||||
permutation,
|
||||
paddle.arange(
|
||||
0,
|
||||
permutation.numel(),
|
||||
dtype=permutation.dtype,
|
||||
device=permutation.place,
|
||||
),
|
||||
overwrite=True,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
class PackedSequence:
|
||||
"""Holds the data and batch sizes of a packed sequence.
|
||||
|
||||
PackedSequence is used to represent a packed sequence, which is typically
|
||||
produced by ``pack_padded_sequence`` and consumed by ``pad_packed_sequence``.
|
||||
|
||||
Args:
|
||||
data (Tensor): The packed data tensor.
|
||||
batch_sizes (Tensor): A tensor containing the batch size at each step.
|
||||
sorted_indices (Tensor|None, optional): The indices used to sort the sequences.
|
||||
unsorted_indices (Tensor|None, optional): The indices to restore the original order.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data: Tensor,
|
||||
batch_sizes: Tensor,
|
||||
sorted_indices: Tensor | None = None,
|
||||
unsorted_indices: Tensor | None = None,
|
||||
):
|
||||
self.data = data
|
||||
self.batch_sizes = batch_sizes
|
||||
self.sorted_indices = sorted_indices
|
||||
self.unsorted_indices = unsorted_indices
|
||||
|
||||
@property
|
||||
def is_pinned(self) -> bool:
|
||||
return (
|
||||
self.data.place.is_cuda_pinned_place()
|
||||
or self.data.place.is_xpu_pinned_place()
|
||||
)
|
||||
|
||||
def to(self, *args, **kwargs) -> PackedSequence:
|
||||
data = self.data.to(*args, **kwargs)
|
||||
if data is self.data:
|
||||
return self
|
||||
# Only convert indices to same device as data, not dtype
|
||||
target_device = data.place
|
||||
sorted_indices = (
|
||||
self.sorted_indices.to(target_device)
|
||||
if self.sorted_indices is not None
|
||||
else None
|
||||
)
|
||||
unsorted_indices = (
|
||||
self.unsorted_indices.to(target_device)
|
||||
if self.unsorted_indices is not None
|
||||
else None
|
||||
)
|
||||
return PackedSequence(
|
||||
data, self.batch_sizes, sorted_indices, unsorted_indices
|
||||
)
|
||||
|
||||
def cuda(self) -> PackedSequence:
|
||||
return self.to(device="gpu")
|
||||
|
||||
def cpu(self) -> PackedSequence:
|
||||
return self.to(device="cpu")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"PackedSequence(data={self.data}, batch_sizes={self.batch_sizes}, "
|
||||
f"sorted_indices={self.sorted_indices}, unsorted_indices={self.unsorted_indices})"
|
||||
)
|
||||
|
||||
def pin_memory(self) -> PackedSequence:
|
||||
return PackedSequence(
|
||||
self.data.pin_memory(),
|
||||
self.batch_sizes,
|
||||
self.sorted_indices.pin_memory()
|
||||
if self.sorted_indices is not None
|
||||
else None,
|
||||
self.unsorted_indices.pin_memory()
|
||||
if self.unsorted_indices is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
@property
|
||||
def is_cuda(self) -> bool:
|
||||
return self.data.is_cuda
|
||||
|
||||
def double(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.float64)
|
||||
|
||||
def float(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.float32)
|
||||
|
||||
def half(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.float16)
|
||||
|
||||
def long(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.int64)
|
||||
|
||||
def int(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.int32)
|
||||
|
||||
def short(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.int16)
|
||||
|
||||
def char(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.int8)
|
||||
|
||||
def byte(self) -> PackedSequence:
|
||||
return self.to(dtype=paddle.uint8)
|
||||
|
||||
|
||||
def pack_padded_sequence(
|
||||
input: Tensor,
|
||||
lengths: Tensor | list[int],
|
||||
batch_first: bool = False,
|
||||
enforce_sorted: bool = True,
|
||||
) -> PackedSequence:
|
||||
r"""Packs a Tensor containing padded sequences of variable length.
|
||||
|
||||
This function packs a Tensor containing padded sequences into a PackedSequence
|
||||
object, which can be used as input to a recurrent neural network.
|
||||
|
||||
Args:
|
||||
input (Tensor): The padded sequence tensor. Shape is ``T x B x *`` if
|
||||
``batch_first`` is False, or ``B x T x *`` if ``batch_first`` is True,
|
||||
where ``T`` is the length of the longest sequence, ``B`` is the batch size.
|
||||
lengths (Tensor|list[int]): The lengths of each sequence in the batch.
|
||||
batch_first (bool, optional): If True, the input is expected to be in
|
||||
``B x T x *`` format. Default: False.
|
||||
enforce_sorted (bool, optional): If True, the input is expected to contain
|
||||
sequences sorted by length in descending order. Default: True.
|
||||
|
||||
Returns:
|
||||
PackedSequence: A PackedSequence object containing the packed data.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
"""
|
||||
if batch_first:
|
||||
input = input.transpose([1, 0, *range(2, len(input.shape))])
|
||||
|
||||
if isinstance(lengths, paddle.Tensor):
|
||||
lengths = lengths.tolist()
|
||||
|
||||
batch_size = input.shape[1]
|
||||
|
||||
if len(lengths) != batch_size:
|
||||
raise ValueError(
|
||||
f"Length of lengths ({len(lengths)}) does not match batch size ({batch_size})"
|
||||
)
|
||||
|
||||
sorted_indices = None
|
||||
unsorted_indices = None
|
||||
|
||||
if not enforce_sorted:
|
||||
sorted_lengths = sorted(
|
||||
enumerate(lengths), key=lambda x: x[1], reverse=True
|
||||
)
|
||||
sorted_indices = paddle.to_tensor(
|
||||
[i for i, _ in sorted_lengths], place=input.place
|
||||
)
|
||||
unsorted_indices = paddle.argsort(sorted_indices)
|
||||
lengths = [l for _, l in sorted_lengths]
|
||||
# Use index_select to reorder along batch dimension (axis=1)
|
||||
input = paddle.index_select(input, sorted_indices, axis=1)
|
||||
|
||||
packed_data_list = []
|
||||
batch_sizes_list = []
|
||||
|
||||
# num_steps may be different from actual input shape[0] after sorting
|
||||
# We need to iterate over the actual sequence length
|
||||
actual_num_steps = input.shape[0]
|
||||
for step in range(actual_num_steps):
|
||||
batch_size_at_step = sum(1 for l in lengths if l > step)
|
||||
if batch_size_at_step > 0:
|
||||
packed_data_list.append(input[step, :batch_size_at_step])
|
||||
batch_sizes_list.append(batch_size_at_step)
|
||||
|
||||
packed_data = paddle.concat(packed_data_list, axis=0)
|
||||
batch_sizes = paddle.to_tensor(batch_sizes_list, dtype="int64")
|
||||
|
||||
return PackedSequence(
|
||||
packed_data, batch_sizes, sorted_indices, unsorted_indices
|
||||
)
|
||||
|
||||
|
||||
def pad_packed_sequence(
|
||||
sequence: PackedSequence,
|
||||
batch_first: bool = False,
|
||||
padding_value: float = 0.0,
|
||||
total_length: int | None = None,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
r"""Pads a packed sequence to a Tensor of padded sequences.
|
||||
|
||||
This function is the inverse of ``pack_padded_sequence``. It takes a PackedSequence
|
||||
and returns a padded Tensor and a list of lengths.
|
||||
|
||||
Args:
|
||||
sequence (PackedSequence): The packed sequence to pad.
|
||||
batch_first (bool, optional): If True, the output will be in ``B x T x *``
|
||||
format. Default: False.
|
||||
padding_value (float, optional): The value to use for padding. Default: 0.0.
|
||||
total_length (int|None, optional): If not None, the output will be padded to
|
||||
this length. Default: None.
|
||||
|
||||
Returns:
|
||||
tuple[Tensor, Tensor]: A tuple containing:
|
||||
- The padded sequence tensor.
|
||||
- A tensor of sequence lengths.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
"""
|
||||
if not isinstance(sequence, PackedSequence):
|
||||
raise TypeError(f"Expected PackedSequence, got {type(sequence)}")
|
||||
|
||||
data = sequence.data
|
||||
batch_sizes = sequence.batch_sizes.tolist()
|
||||
unsorted_indices = sequence.unsorted_indices
|
||||
|
||||
max_seq_len = len(batch_sizes)
|
||||
max_batch_size = batch_sizes[0]
|
||||
|
||||
if total_length is not None:
|
||||
if total_length < max_seq_len:
|
||||
raise ValueError(
|
||||
f"total_length ({total_length}) must be >= max sequence length ({max_seq_len})"
|
||||
)
|
||||
|
||||
trailing_dims = list(data.shape[1:])
|
||||
|
||||
if total_length is not None and total_length > max_seq_len:
|
||||
output = paddle.full(
|
||||
[total_length, max_batch_size, *trailing_dims],
|
||||
padding_value,
|
||||
dtype=data.dtype,
|
||||
device=data.place,
|
||||
)
|
||||
else:
|
||||
output = paddle.full(
|
||||
[max_seq_len, max_batch_size, *trailing_dims],
|
||||
padding_value,
|
||||
dtype=data.dtype,
|
||||
device=data.place,
|
||||
)
|
||||
|
||||
data_offset = 0
|
||||
for step, batch_size in enumerate(batch_sizes):
|
||||
output[step, :batch_size] = data[data_offset : data_offset + batch_size]
|
||||
data_offset += batch_size
|
||||
|
||||
# Calculate lengths from batch_sizes
|
||||
# batch_sizes is in descending order, e.g., [3, 2, 1] means:
|
||||
# - First time step has 3 sequences
|
||||
# - Second time step has 2 sequences
|
||||
# - Third time step has 1 sequence
|
||||
# This means sequence lengths are [3, 2, 1] in sorted order
|
||||
lengths_list = []
|
||||
for i in range(max_batch_size):
|
||||
# Find the length of the i-th sequence (in sorted order)
|
||||
# It's the number of time steps where batch_sizes > i
|
||||
seq_len = sum(1 for bs in batch_sizes if bs > i)
|
||||
lengths_list.append(seq_len)
|
||||
lengths = paddle.to_tensor(lengths_list, dtype="int64", place=data.place)
|
||||
|
||||
if unsorted_indices is not None:
|
||||
output = output[:, unsorted_indices]
|
||||
lengths = lengths[unsorted_indices]
|
||||
|
||||
if batch_first:
|
||||
output = output.transpose([1, 0, *range(2, len(output.shape))])
|
||||
|
||||
return output, lengths
|
||||
|
||||
|
||||
def pad_sequence(
|
||||
sequences: Iterable[Tensor],
|
||||
batch_first: bool = False,
|
||||
padding_value: float = 0.0,
|
||||
padding_side: str = 'right',
|
||||
) -> Tensor:
|
||||
r"""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. ``sequences`` can be a list of sequences with size
|
||||
``L x *``, where ``L`` is the length of the sequence and ``*`` is any
|
||||
number of dimensions (including 0). If ``batch_first`` is ``False``, the
|
||||
output is of size ``T x B x *``, and ``B x T x *`` otherwise, where ``B``
|
||||
is the batch size (the number of elements in ``sequences``), ``T`` is the
|
||||
length of the longest sequence.
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
sequences (list[Tensor]): list of variable length sequences.
|
||||
batch_first (bool, optional): if ``True``, the output will be in
|
||||
``B x T x *`` format, ``T x B x *`` otherwise. Default: ``False``.
|
||||
padding_value (float, optional): value for padded elements.
|
||||
Default: ``0.0``.
|
||||
padding_side (str, optional): the side to pad the sequences on,
|
||||
either ``'right'`` or ``'left'``. Default: ``'right'``.
|
||||
|
||||
Returns:
|
||||
Tensor: Tensor of size ``T x B x *`` if ``batch_first`` is ``False``,
|
||||
or ``B x T x *`` otherwise.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> a = paddle.ones([25, 300])
|
||||
>>> b = paddle.ones([22, 300])
|
||||
>>> c = paddle.ones([15, 300])
|
||||
>>> padded = paddle.nn.utils.rnn.pad_sequence([a, b, c])
|
||||
>>> print(padded.shape)
|
||||
paddle.Size([25, 3, 300])
|
||||
>>> padded = paddle.nn.utils.rnn.pad_sequence([a, b, c], batch_first=True)
|
||||
>>> print(padded.shape)
|
||||
paddle.Size([3, 25, 300])
|
||||
|
||||
"""
|
||||
if not isinstance(sequences, Iterable):
|
||||
raise TypeError(
|
||||
f"pad_sequence expects an iterable of Tensors, but got {type(sequences)}"
|
||||
)
|
||||
sequences = tuple(sequences)
|
||||
for seq in sequences:
|
||||
if not isinstance(seq, paddle.Tensor):
|
||||
raise TypeError(
|
||||
f"pad_sequence expects an iterable of Tensors, but got element of type {type(seq)}"
|
||||
)
|
||||
if padding_side not in ('right', 'left'):
|
||||
raise ValueError(
|
||||
f"padding_side must be 'right' or 'left', but got '{padding_side}'"
|
||||
)
|
||||
|
||||
max_len = max(seq.shape[0] for seq in sequences)
|
||||
trailing_dims = sequences[0].shape[1:]
|
||||
dtype = sequences[0].dtype
|
||||
|
||||
padded_seqs = []
|
||||
for seq in sequences:
|
||||
length = seq.shape[0]
|
||||
if length == max_len:
|
||||
padded_seqs.append(seq)
|
||||
else:
|
||||
pad_size = [max_len - length, *list(trailing_dims)]
|
||||
padding = paddle.full(pad_size, padding_value, dtype=dtype)
|
||||
if padding_side == 'right':
|
||||
padded_seqs.append(paddle.concat([seq, padding], axis=0))
|
||||
else:
|
||||
padded_seqs.append(paddle.concat([padding, seq], axis=0))
|
||||
|
||||
out = paddle.stack(padded_seqs, axis=0)
|
||||
|
||||
if not batch_first:
|
||||
# Transpose from B x T x * to T x B x *
|
||||
perm = [1, 0, *list(range(2, len(out.shape)))]
|
||||
out = out.transpose(perm)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def unpad_sequence(
|
||||
padded_sequences: Tensor,
|
||||
lengths: Tensor,
|
||||
batch_first: bool = False,
|
||||
) -> list[Tensor]:
|
||||
r"""Unpad a padded Tensor into a list of variable length Tensors.
|
||||
|
||||
``unpad_sequence`` unstacks a padded Tensor into a list of variable length
|
||||
Tensors.
|
||||
|
||||
Args:
|
||||
padded_sequences (Tensor): padded sequences.
|
||||
lengths (Tensor): length of original (unpadded) sequences.
|
||||
batch_first (bool, optional): whether batch dimension is first or not.
|
||||
Default: ``False``.
|
||||
|
||||
Returns:
|
||||
list[Tensor]: a list of Tensor objects with original lengths.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> a = paddle.ones([25, 300])
|
||||
>>> b = paddle.ones([22, 300])
|
||||
>>> c = paddle.ones([15, 300])
|
||||
>>> sequences = [a, b, c]
|
||||
>>> padded = paddle.nn.utils.rnn.pad_sequence(sequences)
|
||||
>>> lengths = paddle.to_tensor([v.shape[0] for v in sequences])
|
||||
>>> unpadded = paddle.nn.utils.rnn.unpad_sequence(padded, lengths)
|
||||
>>> paddle.allclose(sequences[0], unpadded[0]).item()
|
||||
True
|
||||
>>> paddle.allclose(sequences[1], unpadded[1]).item()
|
||||
True
|
||||
>>> paddle.allclose(sequences[2], unpadded[2]).item()
|
||||
True
|
||||
|
||||
"""
|
||||
if not batch_first:
|
||||
# Transpose from T x B x * to B x T x *
|
||||
perm = [1, 0, *list(range(2, len(padded_sequences.shape)))]
|
||||
padded_sequences = padded_sequences.transpose(perm)
|
||||
|
||||
unpadded = []
|
||||
for seq, length in zip(padded_sequences, lengths):
|
||||
length_val = length.item()
|
||||
unpadded.append(seq[:length_val])
|
||||
|
||||
return unpadded
|
||||
|
||||
|
||||
def pack_sequence(
|
||||
sequences: list[Tensor],
|
||||
enforce_sorted: bool = True,
|
||||
) -> PackedSequence:
|
||||
r"""Packs a list of variable length Tensors.
|
||||
|
||||
Consecutive call of the next functions: ``pad_sequence``, ``pack_padded_sequence``.
|
||||
|
||||
``sequences`` should be a list of Tensors of size ``L x *``, where `L` is
|
||||
the length of a sequence and `*` is any number of trailing dimensions,
|
||||
including ``0``.
|
||||
|
||||
For unsorted sequences, use `enforce_sorted = False`. If ``enforce_sorted``
|
||||
is ``True``, the sequences should be sorted in the order of decreasing length.
|
||||
``enforce_sorted = True`` is only necessary for ONNX export.
|
||||
|
||||
Args:
|
||||
sequences (list[Tensor]): A list of sequences of decreasing length.
|
||||
enforce_sorted (bool, optional): if ``True``, checks that the input
|
||||
contains sequences sorted by length in a decreasing order. If
|
||||
``False``, this condition is not checked. Default: ``True``.
|
||||
|
||||
Returns:
|
||||
PackedSequence: a PackedSequence object.
|
||||
|
||||
Examples:
|
||||
>>> import paddle
|
||||
|
||||
"""
|
||||
lengths = paddle.to_tensor([v.shape[0] for v in sequences])
|
||||
return pack_padded_sequence(
|
||||
pad_sequence(sequences), lengths, enforce_sorted=enforce_sorted
|
||||
)
|
||||
|
||||
|
||||
def unpack_sequence(packed_sequences: PackedSequence) -> list[Tensor]:
|
||||
r"""Unpack PackedSequence into a list of variable length Tensors.
|
||||
|
||||
``packed_sequences`` should be a PackedSequence object.
|
||||
|
||||
Args:
|
||||
packed_sequences (PackedSequence): A PackedSequence object.
|
||||
|
||||
Returns:
|
||||
list[Tensor]: a list of Tensor objects.
|
||||
|
||||
Examples:
|
||||
>>> import paddle
|
||||
|
||||
"""
|
||||
padded_sequences, lengths = pad_packed_sequence(
|
||||
packed_sequences, batch_first=True
|
||||
)
|
||||
unpacked_sequences = unpad_sequence(
|
||||
padded_sequences, lengths, batch_first=True
|
||||
)
|
||||
return unpacked_sequences
|
||||
Reference in New Issue
Block a user