chore: import upstream snapshot with attribution

python/paddle/sparse/nn/functional/__init__.py

@@ -0,0 +1,40 @@
+#   Copyright (c) 2022 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 .activation import leaky_relu, relu, relu6, softmax
+from .conv import (
+    conv2d,
+    conv3d,
+    subm_conv2d,
+    subm_conv2d_igemm,
+    subm_conv3d,
+    subm_conv3d_igemm,
+)
+from .pooling import max_pool3d
+from .transformer import attention
+
+__all__ = [
+    'conv2d',
+    'conv3d',
+    'subm_conv2d',
+    'subm_conv2d_igemm',
+    'subm_conv3d',
+    'subm_conv3d_igemm',
+    'max_pool3d',
+    'relu',
+    'relu6',
+    'leaky_relu',
+    'softmax',
+    'attention',
+]

python/paddle/sparse/nn/functional/activation.py

@@ -0,0 +1,223 @@
+#   Copyright (c) 2022 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
+
+__all__ = []
+
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def relu(x: Tensor, name: str | None = None) -> Tensor:
+    """
+    sparse relu activation, requiring x to be a SparseCooTensor or SparseCsrTensor.
+
+    .. math::
+
+        out = max(x, 0)
+
+    Parameters:
+        x (Tensor): The input Sparse Tensor with data type float32, float64.
+        name (str|None, optional): Name for the operation (optional, default is None).
+            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Sparse Tensor with the same data type and shape as ``x`` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> dense_x = paddle.to_tensor([-2.0, 0.0, 1.0])
+            >>> sparse_x = dense_x.to_sparse_coo(1)
+            >>> out = paddle.sparse.nn.functional.relu(sparse_x)
+            >>> print(out)
+            Tensor(shape=[3], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
+                   indices=[[0, 2]],
+                   values=[0., 1.])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.sparse_relu(x)
+    else:
+        op_type = 'sparse_relu'
+        helper = LayerHelper(op_type)
+        out = helper.create_sparse_variable_for_type_inference(x.dtype)
+        helper.append_op(
+            type=op_type, inputs={'x': x}, outputs={'out': out}, attrs={}
+        )
+        return out
+
+
+def softmax(x: Tensor, axis: int = -1, name: str | None = None) -> Tensor:
+    r"""
+    sparse softmax activation, requiring x to be a SparseCooTensor or SparseCsrTensor.
+
+    Note:
+        Only support axis=-1 for SparseCsrTensor, which is faster when read data
+        by row (axis=-1).
+
+    From the point of view of dense matrix, for each row :math:`i` and each column :math:`j`
+    in the matrix, we have:
+
+    .. math::
+
+        softmax_ij = \frac{\exp(x_ij - max_j(x_ij))}{\sum_j(exp(x_ij - max_j(x_ij))}
+
+    Parameters:
+        x (Tensor): The input tensor. It can be SparseCooTensor/SparseCsrTensor. The data type can be float32 or float64.
+        axis (int, optional): The axis along which to perform softmax calculations. Only support -1 for SparseCsrTensor.
+        name (str|None, optional): Name for the operation (optional, default is None).
+            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        Tensor: SparseCoo or SparseCsr, whose layout is the same with `x` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.seed(100)
+
+            >>> mask = paddle.rand((3, 4)) < 0.5
+            >>> x = paddle.rand((3, 4)) * mask.astype('float32')
+            >>> print(x)
+            Tensor(shape=[3, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[0.        , 0.95717543, 0.43864486, 0.        ],
+             [0.84765935, 0.45680618, 0.39412445, 0.        ],
+             [0.59444654, 0.        , 0.78364515, 0.        ]])
+
+            >>> csr = x.to_sparse_csr()
+            >>> print(csr)
+            Tensor(shape=[3, 4], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
+                   crows=[0, 2, 5, 7],
+                   cols=[1, 2, 0, 1, 2, 0, 2],
+                   values=[0.95717543, 0.43864486, 0.84765935, 0.45680618, 0.39412445,
+                           0.59444654, 0.78364515])
+
+            >>> out = paddle.sparse.nn.functional.softmax(csr)
+            >>> print(out)
+            Tensor(shape=[3, 4], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
+                   crows=[0, 2, 5, 7],
+                   cols=[1, 2, 0, 1, 2, 0, 2],
+                   values=[0.62680405, 0.37319586, 0.43255258, 0.29261294, 0.27483448,
+                           0.45284089, 0.54715902])
+
+            >>> coo = x.to_sparse_coo(sparse_dim=2)
+            >>> print(coo)
+            Tensor(shape=[3, 4], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
+                   indices=[[0, 0, 1, 1, 1, 2, 2],
+                            [1, 2, 0, 1, 2, 0, 2]],
+                   values=[0.95717543, 0.43864486, 0.84765935, 0.45680618, 0.39412445,
+                           0.59444654, 0.78364515])
+
+            >>> out = paddle.sparse.nn.functional.softmax(coo)
+            >>> print(out)
+            Tensor(shape=[3, 4], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
+                   indices=[[0, 0, 1, 1, 1, 2, 2],
+                            [1, 2, 0, 1, 2, 0, 2]],
+                   values=[0.62680405, 0.37319589, 0.43255258, 0.29261294, 0.27483445,
+                           0.45284092, 0.54715902])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.sparse_softmax(x, axis)
+    else:
+        op_type = 'sparse_softmax'
+        helper = LayerHelper(op_type)
+        out = helper.create_sparse_variable_for_type_inference(x.dtype)
+        helper.append_op(
+            type=op_type,
+            inputs={'x': x},
+            outputs={'out': out},
+            attrs={'axis': axis},
+        )
+        return out
+
+
+def relu6(x: Tensor, name: str | None = None) -> Tensor:
+    """
+    sparse relu6 activation, requiring x to be a SparseCooTensor or SparseCsrTensor.
+
+    .. math::
+
+        relu6(x) = min(max(0, x), 6)
+
+    Parameters:
+        x (Tensor): The input Sparse Tensor with data type float32, float64.
+        name (str|None, optional): Name for the operation (optional, default is None).
+            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Sparse Tensor with the same data type and shape as ``x`` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> dense_x = paddle.to_tensor([-2.0, 0.0, 8.0])
+            >>> sparse_x = dense_x.to_sparse_coo(1)
+            >>> out = paddle.sparse.nn.functional.relu6(sparse_x)
+    """
+    assert in_dynamic_or_pir_mode(), (
+        "Currently, Sparse API only support dynamic mode or pir mode."
+    )
+    return _C_ops.sparse_relu6(x)
+
+
+def leaky_relu(
+    x: Tensor, negative_slope: float = 0.01, name: str | None = None
+) -> Tensor:
+    r"""
+    sparse leaky_relu activation, requiring x to be a SparseCooTensor or SparseCsrTensor.
+
+    .. math::
+        leaky\_relu(x)=
+        \left\{
+            \begin{array}{rcl}
+                x, & & if \ x >= 0 \\
+                negative\_slope * x, & & otherwise \\
+            \end{array}
+        \right.
+
+    Parameters:
+        x (Tensor): The input Sparse Tensor with data type float32, float64.
+        negative_slope (float, optional): Slope of the activation function at
+            :math:`x < 0` . Default is 0.01.
+        name (str|None, optional): Name for the operation (optional, default is None).
+            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Sparse Tensor with the same data type and shape as ``x`` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> dense_x = paddle.to_tensor([-2., 0., 5.])
+            >>> sparse_x = dense_x.to_sparse_coo(1)
+            >>> out = paddle.sparse.nn.functional.leaky_relu(sparse_x, 0.5)
+    """
+    assert in_dynamic_or_pir_mode(), (
+        "Currently, Sparse API only support dynamic mode or pir mode."
+    )
+    return _C_ops.sparse_leaky_relu(x, negative_slope)

python/paddle/sparse/nn/functional/pooling.py

@@ -0,0 +1,117 @@
+#   Copyright (c) 2022 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, Literal
+
+from paddle import _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+from paddle.nn.functional.pooling import _update_padding_nd
+from paddle.utils import convert_to_list
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import (
+        Size3,
+        Size6,
+    )
+    from paddle.nn.functional.common import _PaddingSizeMode
+
+__all__ = []
+
+
+def max_pool3d(
+    x: Tensor,
+    kernel_size: Size3,
+    stride: Size3 | None = None,
+    padding: _PaddingSizeMode | Size3 | Size6 = 0,
+    ceil_mode: bool = False,
+    data_format: Literal['NDHWC'] = "NDHWC",
+    name: str | None = None,
+) -> Tensor:
+    """
+    Implements sparse max pooling 3d operation.
+    See more details in :ref:`api_paddle_sparse_nn_MaxPool3D` .
+
+    Args:
+        x (Tensor): The input SparseCooTensor of pooling operator, which is a 5-D tensor with
+                          shape [N, D, H, W, C]. The format of input tensor `"NDHWC"`, where N represents batch size, C represents the number of channels, D, H and W represent the depth, height and width of the feature respectively.
+        kernel_size (int|list|tuple): The pool kernel size. If the 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, 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.
+        padding (string|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}
+        data_format (string, optional): The data format of the input and output data. An optional string from: `"NCDHW"`, `"NDHWC"`.
+                        The default is `"NCDHW"`. When it is `"NCDHW"`, the data is stored in the order of:
+                        `[batch_size, input_channels, input_depth, input_height, input_width]`. Currently only support `"NDHWC"` .
+        name(str|None, optional): For detailed information, please refer
+                             to :ref:`api_guide_Name`. Usually name is no need to set and
+                             None by default.
+
+    Returns:
+        Tensor: The output tensor of pooling result. The data type is same as input tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> dense_x = paddle.randn((1, 4, 4, 4, 3))
+            >>> sparse_x = dense_x.to_sparse_coo(4)
+            >>> kernel_sizes = [3, 3, 3]
+            >>> paddings = [0, 0, 0]
+            >>> strides = [1, 1, 1]
+            >>> out = paddle.sparse.nn.functional.max_pool3d(sparse_x, kernel_sizes, stride=strides, padding=paddings)
+            >>> print(out.shape)
+            paddle.Size([1, 2, 2, 2, 3])
+    """
+
+    assert in_dynamic_or_pir_mode(), (
+        "Currently, Sparse API only support dynamic mode or pir mode."
+    )
+    assert x.is_sparse_coo(), (
+        "Currently, sparse.relu only support the input of SparseCooTensor"
+    )
+    assert data_format == 'NDHWC', (
+        "Currently, sparse.max_pool3d only support data format of 'NDHWC'"
+    )
+
+    kernel_size = convert_to_list(kernel_size, 3, 'pool_size')
+    if stride is None:
+        stride = kernel_size
+    else:
+        stride = convert_to_list(stride, 3, 'pool_stride')
+
+    channel_last = True
+
+    padding, padding_algorithm = _update_padding_nd(
+        padding, 3, channel_last=channel_last, ceil_mode=ceil_mode
+    )
+
+    # TODO(zkh2016): remove the dependency on dilation from the backend
+    dilation = [1, 1, 1]
+
+    return _C_ops.sparse_maxpool(x, kernel_size, padding, dilation, stride)

python/paddle/sparse/nn/functional/transformer.py

@@ -0,0 +1,105 @@
+#   Copyright (c) 2022 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
+
+__all__ = []
+
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    sparse_mask: Tensor,
+    key_padding_mask: Tensor | None = None,
+    attn_mask: Tensor | None = None,
+    name: str | None = None,
+) -> Tensor:
+    r"""
+    Note:
+        This API is only used from ``CUDA 11.8`` .
+
+    SparseCsrTensor is used to store the intermediate result of Attention matrix
+    in Transformer module, which can reduce memory usage and improve performance.
+    ``sparse_mask`` express the sparse layout in CSR format.
+    The calculation 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 shape of the three parameters are: `[batch_size, num_heads, seq_len, head_dim]`, and
+    ``d`` represents ``head_dim`` .
+
+    Args:
+        query (DenseTensor): `query` in the Attention module. 4D Tensor with float32 or float64.
+        key (DenseTensor): `key` in the Attention module. 4D Tensor with float32 or float64.
+        value (DenseTensor): `value` in the Attention module. 4D Tensor with float32 or float64.
+        sparse_mask (SparseCsrTensor): The sparse layout in the Attention module. Its dense shape
+            is `[batch_size*num_heads, seq_len, seq_len]`. `nnz` of each batch must be the same.
+            dtype of `crows` and `cols` must be int64, dtype of `values` can be float32 or float64.
+        key_padding_mask (DenseTensor|None, optional): The key padding mask tensor in the Attention module.
+            2D tensor with shape: [batch_size, seq_len]. dtype can be float32 or float64. Default: None.
+        attn_mask (DenseTensor|None, optional): The attention mask tensor in the Attention module.
+            2D tensor with shape: [seq_len, seq_len]. dtype can be float32 or float64. Default: None.
+        name (str|None, optional): The default value is None. Normally there is no need for user
+            to set this property. For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        4D tensor with shape: [batch_size, num_heads, seq_len, head_dim]. dtype is same with input.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+
+            >>> batch_size = 16
+            >>> num_heads = 16
+            >>> seq_len = 512
+            >>> head_dim = 32
+
+            >>> query = paddle.rand([batch_size, num_heads, seq_len, head_dim])
+            >>> key = paddle.rand([batch_size, num_heads, seq_len, head_dim])
+            >>> value = paddle.rand([batch_size, num_heads, seq_len, head_dim])
+
+            >>> query.stop_gradient = False
+            >>> key.stop_gradient = False
+            >>> value.stop_gradient = False
+
+            >>> mask = paddle.nn.functional.dropout(paddle.ones([seq_len, seq_len])).expand([batch_size, num_heads, seq_len, seq_len])
+            >>> sp_mask = mask.reshape([-1, seq_len, seq_len]).to_sparse_csr()
+
+            >>> kp_mask = paddle.randint(0, 2, [batch_size, seq_len]).astype(paddle.float32)
+            >>> attn_mask = paddle.randint(0, 2, [seq_len, seq_len]).astype(paddle.float32)
+
+            >>> output = paddle.sparse.nn.functional.attention(query, key, value, sp_mask, kp_mask, attn_mask)
+            >>> output.backward()
+    """
+    assert in_dynamic_or_pir_mode(), (
+        "Currently, Sparse API only support dynamic mode or pir mode."
+    )
+    return _C_ops.sparse_fused_attention(
+        query, key, value, sparse_mask, key_padding_mask, attn_mask
+    )