Numba sparse: Implement basic functionality

pytensor/link/numba/dispatch/basic.py

@@ -14,14 +14,13 @@ from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.link.numba.cache import compile_numba_function_src, hash_from_pickle_dump
-from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
 from pytensor.link.utils import (
     fgraph_to_python,
 )
 from pytensor.scalar.basic import ScalarType
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.random.type import RandomGeneratorType
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import DenseTensorType
 from pytensor.tensor.utils import hash_from_ndarray
 from pytensor.typed_list import TypedListType
 
@@ -112,7 +111,7 @@ def get_numba_type(
         Return Numba scalars for zero dimensional :class:`TensorType`\s.
     """
 
-    if isinstance(pytensor_type, TensorType):
+    if isinstance(pytensor_type, DenseTensorType):
         dtype = pytensor_type.numpy_dtype
         numba_dtype = numba.from_dtype(dtype)
         if force_scalar or (
@@ -125,18 +124,26 @@ def get_numba_type(
         numba_dtype = numba.from_dtype(dtype)
         return numba_dtype
     elif isinstance(pytensor_type, SparseTensorType):
-        dtype = pytensor_type.numpy_dtype
-        numba_dtype = numba.from_dtype(dtype)
+        from pytensor.link.numba.dispatch.sparse.variable import (
+            CSCMatrixType,
+            CSRMatrixType,
+        )
+
+        data_array = numba.types.Array(
+            numba.from_dtype(pytensor_type.numpy_dtype), 1, layout
+        )
+        indices_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
+        indptr_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
         if pytensor_type.format == "csr":
-            return CSRMatrixType(numba_dtype)
+            return CSRMatrixType(data_array, indices_array, indptr_array)
         if pytensor_type.format == "csc":
-            return CSCMatrixType(numba_dtype)
+            return CSCMatrixType(data_array, indices_array, indptr_array)
     elif isinstance(pytensor_type, RandomGeneratorType):
         return numba.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType")
     elif isinstance(pytensor_type, TypedListType):
         return numba.types.List(get_numba_type(pytensor_type.ttype))
-    else:
-        raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
+
+    raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
 
 
 def create_numba_signature(

pytensor/link/numba/dispatch/sparse/__init__.py

+from pytensor.link.numba.dispatch.sparse import variable

pytensor/link/numba/dispatch/sparse.py → pytensor/link/numba/dispatch/sparse/variable.py

 import numpy as np
 import scipy as sp
-import scipy.sparse
 from numba.core import cgutils, types
-from numba.core.imputils import impl_ret_borrowed
+from numba.core.imputils import impl_ret_borrowed, lower_constant
 from numba.extending import (
     NativeValue,
     box,
@@ -27,17 +26,17 @@ class CSMatrixType(types.Type):
     def instance_class(data, indices, indptr, shape):
         raise NotImplementedError()
 
-    def __init__(self, dtype):
-        self.dtype = dtype
-        self.data = types.Array(dtype, 1, "A")
-        self.indices = types.Array(types.int32, 1, "A")
-        self.indptr = types.Array(types.int32, 1, "A")
-        self.shape = types.UniTuple(types.int64, 2)
+    def __init__(self, data_type, indices_type, indptr_type):
+        self._key = (data_type, indices_type, indptr_type)
+        self.data = data_type
+        self.indices = indices_type
+        self.indptr = indptr_type
+        self.shape = types.UniTuple(types.int32, 2)
         super().__init__(self.name)
 
     @property
     def key(self):
-        return (self.name, self.dtype)
+        return self._key
 
 
 make_attribute_wrapper(CSMatrixType, "data", "data")
@@ -63,31 +62,25 @@ class CSCMatrixType(CSMatrixType):
 
 
 @typeof_impl.register(sp.sparse.csc_matrix)
-def typeof_csc_matrix(val, c):
-    data = typeof_impl(val.data, c)
-    return CSCMatrixType(data.dtype)
-
-
 @typeof_impl.register(sp.sparse.csr_matrix)
-def typeof_csr_matrix(val, c):
-    data = typeof_impl(val.data, c)
-    return CSRMatrixType(data.dtype)
-
-
-@register_model(CSRMatrixType)
-class CSRMatrixModel(models.StructModel):
-    def __init__(self, dmm, fe_type):
-        members = [
-            ("data", fe_type.data),
-            ("indices", fe_type.indices),
-            ("indptr", fe_type.indptr),
-            ("shape", fe_type.shape),
-        ]
-        super().__init__(dmm, fe_type, members)
+def typeof_cs_matrix(val, ctx):
+    match val:
+        case sp.sparse.csc_matrix():
+            numba_type = CSCMatrixType
+        case sp.sparse.csr_matrix():
+            numba_type = CSRMatrixType
+        case _:
+            raise ValueError(f"val of type {type(val)} not recognized")
+    return numba_type(
+        typeof_impl(val.data, ctx),
+        typeof_impl(val.indices, ctx),
+        typeof_impl(val.indptr, ctx),
+    )
 
 
 @register_model(CSCMatrixType)
-class CSCMatrixModel(models.StructModel):
+@register_model(CSRMatrixType)
+class CSMatrixModel(models.StructModel):
     def __init__(self, dmm, fe_type):
         members = [
             ("data", fe_type.data),
@@ -98,21 +91,23 @@ class CSCMatrixModel(models.StructModel):
         super().__init__(dmm, fe_type, members)
 
 
-@unbox(CSCMatrixType)
-@unbox(CSRMatrixType)
-def unbox_matrix(typ, obj, c):
+@unbox(CSMatrixType)
+def unbox_cs_matrix(typ, obj, c):
     struct_ptr = cgutils.create_struct_proxy(typ)(c.context, c.builder)
 
+    # Get attributes from python object
     data = c.pyapi.object_getattr_string(obj, "data")
     indices = c.pyapi.object_getattr_string(obj, "indices")
     indptr = c.pyapi.object_getattr_string(obj, "indptr")
     shape = c.pyapi.object_getattr_string(obj, "shape")
 
+    # Unbox them into llvm struct
     struct_ptr.data = c.unbox(typ.data, data).value
     struct_ptr.indices = c.unbox(typ.indices, indices).value
     struct_ptr.indptr = c.unbox(typ.indptr, indptr).value
     struct_ptr.shape = c.unbox(typ.shape, shape).value
 
+    # Decref created attributes
     c.pyapi.decref(data)
     c.pyapi.decref(indices)
     c.pyapi.decref(indptr)
@@ -120,15 +115,13 @@ def unbox_matrix(typ, obj, c):
 
     is_error_ptr = cgutils.alloca_once_value(c.builder, cgutils.false_bit)
     is_error = c.builder.load(is_error_ptr)
-
     res = NativeValue(struct_ptr._getvalue(), is_error=is_error)
 
     return res
 
 
-@box(CSCMatrixType)
-@box(CSRMatrixType)
-def box_matrix(typ, val, c):
+@box(CSMatrixType)
+def box_cs_matrix(typ, val, c):
     struct_ptr = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
 
     data_obj = c.box(typ.data, struct_ptr.data)
@@ -136,11 +129,7 @@ def box_matrix(typ, val, c):
     indptr_obj = c.box(typ.indptr, struct_ptr.indptr)
     shape_obj = c.box(typ.shape, struct_ptr.shape)
 
-    c.pyapi.incref(data_obj)
-    c.pyapi.incref(indices_obj)
-    c.pyapi.incref(indptr_obj)
-    c.pyapi.incref(shape_obj)
-
+    # Call scipy.sparse.cs[c|r]_matrix
     cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(typ.instance_class))
     obj = c.pyapi.call_function_objargs(
         cls_obj, (data_obj, indices_obj, indptr_obj, shape_obj)
@@ -154,53 +143,125 @@ def box_matrix(typ, val, c):
     return obj
 
 
-@overload(np.shape)
-def overload_sparse_shape(x):
-    if isinstance(x, CSMatrixType):
-        return lambda x: x.shape
+def _intrinsic_cs_codegen(context, builder, sig, args):
+    matrix_type = sig.return_type
+    struct = cgutils.create_struct_proxy(matrix_type)(context, builder)
+    data, indices, indptr, shape = args
+    struct.data = data
+    struct.indices = indices
+    struct.indptr = indptr
+    struct.shape = shape
+    return impl_ret_borrowed(
+        context,
+        builder,
+        matrix_type,
+        struct._getvalue(),
+    )
 
 
-@overload_attribute(CSMatrixType, "ndim")
-def overload_sparse_ndim(inst):
-    if not isinstance(inst, CSMatrixType):
-        return
+@intrinsic
+def csr_matrix_from_components(typingctx, data, indices, indptr, shape):
+    sig = CSRMatrixType(data, indices, indptr)(data, indices, indptr, shape)
+    return sig, _intrinsic_cs_codegen
 
-    def ndim(inst):
-        return 2
 
-    return ndim
+@intrinsic
+def csc_matrix_from_components(typingctx, data, indices, indptr, shape):
+    sig = CSCMatrixType(data, indices, indptr)(data, indices, indptr, shape)
+    return sig, _intrinsic_cs_codegen
 
 
-@intrinsic
-def _sparse_copy(typingctx, inst, data, indices, indptr, shape):
-    def _construct(context, builder, sig, args):
-        typ = sig.return_type
-        struct = cgutils.create_struct_proxy(typ)(context, builder)
-        _, data, indices, indptr, shape = args
-        struct.data = data
-        struct.indices = indices
-        struct.indptr = indptr
-        struct.shape = shape
-        return impl_ret_borrowed(
-            context,
-            builder,
-            sig.return_type,
-            struct._getvalue(),
-        )
+@lower_constant(CSRMatrixType)
+@lower_constant(CSCMatrixType)
+def cs_matrix_constant(context, builder, ty, pyval):
+    data_const = context.make_constant_array(builder, ty.data, pyval.data)
+    indices_const = context.make_constant_array(builder, ty.indices, pyval.indices)
+    indptr_const = context.make_constant_array(builder, ty.indptr, pyval.indptr)
+    shape = context.get_constant_generic(builder, ty.shape, pyval.shape)
+    args = (data_const, indices_const, indptr_const, shape)
 
-    sig = inst(inst, inst.data, inst.indices, inst.indptr, inst.shape)
+    sig = ty(*args)
+    return _intrinsic_cs_codegen(context, builder, sig, args)
 
-    return sig, _construct
 
+@overload(sp.sparse.csr_matrix)
+def overload_csr_matrix(arg1, shape, dtype=None):
+    if not isinstance(arg1, types.BaseAnonymousTuple) or len(arg1) != 3:
+        return None
+    if isinstance(shape, types.NoneType):
+        return None
+
+    def impl(arg1, shape, dtype=None):
+        data, indices, indptr = arg1
+        int32_shape = (types.int32(shape[0]), types.int32(shape[1]))
+        return csr_matrix_from_components(data, indices, indptr, int32_shape)
+
+    return impl
 
-@overload_method(CSMatrixType, "copy")
-def overload_sparse_copy(inst):
-    if not isinstance(inst, CSMatrixType):
-        return
 
-    def copy(inst):
-        return _sparse_copy(
-            inst, inst.data.copy(), inst.indices.copy(), inst.indptr.copy(), inst.shape
+@overload(sp.sparse.csc_matrix)
+def overload_csc_matrix(arg1, shape, dtype=None):
+    if not isinstance(arg1, types.BaseAnonymousTuple) or len(arg1) != 3:
+        return None
+    if isinstance(shape, types.NoneType):
+        return None
+
+    def impl(arg1, shape, dtype=None):
+        data, indices, indptr = arg1
+        int32_shape = (types.int32(shape[0]), types.int32(shape[1]))
+        return csc_matrix_from_components(data, indices, indptr, int32_shape)
+
+    return impl
+
+
+@overload(np.shape)
+def overload_sparse_shape(matrix):
+    if isinstance(matrix, CSMatrixType):
+        return lambda matrix: matrix.shape
+
+
+@overload_attribute(CSMatrixType, "ndim")
+def overload_sparse_ndim(matrix):
+    return lambda matrix: 2
+
+
+@overload_method(CSMatrixType, "copy")
+def overload_sparse_copy(matrix):
+    match matrix:
+        case CSRMatrixType():
+            builder = csr_matrix_from_components
+        case CSCMatrixType():
+            builder = csc_matrix_from_components
+        case _:
+            return
+
+    def copy(matrix):
+        return builder(
+            matrix.data.copy(),
+            matrix.indices.copy(),
+            matrix.indptr.copy(),
+            matrix.shape,
         )
 
     return copy
+
+
+@overload_method(CSMatrixType, "astype")
+def overload_sparse_astype(matrix, dtype):
+    match matrix:
+        case CSRMatrixType():
+            builder = csr_matrix_from_components
+        case CSCMatrixType():
+            builder = csc_matrix_from_components
+        case _:
+            return
+
+    def astype(matrix, dtype):
+        return builder(
+            matrix.data.astype(dtype),
+            matrix.indices.copy(),
+            matrix.indptr.copy(),
+            matrix.shape,
+        )
+
+    return astype

tests/link/numba/sparse/test_basic.py

+from functools import partial
+from sys import getrefcount
+
+import numpy as np
+import pytest
+import scipy
+import scipy as sp
+
+import pytensor.sparse as ps
+import pytensor.tensor as pt
+from pytensor.graph import Apply, Op
+from pytensor.sparse.variable import SparseConstant
+from pytensor.tensor.type import DenseTensorType
+
+
+numba = pytest.importorskip("numba")
+
+
+# Make sure the Numba customizations are loaded
+import pytensor.link.numba.dispatch.sparse  # noqa: F401
+from pytensor import config
+from pytensor.sparse import SparseTensorType
+from tests.link.numba.test_basic import compare_numba_and_py
+
+
+pytestmark = pytest.mark.filterwarnings("error")
+
+
+def sparse_assert_fn(a, b):
+    a_is_sparse = sp.sparse.issparse(a)
+    assert a_is_sparse == sp.sparse.issparse(b)
+    if a_is_sparse:
+        assert a.format == b.format
+        assert a.dtype == b.dtype
+        assert a.shape == b.shape
+        np.testing.assert_allclose(a.data, b.data, strict=True)
+        np.testing.assert_allclose(a.indices, b.indices, strict=True)
+        np.testing.assert_allclose(a.indptr, b.indptr, strict=True)
+    else:
+        np.testing.assert_allclose(a, b, strict=True)
+
+
+compare_numba_and_py_sparse = partial(compare_numba_and_py, assert_fn=sparse_assert_fn)
+
+
+def test_sparse_boxing():
+    @numba.njit
+    def boxing_fn(x, y):
+        return x, y, y.data.sum()
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+    y_val = sp.sparse.csc_matrix(np.eye(101))
+
+    res_x_val, res_y_val, res_y_sum = boxing_fn(x_val, y_val)
+
+    assert np.array_equal(res_x_val.data, x_val.data)
+    assert np.array_equal(res_x_val.indices, x_val.indices)
+    assert np.array_equal(res_x_val.indptr, x_val.indptr)
+    assert res_x_val.shape == x_val.shape
+
+    assert np.array_equal(res_y_val.data, y_val.data)
+    assert np.array_equal(res_y_val.indices, y_val.indices)
+    assert np.array_equal(res_y_val.indptr, y_val.indptr)
+    assert res_y_val.shape == y_val.shape
+
+    np.testing.assert_allclose(res_y_sum, y_val.sum())
+
+
+def test_sparse_creation_refcount():
+    @numba.njit
+    def create_csr_matrix(data, indices, ind_ptr):
+        return scipy.sparse.csr_matrix((data, indices, ind_ptr), shape=(5, 5))
+
+    x = scipy.sparse.random(5, 5, density=0.5, format="csr")
+
+    x_data = x.data
+    x_indptr = x.indptr
+    assert getrefcount(x_data) == 3
+    assert getrefcount(x_indptr) == 3
+
+    for i in range(5):
+        a = create_csr_matrix(x.data, x.indices, x.indptr)
+
+    # a.data is a view of the underlying data under x.data, but doesn't reference it directly
+    assert getrefcount(x_data) == 3
+    # x_indptr is reused directly
+    assert getrefcount(x_indptr) == 4
+
+    del a
+    assert getrefcount(x_data) == 3
+    assert getrefcount(x_indptr) == 3
+
+
+def test_sparse_passthrough_refcount():
+    @numba.njit
+    def identity(a):
+        return a
+
+    x = scipy.sparse.random(5, 5, density=0.5, format="csr")
+
+    x_data = x.data
+    assert getrefcount(x_data) == 3
+
+    for i in range(5):
+        identity(x)
+
+    assert getrefcount(x_data) == 3
+
+
+def test_sparse_shape():
+    @numba.njit
+    def test_fn(x):
+        return np.shape(x)
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+
+    res = test_fn(x_val)
+
+    assert res == (100, 100)
+
+
+def test_sparse_ndim():
+    @numba.njit
+    def test_fn(x):
+        return x.ndim
+
+    x_val = sp.sparse.csr_matrix(np.eye(100))
+
+    res = test_fn(x_val)
+
+    assert res == 2
+
+
+def test_sparse_copy():
+    @numba.njit
+    def test_fn(x):
+        return x.copy()
+
+    x = sp.sparse.csr_matrix(np.eye(100))
+
+    y = test_fn(x)
+    assert y is not x
+    for attr in ("data", "indices", "indptr"):
+        y_data = getattr(y, attr)
+        x_data = getattr(x, attr)
+        assert y_data is not x_data
+        assert not np.shares_memory(y_data, x_data)
+        assert (y_data == x_data).all()
+
+
+@pytest.mark.parametrize(
+    "func", [sp.sparse.csr_matrix, sp.sparse.csc_matrix], ids=["csr", "csc"]
+)
+def test_sparse_constructor(func):
+    @numba.njit
+    def csr_matrix_constructor(data, indices, indptr):
+        return func((data, indices, indptr), shape=(3, 3))
+
+    inp = sp.sparse.random(3, 3, density=0.5, format="csr")
+
+    # Test with pure scipy constructor
+    out = func((inp.data, inp.indices, inp.indptr), copy=False)
+    # Scipy does a useless slice on data and indices to trim away useless zeros
+    # which means these attributes are views of the original arrays.
+    assert out.data is not inp.data
+    assert not out.data.flags.owndata
+
+    assert out.indices is not inp.indices
+    assert not out.indices.flags.owndata
+
+    assert out.indptr is inp.indptr
+
+    # Test numba impl
+    out_pt = csr_matrix_constructor(inp.data, inp.indices, inp.indptr)
+    # Should work the same as Scipy's constructor, because it's ultimately used
+    assert type(out_pt) is type(out)
+
+    assert out_pt.data is not inp.data
+    assert not out_pt.data.flags.owndata
+    assert np.shares_memory(out_pt.data, inp.data)
+    assert (out_pt.data == inp.data).all()
+
+    assert out_pt.indices is not inp.indices
+    assert not out_pt.indices.flags.owndata
+    assert np.shares_memory(out_pt.indices, inp.indices)
+    assert (out_pt.indices == inp.indices).all()
+
+    assert out_pt.indptr is inp.indptr
+
+
+@pytest.mark.parametrize("cache", [True, False])
+@pytest.mark.parametrize("format", ["csr", "csc"])
+def test_sparse_constant(format, cache):
+    x = sp.sparse.random(3, 3, density=0.5, format=format, random_state=166)
+    x = ps.as_sparse(x)
+    assert isinstance(x, SparseConstant)
+    assert x.type.format == format
+    y = pt.vector("y", shape=(3,))
+    out = x * y
+
+    y_test = np.array([np.pi, np.e, np.euler_gamma])
+    with config.change_flags(numba__cache=cache):
+        with pytest.warns(
+            UserWarning,
+            match=r"Numba will use object mode to run SparseDenseVectorMultiply's perform method",
+        ):
+            compare_numba_and_py_sparse(
+                [y],
+                [out],
+                [y_test],
+                eval_obj_mode=False,
+            )
+
+
+@pytest.mark.parametrize("format", ["csc", "csr"])
+@pytest.mark.parametrize("dense_out", [True, False])
+def test_sparse_objmode(format, dense_out):
+    class SparseTestOp(Op):
+        def make_node(self, x):
+            out = x.type.clone(shape=(1, x.type.shape[-1]))()
+            if dense_out:
+                out = out.todense().type()
+            return Apply(self, [x], [out])
+
+        def perform(self, node, inputs, output_storage):
+            [x] = inputs
+            [out] = output_storage
+            out[0] = x[0]
+            if dense_out:
+                out[0] = out[0].todense()
+
+    x = ps.matrix(format, dtype=config.floatX, shape=(5, 5), name="x")
+
+    out = SparseTestOp()(x)
+    assert out.type.shape == (1, 5)
+    assert isinstance(out.type, DenseTensorType if dense_out else SparseTensorType)
+
+    x_val = sp.sparse.random(5, 5, density=0.25, dtype=config.floatX, format=format)
+
+    with pytest.warns(
+        UserWarning,
+        match="Numba will use object mode to run SparseTestOp's perform method",
+    ):
+        compare_numba_and_py_sparse([x], out, [x_val])
+
+
+@pytest.mark.parametrize("format", ["csr", "csc"])
+def test_simple_graph(format):
+    x = ps.matrix(format, name="x", shape=(3, 3))
+    y = pt.tensor("y", shape=(3,))
+    z = ps.math.sin(x * y)
+
+    rng = np.random.default_rng((155, format == "csr"))
+    x_test = sp.sparse.random(3, 3, density=0.5, format=format, random_state=rng)
+    y_test = rng.normal(size=(3,))
+
+    with pytest.warns(
+        UserWarning, match=r"Numba will use object mode to run .* perform method"
+    ):
+        compare_numba_and_py_sparse(
+            [x, y],
+            z,
+            [x_test, y_test],
+        )

tests/link/numba/test_sparse.py

-from functools import partial
-
-import numpy as np
-import pytest
-import scipy as sp
-
-
-numba = pytest.importorskip("numba")
-
-
-# Make sure the Numba customizations are loaded
-import pytensor.link.numba.dispatch.sparse  # noqa: F401
-from pytensor import config
-from pytensor.sparse import Dot, SparseTensorType
-from tests.link.numba.test_basic import compare_numba_and_py
-
-
-pytestmark = pytest.mark.filterwarnings("error")
-
-
-def sparse_assert_fn(a, b):
-    a_is_sparse = sp.sparse.issparse(a)
-    assert a_is_sparse == sp.sparse.issparse(b)
-    if a_is_sparse:
-        assert a.format == b.format
-        assert a.dtype == b.dtype
-        assert a.shape == b.shape
-        np.testing.assert_allclose(a.data, b.data, strict=True)
-        np.testing.assert_allclose(a.indices, b.indices, strict=True)
-        np.testing.assert_allclose(a.indptr, b.indptr, strict=True)
-    else:
-        np.testing.assert_allclose(a, b, strict=True)
-
-
-compare_numba_and_py_sparse = partial(compare_numba_and_py, assert_fn=sparse_assert_fn)
-
-
-def test_sparse_unboxing():
-    @numba.njit
-    def test_unboxing(x, y):
-        return x.shape, y.shape
-
-    x_val = sp.sparse.csr_matrix(np.eye(100))
-    y_val = sp.sparse.csc_matrix(np.eye(101))
-
-    res = test_unboxing(x_val, y_val)
-
-    assert res == (x_val.shape, y_val.shape)
-
-
-def test_sparse_boxing():
-    @numba.njit
-    def test_boxing(x, y):
-        return x, y
-
-    x_val = sp.sparse.csr_matrix(np.eye(100))
-    y_val = sp.sparse.csc_matrix(np.eye(101))
-
-    res_x_val, res_y_val = test_boxing(x_val, y_val)
-
-    assert np.array_equal(res_x_val.data, x_val.data)
-    assert np.array_equal(res_x_val.indices, x_val.indices)
-    assert np.array_equal(res_x_val.indptr, x_val.indptr)
-    assert res_x_val.shape == x_val.shape
-
-    assert np.array_equal(res_y_val.data, y_val.data)
-    assert np.array_equal(res_y_val.indices, y_val.indices)
-    assert np.array_equal(res_y_val.indptr, y_val.indptr)
-    assert res_y_val.shape == y_val.shape
-
-
-def test_sparse_shape():
-    @numba.njit
-    def test_fn(x):
-        return np.shape(x)
-
-    x_val = sp.sparse.csr_matrix(np.eye(100))
-
-    res = test_fn(x_val)
-
-    assert res == (100, 100)
-
-
-def test_sparse_ndim():
-    @numba.njit
-    def test_fn(x):
-        return x.ndim
-
-    x_val = sp.sparse.csr_matrix(np.eye(100))
-
-    res = test_fn(x_val)
-
-    assert res == 2
-
-
-def test_sparse_copy():
-    @numba.njit
-    def test_fn(x):
-        y = x.copy()
-        return (
-            y is not x and np.all(x.data == y.data) and np.all(x.indices == y.indices)
-        )
-
-    x_val = sp.sparse.csr_matrix(np.eye(100))
-
-    assert test_fn(x_val)
-
-
-def test_sparse_objmode():
-    x = SparseTensorType("csc", dtype=config.floatX)()
-    y = SparseTensorType("csc", dtype=config.floatX)()
-
-    out = Dot()(x, y)
-
-    x_val = sp.sparse.random(2, 2, density=0.25, dtype=config.floatX, format="csc")
-    y_val = sp.sparse.random(2, 2, density=0.25, dtype=config.floatX, format="csc")
-
-    with pytest.warns(
-        UserWarning,
-        match="Numba will use object mode to run SparseDot's perform method",
-    ):
-        compare_numba_and_py_sparse(
-            [x, y],
-            out,
-            [x_val, y_val],
-        )