core/maxframe/lib/sparse/array.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2025 Alibaba Group Holding Ltd.
#
# 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 functools import partialmethod

from ...utils import ceildiv
from .core import (
    cp,
    cps,
    get_array_module,
    get_sparse_module,
    is_cupy,
    issparse,
    naked,
    np,
)


class SparseNDArray:
    __slots__ = ("__weakref__",)
    __array_priority__ = 21

    def __new__(cls, *args, **kwargs):
        shape = kwargs.get("shape", None)
        if shape is not None and len(shape) == 1:
            from .vector import SparseVector

            return object.__new__(SparseVector)
        if len(args) == 1 and issparse(args[0]) and args[0].ndim == 2:
            from .matrix import SparseMatrix

            return object.__new__(SparseMatrix)

        else:
            if cls is not SparseNDArray:
                return object.__new__(cls)
            else:
                raise ValueError(
                    f"The construct params of {cls.__name__} are invalid: "
                    f"args={args}, kwargs={kwargs}"
                )

    @property
    def raw(self):
        raise NotImplementedError


def call_sparse(method, *args, **kwargs):
    new_args = []
    make_dense = False
    matrix = None
    for arg in args:
        if hasattr(arg, "spmatrix"):
            # todo add support for multiple sparse arrays
            if make_dense or matrix is not None:
                make_dense = True
            matrix = arg
            new_args.append(matrix.spmatrix.data)
        else:
            if isinstance(arg, np.ndarray):
                make_dense = True
            new_args.append(arg)

    spmatrix = matrix.spmatrix
    if make_dense:
        new_args = [arg.toarray() if hasattr(arg, "spmatrix") else arg for arg in args]

    xp = get_array_module(spmatrix)
    try:
        new_data = getattr(xp, method)(*new_args, **kwargs)
    except AttributeError:
        if xp is np:
            from scipy import special
        else:
            from cupyx.scipy import special
        new_data = getattr(special, method)(*new_args, **kwargs)

    if not make_dense:
        new_spmatrix = get_sparse_module(spmatrix).csr_matrix(
            (new_data, spmatrix.indices, spmatrix.indptr), spmatrix.shape
        )
    else:
        new_spmatrix = get_sparse_module(spmatrix).csr_matrix(new_data)
    return SparseNDArray(new_spmatrix, shape=matrix.shape)


class SparseArray(SparseNDArray):
    __slots__ = ("spmatrix",)

    @property
    def ndim(self):
        return len(self.shape)

    def tocsr(self):
        return self

    def toarray(self):
        if self.shape != self.spmatrix.shape:
            return self.spmatrix.toarray().reshape(self.shape)
        else:
            return self.spmatrix.toarray()

    def todense(self):
        return self.toarray()

    def ascupy(self):
        is_cp = get_array_module(self.spmatrix) is cp
        if is_cp:
            return self
        mat_tuple = (
            cp.asarray(self.data),
            cp.asarray(self.indices),
            cp.asarray(self.indptr),
        )
        return SparseNDArray(
            cps.csr_matrix(mat_tuple, shape=self.spmatrix.shape), shape=self.shape
        )

    def asscipy(self):
        is_cp = get_array_module(self.spmatrix) is cp
        if not is_cp:
            return self
        return SparseNDArray(self.spmatrix.get(), shape=self.shape)

    def __array__(self, dtype=None):
        x = self.toarray()
        if dtype and x.dtype != dtype:
            return x.astype(dtype)
        return x

    @property
    def nbytes(self):
        return (
            self.spmatrix.data.nbytes
            + self.spmatrix.indptr.nbytes
            + self.spmatrix.indices.nbytes
        )

    @property
    def raw(self):
        return self.spmatrix

    @property
    def data(self):
        return self.spmatrix.data

    @property
    def indptr(self):
        return self.spmatrix.indptr

    @property
    def indices(self):
        return self.spmatrix.indices

    @property
    def nnz(self):
        return self.spmatrix.nnz

    @property
    def shape(self):
        raise self.spmatrix.shape

    @property
    def dtype(self):
        return self.spmatrix.dtype

    def copy(self):
        return SparseNDArray(self.spmatrix.copy(), shape=self.shape)

    @property
    def real(self):
        xps = get_sparse_module(self.spmatrix)
        return SparseNDArray(
            xps.csr_matrix(
                (self.spmatrix.data.real, self.spmatrix.indices, self.spmatrix.indptr),
                self.spmatrix.shape,
            ),
            shape=self.shape,
        )

    @real.setter
    def real(self, r):
        xps = get_sparse_module(self.spmatrix)
        x = self.spmatrix.toarray()
        if issparse(r):
            r = r.toarray()
        x.real = r
        self.spmatrix = xps.csr_matrix(x)

    @property
    def imag(self):
        xps = get_sparse_module(self.spmatrix)
        return SparseNDArray(
            xps.csr_matrix(
                (self.spmatrix.data.imag, self.spmatrix.indices, self.spmatrix.indptr),
                self.spmatrix.shape,
            ),
            shape=self.shape,
        )

    @imag.setter
    def imag(self, imag):
        xps = get_sparse_module(self.spmatrix)
        x = self.spmatrix.toarray()
        if issparse(imag):
            imag = imag.toarray()
        x.imag = imag
        self.spmatrix = xps.csr_matrix(x)

    def __getattr__(self, attr):
        is_cp = get_array_module(self.spmatrix) is cp
        if attr == "device" and is_cp:
            try:
                return self.spmatrix.device
            except NotImplementedError:
                return cp.cuda.Device(0)
        if attr == "get" and is_cp:
            return lambda: SparseNDArray(self.spmatrix.get(), shape=self.shape)

        return super().__getattribute__(attr)

    def __getstate__(self):
        return self.spmatrix

    def __setstate__(self, state):
        self.spmatrix = state

    def astype(self, dtype, **_):
        dtype = np.dtype(dtype)
        if self.dtype == dtype:
            return self
        return SparseNDArray(self.spmatrix.astype(dtype), shape=self.shape)

    def transpose(self, axes=None):
        raise NotImplementedError

    def swapaxes(self, axis1, axis2):
        if axis1 == 0 and axis2 == 1:
            return self

        assert axis1 == 1 and axis2 == 0
        return self.transpose()

    def reshape(self, shape, **_):
        sp_shape = shape if len(shape) == 2 else (1, shape[0])
        spmatrix = self.spmatrix.tolil().reshape(sp_shape)
        return SparseNDArray(spmatrix, shape=shape)

    def broadcast_to(self, shape):
        # TODO(jisheng): implement broadcast_to
        raise NotImplementedError

    def squeeze(self, axis=None):
        # TODO(jisheng): implement squeeze
        raise NotImplementedError

    @property
    def T(self):
        raise NotImplementedError

    # ---------------- arithmetic ----------------------

    def __add__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("add", self, naked_other)
        if issparse(naked_other):
            x = self.spmatrix + naked_other
        else:
            x = self.toarray() + naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __radd__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("add", naked_other, self)
        if issparse(naked_other):
            x = self.spmatrix + naked_other
        else:
            x = self.toarray() + naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __sub__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("subtract", self, naked_other)
        if issparse(naked_other):
            x = self.spmatrix - naked_other
        else:
            x = self.toarray() - naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rsub__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("subtract", naked_other, self)
        if issparse(naked_other):
            x = naked_other - self.spmatrix
        else:
            x = naked_other - self.toarray()
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __mul__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        if is_cupy(self.spmatrix):
            if not cp.isscalar(naked_other):
                # TODO(jisheng): cupy does not implement multiply method
                is_other_sparse = issparse(naked_other)
                if (
                    is_other_sparse
                    and self.spmatrix.nnz == naked_other.nnz
                    and cp.all(self.spmatrix.indptr == naked_other.indptr)
                    and cp.all(self.spmatrix.indices == naked_other.indices)
                ):
                    x = cps.csr_matrix(
                        (
                            self.spmatrix.data * naked_other.data,
                            self.spmatrix.indices,
                            self.spmatrix.indptr,
                        ),
                        self.spmatrix.shape,
                    )
                else:
                    if is_other_sparse:
                        naked_other = other.toarray()
                    dense = self.spmatrix.toarray()
                    res = cp.multiply(dense, naked_other, out=dense)
                    x = cps.csr_matrix(res)
            else:
                x = self.spmatrix * naked_other
        else:
            x = self.spmatrix.multiply(naked_other)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rmul__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        if is_cupy(self.spmatrix):
            if not cp.isscalar(naked_other):
                # TODO(jisheng): cupy does not implement multiply method
                is_other_sparse = issparse(naked_other)
                if (
                    is_other_sparse
                    and self.spmatrix.nnz == naked_other.nnz
                    and cp.all(self.spmatrix.indptr == naked_other.indptr)
                    and cp.all(self.spmatrix.indices == naked_other.indices)
                ):
                    x = cps.csr_matrix(
                        (
                            naked_other.data * self.spmatrix.data,
                            self.spmatrix.indices,
                            self.spmatrix.indptr,
                        ),
                        self.spmatrix.shape,
                    )
                else:
                    if is_other_sparse:
                        naked_other = other.toarray()
                    dense = self.spmatrix.toarray()
                    res = cp.multiply(naked_other, dense, out=dense)
                    x = cps.csr_matrix(res)
            else:
                x = naked_other * self.spmatrix
        else:
            x = self.spmatrix.multiply(naked_other)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __matmul__(self, other):
        from . import matmul

        return matmul(self, other)

    def __rmatmul__(self, other):
        from . import matmul

        return matmul(other, self)

    def __div__(self, other):
        return self.__truediv__(other)

    def __rdiv__(self, other):
        return self.__rtruediv__(other)

    def __truediv__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        x = self.spmatrix / naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rtruediv__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        try:
            x = naked_other / self.spmatrix
        except TypeError:
            x = naked_other / self.spmatrix.toarray()
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __floordiv__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("floor_divide", self, naked_other)
        else:
            if issparse(naked_other):
                naked_other = other.toarray()
                x = get_sparse_module(self.spmatrix).csr_matrix(
                    self.toarray() // naked_other
                )
            else:
                x = self.toarray() // naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rfloordiv__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("floor_divide", naked_other, self)
        else:
            if issparse(naked_other):
                naked_other = other.toarray()
                x = get_sparse_module(self.spmatrix).csr_matrix(
                    naked_other // self.toarray()
                )
            else:
                x = naked_other // self.toarray()
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __pow__(self, other, modulo=None):
        if modulo is not None:
            return NotImplemented

        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        if get_array_module(naked_other).isscalar(naked_other):
            try:
                x = self.spmatrix.power(naked_other)
            except ValueError as e:  # pragma: no cover
                # https://github.com/scipy/scipy/issues/8678
                assert "WRITEBACKIFCOPY" in e.args[0]
                self.spmatrix = self.spmatrix.copy()
                x = self.spmatrix.power(naked_other)
        else:
            if issparse(naked_other):
                naked_other = other.toarray()
            x = self.toarray() ** naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rpow__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        if issparse(naked_other):
            naked_other = other.toarray()
        x = naked_other ** self.toarray()
        return get_array_module(x).asarray(x)

    def float_power(self, other):
        ret = self.__pow__(other)
        ret = naked(ret).astype(float)
        if issparse(ret):
            return SparseNDArray(ret, shape=self.shape)
        return ret

    def __mod__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        if get_array_module(naked_other).isscalar(naked_other):
            data = self.spmatrix.data % naked_other
            x = get_sparse_module(self.spmatrix).csr_matrix(
                (data, self.spmatrix.indices, self.spmatrix.indptr), self.spmatrix.shape
            )
        else:
            if issparse(naked_other):
                naked_other = other.toarray()
            x = get_sparse_module(self.spmatrix).csr_matrix(
                self.toarray() % naked_other
            )
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __rmod__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented
        is_sparse = issparse(naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        if get_array_module(naked_other).isscalar(naked_other):
            data = naked_other % self.spmatrix.data
            x = get_sparse_module(self.spmatrix).csr_matrix(
                (data, self.spmatrix.indices, self.spmatrix.indptr), self.spmatrix.shape
            )
        else:
            x = naked_other % self.toarray()
            if is_sparse:
                x = get_sparse_module(self.spmatrix).csr_matrix(x)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def fmod(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(self.spmatrix)
        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("fmod", self, naked_other)
        else:
            if issparse(naked_other):
                naked_other = other.toarray()
            x = get_sparse_module(self.spmatrix).csr_matrix(
                xp.fmod(self.toarray(), naked_other)
            )
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def logaddexp(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(self.spmatrix)
        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("logaddexp", self, naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        return xp.logaddexp(self.toarray(), naked_other)

    def logaddexp2(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(self.spmatrix)
        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("logaddexp2", self, naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        return xp.logaddexp2(self.toarray(), naked_other)

    def __neg__(self):
        return SparseNDArray(-self.spmatrix, shape=self.shape)

    def __pos__(self):
        return SparseNDArray(self.spmatrix.copy(), shape=self.shape)

    def __abs__(self):
        return SparseNDArray(abs(self.spmatrix), shape=self.shape)

    def fabs(self):
        xp = get_array_module(self.spmatrix)
        return SparseNDArray(
            get_sparse_module(self.spmatrix).csr_matrix(
                xp.abs(self.spmatrix), dtype="f8"
            ),
            shape=self.shape,
        )

    def rint(self):
        return SparseNDArray(self.spmatrix.rint(), shape=self.shape)

    def sign(self):
        return SparseNDArray(self.spmatrix.sign(), shape=self.shape)

    def conj(self):
        return SparseNDArray(self.spmatrix.conj(), shape=self.shape)

    def exp(self):
        return call_sparse("exp", self)

    def exp2(self):
        return call_sparse("exp2", self)

    def log(self):
        return call_sparse("log", self)

    def log2(self):
        return call_sparse("log2", self)

    def log10(self):
        return call_sparse("log10", self)

    def expm1(self):
        return SparseNDArray(self.spmatrix.expm1(), shape=self.shape)

    def log1p(self):
        return SparseNDArray(self.spmatrix.log1p(), shape=self.shape)

    def sqrt(self):
        return SparseNDArray(self.spmatrix.sqrt(), shape=self.shape)

    def square(self):
        return call_sparse("square", self)

    def cbrt(self):
        return call_sparse("cbrt", self)

    def reciprocal(self):
        return call_sparse("reciprocal", self)

    def _scipy_unary(self, func_name):
        spmatrix = self.spmatrix
        xp = get_array_module(spmatrix)
        if xp is np:
            from scipy import special
        else:
            from cupyx.scipy import special

        new_data = getattr(special, func_name)(spmatrix.data)
        new_spmatrix = get_sparse_module(spmatrix).csr_matrix(
            (new_data, spmatrix.indices, spmatrix.indptr), spmatrix.shape
        )
        return SparseNDArray(new_spmatrix, shape=self.shape)

    def _scipy_binary(self, func_name, other):
        try:
            naked_other = naked(other)
        except TypeError:  # pragma: no cover
            return NotImplemented

        xp = get_array_module(self.spmatrix)

        if xp is np:
            from scipy import special
        else:  # pragma: no cover
            from cupyx.scipy import special

        func = getattr(special, func_name)

        if get_array_module(naked_other).isscalar(naked_other):  # pragma: no cover
            return call_sparse(func, self, naked_other)
        else:
            if issparse(naked_other):  # pragma: no cover
                naked_other = other.toarray()
            x = get_sparse_module(self.spmatrix).csr_matrix(
                func(self.toarray(), naked_other)
            )
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    gamma = partialmethod(_scipy_unary, "gamma")
    gammaln = partialmethod(_scipy_unary, "gammaln")
    loggamma = partialmethod(_scipy_unary, "loggamma")
    gammasgn = partialmethod(_scipy_unary, "gammasgn")
    gammainc = partialmethod(_scipy_binary, "gammainc")
    gammaincinv = partialmethod(_scipy_binary, "gammaincinv")
    gammaincc = partialmethod(_scipy_binary, "gammaincc")
    gammainccinv = partialmethod(_scipy_binary, "gammainccinv")
    beta = partialmethod(_scipy_binary, "beta")
    betaln = partialmethod(_scipy_binary, "betaln")
    psi = partialmethod(_scipy_unary, "psi")
    rgamma = partialmethod(_scipy_unary, "rgamma")
    polygamma = partialmethod(_scipy_binary, "polygamma")
    multigammaln = partialmethod(_scipy_binary, "multigammaln")
    digamma = partialmethod(_scipy_unary, "digamma")
    poch = partialmethod(_scipy_binary, "poch")

    erf = partialmethod(_scipy_unary, "erf")
    erfc = partialmethod(_scipy_unary, "erfc")
    erfcx = partialmethod(_scipy_unary, "erfcx")
    erfi = partialmethod(_scipy_unary, "erfi")
    erfinv = partialmethod(_scipy_unary, "erfinv")
    erfcinv = partialmethod(_scipy_unary, "erfcinv")
    wofz = partialmethod(_scipy_unary, "wofz")
    dawsn = partialmethod(_scipy_unary, "dawsn")
    entr = partialmethod(_scipy_unary, "entr")

    ellipk = partialmethod(_scipy_unary, "ellipk")
    ellipkm1 = partialmethod(_scipy_unary, "ellipkm1")
    ellipkinc = partialmethod(_scipy_binary, "ellipkinc")
    ellipe = partialmethod(_scipy_unary, "ellipe")
    ellipeinc = partialmethod(_scipy_binary, "ellipeinc")
    elliprc = partialmethod(_scipy_binary, "elliprc")

    rel_entr = partialmethod(_scipy_binary, "rel_entr")
    kl_div = partialmethod(_scipy_binary, "kl_div")
    xlogy = partialmethod(_scipy_binary, "xlogy")

    jv = partialmethod(_scipy_binary, "jv")
    jve = partialmethod(_scipy_binary, "jve")
    yn = partialmethod(_scipy_binary, "yn")
    yv = partialmethod(_scipy_binary, "yv")
    yve = partialmethod(_scipy_binary, "yve")
    kn = partialmethod(_scipy_binary, "kn")
    kv = partialmethod(_scipy_binary, "kv")
    kve = partialmethod(_scipy_binary, "kve")
    iv = partialmethod(_scipy_binary, "iv")
    ive = partialmethod(_scipy_binary, "ive")
    hankel1 = partialmethod(_scipy_binary, "hankel1")
    hankel1e = partialmethod(_scipy_binary, "hankel1e")
    hankel2 = partialmethod(_scipy_binary, "hankel2")
    hankel2e = partialmethod(_scipy_binary, "hankel2e")

    hyp0f1 = partialmethod(_scipy_binary, "hyp0f1")

    airy = partialmethod(_scipy_unary, "airy")
    airye = partialmethod(_scipy_unary, "airye")
    itairy = partialmethod(_scipy_unary, "itairy")

    def __eq__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("equal", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix == naked_other
            else:
                x = self.toarray() == other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __ne__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("not_equal", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix != naked_other
            else:
                x = self.toarray() != other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __lt__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("less", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix < naked_other
            else:
                x = self.toarray() < other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __le__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("less_equal", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix <= naked_other
            else:
                x = self.toarray() <= other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __gt__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("greater", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix > naked_other
            else:
                x = self.toarray() > other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __ge__(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("greater_equal", self, naked_other)
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            if issparse(naked_other):
                x = self.spmatrix >= naked_other
            else:
                x = self.toarray() >= other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def logical_and(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            other_xp = get_array_module(naked_other)
            if other_xp.isscalar(naked_other):
                naked_other = other_xp.array(naked_other).astype(bool)
            else:
                naked_other = naked_other.astype(bool)
            x = self.spmatrix.astype(bool).multiply(naked_other)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def logical_or(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            other_xp = get_array_module(naked_other)
            if other_xp.isscalar(naked_other):
                if naked_other != 0:
                    x = np.logical_and(self.toarray(), naked_other)
                else:
                    x = self.spmatrix.astype(bool)
            else:
                naked_other = naked_other.astype(bool)
                x = (self.spmatrix.astype(bool) + naked_other).astype(bool)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def logical_xor(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            other_xp = get_array_module(naked_other)
            if other_xp.isscalar(naked_other):
                naked_other = other_xp.array(naked_other).astype(bool)
            else:
                naked_other = naked_other.astype(bool)
            x = self.spmatrix.astype(bool) != naked_other
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def logical_not(self):
        return call_sparse("logical_not", self)

    @staticmethod
    def _bitwise(this, other, method_name):
        try:
            naked_this = naked(this)
        except TypeError:
            return NotImplemented
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if not issparse(naked_this):
            return SparseArray._bitwise(naked_other, naked_this, method_name)

        if issparse(naked_other):
            naked_other = other.toarray()

        xp = get_array_module(naked_this)
        xps = get_sparse_module(naked_this)
        return SparseNDArray(
            xps.csr_matrix(getattr(xp, method_name)(this.toarray(), naked_other)),
            shape=naked_this.shape,
        )

    def __and__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_and", self, other)
        return self._bitwise(self.spmatrix, other, "bitwise_and")

    def __rand__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_and", other, self)
        return self._bitwise(other, self.spmatrix, "bitwise_and")

    def __or__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_or", self, other)
        return self._bitwise(self.spmatrix, other, "bitwise_or")

    def __ror__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_or", other, self)
        return self._bitwise(other, self.spmatrix, "bitwise_or")

    def __xor__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_xor", self, other)
        return self._bitwise(self.spmatrix, other, "bitwise_xor")

    def __rxor__(self, other):
        if get_array_module(other).isscalar(other):
            return call_sparse("bitwise_xor", other, self)
        return self._bitwise(other, self.spmatrix, "bitwise_xor")

    def isclose(self, other, **kw):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        return xp.isclose(self.toarray(), naked_other, **kw)

    def __invert__(self):
        return call_sparse("invert", self)

    @staticmethod
    def _shift(this, other, method_name):
        try:
            naked_this = naked(this)
        except TypeError:
            return NotImplemented
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xps = get_sparse_module(naked_this)
        xp = get_array_module(naked_this)

        if xp.isscalar(naked_this):
            other_xp = get_array_module(naked_other)
            data = getattr(other_xp, method_name)(naked_this, naked_other.data)
            indices, indptr, shape = (
                naked_other.indices,
                naked_other.indptr,
                naked_other.shape,
            )
        elif isinstance(naked_this, xp.ndarray):
            # dense
            return getattr(xp, method_name)(naked_this, other.toarray())
        else:
            tp = (
                np.int32 if is_cupy(naked_this) else np.bool_
            )  # cupy.sparse does not support bool
            mask = xps.csr_matrix(
                (
                    (naked_this.data > 0).astype(tp),
                    naked_this.indices,
                    naked_this.indptr,
                ),
                naked_this.shape,
            )
            naked_other = mask.multiply(naked_other)
            indices, indptr, shape = (
                naked_this.indices,
                naked_this.indptr,
                naked_this.shape,
            )
            data = getattr(xp, method_name)(naked_this.data, naked_other.data)

        return SparseNDArray(
            xps.csr_matrix((data, indices, indptr), shape), shape=shape
        )

    def __lshift__(self, other):
        return self._shift(self.spmatrix, other, "left_shift")

    def __rlshift__(self, other):
        return self._shift(other, self.spmatrix, "left_shift")

    def __rshift__(self, other):
        return self._shift(self.spmatrix, other, "right_shift")

    def __rrshift__(self, other):
        return self._shift(other, self.spmatrix, "right_shift")

    def sin(self):
        return SparseNDArray(self.spmatrix.sin(), shape=self.shape)

    def cos(self):
        return call_sparse("cos", self)

    def tan(self):
        return SparseNDArray(self.spmatrix.tan(), shape=self.shape)

    def arcsin(self):
        return SparseNDArray(self.spmatrix.arcsin(), shape=self.shape)

    def arccos(self):
        return call_sparse("arccos", self)

    def arctan(self):
        return SparseNDArray(self.spmatrix.arctan(), shape=self.shape)

    def arctan2(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(self.spmatrix)
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("arctan2", self, naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        x = xp.arctan2(self.toarray(), naked_other)
        return SparseNDArray(get_sparse_module(x).csr_matrix(x), shape=self.shape)

    def hypot(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        xp = get_array_module(self.spmatrix)
        other_xp = get_array_module(naked_other)
        if other_xp.isscalar(naked_other):
            return call_sparse("hypot", self, naked_other)
        if issparse(naked_other):
            naked_other = other.toarray()
        x = xp.hypot(self.toarray(), naked_other)
        return SparseNDArray(get_sparse_module(x).csr_matrix(x), shape=self.shape)

    def sinh(self):
        return SparseNDArray(self.spmatrix.sinh(), shape=self.shape)

    def cosh(self):
        xp = get_array_module(self.spmatrix)
        return xp.cosh(self.toarray())

    def tanh(self):
        return SparseNDArray(self.spmatrix.tanh(), shape=self.shape)

    def arcsinh(self):
        return SparseNDArray(self.spmatrix.arcsinh(), shape=self.shape)

    def arccosh(self):
        return call_sparse("arccosh", self)

    def arctanh(self):
        return SparseNDArray(self.spmatrix.arctanh(), shape=self.shape)

    def around(self, decimals=0):
        return call_sparse("around", self, decimals=decimals)

    def deg2rad(self):
        return SparseNDArray(self.spmatrix.deg2rad(), shape=self.shape)

    def rad2deg(self):
        return SparseNDArray(self.spmatrix.rad2deg(), shape=self.shape)

    def angle(self, deg=0):
        return call_sparse("angle", self, deg=deg)

    def dot(self, other, sparse=True):
        raise NotImplementedError

    def concatenate(self, other, axis=0):
        raise NotImplementedError

    def _reduction(
        self, method_name, axis=None, dtype=None, keepdims=None, todense=False, **kw
    ):
        raise NotImplementedError

    def sum(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("sum", axis=axis, dtype=dtype, keepdims=keepdims)

    def prod(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "sum", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def amax(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("max", axis=axis, dtype=dtype, keepdims=keepdims)

    def amin(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("min", axis=axis, dtype=dtype, keepdims=keepdims)

    def all(self, axis=None, dtype=None, keepdims=None):
        ret = self._reduction(
            "all", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )
        if not issparse(ret):
            if get_array_module(ret).isscalar(ret):
                return ret
            xps = get_sparse_module(self.spmatrix)
            ret = SparseNDArray(xps.csr_matrix(ret))
            return ret
        return ret

    def any(self, axis=None, dtype=None, keepdims=None):
        ret = self._reduction(
            "any", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )
        if not issparse(ret):
            if get_array_module(ret).isscalar(ret):
                return ret
            xps = get_sparse_module(self.spmatrix)
            ret = SparseNDArray(xps.csr_matrix(ret))
            return ret
        return ret

    def mean(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("mean", axis=axis, dtype=dtype, keepdims=keepdims)

    def nansum(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "nansum", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def nanprod(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "nanprod", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def nanmax(self, axis=None, dtype=None, keepdims=None):
        ret = self._reduction(
            "nanmax", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )
        if not issparse(ret):
            if get_array_module(ret).isscalar(ret):
                return ret
            xps = get_sparse_module(self.spmatrix)
            ret = SparseNDArray(xps.csr_matrix(ret))
            return ret
        return ret

    def nanmin(self, axis=None, dtype=None, keepdims=None):
        ret = self._reduction(
            "nanmin", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )
        if not issparse(ret):
            if get_array_module(ret).isscalar(ret):
                return ret
            xps = get_sparse_module(self.spmatrix)
            ret = SparseNDArray(xps.csr_matrix(ret))
            return ret
        return ret

    def nanmean(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "nanmean", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def argmax(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("argmax", axis=axis, dtype=dtype, keepdims=keepdims)

    def nanargmax(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "nanargmax", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def argmin(self, axis=None, dtype=None, keepdims=None):
        return self._reduction("argmin", axis=axis, dtype=dtype, keepdims=keepdims)

    def nanargmin(self, axis=None, dtype=None, keepdims=None):
        return self._reduction(
            "nanargmin", axis=axis, dtype=dtype, keepdims=keepdims, todense=True
        )

    def var(self, axis=None, dtype=None, ddof=0, keepdims=None):
        return self._reduction(
            "var", axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, todense=True
        )

    def cumsum(self, axis=None, dtype=None):
        return self.toarray().cumsum(axis=axis)

    def cumprod(self, axis=None, dtype=None):
        return self.toarray().cumprod(axis=axis)

    def nancumsum(self, axis=None, dtype=None):
        xp = get_array_module(self.spmatrix)
        return xp.nancumsum(self.toarray(), axis=axis)

    def nancumprod(self, axis=None, dtype=None):
        xp = get_array_module(self.spmatrix)
        return xp.nancumprod(self.toarray(), axis=axis)

    def count_nonzero(self, axis=None, dtype=None, keepdims=None):
        if axis is None:
            return get_array_module(self.spmatrix).array(
                [self.spmatrix.count_nonzero()]
            )[0]
        else:
            return get_array_module(self.spmatrix).count_nonzero(
                self.toarray(), axis=axis
            )

    def __getitem__(self, item):
        if isinstance(item, SparseArray):
            item = item.spmatrix
        if isinstance(item, list):
            item = tuple(item)

        x = self.spmatrix[item]
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def __setitem__(self, key, value):
        if is_cupy(self.spmatrix):
            return NotImplemented
        else:
            x = self.spmatrix.tolil()
            x[key] = value
            x = x.tocsr()
        self.spmatrix = x

    def _maximum_minimum(self, other, method_name):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if is_cupy(self.spmatrix):
            # TODO(jisheng): cupy does not implement sparse maximum and minimum
            return NotImplemented

        xps = get_sparse_module(self.spmatrix)
        xp = get_array_module(self.spmatrix)
        has_nan = xps.csr_matrix(
            (xp.isnan(self.spmatrix.data), self.spmatrix.indices, self.spmatrix.indptr),
            self.spmatrix.shape,
        )
        if issparse(naked_other):
            has_nan += xps.csr_matrix(
                (xp.isnan(naked_other.data), naked_other.indices, naked_other.indptr),
                naked_other.shape,
            )

        if issparse(naked_other):
            x = getattr(self.spmatrix, method_name)(naked_other)
        else:
            x = getattr(xp, method_name)(self.toarray(), naked_other)

        if has_nan.sum() > 0:
            x = x + (has_nan * np.nan)

        if issparse(x):
            return SparseNDArray(x, shape=self.shape)

        return get_array_module(x).asarray(x)

    def maximum(self, other):
        return self._maximum_minimum(other, "maximum")

    def minimum(self, other):
        return self._maximum_minimum(other, "minimum")

    def fmax(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        x = self.spmatrix.maximum(naked_other)
        if issparse(x):
            return SparseArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def fmin(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        x = self.spmatrix.minimum(naked_other)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def isinf(self):
        return call_sparse("isinf", self)

    def isnan(self):
        return call_sparse("isnan", self)

    def signbit(self):
        return call_sparse("signbit", self)

    def floor(self):
        return SparseNDArray(self.spmatrix.floor(), shape=self.shape)

    def ceil(self):
        return SparseNDArray(self.spmatrix.ceil(), shape=self.shape)

    def trunc(self):
        return SparseNDArray(self.spmatrix.trunc(), shape=self.shape)

    def degrees(self):
        return call_sparse("degrees", self)

    def radians(self):
        return call_sparse("radians", self)

    def clip(self, a_min, a_max):
        try:
            a_min = naked(a_min)
        except TypeError:
            return NotImplemented

        try:
            a_max = naked(a_max)
        except TypeError:
            return NotImplemented

        x = self.spmatrix.maximum(a_min)
        if issparse(x):
            x = x.minimum(a_max)
        elif issparse(a_max):
            x = a_max.minimum(x)
        else:
            xp = get_array_module(x)
            x = xp.minimum(x, a_max)
        if issparse(x):
            return SparseNDArray(x, shape=self.shape)
        return get_array_module(x).asarray(x)

    def iscomplex(self):
        return call_sparse("iscomplex", self)

    def fix(self):
        return call_sparse("fix", self)

    def i0(self):
        xp = get_array_module(self.spmatrix)
        data = xp.i0(self.spmatrix.data).reshape(self.spmatrix.data.shape)
        x = get_sparse_module(self.spmatrix).csr_matrix(
            (data, self.spmatrix.indices, self.spmatrix.indptr), self.spmatrix.shape
        )
        return SparseNDArray(x, shape=self.shape)

    def nan_to_num(self):
        return call_sparse("nan_to_num", self)

    def copysign(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("copysign", self, naked_other)

        if issparse(naked_other):
            naked_other = other.toarray()

        xp = get_array_module(self.spmatrix)
        return xp.copysign(self.toarray(), naked_other)

    def nextafter(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        ret_sparse = False
        if issparse(naked_other):
            ret_sparse = True
            naked_other = other.toarray()

        xp = get_array_module(self.spmatrix)
        xps = get_sparse_module(self.spmatrix)

        x = xp.nextafter(self.toarray(), naked_other)
        if ret_sparse:
            return SparseNDArray(xps.csr_matrix(x), shape=self.shape)
        return x

    def spacing(self):
        if is_cupy(self.spmatrix):
            raise NotImplementedError
        return call_sparse("spacing", self)

    def ldexp(self, other):
        try:
            naked_other = naked(other)
        except TypeError:
            return NotImplemented

        if get_array_module(naked_other).isscalar(naked_other):
            return call_sparse("ldexp", self, naked_other)

        if issparse(naked_other):
            naked_other = other.toarray()

        return SparseNDArray(self.spmatrix.multiply(2**naked_other))

    def frexp(self, **kw):
        xp = get_array_module(self.spmatrix)
        xps = get_sparse_module(self.spmatrix)
        x, y = xp.frexp(self.toarray(), **kw)
        return (
            SparseNDArray(xps.csr_matrix(x), shape=self.shape),
            SparseNDArray(xps.csr_matrix(y), shape=self.shape),
        )

    def modf(self, **kw):
        xp = get_array_module(self.spmatrix)
        xps = get_sparse_module(self.spmatrix)
        x, y = xp.modf(self.toarray(), **kw)
        return (
            SparseNDArray(xps.csr_matrix(x), shape=self.shape),
            SparseNDArray(xps.csr_matrix(y), shape=self.shape),
        )

    def sinc(self):
        return call_sparse("sinc", self)

    def isfinite(self):
        return call_sparse("isfinite", self)

    def isreal(self):
        return call_sparse("isreal", self)

    def digitize(self, bins, right=False):
        return call_sparse("digitize", self, bins=bins, right=right)

    def repeat(self, repeats, axis=None):
        if axis is None:
            raise NotImplementedError

        xp = get_array_module(self.spmatrix)
        xps = get_sparse_module(self.spmatrix)
        r = xp.repeat(self.toarray(), repeats, axis=axis)
        x = xps.csr_matrix(r)
        return SparseNDArray(x, shape=r.shape)

    @staticmethod
    def _expand_val(val, expect_val_size, xp):
        if val.size > expect_val_size:
            val = val[:expect_val_size]
        elif val.size < expect_val_size:
            n_repeat = ceildiv(expect_val_size, val.size)
            val = xp.tile(val, n_repeat)[:expect_val_size]
        return val

    def fill_diagonal(self, val, wrap=False):
        lil_matrix = self.spmatrix.tolil()

        xp = get_array_module(self.spmatrix)
        val = xp.asarray(val)
        if val.ndim > 1:
            val = val.ravel()
        is_tall_matrix = lil_matrix.shape[0] > lil_matrix.shape[1] + 1
        n_rows, n_cols = lil_matrix.shape

        if not wrap or not is_tall_matrix:
            if val.ndim > 0:
                # check if val is long enough
                expect_val_size = min(n_rows, n_cols)
                val = self._expand_val(val, expect_val_size, xp)
            lil_matrix.setdiag(val)
            matrix = lil_matrix
        else:
            block_size = n_cols + 1

            n_block = n_rows // block_size
            n_vals = n_cols * n_block
            if n_rows % block_size > 0:
                # 1 chunk left
                n_block += 1
                n_vals += min(n_rows % block_size, n_cols)

            if val.ndim > 0:
                val = self._expand_val(val, n_vals, xp)

            sub_matrices = []
            for i in range(n_block):
                sub_lil_matrix = lil_matrix[i * block_size : (i + 1) * block_size]
                if val.ndim > 0:
                    sub_val = val[i * n_cols : (i + 1) * n_cols]
                else:
                    sub_val = val
                sub_lil_matrix.setdiag(sub_val)
                sub_matrices.append(sub_lil_matrix)

            xps = get_sparse_module(self.spmatrix)
            matrix = SparseArray(xps.vstack(sub_matrices, format="csr"))

        self.spmatrix = matrix.tocsr()

    def unique(
        self, return_index=False, return_inverse=False, return_counts=False, axis=None
    ):
        if return_inverse or return_index:  # pragma: no cover
            raise NotImplementedError
        if self.ndim == 2 and axis is not None:  # pragma: no cover
            raise NotImplementedError

        xp = get_array_module(self.spmatrix)
        return xp.unique(self.spmatrix.data, return_counts=return_counts)