paddlepaddle--paddle/python/paddle/sparse/creation.py

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

import numpy as np

import paddle
from paddle import _C_ops, in_dynamic_mode
from paddle.base.data_feeder import convert_dtype
from paddle.base.framework import (
    _current_expected_place,
    _get_paddle_place,
    core,
    dygraph_only,
    in_pir_mode,
)
from paddle.base.layer_helper import LayerHelper
from paddle.tensor import max, to_tensor

if TYPE_CHECKING:
    import numpy.typing as npt

    from paddle import CPUPlace, CUDAPinnedPlace, CUDAPlace, Tensor
    from paddle._typing import DTypeLike, NumericSequence, ShapeLike

__all__ = [
    'sparse_coo_tensor',
    'sparse_csr_tensor',
]


def _handle_dtype(data, dtype):
    if dtype:
        if convert_dtype(dtype) != convert_dtype(data.dtype):
            return data.astype(convert_dtype(dtype))
    return data


def _infer_dense_shape(indices, values):
    assert len(indices.shape) == 2
    lens = max(indices, axis=1)
    lens = lens + 1
    lens = lens.numpy()
    if len(values.shape) > 1:
        lens = np.append(lens, values.shape[1:])
    return list(lens)


def _get_place(place):
    place = _get_paddle_place(place)
    if place is None:
        place = _current_expected_place()
    elif not isinstance(
        place, (core.Place, core.CPUPlace, core.CUDAPinnedPlace, core.CUDAPlace)
    ):
        raise ValueError(
            "'place' must be any of paddle.Place, paddle.CPUPlace, paddle.CUDAPinnedPlace, paddle.CUDAPlace"
        )
    return place


def _check_indices_dtype(dtype):
    if dtype not in [paddle.int8, paddle.int16, paddle.int32, paddle.int64]:
        raise TypeError(
            "the dtype of indices must be 'int8' or 'int16' or 'int32' or 'int64'"
        )


def sparse_coo_tensor(
    indices: (
        list[list[int]]
        | tuple[tuple[int, ...], ...]
        | npt.NDArray[np.int_]
        | Tensor
    ),
    values: NumericSequence | npt.NDArray[Any] | Tensor,
    shape: ShapeLike | None = None,
    dtype: DTypeLike | None = None,
    place: CPUPlace | CUDAPinnedPlace | CUDAPlace | str | None = None,
    stop_gradient: bool = True,
) -> Tensor:
    r"""
    Constructs a sparse ``paddle.Tensor`` in coordinate format according to the indices
    and values of the specified non-zero elements.

    Args:
        indices(list|tuple|ndarray|Tensor): the indices of non-zero elements.
            Can be a list, tuple, numpy\.ndarray, paddle\.Tensor. The indices must be 2-D.
        values(list|tuple|ndarray|Tensor): Initial values for the tensor.
            Can be a scalar, list, tuple, numpy\.ndarray, paddle\.Tensor.
        shape(list|tuple|None, optional): The shape of the sparse tensor also represents the shape of
            original dense tensor. If not provided the smallest shape will be inferred to
            hold all elements.
        dtype(str|paddle.dtype|np.dtype, optional): The desired data type of returned tensor. Can be 'bool' , 'float16' ,
            'float32' , 'float64' , 'int8' , 'int16' , 'int32' , 'int64' , 'uint8',
            'complex64' , 'complex128'. Default: None, infers dtype from ``data``
            except for python float number which gets dtype from ``get_default_type`` .
        place(CPUPlace|CUDAPinnedPlace|CUDAPlace|str|None, optional): The place to allocate Tensor. Can be
            CPUPlace, CUDAPinnedPlace, CUDAPlace. Default: None, means global place. If ``place`` is
            string, It can be ``cpu``, ``gpu:x`` and ``gpu_pinned``, where ``x`` is the index of the GPUs.
        stop_gradient(bool, optional): Whether to block the gradient propagation of Autograd. Default: True.

    Returns:
        Tensor: A Tensor constructed from ``indices`` and ``values`` .

    Examples:

        .. code-block:: pycon

            >>> import paddle

            >>> indices = [[0, 1, 2], [1, 2, 0]]
            >>> values = [1.0, 2.0, 3.0]
            >>> dense_shape = [3, 3]
            >>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
            >>> print(coo)
            Tensor(shape=[3, 3], dtype=paddle.float32, place=Place(cpu), stop_gradient=True,
                   indices=[[0, 1, 2],
                            [1, 2, 0]],
                   values=[1., 2., 3.])
    """

    if in_dynamic_mode():
        place = _get_place(place)
        if not isinstance(indices, core.eager.Tensor):
            indices = to_tensor(
                indices, dtype=None, place=place, stop_gradient=True
            )
        if not isinstance(values, core.eager.Tensor):
            values = to_tensor(values, dtype, place, stop_gradient)
        if len(indices.shape) != 2:
            raise ValueError("'indices' must be 2-D.")

        nnz = indices.shape[1]
        sparse_dim = indices.shape[0]

        _check_indices_dtype(indices.dtype)

        if nnz != values.shape[0]:
            raise ValueError(
                f"the indices and values must have same number of non-zero, but get {nnz} and {values.shape[0]}"
            )

        dense_dim = len(values.shape) - 1

        if not indices.place._equals(place):
            indices = indices._copy_to(place, False)

        if not values.place._equals(place):
            values = values._copy_to(place, False)
        values = _handle_dtype(values, dtype)
        values.stop_gradient = stop_gradient

        min_shape = _infer_dense_shape(indices, values)

        if shape is None:
            shape = min_shape
        else:
            shape = list(shape)
            if shape < min_shape:
                raise ValueError(
                    f"the minimum shape required is {min_shape}, but get {shape}"
                )
            if len(shape) != sparse_dim + dense_dim:
                raise ValueError(
                    f"the number of dimensions(len(shape) must be sparse_dim({sparse_dim}) + dense_dim({dense_dim}), but get {len(shape)}"
                )

        return _C_ops.sparse_sparse_coo_tensor(values, indices, shape)
    elif in_pir_mode():
        return _C_ops.sparse_sparse_coo_tensor(values, indices, shape)
        if shape[0] is None:
            shape[0] = -1
    else:
        op_type = 'sparse_sparse_coo_tensor'
        inputs = {'values': values, 'indices': indices}
        if shape[0] is None:
            shape[0] = -1
        attrs = {'shape': shape}
        helper = LayerHelper(op_type)
        out = helper.create_sparse_variable_for_type_inference(dtype)
        helper.append_op(
            type=op_type, inputs=inputs, outputs={'out': out}, attrs=attrs
        )
        return out


def _infer_dense_csr_shape(crows, cols):
    crows_numpy = crows.numpy()
    cols_numpy = cols.numpy()
    batchs = np.sum(crows_numpy[:-1] > crows_numpy[1:]) + 1
    if (int(len(crows_numpy) / batchs) * batchs) != len(crows_numpy):
        raise ValueError(
            f"The calculated original matrix batch size is {batchs}, but it cannot correctly split the row data. Please carefully check the data or the input shape."
        )
    col = cols_numpy.max() + 1
    row = int(len(crows_numpy) / batchs) - 1
    if batchs == 1:
        return [row, col]
    else:
        return [batchs, row, col]


# TODO: need to support shape is None
@dygraph_only
def sparse_csr_tensor(
    crows: list[int] | tuple[int, ...] | npt.NDArray[np.int_] | Tensor,
    cols: list[int] | tuple[int, ...] | npt.NDArray[np.int_] | Tensor,
    values: NumericSequence | npt.NDArray[Any] | Tensor,
    shape: ShapeLike | None = None,
    dtype: DTypeLike | None = None,
    place: CPUPlace | CUDAPinnedPlace | CUDAPlace | str | None = None,
    stop_gradient: bool = True,
) -> Tensor:
    r"""
    Constructs a sparse ``paddle.Tensor`` in CSR(Compressed Sparse Row) format according to the
    ``crows``, ``cols`` and ``values``.
    Currently, the crows and cols of each batch must be incremented.

    Args:
        crows(list|tuple|ndarray|Tensor): 1-D array, each element in the rows represents the
            starting position of the first non-zero element of each row in values.
            Can be a list, tuple, numpy\.ndarray, paddle\.Tensor.
        cols(list|tuple|ndarray|Tensor): 1-D array, the column of non-zero elements.
            Can be a list, tuple, numpy\.ndarray, paddle\.Tensor.
        values(list|tuple|ndarray|Tensor): 1-D array, the non-zero elements.
            Can be a scalar, list, tuple, numpy\.ndarray, paddle\.Tensor.
        shape(list|tuple, optional): The shape of the sparse tensor also represents the shape of
            original dense tensor.
            hold all elements.
        dtype(str|paddle.dtype|np.dtype, optional): The desired data type of returned tensor. Can be 'bool' , 'float16' ,
            'float32' , 'float64' , 'int8' , 'int16' , 'int32' , 'int64' , 'uint8',
            'complex64' , 'complex128'. Default: None, infers dtype from ``data``
            except for python float number which gets dtype from ``get_default_type`` .
        place(CPUPlace|CUDAPinnedPlace|CUDAPlace|str|None, optional): The place to allocate Tensor. Can be
            CPUPlace, CUDAPinnedPlace, CUDAPlace. Default: None, means global place. If ``place`` is
            string, It can be ``cpu``, ``gpu:x`` and ``gpu_pinned``, where ``x`` is the index of the GPUs.
        stop_gradient(bool, optional): Whether to block the gradient propagation of Autograd. Default: True.

    Returns:
        Tensor: A Tensor constructed from ``crows``, ``cols`` and ``values`` .

    Examples:

        .. code-block:: pycon

            >>> import paddle

            >>> crows = [0, 2, 3, 5]
            >>> cols = [1, 3, 2, 0, 1]
            >>> values = [1, 2, 3, 4, 5]
            >>> dense_shape = [3, 4]
            >>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
            >>> print(csr)
            Tensor(shape=[3, 4], dtype=paddle.int64, place=Place(cpu), stop_gradient=True,
                   crows=[0, 2, 3, 5],
                   cols=[1, 3, 2, 0, 1],
                   values=[1, 2, 3, 4, 5])
    """

    place = _get_place(place)

    if not isinstance(crows, core.eager.Tensor):
        crows = to_tensor(crows, dtype=None, place=place, stop_gradient=True)
    if not isinstance(cols, core.eager.Tensor):
        cols = to_tensor(cols, dtype=None, place=place, stop_gradient=True)
    if not isinstance(values, core.eager.Tensor):
        values = to_tensor(values, dtype, place, stop_gradient)

    _check_indices_dtype(crows.dtype)
    _check_indices_dtype(cols.dtype)

    if shape is not None:
        if len(shape) != 2 and len(shape) != 3:
            raise ValueError(
                f"SparseCsrTensor only support 2-D or 3-D matrix. but get shape {shape}"
            )
    else:
        shape = _infer_dense_csr_shape(crows, cols)

    rows = shape[len(shape) - 2]

    if not crows.place._equals(place):
        crows = crows._copy_to(place, False)

    if not cols.place._equals(place):
        cols = cols._copy_to(place, False)

    if not values.place._equals(place):
        values = values._copy_to(place, False)
    values = _handle_dtype(values, dtype)
    values.stop_gradient = stop_gradient

    if len(crows.shape) != 1 or len(cols.shape) != 1 or len(values.shape) != 1:
        raise ValueError("The 'crows', 'cols' and 'values' must be 1-D.")

    if len(cols) != len(values):
        raise ValueError("the length of cols must be same as length of values")

    if len(shape) == 2:
        if crows.shape[0] != rows + 1:
            raise ValueError(
                f"The length({crows.shape[0]}) of crows must be equal to the rows({rows})+1 of matrix."
            )
        if crows[0] != 0:
            raise ValueError("the 0th value of crows must be 0")

        if crows[-1] != values.shape[0]:
            raise ValueError(
                "the last value of crows must be equal the number of non-zero"
            )
    else:
        if crows.shape[0] % (rows + 1) != 0:
            raise ValueError(
                f"The length({crows.shape[0]}) of crows must be divisible the rows({rows})+1 of matrix."
            )
    # TODO(zkh2016): check whether the value in crows and cols is legal

    return core.eager.sparse_csr_tensor(
        crows, cols, values, shape, stop_gradient
    )