Add SciPy conversions.

pixi.toml

@@ -27,8 +27,8 @@ mkdocs-jupyter = "*"
 
 [feature.tests.tasks]
 test = "pytest --pyargs sparse -n auto"
-test-mlir = { cmd = "pytest --pyargs sparse/mlir_backend -vvv", env = { SPARSE_BACKEND = "MLIR" } }
-test-finch = { cmd = "pytest --pyargs sparse/tests -n auto -vvv", env = { SPARSE_BACKEND = "Finch", PYTHONFAULTHANDLER = "${HOME}/faulthandler.log" }, depends-on = ["precompile"] }
+test-mlir = { cmd = "pytest --pyargs sparse/mlir_backend -v", env = { SPARSE_BACKEND = "MLIR" } }
+test-finch = { cmd = "pytest --pyargs sparse/tests -n auto -v", env = { SPARSE_BACKEND = "Finch", PYTHONFAULTHANDLER = "${HOME}/faulthandler.log" }, depends-on = ["precompile"] }
 
 [feature.tests.dependencies]
 pytest = ">=3.5"

sparse/mlir_backend/__init__.py

@@ -7,24 +7,20 @@
         "to enable MLIR backend."
     ) from e
 
-from ._constructors import (
-    PackedArgumentTuple,
-    asarray,
-)
+from ._common import PackedArgumentTuple
+from ._conversions import asarray, to_numpy, to_scipy
 from ._dtypes import (
     asdtype,
 )
 from ._ops import (
     add,
-    broadcast_to,
-    reshape,
 )
 
 __all__ = [
     "add",
-    "broadcast_to",
     "asarray",
     "asdtype",
-    "reshape",
     "PackedArgumentTuple",
+    "to_numpy",
+    "to_scipy",
 ]

sparse/mlir_backend/_array.py

@@ -0,0 +1,45 @@
+import dataclasses
+
+from ._common import _hold_ref, numpy_to_ranked_memref, ranked_memref_to_numpy
+from ._levels import StorageFormat
+
+
+class Array:
+    def __init__(self, *, storage, shape: tuple[int, ...]) -> None:
+        storage_rank = storage.get_storage_format().rank
+        if len(shape) != storage_rank:
+            raise ValueError(f"Mismatched rank, `{storage_rank=}`, `{shape=}`")
+
+        self._storage = storage
+        self._shape = shape
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return self._shape
+
+    @property
+    def ndim(self) -> int:
+        return len(self.shape)
+
+    @property
+    def dtype(self):
+        return self._storage.get_storage_format().dtype
+
+    def _get_storage_format(self) -> StorageFormat:
+        return self._storage.get_storage_format()
+
+    def _get_mlir_type(self):
+        return self._get_storage_format().get_mlir_type(shape=self.shape)
+
+    def _to_module_arg(self):
+        return self._storage.to_module_arg()
+
+    def copy(self):
+        storage_format: StorageFormat = dataclasses.replace(self._get_storage_format(), owns_memory=False)
+
+        fields = self._storage.get__fields_()
+        arrs = [ranked_memref_to_numpy(f).copy() for f in fields]
+        memrefs = [numpy_to_ranked_memref(arr) for arr in arrs]
+        arr = Array(storage=storage_format.get_ctypes_type()(*memrefs), shape=self.shape)
+        for carr in arrs:
+            _hold_ref(arr, carr)

sparse/mlir_backend/_common.py

@@ -1,11 +1,9 @@
-import abc
 import ctypes
 import functools
 import weakref
 from dataclasses import dataclass
 
 import mlir.runtime as rt
-from mlir import ir
 
 import numpy as np
 
@@ -17,8 +15,12 @@ def fn_cache(f, maxsize: int | None = None):
     return functools.wraps(f)(functools.lru_cache(maxsize=maxsize)(f))
 
 
-@fn_cache
 def get_nd_memref_descr(rank: int, dtype: type[DType]) -> ctypes.Structure:
+    return _get_nd_memref_descr(int(rank), asdtype(dtype))
+
+
+@fn_cache
+def _get_nd_memref_descr(rank: int, dtype: type[DType]) -> ctypes.Structure:
     return rt.make_nd_memref_descriptor(rank, dtype.to_ctype())
 
 
@@ -43,12 +45,6 @@ def free_memref(obj: ctypes.Structure) -> None:
     libc.free(ctypes.cast(obj.allocated, ctypes.c_void_p))
 
 
-class MlirType(abc.ABC):
-    @classmethod
-    @abc.abstractmethod
-    def get_mlir_type(cls) -> ir.Type: ...
-
-
 @dataclass
 class PackedArgumentTuple:
     contents: tuple
@@ -67,13 +63,13 @@ def _hold_self_ref_in_ret(fn):
     @functools.wraps(fn)
     def wrapped(self, *a, **kw):
         ret = fn(self, *a, **kw)
-        _take_owneship(ret, self)
+        _hold_ref(ret, self)
         return ret
 
     return wrapped
 
 
-def _take_owneship(owner, obj):
+def _hold_ref(owner, obj):
     ptr = ctypes.py_object(obj)
     ctypes.pythonapi.Py_IncRef(ptr)