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/__init__.py
+1177
View File
File diff suppressed because it is too large Load Diff
python/paddle/compat/distributions/__init__.py
+22
View File
@@ -0,0 +1,22 @@
# Copyright (c) 2026 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 . import categorical as categorical
from .categorical import Categorical
__all__ = [
'Categorical',
]
python/paddle/compat/distributions/categorical.py
+153
View File
@@ -0,0 +1,153 @@
# Copyright (c) 2026 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 numpy as np
import paddle
from paddle.distribution import constraint, distribution
from paddle.tensor import multinomial
class Categorical(distribution.Distribution):
arg_constraints = {
"probs": constraint.simplex,
"logits": constraint.real_vector,
}
has_enumerate_support = True
def __init__(
self,
probs=None,
logits=None,
validate_args: bool | None = None,
) -> None:
if (probs is None) == (logits is None):
raise ValueError(
"Either `probs` or `logits` must be specified, but not both."
)
if probs is not None:
if probs.dim() < 1:
raise ValueError(
"`probs` parameter must be at least one-dimensional."
)
self._probs = probs / probs.sum(-1, keepdim=True)
self._logits = None
self._param = self._probs
else:
if logits.dim() < 1:
raise ValueError(
"`logits` parameter must be at least one-dimensional."
)
self._logits = logits - paddle.logsumexp(
logits, axis=-1, keepdim=True
)
self._probs = None
self._param = self._logits
self._num_events = self._param.shape[-1]
batch_shape = (
tuple(self._param.shape[:-1]) if self._param.dim() > 1 else ()
)
super().__init__(batch_shape, validate_args=validate_args)
def expand(self, batch_shape, _instance=None):
new = (
self.__class__.__new__(self.__class__)
if _instance is None
else _instance
)
batch_shape = tuple(batch_shape)
param_shape = (*batch_shape, self._num_events)
new._probs = (
self._probs.expand(param_shape) if self._probs is not None else None
)
new._logits = (
self._logits.expand(param_shape)
if self._logits is not None
else None
)
new._param = new._logits if new._logits is not None else new._probs
new._num_events = self._num_events
super(Categorical, new).__init__(batch_shape, validate_args=False)
new._validate_args_enabled = self._validate_args_enabled
return new
@property
def support(self):
return constraint.integer_interval(0, self._num_events - 1)
@property
def logits(self):
if self._logits is None:
eps = paddle.finfo(self.probs.dtype).eps
probs = paddle.clip(self.probs, min=eps, max=1 - eps)
self._logits = paddle.log(probs)
return self._logits
@property
def probs(self):
if self._probs is None:
self._probs = paddle.nn.functional.softmax(self.logits, axis=-1)
return self._probs
@property
def param_shape(self):
return self._param.shape
@property
def mean(self):
return paddle.full_like(self.probs[..., 0], float('nan'))
@property
def mode(self):
return paddle.argmax(self.probs, axis=-1)
@property
def variance(self):
return paddle.full_like(self.probs[..., 0], float('nan'))
def sample(self, sample_shape=()):
sample_shape = tuple(sample_shape)
probs_2d = self.probs.reshape([-1, self._num_events])
samples_2d = multinomial(probs_2d, int(np.prod(sample_shape)), True).T
return samples_2d.reshape(self._extend_shape(sample_shape))
def log_prob(self, value):
if self._validate_args_enabled and paddle.in_dynamic_mode():
self._validate_sample(value)
value = paddle.cast(value, dtype='int64').unsqueeze(-1)
log_pmf = self.logits
output_shape = paddle.broadcast_shape(value.shape, log_pmf.shape)
value = paddle.broadcast_to(value, [*output_shape[:-1], 1])
log_pmf = paddle.broadcast_to(log_pmf, output_shape)
return paddle.take_along_axis(log_pmf, value, axis=-1).squeeze(-1)
def entropy(self):
min_real = paddle.finfo(self.logits.dtype).min
logits = paddle.clip(self.logits, min=min_real)
p_log_p = logits * self.probs
return -p_log_p.sum(-1)
def enumerate_support(self, expand=True):
values = paddle.arange(self._num_events, dtype='int64')
values = values.reshape(
[self._num_events] + [1] * len(self._batch_shape)
)
if expand:
values = paddle.broadcast_to(
values, [self._num_events, *self._batch_shape]
)
return values
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
python/paddle/compat/proxy.py
+682
View File
@@ -0,0 +1,682 @@
# 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 importlib
import importlib.abc
import importlib.util
import inspect
import pkgutil
import sys
import types
import warnings
from contextlib import contextmanager
from functools import cache
from typing import TYPE_CHECKING, Any, Literal
if TYPE_CHECKING:
from collections.abc import Generator, Iterable
from typing import TypeAlias
_ScopeType: TypeAlias = str | Iterable[str] | None
def warning_about_fake_interface(name: str):
warnings.warn(
f"The interface '{name}' is a fake implementation for torch compatibility. "
"It does not have the actual functionality of PyTorch. "
"Please refer to the PaddlePaddle documentation for equivalent functionality.",
category=UserWarning,
stacklevel=2,
)
def create_fake_class(name, attrs: dict[str, Any]):
"""Create a fake class with the given name and attributes."""
new_fn = lambda *args, **kwargs: warning_about_fake_interface(name)
attrs["__init__"] = new_fn
return type(name, (), attrs)
def create_fake_function(name):
"""Create a fake function with the given name and implementation."""
fn = lambda *args, **kwargs: warning_about_fake_interface(name)
fn.__name__ = name
return fn
class OverriddenAttribute:
def get_value(self):
raise NotImplementedError
class LazyImportOverriddenAttribute(OverriddenAttribute):
def __init__(self, full_name: str):
self._full_name = full_name
def get_value(self):
parts = self._full_name.split(".")
root_module = importlib.import_module(parts[0])
result = root_module
for part in parts[1:]:
result = getattr(result, part)
return result
class RawOverriddenAttribute(OverriddenAttribute):
def __init__(self, value: Any):
self._value = value
def get_value(self):
return self._value
class ProxyModule(types.ModuleType):
def __init__(
self,
original_module: types.ModuleType,
proxy_name: str,
overrides: dict[str, OverriddenAttribute],
):
super().__init__(proxy_name)
self._original_module = original_module
self._proxy_name = proxy_name
self._overrides = overrides
def __getattr__(self, name: str) -> Any:
if name in self._overrides:
return self._overrides[name].get_value()
return getattr(self._original_module, name)
class CallableProxyModule(ProxyModule):
"""
Preserve callability for modules whose type defines ``__call__``.
``callable(obj)`` does not consult ``obj.__getattr__("__call__")``. It checks the
type-level call slot instead, so callable modules need a dedicated proxy subtype.
"""
def __call__(self, *args, **kwargs):
return self._original_module(*args, **kwargs)
def _create_proxy_module(
original_module: types.ModuleType,
proxy_name: str,
overrides: dict[str, OverriddenAttribute],
) -> ProxyModule:
"""Wrap callable modules with a callable proxy and plain modules with ProxyModule."""
if callable(original_module):
return CallableProxyModule(original_module, proxy_name, overrides)
return ProxyModule(original_module, proxy_name, overrides)
GLOBAL_OVERRIDES: dict[str, OverriddenAttribute] = {
"torch.relu": LazyImportOverriddenAttribute("paddle.nn.functional.relu"),
"torch.TorchVersion": LazyImportOverriddenAttribute(
"paddle.paddle_version.PaddleVersion"
),
}
TORCH_PROXY_BLOCKED_MODULES = {
"tvm_ffi",
}
MAGIC_DISABLED_MODULE_ATTR: str = "__disable_torch_proxy__"
MAGIC_ENABLED_MODULE_ATTR: str = "__enable_torch_proxy__"
def _extend_torch_proxy_overrides(
overrides: dict[str, OverriddenAttribute],
) -> None:
GLOBAL_OVERRIDES.update(overrides)
@cache
def _register_compat_override():
import paddle.compat
PADDLE_PREFIX = "paddle.compat"
TORCH_PREFIX = "torch"
PUBLIC_ATTR_DECLARATION = "__all__"
compat_overrides = {}
for module_info in pkgutil.walk_packages(
paddle.compat.__path__,
paddle.compat.__name__ + ".",
):
module = importlib.import_module(module_info.name)
if hasattr(module, PUBLIC_ATTR_DECLARATION):
public_attrs = getattr(module, PUBLIC_ATTR_DECLARATION)
torch_module_name = module_info.name.replace(
PADDLE_PREFIX, TORCH_PREFIX, 1
)
for attr_name in public_attrs:
if attr_name.startswith("_"):
continue
paddle_attr = getattr(module, attr_name)
torch_attr_name = f"{torch_module_name}.{attr_name}"
compat_overrides[torch_attr_name] = RawOverriddenAttribute(
paddle_attr
)
_extend_torch_proxy_overrides(compat_overrides)
def _is_specific_module_or_its_submodule(name: str, module: str) -> bool:
return name == module or name.startswith(f"{module}.")
def _is_torch_module(name: str) -> bool:
return _is_specific_module_or_its_submodule(name, "torch")
def _is_torch_proxy_local_enabled_module(name: str, scope: set[str]) -> bool:
for enabled_module in scope:
if _is_specific_module_or_its_submodule(name, enabled_module):
return True
return False
def _is_torch_proxy_blocked_module(name: str) -> bool:
for blocked_module in TORCH_PROXY_BLOCKED_MODULES:
if _is_specific_module_or_its_submodule(name, blocked_module):
return True
return False
def _is_called_by_module_with_specific_dunder_attr(dunder_attr: str) -> bool:
frame = inspect.currentframe()
while frame is not None:
if frame.f_globals.get(dunder_attr):
return True
frame = frame.f_back
return False
def _is_called_by_torch_proxy_blocked_module():
return _is_called_by_module_with_specific_dunder_attr(
MAGIC_DISABLED_MODULE_ATTR
)
def _is_called_by_torch_proxy_local_enabled_module():
return _is_called_by_module_with_specific_dunder_attr(
MAGIC_ENABLED_MODULE_ATTR
)
class TorchProxyMetaFinder:
"""
PyTorch compatibility layer for PaddlePaddle.
This class provides a way to `import torch` but actually loads PaddlePaddle.
Inspired by the setuptools _distutils_hack.
"""
_local_enabled_scope: set[str]
_globally_enabled: bool
def __init__(self, scope: set[str] | None = None):
self._set_scope(scope)
def _set_scope(self, scope: set[str] | None):
self._local_enabled_scope = scope or set()
self._globally_enabled = scope is None
def find_spec(self, fullname, path, target=None):
if _is_torch_proxy_blocked_module(fullname):
return self._find_spec_for_torch_proxy_blocked_module(fullname)
if _is_torch_proxy_local_enabled_module(
fullname, self._local_enabled_scope
):
return self._find_spec_for_torch_proxy_local_enabled_module(
fullname
)
if not _is_torch_module(fullname):
return None
if _is_called_by_torch_proxy_blocked_module():
if fullname in TORCH_MODULES_CACHE:
return self._find_spec_for_cached_torch_module(fullname)
return None
if (
not self._globally_enabled
and not _is_called_by_torch_proxy_local_enabled_module()
):
if fullname in TORCH_MODULES_CACHE:
return self._find_spec_for_cached_torch_module(fullname)
return None
return self._find_spec_for_torch_module(fullname)
def _find_spec_for_specific_module(
self,
fullname: str,
enable_proxy_when_exec_module: bool,
patched_dunder_attr: str,
):
# Return a special loader that imports the blocked module without torch proxy
with use_compat_guard(enable=False):
spec = importlib.util.find_spec(fullname)
if spec is None:
return None
original_loader = spec.loader
if original_loader is None:
return None
class SpecificModuleLoader(importlib.abc.Loader):
def create_module(self, spec):
mod = original_loader.create_module(spec)
if mod is None:
# If original loader returns None, create default module
# and ensure it has necessary attributes from spec
mod = types.ModuleType(spec.name)
mod.__spec__ = spec
mod.__loader__ = self
if spec.origin is not None:
mod.__file__ = spec.origin
if spec.submodule_search_locations is not None:
mod.__path__ = list(spec.submodule_search_locations)
return mod
def exec_module(self, module):
# Import the real module with torch proxy disabled
with use_compat_guard(
enable=enable_proxy_when_exec_module, silent=True
):
original_loader.exec_module(module)
# Mark module as torch proxy disabled/local enabled
module.__dict__[patched_dunder_attr] = True
spec.loader = SpecificModuleLoader()
return spec
def _find_spec_for_torch_proxy_local_enabled_module(self, fullname: str):
return self._find_spec_for_specific_module(
fullname,
enable_proxy_when_exec_module=True,
patched_dunder_attr=MAGIC_ENABLED_MODULE_ATTR,
)
def _find_spec_for_torch_proxy_blocked_module(self, fullname: str):
return self._find_spec_for_specific_module(
fullname,
enable_proxy_when_exec_module=False,
patched_dunder_attr=MAGIC_DISABLED_MODULE_ATTR,
)
def _find_spec_for_cached_torch_module(self, fullname: str):
module = TORCH_MODULES_CACHE[fullname]
# Return cached module before enable proxy
class CachedTorchModuleLoader(importlib.abc.Loader):
def create_module(self, spec):
return module
def exec_module(self, module):
pass
# Always treat cached modules as packages to allow submodules to be loaded.
# This is necessary because some modules (e.g. torch._C) are not packages
# but have submodules (e.g. torch._C._dynamo) attached to them.
spec = importlib.util.spec_from_loader(
fullname,
CachedTorchModuleLoader(),
origin=getattr(module, "__file__", None),
is_package=True,
)
spec.submodule_search_locations = list(getattr(module, "__path__", []))
return spec
def _find_spec_for_torch_module(self, fullname: str):
# Map the requested torch fullname to the corresponding paddle fullname.
module_name = fullname.replace("torch", "paddle", 1)
source_module = importlib.import_module(module_name)
overrides = {
k.removeprefix(f"{fullname}."): v
for k, v in GLOBAL_OVERRIDES.items()
if k.startswith(f"{fullname}.")
}
is_pkg = hasattr(source_module, "__path__")
class TorchProxyLoader(importlib.abc.Loader):
def __init__(self, source, target_name):
self._source = source
self._target_name = target_name
def create_module(self, spec):
# Create a new module object that will act as the "torch..." module.
mod = _create_proxy_module(
self._source, self._target_name, overrides
)
# Preserve file/path information for tooling/debugging.
mod.__file__ = getattr(self._source, "__file__", None)
if is_pkg:
# package must expose __path__ so import machinery can find submodules
mod.__path__ = list(getattr(self._source, "__path__", []))
mod.__package__ = self._target_name
else:
mod.__package__ = self._target_name.rpartition('.')[0]
return mod
def exec_module(self, module):
# Populate the new module with attributes from the source paddle module.
# Skip a few special attributes that should reflect the new module name.
for k, v in self._source.__dict__.items():
if k in ("__name__", "__package__", "__path__", "__spec__"):
continue
if k in overrides:
continue
if isinstance(v, types.ModuleType):
v = _create_proxy_module(
v,
f"{self._target_name}.{k}",
{
kk.removeprefix(f"{k}."): vv
for kk, vv in overrides.items()
if kk.startswith(f"{k}.")
},
)
module.__dict__[k] = v
# Use fullname for the spec name and mark as package when appropriate so that
# statements like `import torch.nn.functional` work correctly.
return importlib.util.spec_from_loader(
fullname,
TorchProxyLoader(source_module, fullname),
is_package=is_pkg,
origin=getattr(source_module, "__file__", None),
)
TORCH_PROXY_FINDER = TorchProxyMetaFinder()
TORCH_MODULES_CACHE: dict[str, types.ModuleType] = {}
def _clear_torch_proxy_modules():
for name, module in list(sys.modules.items()):
if _is_torch_module(name) and isinstance(module, ProxyModule):
del sys.modules[name]
def _swap_torch_modules_to_cache():
for name, module in list(sys.modules.items()):
if _is_torch_module(name):
if not isinstance(module, ProxyModule):
TORCH_MODULES_CACHE[name] = sys.modules[name]
del sys.modules[name]
def _copy_torch_modules_from_cache():
for name in list(TORCH_MODULES_CACHE):
assert _is_torch_module(name), f"`{name}` is not a PyTorch module"
sys.modules[name] = TORCH_MODULES_CACHE[name]
def _modify_scope_of_torch_proxy(
scope: set[str] | None,
*,
silent: bool = False,
) -> None:
def _warn_or_not(msg: str):
if silent:
return
warnings.warn(msg)
if TORCH_PROXY_FINDER not in sys.meta_path:
TORCH_PROXY_FINDER._set_scope(scope)
return
if TORCH_PROXY_FINDER._globally_enabled:
if scope is not None:
_warn_or_not(
"PyTorch already enabled globally, scope modification ignored."
)
TORCH_PROXY_FINDER._set_scope(scope)
return
if scope is None:
_warn_or_not(
"Enabling PyTorch compat globally, previous scope will be ignored."
)
TORCH_PROXY_FINDER._globally_enabled = True
return
if scope != TORCH_PROXY_FINDER._local_enabled_scope:
_warn_or_not(
f"Extending PyTorch compat scope, previous scope: {TORCH_PROXY_FINDER._local_enabled_scope}, new scope: {scope}."
)
TORCH_PROXY_FINDER._local_enabled_scope |= scope
def _parse_scope(scope: str | Iterable[str] | None) -> set[str] | None:
if scope is None:
return None
if isinstance(scope, str):
return {scope}
return set(scope)
def enable_compat(
*,
scope: _ScopeType = None,
blocked_modules: _ScopeType = None,
backend: Literal["torch"] = "torch",
silent: bool = False,
) -> None:
"""
Enable the PyTorch compat by adding the TorchProxyMetaFinder to sys.meta_path.
This allows importing 'torch' modules that are actually proxies to PaddlePaddle.
Args:
scope (str or Iterable[str], optional): Specific module or modules to enable
PyTorch compat for. If None, enables PyTorch compat globally. Defaults to None.
blocked_modules (str or Iterable[str], optional): Specific module or modules to
exclude from PyTorch compat. Defaults to None.
backend (str, optional): The backend to enable compat for. Currently only
"torch" is supported. Defaults to "torch".
silent (bool, optional): If True, suppresses warnings about scope changes.
Defaults to False.
Example:
.. code-block:: pycon
:name: enable-compat-in-global-scope
>>> import paddle
>>> paddle.enable_compat() # Enable torch compat globally
>>> import torch # type: ignore[import-not-found] # This will import paddle as torch
>>> assert torch.sin is paddle.sin
>>> paddle.disable_compat() # Disable torch compat
.. code-block:: pycon
:name: enable-compat-in-specific-scope
>>> import paddle
>>> paddle.enable_compat(scope={"triton"}) # Enable torch compat for 'triton' module only
>>> import triton # type: ignore[import-untyped] # All `import torch` inside `triton` will proxy to paddle
>>> try:
... import torch # type: ignore[import-not-found] # This will raise ModuleNotFoundError
... except ModuleNotFoundError:
... print("PyTorch compat is not enabled globally.")
>>> paddle.disable_compat() # Disable torch compat
"""
assert backend == "torch", f"Unsupported backend: {backend}"
blocked_modules = _parse_scope(blocked_modules)
if blocked_modules is not None:
extend_torch_proxy_blocked_modules(blocked_modules)
scope = _parse_scope(scope)
_register_compat_override()
_swap_torch_modules_to_cache()
_modify_scope_of_torch_proxy(scope, silent=silent)
sys.meta_path.insert(0, TORCH_PROXY_FINDER)
def disable_compat() -> None:
"""
Disable the PyTorch proxy by removing the TorchProxyMetaFinder from sys.meta_path.
This prevents 'torch' imports from being proxied to PaddlePaddle.
Example:
.. code-block:: pycon
>>> import paddle
>>> paddle.enable_compat() # Enable torch compat globally
>>> import torch # type: ignore[import-not-found] # This will import paddle as torch
>>> assert torch.sin is paddle.sin
>>> paddle.disable_compat() # Disable torch compat
>>> try:
... import torch # This will raise ModuleNotFoundError
... except ModuleNotFoundError:
... print("PyTorch compat is disabled.")
"""
if TORCH_PROXY_FINDER in sys.meta_path:
sys.meta_path.remove(TORCH_PROXY_FINDER)
_clear_torch_proxy_modules()
_copy_torch_modules_from_cache()
return
warnings.warn("torch compat is not installed.")
@contextmanager
def use_compat_guard(
*,
enable: bool = True,
scope: _ScopeType = None,
silent: bool = False,
) -> Generator[None, None, None]:
"""
Context manager to temporarily enable or disable the PyTorch compat.
When `enable` is True (default), the PyTorch compat is enabled for the duration
of the context and restored to its previous state afterwards. When `enable`
is False, the PyTorch compat is disabled for the duration of the context and
restored afterwards.
Args:
enable (bool, optional): Whether to enable or disable the PyTorch compat
within the context. Defaults to True.
scope (str or Iterable[str], optional): Specific module or modules to enable
PyTorch compat for. If None, uses the global scope. Defaults to None.
silent (bool, optional): If True, suppresses warnings about scope changes.
Defaults to False.
Example:
.. code-block:: pycon
>>> import paddle
>>> with paddle.use_compat_guard():
... # code that requires the Torch compat to be enabled
... import torch # type: ignore[import-not-found]
...
... assert torch.sin is paddle.sin
... # Temporarily disable the Torch compat
... with paddle.use_compat_guard(enable=False):
... try:
... import torch
... except ModuleNotFoundError:
... print("Torch compat is disabled within this block.")
... # Torch compat is re-enabled here
... import torch
...
... assert torch.sin is paddle.sin
"""
scope = _parse_scope(scope)
already_has_torch_proxy = TORCH_PROXY_FINDER in sys.meta_path
original_local_enabled_scope = set(TORCH_PROXY_FINDER._local_enabled_scope)
original_globally_enabled = TORCH_PROXY_FINDER._globally_enabled
if enable == already_has_torch_proxy and (
(original_globally_enabled and scope is None)
or (original_local_enabled_scope == (scope or set()))
):
yield
return
if enable:
enable_compat(scope=scope, silent=silent)
try:
yield
finally:
TORCH_PROXY_FINDER._local_enabled_scope = (
original_local_enabled_scope
)
TORCH_PROXY_FINDER._globally_enabled = original_globally_enabled
disable_compat()
else:
disable_compat()
try:
yield
finally:
enable_compat(scope=None, silent=True)
TORCH_PROXY_FINDER._local_enabled_scope = (
original_local_enabled_scope
)
TORCH_PROXY_FINDER._globally_enabled = original_globally_enabled
def extend_torch_proxy_blocked_modules(modules: Iterable[str]) -> None:
"""Add modules to the PyTorch proxy blocked list.
Modules in the blocked list will not use PyTorch compat when imported,
and their functions will not trigger PyTorch compat when called.
By default, some modules are already in the blocked list, such as 'tvm_ffi'.
Args:
modules(Iterable[str]): An iterable of module names to block from PyTorch compat.
Example:
.. code-block:: pycon
>>> import paddle
>>> paddle.enable_compat() # Enable torch compat globally
>>> # Add 'my_custom_module' to the blocked list
>>> paddle.compat.extend_torch_proxy_blocked_modules(['my_custom_module'])
>>> # doctest: +SKIP('my_custom_module is not available')
>>> import my_custom_module # type: ignore[import-not-found] # This import will not use torch compat
"""
TORCH_PROXY_BLOCKED_MODULES.update(modules)
def paddle_triton_fun():
"""
Enable the triton support and return triton module.
Args: None.
Returns: triton module
Example:
.. code-block:: pycon
>>> # doctest: +REQUIRES(env:GPU)
>>> from paddle.compat import paddle_triton_fun
>>> triton = paddle_triton_fun()
>>> import triton.language as tl
>>> @triton.jit
>>> def add_kernel(X, Y, Z, N, BLOCK: tl.constexpr):
... pid = tl.program_id(0)
... offs = pid * BLOCK + tl.arange(0, BLOCK)
... mask = offs < N
... x = tl.load(X + offs, mask=mask)
... y = tl.load(Y + offs, mask=mask)
... tl.store(Z + offs, x + y, mask=mask)
"""
enable_compat(scope={"triton"})
import triton
return triton
python/paddle/compat/utils.py
+48
View File
@@ -0,0 +1,48 @@
# 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 paddle
from paddle import Tensor
from paddle.framework import (
in_dynamic_mode,
)
def _check_out_status(
out: Tensor | tuple[Tensor, Tensor] | list[Tensor],
expect_multiple: bool = False,
):
if out is None:
return
if not in_dynamic_mode():
raise RuntimeError(
"Using `out` static graph CINN backend is currently not supported. Directly return the tensor tuple instead.\n"
)
if expect_multiple:
if not isinstance(out, (tuple, list)) or len(out) != 2:
raise TypeError(
f"Expected a list or tuple of two tensors, got {type(out)} instead."
)
if not (
isinstance(out[0], paddle.Tensor)
and isinstance(out[1], paddle.Tensor)
):
raise TypeError(
f"Expected Tensor type in the tuple/list, got ({type(out[0])}, {type(out[1])}) instead."
)
else:
if not isinstance(out, paddle.Tensor):
raise TypeError(f"Expected a Tensor, got {type(out)} instead.")