paddlepaddle--paddle/python/paddle/compat/nn/functional/__init__.py

# 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,
    )