Adjust public API, tests, and format definition language.

Author: hameerabbasi

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 -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"] }
+test-mlir = { cmd = "pytest --pyargs sparse/mlir_backend -v" }
+test-finch = { cmd = "pytest --pyargs sparse/tests -n auto -v", depends-on = ["precompile"] }
 
 [feature.tests.dependencies]
 pytest = ">=3.5"
@@ -51,10 +51,19 @@ precompile = "python -c 'import finch'"
 scipy = ">=0.19"
 finch-tensor = ">=0.1.31"
 
+[feature.finch.activation.env]
+SPARSE_BACKEND = "Finch"
+
+[feature.finch.target.osx-arm64.activation.env]
+PYTHONFAULTHANDLER = "${HOME}/faulthandler.log"
+
 [feature.mlir.dependencies]
 scipy = ">=0.19"
 mlir-python-bindings = "19.*"
 
+[feature.mlir.activation.env]
+SPARSE_BACKEND = "MLIR"
+
 [environments]
 tests = ["tests", "extras"]
 docs = ["docs", "extras"]

sparse/mlir_backend/__init__.py

@@ -1,16 +1,17 @@
 try:
     import mlir  # noqa: F401
+
+    del mlir
 except ModuleNotFoundError as e:
     raise ImportError(
         "MLIR Python bindings not installed. Run "
         "`conda install conda-forge::mlir-python-bindings` "
         "to enable MLIR backend."
     ) from e
 
-from ._common import PackedArgumentTuple
-from ._conversions import asarray, to_numpy, to_scipy
+from . import levels
+from ._conversions import asarray, from_constituent_arrays, to_numpy, to_scipy
 from ._dtypes import asdtype
-from ._levels import StorageFormat
 from ._ops import add
 
-__all__ = ["add", "asarray", "asdtype", "to_numpy", "to_scipy", "PackedArgumentTuple", "StorageFormat"]
+__all__ = ["add", "asarray", "asdtype", "to_numpy", "to_scipy", "levels", "from_constituent_arrays"]

sparse/mlir_backend/_array.py

@@ -1,7 +1,9 @@
 import dataclasses
 
+import numpy as np
+
 from ._common import _hold_ref, numpy_to_ranked_memref, ranked_memref_to_numpy
-from ._levels import StorageFormat
+from .levels import StorageFormat
 
 
 class Array:
@@ -43,3 +45,8 @@ def copy(self):
         arr = Array(storage=storage_format._get_ctypes_type()(*memrefs), shape=self.shape)
         for carr in arrs:
             _hold_ref(arr, carr)
+
+        return arr
+
+    def get_constituent_arrays(self) -> tuple[np.ndarray, ...]:
+        return self._storage.get_constituent_arrays()

sparse/mlir_backend/_common.py

@@ -1,7 +1,6 @@
 import ctypes
 import functools
 import weakref
-from dataclasses import dataclass
 
 import mlir.runtime as rt
 
@@ -45,20 +44,6 @@ def free_memref(obj: ctypes.Structure) -> None:
     libc.free(ctypes.cast(obj.allocated, ctypes.c_void_p))
 
 
-@dataclass
-class PackedArgumentTuple:
-    contents: tuple
-
-    def __getitem__(self, index):
-        return self.contents[index]
-
-    def __iter__(self):
-        yield from self.contents
-
-    def __len__(self):
-        return len(self.contents)
-
-
 def _hold_self_ref_in_ret(fn):
     @functools.wraps(fn)
     def wrapped(self, *a, **kw):

sparse/mlir_backend/_conversions.py

@@ -1,10 +1,11 @@
+import dataclasses
 import functools
 
 import numpy as np
 
 from ._array import Array
 from ._common import _hold_ref, numpy_to_ranked_memref, ranked_memref_to_numpy
-from ._levels import Level, LevelFormat, LevelProperties, StorageFormat, get_storage_format
+from .levels import Level, LevelFormat, LevelProperties, StorageFormat, get_storage_format
 
 try:
     import scipy.sparse as sps
@@ -178,3 +179,9 @@ def asarray(arr, copy: bool | None = None) -> Array:
         return arr
 
     return _from_numpy(np.asarray(arr, copy=copy), copy=None)
+
+
+def from_constituent_arrays(*, format: StorageFormat, arrays: tuple[np.ndarray, ...], shape: tuple[int, ...]) -> Array:
+    storage_format: StorageFormat = dataclasses.replace(format, owns_memory=False)
+    storage = storage_format._get_ctypes_type().from_constituent_arrays(arrays)
+    return Array(storage=storage, shape=shape)

sparse/mlir_backend/_levels.py renamed to sparse/mlir_backend/levels.py

@@ -11,7 +11,6 @@
 import numpy as np
 
 from ._common import (
-    PackedArgumentTuple,
     _hold_ref,
     fn_cache,
     free_memref,
@@ -24,6 +23,8 @@
 
 _CAMEL_TO_SNAKE = [re.compile("(.)([A-Z][a-z]+)"), re.compile("([a-z0-9])([A-Z])")]
 
+__all__ = ["LevelProperties", "LevelFormat", "StorageFormat", "Level", "get_storage_format"]
+
 
 def _camel_to_snake(name: str) -> str:
     for exp in _CAMEL_TO_SNAKE:
@@ -134,8 +135,11 @@ def to_module_arg(self) -> list:
             def get__fields_(self) -> list:
                 return [getattr(self, field[0]) for field in self._fields_]
 
-            def to_constituent_arrays(self) -> PackedArgumentTuple:
-                return PackedArgumentTuple(tuple(ranked_memref_to_numpy(field) for field in self.get__fields_()))
+            def get_constituent_arrays(self) -> tuple[np.ndarray, ...]:
+                arrays = tuple(ranked_memref_to_numpy(field) for field in self.get__fields_())
+                for arr in arrays:
+                    _hold_ref(arr, self)
+                return arrays
 
             def get_storage_format(self) -> StorageFormat:
                 return storage_format
@@ -155,13 +159,6 @@ def __del__(self) -> None:
 
         return Storage
 
-    def from_constituent_arrays(self, arrs: list[np.ndarray], shape: tuple[int, ...]):
-        from ._array import Array
-
-        storage = self._get_ctypes_type().from_constituent_arrays(arrs)
-
-        return Array(storage=storage, shape=shape)
-
 
 def get_storage_format(
     *,

sparse/mlir_backend/tests/test_simple.py

@@ -178,65 +178,95 @@ def test_add(rng, dtype):
 
 @parametrize_dtypes
 def test_csf_format(dtype):
+    format = sparse.levels.get_storage_format(
+        levels=(
+            sparse.levels.Level(sparse.levels.LevelFormat.Dense),
+            sparse.levels.Level(sparse.levels.LevelFormat.Compressed),
+            sparse.levels.Level(sparse.levels.LevelFormat.Compressed),
+        ),
+        order="C",
+        pos_width=64,
+        crd_width=64,
+        dtype=sparse.asdtype(dtype),
+        owns_memory=False,
+    )
+
     SHAPE = (2, 2, 4)
     pos_1, crd_1, pos_2, crd_2, data = get_exampe_csf_arrays(dtype)
-    csf = [pos_1, crd_1, pos_2, crd_2, data]
+    constituent_arrays = (pos_1, crd_1, pos_2, crd_2, data)
 
-    csf_tensor = sparse.asarray(csf, shape=SHAPE, dtype=sparse.asdtype(dtype), format="csf")
-    result = csf_tensor.to_scipy_sparse()
-    for actual, expected in zip(result, csf, strict=False):
+    csf_array = sparse.from_constituent_arrays(format=format, arrays=constituent_arrays, shape=SHAPE)
+    result_arrays = csf_array.get_constituent_arrays()
+    for actual, expected in zip(result_arrays, constituent_arrays, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
-    res_tensor = sparse.add(csf_tensor, csf_tensor).to_scipy_sparse()
-    csf_2 = [pos_1, crd_1, pos_2, crd_2, data * 2]
-    for actual, expected in zip(res_tensor, csf_2, strict=False):
+    res_arrays = sparse.add(csf_array, csf_array).get_constituent_arrays()
+    expected_arrays = (pos_1, crd_1, pos_2, crd_2, data * 2)
+    for actual, expected in zip(res_arrays, expected_arrays, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
 
 @parametrize_dtypes
 def test_coo_3d_format(dtype):
+    format = sparse.levels.get_storage_format(
+        levels=(
+            sparse.levels.Level(sparse.levels.LevelFormat.Compressed, sparse.levels.LevelProperties.NonUnique),
+            sparse.levels.Level(sparse.levels.LevelFormat.Singleton, sparse.levels.LevelProperties.NonUnique),
+            sparse.levels.Level(sparse.levels.LevelFormat.Singleton, sparse.levels.LevelProperties.NonUnique),
+        ),
+        order="C",
+        pos_width=64,
+        crd_width=64,
+        dtype=sparse.asdtype(dtype),
+        owns_memory=False,
+    )
+
     SHAPE = (2, 2, 4)
     pos = np.array([0, 7])
     crd = np.array([[0, 1, 0, 0, 1, 1, 0], [1, 3, 1, 0, 0, 1, 0], [3, 1, 1, 0, 1, 1, 1]])
     data = np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype)
-    coo = [pos, crd, data]
+    carrs = (pos, crd, data)
 
-    coo_tensor = sparse.asarray(coo, shape=SHAPE, dtype=sparse.asdtype(dtype), format="coo")
-    result = coo_tensor.to_scipy_sparse()
-    for actual, expected in zip(result, coo, strict=False):
+    coo_array = sparse.from_constituent_arrays(format=format, arrays=carrs, shape=SHAPE)
+    result = coo_array.get_constituent_arrays()
+    for actual, expected in zip(result, carrs, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
     # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
-    # res_tensor = sparse.add(coo_tensor, coo_tensor).to_scipy_sparse()
-    # coo_2 = [pos, crd, data * 2]
-    # for actual, expected in zip(res_tensor, coo_2, strict=False):
+    # res_arrays = sparse.add(coo_array, coo_array).get_constituent_arrays()
+    # res_expected = (pos, crd, data * 2)
+    # for actual, expected in zip(res_arrays, res_expected, strict=False):
     #     np.testing.assert_array_equal(actual, expected)
 
 
 @parametrize_dtypes
 def test_sparse_vector_format(dtype):
+    format = sparse.levels.get_storage_format(
+        levels=(sparse.levels.Level(sparse.levels.LevelFormat.Compressed),),
+        order="C",
+        pos_width=64,
+        crd_width=64,
+        dtype=sparse.asdtype(dtype),
+        owns_memory=False,
+    )
+
     SHAPE = (10,)
     pos = np.array([0, 6])
     crd = np.array([0, 1, 2, 6, 8, 9])
     data = np.array([1, 2, 3, 4, 5, 6], dtype=dtype)
-    sparse_vector = [pos, crd, data]
+    carrs = (pos, crd, data)
 
-    sv_tensor = sparse.asarray(
-        sparse_vector,
-        shape=SHAPE,
-        dtype=sparse.asdtype(dtype),
-        format="sparse_vector",
-    )
-    result = sv_tensor.to_scipy_sparse()
-    for actual, expected in zip(result, sparse_vector, strict=False):
+    sv_array = sparse.from_constituent_arrays(format=format, arrays=carrs, shape=SHAPE)
+    result = sv_array.get_constituent_arrays()
+    for actual, expected in zip(result, carrs, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
-    res_tensor = sparse.add(sv_tensor, sv_tensor).to_scipy_sparse()
-    sparse_vector_2 = [pos, crd, data * 2]
-    for actual, expected in zip(res_tensor, sparse_vector_2, strict=False):
+    res_arrs = sparse.add(sv_array, sv_array).get_constituent_arrays()
+    sv2_expected = (pos, crd, data * 2)
+    for actual, expected in zip(res_arrs, sv2_expected, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
     dense = np.array([1, 2, 3, 0, 0, 0, 4, 0, 5, 6], dtype=dtype)
-    dense_tensor = sparse.asarray(dense)
-    res_tensor = sparse.add(dense_tensor, sv_tensor).to_scipy_sparse()
-    np.testing.assert_array_equal(res_tensor, dense * 2)
+    dense_array = sparse.asarray(dense)
+    res = sparse.to_numpy(sparse.add(dense_array, sv_array))
+    np.testing.assert_array_equal(res, dense * 2)