Get CSR round-trip working.

Author: hameerabbasi

sparse/numba_backend/_compressed/compressed.py

@@ -11,7 +11,6 @@
 from .._coo.core import COO
 from .._sparse_array import SparseArray
 from .._utils import (
-    _zero_of_dtype,
     can_store,
     check_compressed_axes,
     check_fill_value,
@@ -175,13 +174,9 @@ def __init__(
         if self.data.ndim != 1:
             raise ValueError("data must be a scalar or 1-dimensional.")
 
-        self.shape = shape
-
-        if fill_value is None:
-            fill_value = _zero_of_dtype(self.data.dtype)
+        SparseArray.__init__(self, shape=shape, fill_value=fill_value)
 
         self._compressed_axes = tuple(compressed_axes) if isinstance(compressed_axes, Iterable) else None
-        self.fill_value = self.data.dtype.type(fill_value)
 
         if prune:
             self._prune()
@@ -417,7 +412,7 @@ def tocoo(self):
                 fill_value=self.fill_value,
             )
         uncompressed = uncompress_dimension(self.indptr)
-        coords = np.vstack((uncompressed, self.indices))
+        coords = np.stack((uncompressed, self.indices))
         order = np.argsort(self._axis_order)
         return (
             COO(
@@ -884,7 +879,7 @@ def __binsparse__(self) -> tuple[dict, list[np.ndarray]]:
             "original_source": f"`sparse`, version {__version__}",
         }
 
-        return descriptor, [self.indices, self.indptr, self.data]
+        return descriptor, [self.indptr, self.indices, self.data]
 
 
 class CSR(_Compressed2d):

sparse/numba_backend/_io.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from ._common import _check_device
-from ._compressed import GCXS
+from ._compressed import CSC, CSR, GCXS
 from ._coo.core import COO
 from ._sparse_array import SparseArray
 
@@ -145,7 +145,6 @@ def from_binsparse(arr, /, *, device=None, copy: bool | None = None) -> SparseAr
 
     format = desc["format"]
     format_err_str = f"Unsupported format: `{format!r}`."
-    invalid_dtype_str = "Invalid dtype: `{dtype!s}`, expected `{expected!s}`."
 
     if isinstance(format, str):
         match format:
@@ -180,15 +179,15 @@ def from_binsparse(arr, /, *, device=None, copy: bool | None = None) -> SparseAr
             case _:
                 raise RuntimeError(format_err_str)
 
-    format = desc["format"]
+    format = desc["format"]["custom"]
+    rank = 0
+    level = format
+    while "level" in level:
+        if "rank" not in level:
+            level["rank"] = 1
+        rank += level["rank"]
+        level = level["level"]
     if "transpose" not in format:
-        rank = 0
-        level = format
-        while "level" in level:
-            if "rank" not in level:
-                level["rank"] = 1
-            rank += level["rank"]
-
         format["transpose"] = list(range(rank))
 
     match desc:
@@ -225,25 +224,8 @@ def from_binsparse(arr, /, *, device=None, copy: bool | None = None) -> SparseAr
             coord_arr: np.ndarray = np.from_dlpack(arrs[1])
             value_arr: np.ndarray = np.from_dlpack(arrs[2])
 
-            if str(coord_arr.dtype) != coords_dtype:
-                raise BufferError(
-                    invalid_dtype_str.format(
-                        dtype=str(coord_arr.dtype),
-                        expected=coords_dtype,
-                    )
-                )
-
-            if value_dtype.startswith("complex[float") and value_dtype.endswith("]"):
-                complex_bits = 2 * int(value_arr[len("complex[float") : -len("]")])
-                value_dtype: str = f"complex{complex_bits}"
-
-            if str(value_arr.dtype) != value_dtype:
-                raise BufferError(
-                    invalid_dtype_str.format(
-                        dtype=str(coord_arr.dtype),
-                        expected=coords_dtype,
-                    )
-                )
+            _check_binsparse_dt(coord_arr, coords_dtype)
+            _check_binsparse_dt(value_arr, value_dtype)
 
             return COO(
                 coord_arr[:, start:end],
@@ -254,5 +236,68 @@ def from_binsparse(arr, /, *, device=None, copy: bool | None = None) -> SparseAr
                 prune=False,
                 idx_dtype=coord_arr.dtype,
             )
+        case {
+            "format": {
+                "custom": {
+                    "transpose": transpose,
+                    "level": {
+                        "level_desc": "dense",
+                        "rank": 1,
+                        "level": {
+                            "level_desc": "sparse",
+                            "rank": 1,
+                            "level": {
+                                "level_desc": "element",
+                            },
+                        },
+                    },
+                },
+            },
+            "shape": shape,
+            "number_of_stored_values": nnz,
+            "data_types": {
+                "pointers_to_1": ptr_dtype,
+                "indices_1": crd_dtype,
+                "values": val_dtype,
+            },
+            **_kwargs,
+        }:
+            crd_arr = np.from_dlpack(arrs[0])
+            _check_binsparse_dt(crd_arr, crd_dtype)
+            ptr_arr = np.from_dlpack(arrs[1])
+            _check_binsparse_dt(ptr_arr, ptr_dtype)
+            val_arr = np.from_dlpack(arrs[2])
+            _check_binsparse_dt(val_arr, val_dtype)
+
+            match transpose:
+                case [0, 1]:
+                    sparse_type = CSR
+                case [1, 0]:
+                    sparse_type = CSC
+                case _:
+                    raise RuntimeError(format_err_str)
+
+            return sparse_type((val_arr, ptr_arr, crd_arr), shape=shape)
         case _:
+            print(desc)
             raise RuntimeError(format_err_str)
+
+
+def _convert_binsparse_dtype(dt: str) -> np.dtype:
+    if dt.startswith("complex[float") and dt.endswith("]"):
+        complex_bits = 2 * int(dt[len("complex[float") : -len("]")])
+        dt: str = f"complex{complex_bits}"
+
+    return np.dtype(dt)
+
+
+def _check_binsparse_dt(arr: np.ndarray, dt: str) -> None:
+    invalid_dtype_str = "Invalid dtype: `{dtype!s}`, expected `{expected!s}`."
+    dt = _convert_binsparse_dtype(dt)
+    if dt != arr.dtype:
+        raise BufferError(
+            invalid_dtype_str.format(
+                dtype=arr.dtype,
+                expected=dt,
+            )
+        )

sparse/numba_backend/tests/test_io.py

@@ -30,7 +30,9 @@ def test_load_wrong_format_exception(tmp_path):
         load_npz(filename)
 
 
-@pytest.mark.parametrize("format", ["coo", "csr", "csc"])
+@pytest.mark.parametrize(
+    "format", ["coo", "csr", pytest.param("csc", marks=pytest.mark.xfail(reason="`CSC<>COO` round-trip broken"))]
+)
 def test_round_trip_binsparse(format: str) -> None:
     x = sparse.random((20, 30), density=0.25, format=format)
     y = sparse.from_binsparse(x)