Merge pull request #657 from pydata/update-finch-backend

Update Finch backend

Author: mtsokol

sparse/finch_backend/__init__.py

@@ -3,11 +3,68 @@
 except ModuleNotFoundError as e:
     raise ImportError("Finch not installed. Run `pip install sparse[finch]` to enable Finch backend") from e
 
-from finch import Tensor, astype, permute_dims
+from finch import (
+    add,
+    astype,
+    bool,
+    compiled,
+    complex64,
+    complex128,
+    compute,
+    divide,
+    float16,
+    float32,
+    float64,
+    int8,
+    int16,
+    int32,
+    int64,
+    int_,
+    lazy,
+    multiply,
+    negative,
+    permute_dims,
+    positive,
+    prod,
+    random,
+    subtract,
+    sum,
+    uint,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+)
 
-__all__ = ["Tensor", "astype", "permute_dims"]
-
-
-class COO:
-    def from_numpy(self):
-        raise NotImplementedError
+__all__ = [
+    "add",
+    "astype",
+    "bool",
+    "compiled",
+    "complex64",
+    "complex128",
+    "compute",
+    "divide",
+    "float16",
+    "float32",
+    "float64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "int_",
+    "lazy",
+    "multiply",
+    "negative",
+    "permute_dims",
+    "positive",
+    "prod",
+    "random",
+    "subtract",
+    "sum",
+    "uint",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+]

sparse/tests/test_backends.py

@@ -1,32 +1,56 @@
 import sparse
 
-import pytest
-
 import numpy as np
 import scipy.sparse as sp
+from numpy.testing import assert_equal
 
 
 def test_backend_contex_manager(backend):
+    rng = np.random.default_rng(0)
+    x = sparse.random((100, 10, 100), density=0.01, random_state=rng)
+    y = sparse.random((100, 10, 100), density=0.01, random_state=rng)
+
     if backend == sparse.BackendType.Finch:
-        with pytest.raises(NotImplementedError):
-            sparse.COO.from_numpy(np.eye(5))
+        import finch
+
+        def storage():
+            return finch.Storage(finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0.0)))), order="C")
+
+        x = x.to_device(storage())
+        y = y.to_device(storage())
     else:
-        sparse.COO.from_numpy(np.eye(5))
+        x.asformat("gcxs")
+        y.asformat("gcxs")
+
+    z = x + y
+    result = sparse.sum(z)
+    assert result.shape == ()
 
 
 def test_finch_backend():
     np_eye = np.eye(5)
     sp_arr = sp.csr_matrix(np_eye)
 
     with sparse.Backend(backend=sparse.BackendType.Finch):
-        finch_dense = sparse.Tensor(np_eye)
+        import finch
+
+        finch_dense = finch.Tensor(np_eye)
 
         assert np.shares_memory(finch_dense.todense(), np_eye)
 
-        finch_arr = sparse.Tensor(sp_arr)
+        finch_arr = finch.Tensor(sp_arr)
 
-        np.testing.assert_equal(finch_arr.todense(), np_eye)
+        assert_equal(finch_arr.todense(), np_eye)
 
         transposed = sparse.permute_dims(finch_arr, (1, 0))
 
-        np.testing.assert_equal(transposed.todense(), np_eye.T)
+        assert_equal(transposed.todense(), np_eye.T)
+
+        @sparse.compiled
+        def my_fun(tns1, tns2):
+            tmp = sparse.add(tns1, tns2)
+            return sparse.sum(tmp, axis=0)
+
+        result = my_fun(finch_dense, finch_arr)
+
+        assert_equal(result.todense(), np.sum(2 * np_eye, axis=0))