adapt tests

Author: flying-sheep

src/fast_array_utils/_plugins/numba_sparse.py

@@ -62,24 +62,18 @@ def instance_class(
     ) -> CSBase:
         return cls.cls((data, indices, indptr), shape, copy=False)
 
-    def __init__(
-        self,
-        ndim: int,
-        *,
-        dtype_data: nbtypes.Type,
-        dtype_indices: nbtypes.Type,
-        dtype_indptr: nbtypes.Type,
-    ) -> None:
-        self.dtype = nbtypes.DType(dtype_data)
-        self.data = nbtypes.Array(dtype_data, 1, "A")
-        self.indices = nbtypes.Array(dtype_indices, 1, "A")
-        self.indptr = nbtypes.Array(dtype_indptr, 1, "A")
+    def __init__(self, ndim: int, *, dtype: nbtypes.Type, dtype_ind: nbtypes.Type) -> None:
+        self.dtype = nbtypes.DType(dtype)
+        self.dtype_ind = nbtypes.DType(dtype_ind)
+        self.data = nbtypes.Array(dtype, 1, "A")
+        self.indices = nbtypes.Array(dtype_ind, 1, "A")
+        self.indptr = nbtypes.Array(dtype_ind, 1, "A")
         self.shape = nbtypes.UniTuple(nbtypes.intp, ndim)
         super().__init__(self.name)
 
     @property
     def key(self) -> tuple[str | nbtypes.Type, ...]:
-        return (self.name, self.dtype, self.indices.dtype, self.indptr.dtype)
+        return (self.name, self.dtype, self.dtype_ind)
 
 
 # make data model attributes available in numba functions
@@ -88,13 +82,15 @@ def key(self) -> tuple[str | nbtypes.Type, ...]:
 
 
 def make_typeof_fn(typ: type[CSType]) -> Callable[[CSBase, _TypeofContext], CSType]:
+    """Create a `typeof` function that maps a scipy matrix/array type to a numba `Type`."""
+
     def typeof(val: CSBase, c: _TypeofContext) -> CSType:
+        if val.indptr.dtype != val.indices.dtype:
+            msg = "indptr and indices must have the same dtype"
+            raise TypeError(msg)
         data = cast("nbtypes.Array", typeof_impl(val.data, c))
-        indices = cast("nbtypes.Array", typeof_impl(val.indices, c))
         indptr = cast("nbtypes.Array", typeof_impl(val.indptr, c))
-        return typ(
-            val.ndim, dtype_data=data.dtype, dtype_indices=indices.dtype, dtype_indptr=indptr.dtype
-        )
+        return typ(val.ndim, dtype=data.dtype, dtype_ind=indptr.dtype)
 
     return typeof
 
@@ -106,6 +102,11 @@ def typeof(val: CSBase, c: _TypeofContext) -> CSType:
 
 
 class CSModel(_Base):
+    """Numba data model for compressed sparse matrices.
+
+    This is the class that is used by numba to lower the array types.
+    """
+
     def __init__(self, dmm: DataModelManager, fe_type: CSType) -> None:
         members = [
             ("data", fe_type.data),
@@ -116,6 +117,7 @@ def __init__(self, dmm: DataModelManager, fe_type: CSType) -> None:
         super().__init__(dmm, fe_type, members)
 
 
+# create all the actual types and data models
 CLASSES: Sequence[type[CSBase]] = [
     sparse.csr_matrix,
     sparse.csc_matrix,

tests/test_numpy_scipy_sparse.py

@@ -47,22 +47,24 @@ def copy_mat(mat: CSBase) -> CSBase:
 
 
 @pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk | Flags.Gpu)
-@pytest.mark.parametrize("dtype_indptr", [np.int32, np.int64], ids=["p=32", "p=64"])
-@pytest.mark.parametrize("dtype_index", [np.int32, np.int64], ids=["i=32", "i=64"])
+@pytest.mark.parametrize("dtype_ind", [np.int32, np.int64], ids=["i=32", "i=64"])
 @pytest.mark.parametrize("dtype_data", [np.int64, np.float64], ids=["d=i64", "d=f64"])
 def test_copy(
     array_type: ArrayType[CSBase, None],
     dtype_data: type[np.int64 | np.float64],
-    dtype_index: type[np.int32 | np.int64],
-    dtype_indptr: type[np.int32 | np.int64],
+    dtype_ind: type[np.int32 | np.int64],
 ) -> None:
     mat = array_type.random((10, 10), density=0.1, dtype=dtype_data)
-    mat.indices = mat.indices.astype(dtype_index)
-    mat.indptr = mat.indptr.astype(dtype_indptr)
+    mat.indptr = mat.indptr.astype(dtype_ind)
+    mat.indices = mat.indices.astype(dtype_ind)
+
     copied = copy_mat(mat)
-    assert mat.data is not copied.data
-    assert mat.indices is not copied.indices
-    assert mat.indptr is not copied.indptr
+
+    # check that the copied arrays point to different memory locations
+    assert mat.data.ctypes.data != copied.data.ctypes.data
+    assert mat.indices.ctypes.data != copied.indices.ctypes.data
+    assert mat.indptr.ctypes.data != copied.indptr.ctypes.data
+    # check that the array contents and dtypes are the same
     assert mat.shape == copied.shape
     np.testing.assert_equal(mat.toarray(), copied.toarray(), strict=True)
     np.testing.assert_equal(mat.data, copied.data, strict=True)