wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/paddle/compat/nn/__init__.py
+920
View File
@@ -0,0 +1,920 @@
# 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
import warnings
from math import sqrt
from typing import TYPE_CHECKING
import paddle
from paddle import nn
from paddle.nn.modules.utils import _single
from paddle.utils.decorator_utils import ForbidKeywordsDecorator
from . import functional
from .transformer import MultiheadAttention
if TYPE_CHECKING:
from paddle import Tensor
from paddle._typing import (
DTypeLike,
PlaceLike,
Size1,
Size2,
Size3,
)
__all__ = [
'Unfold',
'Linear',
'Softmax',
'AvgPool1D',
'AvgPool2D',
'AvgPool3D',
'AvgPool1d',
'AvgPool2d',
'AvgPool3d',
'BatchNorm1D',
'BatchNorm2D',
'BatchNorm3D',
'BatchNorm1d',
'BatchNorm2d',
'BatchNorm3d',
'MultiheadAttention',
'SmoothL1Loss',
]
class BatchNorm1D(nn.BatchNorm1D):
def __init__(
self,
num_features: int,
eps: float = 1e-5,
momentum: float | None = 0.1,
affine: bool = True,
track_running_stats: bool = True,
device: PlaceLike | None = None,
dtype: DTypeLike | None = None,
) -> None:
if momentum is None:
paddle_momentum = None
else:
paddle_momentum = 1.0 - momentum
super().__init__(
num_features=num_features,
momentum=paddle_momentum,
epsilon=eps,
use_global_stats=None if track_running_stats else False,
affine=affine,
device=device,
dtype=dtype,
)
self.momentum = momentum
class BatchNorm2D(nn.BatchNorm2D):
def __init__(
self,
num_features: int,
eps: float = 1e-5,
momentum: float | None = 0.1,
affine: bool = True,
track_running_stats: bool = True,
device: PlaceLike | None = None,
dtype: DTypeLike | None = None,
) -> None:
if momentum is None:
paddle_momentum = None
else:
paddle_momentum = 1.0 - momentum
super().__init__(
num_features=num_features,
momentum=paddle_momentum,
epsilon=eps,
use_global_stats=None if track_running_stats else False,
affine=affine,
device=device,
dtype=dtype,
)
self.momentum = momentum
class BatchNorm3D(nn.BatchNorm3D):
def __init__(
self,
num_features: int,
eps: float = 1e-5,
momentum: float | None = 0.1,
affine: bool = True,
track_running_stats: bool = True,
device: PlaceLike | None = None,
dtype: DTypeLike | None = None,
) -> None:
if momentum is None:
paddle_momentum = None
else:
paddle_momentum = 1.0 - momentum
super().__init__(
num_features=num_features,
momentum=paddle_momentum,
epsilon=eps,
use_global_stats=None if track_running_stats else False,
affine=affine,
device=device,
dtype=dtype,
)
self.momentum = momentum
BatchNorm1d = BatchNorm1D
BatchNorm2d = BatchNorm2D
BatchNorm3d = BatchNorm3D
class AvgPool1D(nn.Layer):
r"""
This operation applies a 1D average pooling over an input signal composed
of several input planes, based on the input, output_size, return_mask parameters.
Input(X) and output(Out) are in NCL format, where N is batch
size, C is the number of channels, L is the length of the feature.
The output tensor shape will be [N, C, output_size].
The output value of the layer with input size (N, C, L),
output (N, C, :math:`L_{out}`) and kernel_size ksize can be precisely described as
For average pool1d:
.. math::
Output(N_i, C_i, l) = \frac{Input[N_i, C_i, stride \times l:stride \times l+k]}{ksize}
Parameters:
kernel_size(int|list|tuple): The pool kernel size. If pool kernel size is a tuple or list,
it must contain an integer.
stride(int|list|tuple|None, optional): The pool stride size. If pool stride size is a tuple or list,
it must contain an integer. Default None, then stride will be equal to the kernel_size.
padding(str|int|list|tuple, optional): The padding size. Padding could be in one of the following forms.
1. A string in ['valid', 'same'].
2. An int, which means the feature map is zero padded by size of `padding` on every sides.
3. A list[int] or tuple(int) whose length is 1, which means the feature map is zero padded by the size of `padding[0]` on every sides.
4. A list[int] or tuple(int) whose length is 2. It has the form [pad_before, pad_after].
5. A list or tuple of pairs of integers. It has the form [[pad_before, pad_after], [pad_before, pad_after], ...]. Note that, the batch dimension and channel dimension should be [0,0] or (0,0).
The default value is 0.
ceil_mode(bool, optional): ${ceil_mode_comment}Whether to use the ceil function to calculate output height
and width. If it is set to False, the floor function will be used. The default value is False.
count_include_pad(bool, optional): Whether to include padding points in average pooling mode, default is `False`.
Shape:
- x(Tensor): The input tensor of avg pool1d operator, which is a 3-D tensor.
The data type can be float32, float64.
- output(Tensor): The output tensor of avg pool1d operator, which is a 3-D tensor.
The data type is same as input x.
Returns:
A callable object of AvgPool1D.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.compat.nn as nn
>>> data = paddle.uniform([1, 3, 32], dtype="float32", min=-1, max=1)
>>> AvgPool1D = nn.AvgPool1D(kernel_size=2, stride=2, padding=0)
>>> pool_out = AvgPool1D(data)
>>> print(pool_out.shape)
paddle.Size([1, 3, 16])
"""
__constants__ = [
"kernel_size",
"stride",
"padding",
"ceil_mode",
"count_include_pad",
]
kernel_size: Size1
stride: Size1
padding: Size1
ceil_mode: bool
count_include_pad: bool
@ForbidKeywordsDecorator(
illegal_keys={"exclusive", "name"},
func_name="paddle.compat.nn.AvgPool1D",
correct_name="paddle.nn.AvgPool1D",
)
def __init__(
self,
kernel_size: Size1,
stride: Size1 | None = None,
padding: Size1 = 0,
ceil_mode: bool = False,
count_include_pad: bool = True,
) -> None:
super().__init__()
self.kernel_size = _single(kernel_size)
self.stride = _single(stride if stride is not None else kernel_size)
self.padding = _single(padding)
self.ceil_mode = ceil_mode
self.count_include_pad = count_include_pad
def forward(self, input: Tensor) -> Tensor:
return nn.functional.avg_pool1d(
input,
self.kernel_size,
self.stride,
self.padding,
not self.count_include_pad,
self.ceil_mode,
)
def extra_repr(self) -> str:
return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
class AvgPool2D(nn.Layer):
r"""
This operation applies 2D average pooling over input features based on the input,
and kernel_size, stride, padding parameters. Input(X) and Output(Out) are
in NCHW format, where N is batch size, C is the number of channels,
H is the height of the feature, and W is the width of the feature.
Example:
Input:
X shape: :math:`(N, C, :math:`H_{in}`, :math:`W_{in}`)`
Attr:
kernel_size: ksize
Output:
Out shape: :math:`(N, C, :math:`H_{out}`, :math:`W_{out}`)`
.. math::
Output(N_i, C_j, h, w) = \frac{\sum_{m=0}^{ksize[0]-1} \sum_{n=0}^{ksize[1]-1}
Input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)}{ksize[0] * ksize[1]}
Parameters:
kernel_size(int|list|tuple): The pool kernel size. If pool kernel size is a tuple or list,
it must contain two integers, (pool_size_Height, pool_size_Width).
Otherwise, the pool kernel size will be a square of an int.
stride(int|list|tuple|None, optional): The pool stride size. If pool stride size is a tuple or list,
it must contain two integers, (pool_stride_Height, pool_stride_Width).
Otherwise, the pool stride size will be a square of an int.
Default None, then stride will be equal to the kernel_size.
padding(str|int|list|tuple, optional): The padding size. Padding could be in one of the following forms.
1. A string in ['valid', 'same'].
2. An int, which means the feature map is zero padded by size of `padding` on every sides.
3. A list[int] or tuple(int) whose length is 2, [pad_height, pad_weight] whose value means the padding size of each dimension.
4. A list[int] or tuple(int) whose length is 4. [pad_height_top, pad_height_bottom, pad_width_left, pad_width_right] whose value means the padding size of each side.
5. A list or tuple of pairs of integers. It has the form [[pad_before, pad_after], [pad_before, pad_after], ...]. Note that, the batch dimension and channel dimension should be [0,0] or (0,0).
The default value is 0.
ceil_mode(bool, optional): When True, will use `ceil` instead of `floor` to compute the output shape.
count_include_pad(bool, optional): Whether to include padding points in average pooling
mode, default is `False`.
divisor_override(float, optional): If specified, it will be used as divisor, otherwise kernel_size will be
used. Default None.
Shape:
- x(Tensor): The input tensor of avg pool2d operator, which is a 4-D tensor.
The data type can be float32, float64.
- output(Tensor): The output tensor of avg pool2d operator, which is a 4-D tensor.
The data type is same as input x.
Returns:
A callable object of AvgPool2D.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.compat.nn as nn
>>> # max pool2d
>>> input = paddle.uniform([1, 3, 32, 32], dtype="float32", min=-1, max=1)
>>> AvgPool2D = nn.AvgPool2D(kernel_size=2, stride=2, padding=0)
>>> output = AvgPool2D(input)
>>> print(output.shape)
paddle.Size([1, 3, 16, 16])
"""
__constants__ = [
"kernel_size",
"stride",
"padding",
"ceil_mode",
"count_include_pad",
"divisor_override",
]
kernel_size: Size2
stride: Size2
padding: Size2
ceil_mode: bool
count_include_pad: bool
divisor_override: int | None
@ForbidKeywordsDecorator(
illegal_keys={"exclusive", "data_format", "name"},
func_name="paddle.compat.nn.AvgPool2D",
correct_name="paddle.nn.AvgPool2D",
)
def __init__(
self,
kernel_size: Size2,
stride: Size2 | None = None,
padding: Size2 = 0,
ceil_mode: bool = False,
count_include_pad: bool = True,
divisor_override: int | None = None,
):
super().__init__()
self.kernel_size = kernel_size
self.stride = stride if (stride is not None) else kernel_size
self.padding = padding
self.ceil_mode = ceil_mode
self.count_include_pad = count_include_pad
self.divisor_override = divisor_override
def forward(self, input: Tensor) -> Tensor:
return nn.functional.avg_pool2d(
input,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
not self.count_include_pad,
self.divisor_override,
)
def extra_repr(self) -> str:
return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
class AvgPool3D(nn.Layer):
"""
This operation applies 3D max pooling over input features based on the input,
and kernel_size, stride, padding parameters. Input(X) and Output(Out) are
in NCDHW format, where N is batch size, C is the number of channels,
H is the height of the feature, D is the depth of the feature, and W is the width of the feature.
Parameters:
kernel_size(int|list|tuple): The pool kernel size. If pool kernel size
is a tuple or list, it must contain three integers,
(kernel_size_Depth, kernel_size_Height, kernel_size_Width).
Otherwise, the pool kernel size will be the cube of an int.
stride(int|list|tuple|None, optional): The pool stride size. If pool stride size is a tuple or list,
it must contain three integers, [stride_Depth, stride_Height, stride_Width).
Otherwise, the pool stride size will be a cube of an int.
Default None, then stride will be equal to the kernel_size.
padding(str|int|list|tuple, optional): The padding size. Padding could be in one of the following forms.
1. A string in ['valid', 'same'].
2. An int, which means the feature map is zero padded by size of `padding` on every sides.
3. A list[int] or tuple(int) whose length is 3, [pad_depth, pad_height, pad_weight] whose value means the padding size of each dimension.
4. A list[int] or tuple(int) whose length is 6. [pad_depth_front, pad_depth_back, pad_height_top, pad_height_bottom, pad_width_left, pad_width_right] whose value means the padding size of each side.
5. A list or tuple of pairs of integers. It has the form [[pad_before, pad_after], [pad_before, pad_after], ...]. Note that, the batch dimension and channel dimension should be [0,0] or (0,0).
The default value is 0.
ceil_mode(bool, optional): ${ceil_mode_comment}
count_include_pad(bool, optional): Whether to include padding points in average pooling mode, default is True.
divisor_override(int|float, optional): if specified, it will be used as divisor, otherwise kernel_size will
be used. Default None.
Returns:
A callable object of AvgPool3D.
Shape:
- x(Tensor): The input tensor of avg pool3d operator, which is a 5-D tensor.
The data type can be float16, float32, float64.
- output(Tensor): The output tensor of avg pool3d operator, which is a 5-D tensor.
The data type is same as input x.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.compat.nn as nn
>>> # avg pool3d
>>> input = paddle.uniform([1, 2, 3, 32, 32], dtype="float32", min=-1, max=1)
>>> AvgPool3D = nn.AvgPool3D(kernel_size=2, stride=2, padding=0)
>>> output = AvgPool3D(input)
>>> print(output.shape)
paddle.Size([1, 2, 1, 16, 16])
"""
__constants__ = [
"kernel_size",
"stride",
"padding",
"ceil_mode",
"count_include_pad",
"divisor_override",
]
kernel_size: Size3
stride: Size3
padding: Size3
ceil_mode: bool
count_include_pad: bool
divisor_override: int | None
@ForbidKeywordsDecorator(
illegal_keys={"exclusive", "data_format", "name"},
func_name="paddle.compat.nn.AvgPool3D",
correct_name="paddle.nn.AvgPool3D",
)
def __init__(
self,
kernel_size: Size3,
stride: Size3 | None = None,
padding: Size3 = 0,
ceil_mode: bool = False,
count_include_pad: bool = True,
divisor_override: int | None = None,
) -> None:
super().__init__()
self.kernel_size = kernel_size
self.stride = stride if (stride is not None) else kernel_size
self.padding = padding
self.ceil_mode = ceil_mode
self.count_include_pad = count_include_pad
self.divisor_override = divisor_override
def forward(self, input: Tensor) -> Tensor:
return nn.functional.avg_pool3d(
input,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
not self.count_include_pad,
self.divisor_override,
)
def extra_repr(self) -> str:
return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
def __setstate__(self, state):
super().__setstate__(state)
self.__dict__.setdefault("padding", 0)
self.__dict__.setdefault("ceil_mode", False)
self.__dict__.setdefault("count_include_pad", True)
class Unfold(nn.Unfold):
"""
A compatible version of paddle.nn.Unfold:
The keyword arguments are in non-plural forms, example: `kernel_size` instead of `kernel_sizes`. `padding` restricts the size of the input to be 1(int) or 2, Size4 is not allowed.
All the input parameters allow `Tensor` or `pir.Value` as inputs, and will be converted to lists. Other aspects are the same. To use a more input-flexible version of Unfold, please refer to `paddle.nn.Unfold`.
Args:
kernel_size(int|list|tuple|Tensor): The size of convolution kernel, should be [k_h, k_w]
or an integer k treated as [k, k].
stride(int|list|tuple|Tensor, optional): The strides, should be [stride_h, stride_w]
or an integer stride treated as [sride, stride]. For default, strides will be [1, 1].
padding(int|list|tuple|Tensor, optional): The paddings of each dimension, should be
a single integer or [padding_h, padding_w]. If [padding_h, padding_w] was given, it will expanded to
[padding_h, padding_w, padding_h, padding_w]. If an integer padding was given,
[padding, padding, padding, padding] will be used. By default, paddings will be 0.
dilation(int|list|tuple|Tensor, optional): The dilations of convolution kernel, should be
[dilation_h, dilation_w], or an integer dilation treated as [dilation, dilation].
For default, it will be [1, 1].
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.randn((100, 3, 224, 224))
>>> unfold = paddle.compat.nn.Unfold(kernel_size=[3, 3])
>>> result = unfold(x)
>>> print(result.shape)
paddle.Size([100, 27, 49284])
"""
kernel_sizes: Size2
dilations: Size2
paddings: Size2
strides: Size2
@ForbidKeywordsDecorator(
illegal_keys={"kernel_sizes", "dilations", "paddings", "strides"},
func_name="paddle.compat.nn.Unfold",
correct_name="paddle.nn.Unfold",
)
def __init__(
self,
kernel_size: Size2,
dilation: Size2 = 1,
padding: Size2 = 0,
stride: Size2 = 1,
) -> None:
super().__init__(kernel_size, dilation, padding, stride)
def forward(self, input: Tensor) -> Tensor:
def to_list_if_necessary(x):
if isinstance(x, (paddle.pir.Value, paddle.Tensor)):
x = x.tolist()
return x
return nn.functional.unfold(
input,
kernel_sizes=to_list_if_necessary(self.kernel_sizes),
strides=to_list_if_necessary(self.strides),
paddings=to_list_if_necessary(self.paddings),
dilations=to_list_if_necessary(self.dilations),
)
class Linear(nn.Layer):
r"""
Python compatible fully-connected linear transformation layer. For each input :math:`X` ,
the equation is:
.. math::
Out = XW^T + b
where :math:`W` is the weight and :math:`b` is the bias.
Linear layer takes only one multi-dimensional tensor as input with the
shape :math:`[*, in\_features]` , where :math:`*` means any
number of additional dimensions. It multiplies input tensor with the transpose
of weight (a 2-D tensor of shape :math:`[out\_features, in\_features]` ) and
produces an output tensor of shape :math:`[*, out\_features]` .
If ``bias`` is not False, the bias (a 1-D tensor of
shape :math:`[out\_features]` ) will be created and added to the output. At the
end of the initialization, ``reset_parameters`` will be called to initialize
the ``weight`` and ``bias`` (if available) randomly.
Parameters:
in_features (int):
The number of input units.
out_features (int):
The number of output units.
bias (bool): If True, the bias (a 1-D tensor of shape :math:`[out\_features]` ) will be created and
added to the output. Default: True.
device (PlaceLike): The device of the parameters created. Default: None,
representing the default paddle device.
dtype (DTypeLike): The dtype of the parameters created. Default: None, and is set by
the default dtype of Linear (float32).
Variables:
weight (paddle.Tensor): learnable parameters of the module of shape :math:`[out\_features, in\_features]`.
The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k` is :math:`\frac{1}{in\_features}`.
bias (paddle.Tensor): learnable parameters of the module of shape :math:`[out\_features]`. If ``bias`` is True,
the values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k` is :math:`\frac{1}{in\_features}`.
Shape:
- input: Multi-dimensional tensor with shape :math:`[*, in\_features]` . Its data types are float16, float32, float64 ,The default is float32 .
- output: Multi-dimensional tensor with shape :math:`[*, out\_features]` . The data type is the same as the input .
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(100)
>>> # Define the linear layer.
>>> linear = paddle.compat.nn.Linear(2, 4, bias=True)
>>> print(linear.weight)
Parameter containing:
Tensor(shape=[4, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
[[-0.49191639, 0.28120756],
[-0.17887023, 0.40572405],
[ 0.35139430, 0.45717543],
[-0.06135514, -0.21088189]])
>>> print(linear.bias)
Parameter containing:
Tensor(shape=[4], dtype=float32, place=Place(cpu), stop_gradient=False,
[ 0.49166456, -0.06108528, -0.14973064, 0.31168410])
>>> x = paddle.arange(6, dtype="float32").reshape([3, 2])
>>> y = linear(x)
>>> print(y)
Tensor(shape=[3, 4], dtype=float32, place=Place(cpu), stop_gradient=False,
[[ 0.77287209, 0.34463876, 0.30744481, 0.10080221],
[ 0.35145447, 0.79834640, 1.92458415, -0.44367185],
[-0.06996319, 1.25205410, 3.54172373, -0.98814595]])
"""
__constants__ = ["in_features", "out_features"]
in_features: int
out_features: int
weight: Tensor
@ForbidKeywordsDecorator(
illegal_keys={"weight_attr", "bias_attr", "name"},
func_name="paddle.compat.nn.Linear",
correct_name="paddle.nn.Linear",
)
def __init__(
self,
in_features: int,
out_features: int,
bias: bool = True,
device: PlaceLike | None = None,
dtype: DTypeLike | None = None,
) -> None:
super().__init__()
self._dtype = (
self._helper.get_default_dtype() if dtype is None else dtype
)
self.in_features = in_features
self.out_features = out_features
self.weight = self.create_parameter(
shape=[out_features, in_features],
attr=None,
dtype=self._dtype,
is_bias=False,
device=device,
)
self.bias = None
if bias:
self.bias = self.create_parameter(
shape=[out_features],
attr=None,
dtype=self._dtype,
is_bias=True,
device=device,
)
# The same parameter initialization as PyTorch
self.reset_parameters()
def forward(self, input: Tensor) -> Tensor:
return functional.linear.__wrapped__( # bypass ForbidKeywordsDecorator
input=input, weight=self.weight, bias=self.bias
)
def extra_repr(self) -> str:
"""
Return the extra representation of the module.
"""
return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"
def reset_parameters(self) -> None:
"""
Resets parameters based on their initialization used in ``__init__``.
"""
bound = 1 / sqrt(self.in_features) if self.in_features > 0 else 0
if self.in_features > 0 and self.out_features > 0:
nn.init.uniform_(self.weight, -bound, bound)
if self.bias is not None and self.out_features > 0:
nn.init.uniform_(self.bias, -bound, bound)
class Softmax(nn.Layer):
r"""
Softmax Activation.
This operator implements the softmax layer. The calculation process is as follows:
1. The dimension :attr:`dim` of ``input`` will be permuted to the last.
2. Then ``input`` will be logically flattened to a 2-D matrix. The matrix's second
dimension(row length) is the same as the dimension :attr:`dim` of ``input``,
and the first dimension(column length) is the product of all other dimensions
of ``input``. For each row of the matrix, the softmax operator squashes the
K-dimensional(K is the width of the matrix, which is also the size of ``input``'s
dimension :attr:`dim`) vector of arbitrary real values to a K-dimensional
vector of real values in the range [0, 1] that add up to 1.
3. After the softmax operation is completed, the inverse operations of steps 1 and 2
are performed to restore the two-dimensional matrix to the same dimension as the ``input`` .
It computes the exponential of the given dimension and the sum of exponential
values of all the other dimensions in the K-dimensional vector input.
Then the ratio of the exponential of the given dimension and the sum of
exponential values of all the other dimensions is the output of the softmax
operator.
For each row :math:`i` and each column :math:`j` in the matrix, we have:
.. math::
Softmax[i, j] = \frac{\exp(x[i, j])}{\sum_j(exp(x[i, j])}
Example:
.. code-block:: text
Case 1:
Input:
x.shape = [2, 3, 4]
x.data = [[[2.0, 3.0, 4.0, 5.0],
[3.0, 4.0, 5.0, 6.0],
[7.0, 8.0, 8.0, 9.0]],
[[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0],
[6.0, 7.0, 8.0, 9.0]]]
Attrs:
dim = -1
Output:
out.shape = [2, 3, 4]
out.data = [[[0.0320586 , 0.08714432, 0.23688282, 0.64391426],
[0.0320586 , 0.08714432, 0.23688282, 0.64391426],
[0.07232949, 0.19661193, 0.19661193, 0.53444665]],
[[0.0320586 , 0.08714432, 0.23688282, 0.64391426],
[0.0320586 , 0.08714432, 0.23688282, 0.64391426],
[0.0320586 , 0.08714432, 0.23688282, 0.64391426]]]
Case 2:
Input:
x.shape = [2, 3, 4]
x.data = [[[2.0, 3.0, 4.0, 5.0],
[3.0, 4.0, 5.0, 6.0],
[7.0, 8.0, 8.0, 9.0]],
[[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0],
[6.0, 7.0, 8.0, 9.0]]]
Attrs:
dim = 1
Output:
out.shape = [2, 3, 4]
out.data = [[[0.00657326, 0.00657326, 0.01714783, 0.01714783],
[0.01786798, 0.01786798, 0.04661262, 0.04661262],
[0.97555875, 0.97555875, 0.93623955, 0.93623955]],
[[0.00490169, 0.00490169, 0.00490169, 0.00490169],
[0.26762315, 0.26762315, 0.26762315, 0.26762315],
[0.72747516, 0.72747516, 0.72747516, 0.72747516]]]
Parameters:
dim (int, optional): The dim along which to perform log_softmax
calculations. It should be in range [-D, D), where D is the
dimensions of ``input`` . If ``dim`` < 0, it works the same way as
:math:`dim + D` . Default is None.
Shape:
- input: Tensor with any shape.
- output: Tensor with the same shape as input.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.to_tensor(
... [
... [
... [2.0, 3.0, 4.0, 5.0],
... [3.0, 4.0, 5.0, 6.0],
... [7.0, 8.0, 8.0, 9.0],
... ],
... [
... [1.0, 2.0, 3.0, 4.0],
... [5.0, 6.0, 7.0, 8.0],
... [6.0, 7.0, 8.0, 9.0],
... ],
... ],
... dtype='float32',
... )
>>> m = paddle.compat.nn.Softmax()
>>> out = m(x)
>>> print(out)
Tensor(shape=[2, 3, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[0.73105854, 0.73105854, 0.73105854, 0.73105854],
[0.11920292, 0.11920292, 0.11920292, 0.11920292],
[0.73105854, 0.73105854, 0.50000000, 0.50000000]],
[[0.26894143, 0.26894143, 0.26894143, 0.26894143],
[0.88079703, 0.88079703, 0.88079703, 0.88079703],
[0.26894143, 0.26894143, 0.50000000, 0.50000000]]])
"""
@ForbidKeywordsDecorator(
illegal_keys={"axis"},
func_name="paddle.compat.nn.Softmax",
correct_name="paddle.nn.Softmax",
)
def __init__(self, dim: int | None = None) -> None:
super().__init__()
self._dim = dim
self._dtype = None
def forward(self, input: Tensor) -> Tensor:
return functional.softmax(input, self._dim)
def extra_repr(self) -> str:
return f"dim={self._dim}"
class SmoothL1Loss(nn.Layer):
r"""
PyTorch compatible version of :ref:`api_paddle_nn_SmoothL1Loss`, aligned with
``torch.nn.SmoothL1Loss``. The per-element loss is
.. math::
z_i = \left\{\begin{array}{rcl}
0.5 (x_i - y_i)^2 / beta & & {if |x_i - y_i| < beta} \\
|x_i - y_i| - 0.5 * beta & & {otherwise}
\end{array} \right.
which equals Paddle's Huber loss divided by ``beta``. This differs from
:ref:`api_paddle_nn_SmoothL1Loss` whose default ``is_huber=True`` returns the
raw Huber loss.
Parameters:
size_average (bool|None, optional): Deprecated (see ``reduction``). When
``size_average`` or ``reduce`` is not ``None``, it is translated into
``reduction`` with a ``DeprecationWarning``. Default is ``None``.
reduce (bool|None, optional): Deprecated (see ``reduction``). Default is ``None``.
reduction (str, optional): Indicate how to calculate the loss, the candidates
are ``'none'`` | ``'mean'`` | ``'sum'``. Default is ``'mean'``.
beta (float, optional): Non-negative threshold at which to change between L1
and L2 loss. When ``beta == 0`` the loss degrades to the L1 loss, matching
PyTorch. Default is ``1.0``.
Call Parameters:
input (Tensor): Input tensor, the data type is float32 or float64.
target (Tensor): Label tensor with the same shape as ``input``.
Returns:
Tensor, The tensor storing the smooth L1 loss of ``input`` and ``target``.
Examples:
.. code-block:: pycon
>>> import paddle
>>> input = paddle.to_tensor([[0.5, 1.5], [2.0, 0.0]], dtype='float32')
>>> target = paddle.to_tensor([[1.0, 1.0], [1.0, 0.5]], dtype='float32')
>>> loss = paddle.compat.nn.SmoothL1Loss(beta=1.0)
>>> output = loss(input, target)
>>> print(output)
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
0.21875000)
"""
__constants__ = ["reduction", "beta"]
reduction: str
beta: float
@ForbidKeywordsDecorator(
illegal_keys={"delta", "is_huber", "name", "label"},
func_name="paddle.compat.nn.SmoothL1Loss",
correct_name="paddle.nn.SmoothL1Loss",
)
def __init__(
self,
size_average: bool | None = None,
reduce: bool | None = None,
reduction: str = 'mean',
beta: float = 1.0,
) -> None:
super().__init__()
if size_average is not None or reduce is not None:
reduction = (
'none'
if reduce is False
else ('sum' if size_average is False else 'mean')
)
warnings.warn(
"'size_average' and 'reduce' args of 'SmoothL1Loss' will be "
f"deprecated, please use reduction='{reduction}' instead.",
DeprecationWarning,
stacklevel=2,
)
self.reduction = reduction
self.beta = beta
def forward(self, input: Tensor, target: Tensor) -> Tensor:
return functional.smooth_l1_loss.__wrapped__(
input, target, reduction=self.reduction, beta=self.beta
)
def extra_repr(self) -> str:
return f"reduction={self.reduction}, beta={self.beta}"
AvgPool1d = AvgPool1D
AvgPool2d = AvgPool2D
AvgPool3d = AvgPool3D
python/paddle/compat/nn/functional/__init__.py
+595
View File
@@ -0,0 +1,595 @@
# 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
import warnings
from typing import TYPE_CHECKING, Literal
import paddle
from paddle import _C_ops
from paddle.base.framework import Variable
from paddle.framework import (
in_dynamic_mode,
)
from paddle.tensor import log_softmax, softmax
from paddle.utils.decorator_utils import ForbidKeywordsDecorator
from .sdpa import scaled_dot_product_attention
if TYPE_CHECKING:
from typing import TypeAlias
from paddle import Tensor
from paddle._typing import (
ShapeLike,
Size2,
)
_PaddingTensorMode: TypeAlias = Literal[
"zeros", "constant", "reflect", "replicate", "circular"
]
_ReduceMode: TypeAlias = Literal["mean", "sum", "none"]
__all__ = [
'pad',
'softmax',
'log_softmax',
'linear',
'scaled_dot_product_attention',
'unfold',
'smooth_l1_loss',
'batch_norm',
'instance_norm',
]
def _check_valid_pad_len(pad_len, x_dim, is_constant):
if pad_len > 6 or pad_len < 0:
raise ValueError(f"Expect len(pad) <= 6 and not -1, got: {pad_len}")
max_dim = 2 * x_dim - (0 if is_constant else 2)
if pad_len > max_dim:
raise ValueError(
f"len(pad) is bounded by input.ndim: expect len(pad) <= {max_dim}, got: {pad_len}"
)
@ForbidKeywordsDecorator(
illegal_keys={"x", "name", "data_format", "pad_from_left_axis"},
func_name="paddle.compat.nn.functional.pad",
correct_name="paddle.nn.functional.pad",
)
def pad(
input: Tensor,
pad: ShapeLike,
mode: _PaddingTensorMode = 'constant',
value: float = 0.0,
) -> Tensor:
"""
PyTorch compatible version of :ref:`api_paddle_nn_functional_pad`. For the original API, see :ref:`api_paddle_nn_functional_pad` for more details.
Pad tensor according to ``'pad'`` and ``'mode'``. All the padding operations under the hood starts from the **right** (last dim) of the tensor.
Args:
input (Tensor): The input tensor with data type float32, float64, int32, int64, complex64 or complex128.
pad (Tensor|list[int]|tuple[int]): The padding size with data type int. Refer to Note for details.
mode (str, optional): Four modes: ``'constant'`` (default), ``'reflect'``, ``'replicate'``, ``'circular'``. Default is ``'constant'``.
- 'constant' mode, uses a constant value to pad the input tensor.
- 'reflect' mode, uses reflection of the input boundaries to pad the input tensor.
- 'replicate' mode, uses input boundaries to pad the input tensor.
- 'circular' mode, uses circular input to pad the input tensor.
value (float, optional): The value to fill the padded areas in 'constant' mode . Default is :math:`0.0`.
Note:
For non ``'constant'`` mode, padding size can not be greater than ``min(2 * input.ndim - 2, 6)``.
Only 2D, 3D, 4D and 5D tensors are supported with up to the last 3 dims (if ndim >= 3) can be padded.
Returns:
Tensor, a Tensor padded according to pad and mode and data type is same as input.
Examples:
.. code-block:: pycon
>>> import paddle
>>> input_shape = (1, 1, 3)
>>> input_ = paddle.arange(paddle.prod(paddle.to_tensor(input_shape)), dtype="float32").reshape(input_shape) + 1
>>> y = paddle.compat.nn.functional.pad(input_, [1, 0, 0, 1], value=0, mode='constant')
>>> print(y)
Tensor(shape=[1, 2, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[0., 1., 2., 3.],
[0., 0., 0., 0.]]])
>>> # reflect 2D padding
>>> input_ = paddle.arange(6).reshape([2, 3])
>>> y = paddle.compat.nn.functional.pad(input=input_, pad=(1, 1), mode='reflect')
>>> print(y)
Tensor(shape=[2, 5], dtype=int64, place=Place(cpu), stop_gradient=True,
[[1, 0, 1, 2, 1],
[4, 3, 4, 5, 4]])
"""
assert mode in [
'reflect',
'replicate',
'constant',
'circular',
], (
f"mode should be one of constant, reflect, replicate, circular, but got {mode}."
)
x_dim = len(input.shape)
if in_dynamic_mode():
if isinstance(pad, (Variable, paddle.Tensor)) and pad.size == 0:
return input.clone()
if (
mode == "constant"
and isinstance(pad, (list, tuple))
and len(pad) != (x_dim - 2) * 2
):
paddings = pad
pad_value = value
padding_len = len(paddings)
# pad the length of paddings to 2*x_dim
if padding_len < 2 * x_dim:
pad_len_for_paddings = 2 * x_dim - padding_len
paddings = paddings + ([0] if isinstance(pad, list) else (0,)) * (
pad_len_for_paddings
)
# since the kernel pad from left axis, if we want to pad from right axis, we need to reverse the paddings
paddings = [
paddings[i - 1] if i % 2 == 1 else paddings[i + 1]
for i in range(2 * x_dim - 1, -1, -1)
]
pad_val = (
pad_value
if isinstance(pad_value, paddle.pir.Value)
else float(pad_value)
)
return _C_ops.pad(input, paddings, pad_val)
assert x_dim >= 1 and x_dim <= 5, (
f"Input tensor dimension must be in [1-5] but got {x_dim}"
)
is_constant_mode = mode == 'constant'
if (not is_constant_mode) and x_dim < 2:
raise ValueError(
f"Only 2D, 3D, 4D, 5D padding with non-constant padding are supported for now, got ndim: {x_dim}"
)
# pad the `pad` to be length = 6 (right padding), for example [1, 2] -> [1, 2, 0, 0, 0, 0]
if isinstance(pad, (Variable, paddle.pir.Value)):
pad_len = pad.shape[0]
_check_valid_pad_len(pad_len, x_dim, is_constant_mode)
pad = paddle.concat(
[
pad,
paddle.zeros((6 - pad_len,), dtype="int32"),
],
axis=0,
)
else:
pad = list(pad)
pad_len = len(pad)
_check_valid_pad_len(pad_len, x_dim, is_constant_mode)
pad.extend([0] * (6 - pad_len))
ndim_to_unsqueeze = list(range(5 - x_dim))
input = input.unsqueeze(axis=ndim_to_unsqueeze)
out = _C_ops.pad3d(
input,
pad.tolist() if isinstance(pad, Variable) else pad,
mode,
value,
"NCDHW",
)
if ndim_to_unsqueeze:
return out.squeeze(axis=ndim_to_unsqueeze)
return out
@ForbidKeywordsDecorator(
illegal_keys={"x", "name"},
func_name="paddle.compat.nn.functional.linear",
correct_name="paddle.nn.functional.linear",
)
def linear(input: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor:
r"""
Fully-connected linear transformation operator. For each input :math:`x` ,
the equation is:
.. math::
Out = xW^T + b
where :math: `W` is the weight and :math:`b` is the bias.
If the weight is a 2-D tensor of shape :math:`[out\_features, in\_features]` ,
input should be a multi-dimensional tensor of shape
:math:`[*, in\_features]` , where :math:`*` means any number of
additional dimensions. The linear operator multiplies input tensor with
weight and produces an output tensor of shape :math:`[*, out\_features]` ,
If :math:`bias` is not None, the bias should be a 1-D tensor of shape
:math:`[out\_features]` and will be added to the output.
This implementation is aligned with PyTorch's linear function which computes
:math:`y = xW^T + b`.
Parameters:
input (Tensor): Input tensor. The data type should be bfloat16, float16, float32 or float64.
The input tensor should be of shape :math:`[*, in\_features]`, where :math:`*` means any number of additional dimensions, including none
weight (Tensor): Weight tensor. The data type should be float16, float32 or float64.
Shape should be [out_features, in_features].
bias (Tensor, optional): Bias tensor. The data type should be float16, float32 or float64.
If it is set to None, no bias will be added to the output units.
Returns:
Tensor, the shape is :math:`[*, out\_features]` and the
data type is the same with input :math:`x` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(2025)
>>> x = paddle.arange(6, dtype=paddle.float32).reshape([3, 2])
>>> x
Tensor(shape=[3, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0., 1.],
[2., 3.],
[4., 5.]])
>>> weight = paddle.full(shape=[4, 2], fill_value=0.5, dtype="float32", name="weight")
>>> weight
Tensor(shape=[4, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0.50000000, 0.50000000],
[0.50000000, 0.50000000],
[0.50000000, 0.50000000],
[0.50000000, 0.50000000]])
>>> bias = paddle.ones(shape=[4], dtype="float32", name="bias")
>>> y = paddle.compat.nn.functional.linear(x, weight, bias)
>>> print(y)
Tensor(shape=[3, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
[[1.50000000, 1.50000000, 1.50000000, 1.50000000],
[3.50000000, 3.50000000, 3.50000000, 3.50000000],
[5.50000000, 5.50000000, 5.50000000, 5.50000000]])
"""
if (
paddle.get_flags("FLAGS_use_legacy_linear")["FLAGS_use_legacy_linear"]
or not paddle.is_compiled_with_cuda()
or not paddle.framework.in_dynamic_or_pir_mode()
):
# Fallback to old logic when in non-cuda or legacy mode.
out = _C_ops.matmul(input, weight, False, True)
if bias is not None:
out = _C_ops.add(out, bias)
return out
else:
# transpose y is True, since _C_ops.linear(input, weight.T, bias) can introduce more overhead. With CINN, matmul and add can be fused.
# Note(Pan Zhaowu): In accuracy compatible kernel mode, we use linear_v2 op that receives transposed weight, aligning with torch. Note that this will incurs a real transpose op, which might cause performance degradation.
if bias is not None:
return _C_ops.linear_v2(input, weight.contiguous(), bias, True)
else:
return _C_ops.matmul(input, weight.contiguous(), False, True)
@ForbidKeywordsDecorator(
illegal_keys={
"x",
"kernel_sizes",
"dilations",
"paddings",
"strides",
"name",
},
func_name="paddle.compat.nn.functional.unfold",
correct_name="paddle.nn.functional.unfold",
)
def unfold(
input: Tensor,
kernel_size: Size2,
dilation: Size2 = 1,
padding: Size2 = 0,
stride: Size2 = 1,
) -> Tensor:
r"""
Return a col buffer of sliding local blocks of input, also known
as im2col for batched 2D image tensors. For each block under the convolution filter,
all element will be rearranged as a column. While the convolution filter sliding over
the input feature map, a series of such columns will be formed.
For each input :math:`input` with shape [N, C, H, W], the output shape [N, Cout, Lout]
can be calculated as following.
.. math::
dkernel[0] &= dilation[0] \times (kernel\_sizes[0] - 1) + 1
dkernel[1] &= dilation[1] \times (kernel\_sizes[1] - 1) + 1
hout &= \frac{H + padding[0] + padding[2] - dkernel[0]}{stride[0]} + 1
wout &= \frac{W + padding[1] + padding[3] - dkernel[1]}{stride[1]} + 1
Cout &= C \times kernel\_sizes[0] \times kernel\_sizes[1]
Lout &= hout \times wout
Parameters:
input(Tensor): 4-D Tensor, input tensor of format [N, C, H, W],
data type can be float32 or float64
kernel_size(int|list|tuple): The size of convolution kernel, should be [k_h, k_w]
or an integer k treated as [k, k].
dilation(int|list|tuple, optional): the dilation of convolution kernel, should be
[dilation_h, dilation_w], or an integer dilation treated as
[dilation, dilation]. For default, it will be [1, 1].
padding(int|list|tuple, optional): The paddings of each dimension, should be
a single integer or [padding_h, padding_w]. If [padding_h, padding_w] was given, it will expanded to
[padding_h, padding_w, padding_h, padding_w]. If an integer padding was given,
[padding, padding, padding, padding] will be used. By default, paddings will be 0.
strides(int|list|tuple, optional): The strides, should be [stride_h, stride_w]
or an integer stride treated as [stride, stride].
For default, strides will be [1, 1].
Returns:
Tensor, The tensor corresponding to the sliding local blocks.
The output shape is [N, Cout, Lout] as described above.
Cout is the total number of values within each block,
and Lout is the total number of such blocks.
The data type of output is the same as the input :math:`input`
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.compat.nn.functional as F
>>> x = paddle.randn((100, 3, 224, 224))
>>> y = F.unfold(x, [3, 3], 1, 1, 1)
"""
def to_list_if_necessary(x):
if isinstance(x, (paddle.pir.Value, paddle.Tensor)):
x = x.tolist()
return x
return paddle.nn.functional.unfold(
x=input,
kernel_sizes=to_list_if_necessary(kernel_size),
strides=to_list_if_necessary(stride),
paddings=to_list_if_necessary(padding),
dilations=to_list_if_necessary(dilation),
)
@ForbidKeywordsDecorator(
illegal_keys={"label", "delta", "is_huber", "name"},
func_name="paddle.compat.nn.functional.smooth_l1_loss",
correct_name="paddle.nn.functional.smooth_l1_loss",
)
def smooth_l1_loss(
input: Tensor,
target: Tensor,
size_average: bool | None = None,
reduce: bool | None = None,
reduction: _ReduceMode = 'mean',
beta: float = 1.0,
) -> Tensor:
r"""
PyTorch compatible version of :ref:`api_paddle_nn_functional_smooth_l1_loss`.
Computes the Smooth L1 loss, aligned with ``torch.nn.functional.smooth_l1_loss``.
The per-element loss is:
.. math::
z_i = \left\{\begin{array}{rcl}
0.5 (x_i - y_i)^2 / beta & & {if |x_i - y_i| < beta} \\
|x_i - y_i| - 0.5 * beta & & {otherwise}
\end{array} \right.
This equals Paddle's Huber loss divided by ``beta`` (i.e. ``is_huber=False`` with
``delta=beta``), which is the key difference from
:ref:`api_paddle_nn_functional_smooth_l1_loss` whose default ``is_huber=True``
returns the raw Huber loss.
Args:
input (Tensor): Input tensor, the data type is float32 or float64.
target (Tensor): Label tensor with the same shape as ``input``.
size_average (bool|None, optional): Deprecated (see ``reduction``). When
``size_average`` or ``reduce`` is not ``None``, it is translated into
``reduction`` with a ``DeprecationWarning``. Default is ``None``.
reduce (bool|None, optional): Deprecated (see ``reduction``). Default is ``None``.
reduction (str, optional): Indicate how to calculate the loss, the candidates
are ``'none'`` | ``'mean'`` | ``'sum'``. Default is ``'mean'``.
beta (float, optional): Specifies the threshold at which to change between L1
and L2 loss. The value must be non-negative. When ``beta == 0`` the loss
degrades to the L1 loss, matching PyTorch. Default is ``1.0``.
Returns:
Tensor, The tensor storing the smooth L1 loss of ``input`` and ``target``.
Examples:
.. code-block:: pycon
>>> import paddle
>>> input = paddle.to_tensor([[0.5, 1.5], [2.0, 0.0]], dtype='float32')
>>> target = paddle.to_tensor([[1.0, 1.0], [1.0, 0.5]], dtype='float32')
>>> output = paddle.compat.nn.functional.smooth_l1_loss(input, target, beta=1.0)
>>> print(output)
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
0.21875000)
"""
# Translate PyTorch's deprecated size_average / reduce into reduction.
if size_average is not None or reduce is not None:
reduction = (
'none'
if reduce is False
else ('sum' if size_average is False else 'mean')
)
warnings.warn(
"'size_average' and 'reduce' args of 'smooth_l1_loss' will be "
f"deprecated, please use reduction='{reduction}' instead.",
DeprecationWarning,
stacklevel=2,
)
if beta < 0:
raise ValueError(
f"smooth_l1_loss does not accept negative beta, but got beta={beta}."
)
if beta == 0:
return paddle.nn.functional.l1_loss(input, target, reduction=reduction)
return paddle.nn.functional.smooth_l1_loss(
input, target, reduction=reduction, delta=beta, is_huber=False
)
@ForbidKeywordsDecorator(
illegal_keys={"x", "epsilon", "data_format", "use_global_stats", "name"},
func_name="paddle.compat.nn.functional.batch_norm",
correct_name="paddle.nn.functional.batch_norm",
)
def batch_norm(
input: Tensor,
running_mean: Tensor,
running_var: Tensor,
weight: Tensor | None = None,
bias: Tensor | None = None,
training: bool = False,
momentum: float = 0.1,
eps: float = 1e-05,
) -> Tensor:
r"""
PyTorch compatible version of :ref:`api_paddle_nn_functional_batch_norm`.
Aligned with ``torch.nn.functional.batch_norm``.
See :ref:`api_paddle_nn_functional_batch_norm` for more details.
Args:
input (Tensor): Input tensor, the data type is float32 or float64.
running_mean (Tensor|None): Running mean.
running_var (Tensor|None): Running variance.
weight (Tensor|None, optional): The weight tensor. Default: None.
bias (Tensor|None, optional): The bias tensor. Default: None.
training (bool, optional): True means train mode. Default: False.
momentum (float, optional): The value used for the moving_mean and moving_var computation. Default: 0.1.
eps (float, optional): The small value added to variance to prevent division by zero. Default: 1e-05.
Returns:
Tensor, the output of batch normalization.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.arange(12, dtype="float32").reshape([2, 1, 2, 3])
>>> running_mean = paddle.to_tensor([0], dtype="float32")
>>> running_var = paddle.to_tensor([1], dtype="float32")
>>> weight = paddle.to_tensor([2], dtype="float32")
>>> bias = paddle.to_tensor([1], dtype="float32")
>>> out = paddle.compat.nn.functional.batch_norm(x, running_mean, running_var, weight, bias)
>>> print(out)
Tensor(shape=[2, 1, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[1. , 2.99998999 , 4.99997997 ],
[6.99996996 , 8.99995995 , 10.99995041]]],
[[[12.99993992, 14.99992943, 16.99991989],
[18.99991035, 20.99990082, 22.99988937]]]])
"""
return paddle.nn.functional.batch_norm(
x=input,
running_mean=running_mean,
running_var=running_var,
weight=weight,
bias=bias,
training=training,
momentum=1.0 - momentum,
epsilon=eps,
)
@ForbidKeywordsDecorator(
illegal_keys={"x", "data_format", "name"},
func_name="paddle.compat.nn.functional.instance_norm",
correct_name="paddle.nn.functional.instance_norm",
)
def instance_norm(
input: Tensor,
running_mean: Tensor | None = None,
running_var: Tensor | None = None,
weight: Tensor | None = None,
bias: Tensor | None = None,
use_input_stats: bool = True,
momentum: float = 0.1,
eps: float = 1e-05,
) -> Tensor:
r"""
PyTorch compatible version of :ref:`api_paddle_nn_functional_instance_norm`.
Aligned with ``torch.nn.functional.instance_norm``.
See :ref:`api_paddle_nn_functional_instance_norm` for more details.
Args:
input (Tensor): Input tensor, the data type is float32 or float64.
running_mean (Tensor|None, optional): Running mean. Default: None.
running_var (Tensor|None, optional): Running variance. Default: None.
weight (Tensor|None, optional): The weight tensor. Default: None.
bias (Tensor|None, optional): The bias tensor. Default: None.
use_input_stats (bool, optional): Whether to use input statistics. Default: True.
momentum (float, optional): The value used for the moving_mean and moving_var computation. Default: 0.1.
eps (float, optional): The small value added to variance to prevent division by zero. Default: 1e-05.
Returns:
Tensor, the output of instance normalization.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.rand((2, 2, 2, 3))
>>> out = paddle.compat.nn.functional.instance_norm(x)
>>> print(out)
"""
return paddle.nn.functional.instance_norm(
x=input,
running_mean=running_mean,
running_var=running_var,
weight=weight,
bias=bias,
use_input_stats=use_input_stats,
momentum=1.0 - momentum,
eps=eps,
)
python/paddle/compat/nn/functional/sdpa.py
+126
View File
@@ -0,0 +1,126 @@
# 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 typing import TYPE_CHECKING
import paddle.nn.functional as F
if TYPE_CHECKING:
from paddle import Tensor
def scaled_dot_product_attention(
query: Tensor,
key: Tensor,
value: Tensor,
attn_mask: Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
scale: float | None = None,
enable_gqa: bool = False,
) -> Tensor:
r"""
The equation is:
.. math::
result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V
where : ``Q``, ``K``, and ``V`` represent the three input parameters of the attention module.
The dimensions of the three parameters are the same.
``d`` represents the size of the last dimension of the three parameters.
Warning:
This API only verifies inputs with dtype float16 and bfloat16, other dtypes may fall back to math
implementation, which is less optimized.
Note:
This API differs from :ref:`api_paddle_nn_functional_scaled_dot_product_attention` in that:
The QKV layout of this API is [batch_size, num_heads, seq_len, head_dim] or [num_heads, seq_len, head_dim].
Args:
query(Tensor): The query tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
3-D tensor with shape:
[num_heads, seq_len, head_dim].
The dtype can be float16 or bfloat16.
key(Tensor): The key tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
3-D tensor with shape:
[num_heads, seq_len, head_dim].
The dtype can be float16 or bfloat16.
value(Tensor): The value tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
3-D tensor with shape:
[num_heads, seq_len, head_dim].
The dtype can be float16 or bfloat16.
attn_mask(Tensor, optional): The attention mask tensor. The shape should be broadcastable to
[batch_size, num_heads, seq_len_key, seq_len_query]. The dtype can be bool
or same type of query. The bool mask indicates the positions should take part
in attention. The non-bool mask will be added to attention score.
is_causal(bool, optional): Whether enable causal mode. If True, the attention masking is a lower
triangular matrix when the mask is a square matrix. The attention masking has the
form of the upper left causal bias when the mask is a non-square matrix.
An error is thrown if both attn_mask and is_causal are set.
scale(float, optional): The scaling factor used in the calculation of attention weights.
If None, scale = 1 / sqrt(head_dim).
enable_gqa(bool, optional): Whether enable GQA mode. Default False.
Returns:
out(Tensor): The attention tensor.
4-D tensor with shape: [batch_size, num_heads, seq_len, head_dim].
3-D tensor with shape: [num_heads, seq_len, head_dim].
The dtype can be float16 or bfloat16.
Examples:
.. code-block:: pycon
>>> # doctest: +SKIP('bfloat need V100 compile')
>>> import paddle
>>> q = paddle.rand((1, 2, 128, 16), dtype=paddle.bfloat16)
>>> output = paddle.compat.nn.functional.scaled_dot_product_attention(q, q, q, None, 0.9, False)
>>> print(output)
>>> # doctest: -SKIP
"""
if is_causal and attn_mask is not None:
raise RuntimeError(
"Explicit attn_mask should not be set when is_causal=True"
)
query, key, value = (
query.swapaxes(-3, -2),
key.swapaxes(-3, -2),
value.swapaxes(-3, -2),
)
out = F.scaled_dot_product_attention(
query,
key,
value,
attn_mask,
dropout_p,
is_causal,
True, # training
None, # backend
scale,
enable_gqa,
None, # name
)
return out.swapaxes(-3, -2)
python/paddle/compat/nn/transformer.py
+563
View File
@@ -0,0 +1,563 @@
# 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 typing import TYPE_CHECKING
import paddle
import paddle.nn.functional as F
from paddle import nn
from paddle.nn.initializer import XavierNormal, XavierUniform
if TYPE_CHECKING:
from paddle import Tensor
from paddle._typing import DTypeLike, PlaceLike
class MultiheadAttention(nn.Layer):
r"""
Allows the model to jointly attend to information from different representation subspaces.
Multi-Head Attention is defined as:
.. math::
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)W^O
where :math:`\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
Please refer to `Attention Is All You Need <https://arxiv.org/pdf/1706.03762.pdf>`_
for more details.
.. note::
This layer will use the optimized implementation
:func:`paddle.nn.functional.scaled_dot_product_attention` if no need to return the attention weights.
Parameters:
embed_dim (int): Total dimension of the model.
num_heads (int): The number of heads in multi-head attention.
dropout (float, optional): The dropout probability used on attention
weights to drop some attention targets. 0 for no dropout. Default 0.0.
bias (bool, optional): If specified, adds bias to input / output projection layers.
Default: True.
add_bias_kv (bool, optional): If specified, adds bias to the key and value sequences
at axis=0. Default: False.
add_zero_attn (bool, optional): If specified, adds a new batch of zeros to the
key and value sequences at axis=1. Default: False.
kdim (int, optional): Total number of features for keys. If None, assumed equal to
`embed_dim`. Default: None.
vdim (int, optional): Total number of features for values. If None, assumed equal to
`embed_dim`. Default: None.
batch_first (bool, optional): If True, then the input and output tensors are provided
as [batch, seq, feature]. Default: False.
device (PlaceLike|None, optional): The device to initialize parameters on. Default: None.
dtype (DTypeLike|None, optional): The data type of the parameters. Default: None.
Examples:
.. code-block:: pycon
>>> import paddle
>>> from paddle.compat import nn
>>> # Example with batch_first=True
>>> embed_dim, num_heads = 128, 8
>>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
>>> # query: [batch_size, target_seq_len, embed_dim]
>>> query = paddle.randn([32, 10, embed_dim])
>>> # key, value: [batch_size, source_seq_len, embed_dim]
>>> key = paddle.randn([32, 20, embed_dim])
>>> value = paddle.randn([32, 20, embed_dim])
>>> attn_output, attn_output_weights = multihead_attn(query, key, value)
>>> print(attn_output.shape)
paddle.Size([32, 10, 128])
"""
def __init__(
self,
embed_dim: int,
num_heads: int,
dropout: float = 0.0,
bias: bool = True,
add_bias_kv: bool = False,
add_zero_attn: bool = False,
kdim: int | None = None,
vdim: int | None = None,
batch_first: bool = False,
device: PlaceLike | None = None,
dtype: DTypeLike | None = None,
) -> None:
if dtype:
super().__init__(dtype=dtype)
else:
super().__init__()
self.embed_dim = embed_dim
self.kdim = kdim if kdim is not None else embed_dim
self.vdim = vdim if vdim is not None else embed_dim
self._qkv_same_embed_dim = (
self.kdim == embed_dim and self.vdim == embed_dim
)
self.num_heads = num_heads
self.dropout = dropout
self.batch_first = batch_first
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == self.embed_dim
self.in_proj_bias = None
self.q_proj_bias = None
self.k_proj_bias = None
self.v_proj_bias = None
if self._qkv_same_embed_dim:
self.in_proj_weight = self.create_parameter(
shape=[3 * embed_dim, embed_dim],
dtype=self._dtype,
is_bias=False,
device=device,
default_initializer=XavierUniform(),
)
self.q_proj_weight = None
self.k_proj_weight = None
self.v_proj_weight = None
if bias:
self.in_proj_bias = self.create_parameter(
shape=[3 * embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
)
else:
self.q_proj_weight = self.create_parameter(
shape=[embed_dim, embed_dim],
dtype=self._dtype,
is_bias=False,
device=device,
default_initializer=XavierUniform(),
)
self.k_proj_weight = self.create_parameter(
shape=[embed_dim, self.kdim],
dtype=self._dtype,
is_bias=False,
device=device,
default_initializer=XavierUniform(),
)
self.v_proj_weight = self.create_parameter(
shape=[embed_dim, self.vdim],
dtype=self._dtype,
is_bias=False,
device=device,
default_initializer=XavierUniform(),
)
self.in_proj_weight = None
if bias:
self.q_proj_bias = self.create_parameter(
shape=[embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
)
self.k_proj_bias = self.create_parameter(
shape=[embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
)
self.v_proj_bias = self.create_parameter(
shape=[embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
)
self.out_proj = paddle.compat.nn.Linear(
embed_dim, embed_dim, bias=bias, dtype=self._dtype
)
self.add_bias_kv = add_bias_kv
self.add_zero_attn = add_zero_attn
if add_bias_kv:
self.bias_k = self.create_parameter(
shape=[1, 1, embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
default_initializer=XavierNormal(),
)
self.bias_v = self.create_parameter(
shape=[1, 1, embed_dim],
dtype=self._dtype,
is_bias=True,
device=device,
default_initializer=XavierNormal(),
)
else:
self.bias_k = self.bias_v = None
def _convert_bool_mask_to_float(
self, mask: paddle.Tensor, dtype: DTypeLike
) -> paddle.Tensor:
"""
Convert boolean mask to float mask. True -> -inf, False -> 0.0
Args:
mask (paddle.Tensor): boolean mask
dtype (DTypeLike): float dtype
Returns:
paddle.Tensor: float mask
"""
assert mask.dtype == paddle.bool, (
f"mask must be boolean, but got {mask.dtype}"
)
filler = paddle.to_tensor(paddle.finfo(dtype).min, dtype=dtype)
return paddle.where(mask, filler, paddle.zeros_like(mask, dtype=dtype))
def _combine_masks(
self, mask1: paddle.Tensor, mask2: paddle.Tensor, dtype: DTypeLike
) -> paddle.Tensor:
"""
Safely combine two masks, mask can be bool or float.
If both mask are bool, this function equals to
paddle.logical_or(mask1, mask2) and return boolean mask.
Otherwise, the boolean mask will be converted to float and combined with
the float mask using addition.
Args:
mask1 (paddle.Tensor): mask1
mask2 (paddle.Tensor): mask2
Returns:
paddle.Tensor: combined mask
"""
if mask1.dtype == paddle.bool and mask2.dtype == paddle.bool:
return mask1 | mask2
if mask1.dtype == paddle.bool:
mask1 = self._convert_bool_mask_to_float(mask1, dtype=dtype)
if mask2.dtype == paddle.bool:
mask2 = self._convert_bool_mask_to_float(mask2, dtype=dtype)
return mask1 + mask2
def _pad_mask(self, mask: Tensor, pad_amt: int = 1) -> Tensor:
shape = mask.shape
pad_shape = [*shape[:-1], pad_amt]
pad_tensor = paddle.zeros(pad_shape, dtype=mask.dtype)
return paddle.concat([mask, pad_tensor], axis=-1)
def _project_qkv(
self, query: Tensor, key: Tensor, value: Tensor
) -> tuple[Tensor, Tensor, Tensor]:
# in: [batch, seq_len, embed]
# out: [batch, seq_len, embed]
if self._qkv_same_embed_dim:
if id(query) == id(key) and id(key) == id(value):
qkv = F.linear(query, self.in_proj_weight.T, self.in_proj_bias)
q, k, v = qkv.split(3, axis=-1)
else:
q_w, k_w, v_w = self.in_proj_weight.chunk(3, axis=0)
q_b, k_b, v_b = (
self.in_proj_bias.chunk(3, axis=0)
if self.in_proj_bias is not None
else (None,) * 3
)
q = F.linear(query, q_w.T, q_b)
k = F.linear(key, k_w.T, k_b)
v = F.linear(value, v_w.T, v_b)
else:
q = F.linear(query, self.q_proj_weight.T, self.q_proj_bias)
k = F.linear(key, self.k_proj_weight.T, self.k_proj_bias)
v = F.linear(value, self.v_proj_weight.T, self.v_proj_bias)
return q, k, v
def _prepare_qkv_heads(
self,
q: Tensor,
k: Tensor,
v: Tensor,
batch_size: int,
target_seq_len: int,
) -> tuple[Tensor, Tensor, Tensor]:
# in: [batch, seq_len, num_head * dim]
# out: [batch, num_head, seq_len, dim]
if self.add_bias_kv:
k = paddle.concat(
[k, self.bias_k.expand([batch_size, -1, -1])], axis=1
)
v = paddle.concat(
[v, self.bias_v.expand([batch_size, -1, -1])], axis=1
)
q = q.reshape(
[batch_size, target_seq_len, self.num_heads, self.head_dim]
).transpose([0, 2, 1, 3])
current_src_len = k.shape[1]
k = k.reshape(
[batch_size, current_src_len, self.num_heads, self.head_dim]
).transpose([0, 2, 1, 3])
v = v.reshape(
[batch_size, current_src_len, self.num_heads, self.head_dim]
).transpose([0, 2, 1, 3])
if self.add_zero_attn:
zeros = paddle.zeros(
[batch_size, self.num_heads, 1, self.head_dim], dtype=k.dtype
)
k = paddle.concat([k, zeros], axis=2)
v = paddle.concat([v, zeros], axis=2)
return q, k, v
def _prepare_attn_mask(
self,
attn_mask: Tensor | None,
key_padding_mask: Tensor | None,
target_seq_len: int,
src_len_before_bias: int,
dtype: DTypeLike,
batch_size: int,
is_causal: bool,
need_weights: bool,
) -> Tensor | None:
# Do not generate attn_mask if is_causal is True and add_bias_kv is False
# and add_zero_attn is False. In such case, we pass attn_mask as None to
# select efficient implementation backend of sdpa.
if (
is_causal
and not self.add_bias_kv
and not self.add_zero_attn
and key_padding_mask is None
and not need_weights
):
return None
if attn_mask is None and not is_causal and key_padding_mask is None:
return None
if attn_mask is None:
if is_causal:
attn_mask = paddle.triu(
paddle.ones(
[target_seq_len, src_len_before_bias], dtype=paddle.bool
),
diagonal=1,
)
else:
attn_mask = paddle.zeros(
[target_seq_len, src_len_before_bias], dtype=dtype
)
pad_count = int(self.add_zero_attn + self.add_bias_kv)
if pad_count > 0:
attn_mask = self._pad_mask(attn_mask, pad_amt=pad_count)
if key_padding_mask is not None:
key_padding_mask = self._pad_mask(
key_padding_mask, pad_amt=pad_count
)
if attn_mask.dim() == 2:
attn_mask = attn_mask.expand(
[batch_size * self.num_heads, *attn_mask.shape]
)
if attn_mask.dim() == 3:
attn_mask = attn_mask.reshape(
[batch_size, self.num_heads, target_seq_len, -1]
)
if key_padding_mask is not None:
# [N, len_k+pad_count] -> [N, 1, 1, len_k+pad_count]
key_padding_mask = key_padding_mask.unsqueeze(axis=[1, 2])
key_padding_mask = key_padding_mask.repeat(
[1, *attn_mask.shape[1:3], 1]
)
attn_mask = self._combine_masks(attn_mask, key_padding_mask, dtype)
if attn_mask.dtype != dtype:
if attn_mask.dtype == paddle.bool:
attn_mask = self._convert_bool_mask_to_float(attn_mask, dtype)
else:
attn_mask = attn_mask.astype(dtype)
return attn_mask
def _attention_core(
self,
q: Tensor,
k: Tensor,
v: Tensor,
final_mask: Tensor | None,
need_weights: bool,
is_causal: bool,
) -> tuple[Tensor, Tensor | None]:
# in: [batch, num_head, seq_len, head_dim]
# out: [batch, num_head, seq_len, head_dim]
batch_size, _, target_seq_len, _ = q.shape
is_causal = is_causal and final_mask is None
if not need_weights:
attn_output = (
paddle.compat.nn.functional.scaled_dot_product_attention(
q,
k,
v,
attn_mask=final_mask,
dropout_p=self.dropout if self.training else 0.0,
is_causal=is_causal,
)
)
attn_output = attn_output.transpose([0, 2, 1, 3])
attn_output = attn_output.reshape(
[batch_size, target_seq_len, self.embed_dim]
)
return attn_output, None
else:
scores = paddle.matmul(q, k, transpose_y=True)
scores = scores / (self.head_dim**0.5)
if final_mask is not None:
if final_mask.dtype == paddle.bool:
final_mask = self._convert_bool_mask_to_float(
final_mask, scores.dtype
)
scores = scores + final_mask
weights = F.softmax(scores, axis=-1)
weights = F.dropout(weights, self.dropout, training=self.training)
ctx = paddle.matmul(weights, v)
attn_output = ctx.transpose([0, 2, 1, 3]).reshape(
[batch_size, target_seq_len, self.embed_dim]
)
return attn_output, weights if need_weights else None
def forward(
self,
query: paddle.Tensor,
key: paddle.Tensor,
value: paddle.Tensor,
key_padding_mask: paddle.Tensor | None = None,
need_weights: bool = True,
attn_mask: paddle.Tensor | None = None,
average_attn_weights: bool = True,
is_causal: bool = False,
) -> tuple[paddle.Tensor, paddle.Tensor | None]:
r"""
Forward pass of the MultiheadAttention layer.
.. note::
If ``need_weights`` is ``False``, this api will fallback to native math implementation,
otherwise it will call ``paddle.compat.nn.functional.scaled_dot_product_attention`` to
compute the attention score.
To achieve better performance, explicitly set ``need_weights=False``,
and set ``is_causal=True`` if the attn_mask is the causal mask.
Parameters:
query (Tensor): The query embeddings. Shape depends on `batch_first`.
If `batch_first` is False, shape is `[target_seq_len, batch_size, embed_dim]`.
If `batch_first` is True, shape is `[batch_size, target_seq_len, embed_dim]`.
key (Tensor): The key embeddings. Shape depends on `batch_first`.
If `batch_first` is False, shape is `[source_seq_len, batch_size, kdim]`.
If `batch_first` is True, shape is `[batch_size, source_seq_len, kdim]`.
value (Tensor): The value embeddings. Shape depends on `batch_first`.
If `batch_first` is False, shape is `[source_seq_len, batch_size, vdim]`.
If `batch_first` is True, shape is `[batch_size, source_seq_len, vdim]`.
key_padding_mask (Tensor, optional): If specified, a mask indicating which
elements within `key` to ignore for the purpose of attention (i.e. treat as "padding").
Can be a boolean mask (True indicates padding) or a float mask.
Shape is `[batch_size, source_seq_len]`. Default: None.
need_weights (bool, optional): Indicate whether to return the attention
weights. Default: True.
attn_mask (Tensor, optional): 2D or 3D mask that prevents attention to certain positions.
A 2D mask will be broadcasted for all batches while a 3D mask allows different masks
for the entries in the batch. Shape is `[target_seq_len, source_seq_len]` or
`[batch_size * num_heads, target_seq_len, source_seq_len]`. Default: None.
average_attn_weights (bool, optional): If True, indicates that the returned
`attn_weights` should be averaged across heads. Default: True.
is_causal (bool, optional): If True, implies that a causal mask is applied to
the attention implementation. If attn_mask is None and is_causal is True,
a causal mask is automatically created and used in the attention computation.
Default: False.
Returns:
tuple[Tensor, Tensor|None]:
- **attn_output** (Tensor): The output of the attention mechanism.
Shape matches `query` (based on `batch_first`).
- **attn_output_weights** (Tensor|None): The attention weights. Returns None if
`need_weights` is False. Shape is `[batch_size, target_seq_len, source_seq_len]`
if `average_attn_weights` is True.
If `average_attn_weights` is False, shape is
`[batch_size, num_heads, target_seq_len, source_seq_len]`.
"""
is_batched = query.dim() == 3
if not is_batched:
query = query.unsqueeze(0 if self.batch_first else 1)
key = key.unsqueeze(0 if self.batch_first else 1)
value = value.unsqueeze(0 if self.batch_first else 1)
if key_padding_mask is not None and key_padding_mask.dim() != 2:
key_padding_mask = key_padding_mask.unsqueeze(0)
if not self.batch_first:
query = query.transpose([1, 0, 2])
key = key.transpose([1, 0, 2])
value = value.transpose([1, 0, 2])
batch_size, target_seq_len, _ = query.shape
src_len_before_bias = key.shape[1]
if key_padding_mask is not None:
assert key_padding_mask.shape == (batch_size, src_len_before_bias)
q, k, v = self._project_qkv(query, key, value)
q, k, v = self._prepare_qkv_heads(q, k, v, batch_size, target_seq_len)
final_mask = self._prepare_attn_mask(
attn_mask=attn_mask,
key_padding_mask=key_padding_mask,
target_seq_len=target_seq_len,
src_len_before_bias=src_len_before_bias,
dtype=q.dtype,
batch_size=batch_size,
is_causal=is_causal,
need_weights=need_weights,
)
attn_output, attn_weights = self._attention_core(
q, k, v, final_mask, need_weights, is_causal
)
attn_output = self.out_proj(attn_output)
if not self.batch_first:
attn_output = attn_output.transpose([1, 0, 2])
if need_weights and attn_weights is not None:
if average_attn_weights:
attn_weights = attn_weights.mean(axis=1)
if not is_batched:
attn_output = attn_output.squeeze(0 if self.batch_first else 1)
if attn_weights is not None:
attn_weights = attn_weights.squeeze(0)
return attn_output, attn_weights