From a04b56abd94476b02f8bd945cc9e4a9e9890ac13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= <mat646@gmail.com>
Date: Tue, 26 Mar 2024 17:00:40 +0100
Subject: [PATCH 1/2] Update Finch backend

---
 sparse/finch_backend/__init__.py | 71 ++++++++++++++++++++++++++++----
 sparse/tests/test_backends.py    | 40 ++++++++++++++----
 2 files changed, 95 insertions(+), 16 deletions(-)

diff --git a/sparse/finch_backend/__init__.py b/sparse/finch_backend/__init__.py
index 42f2931b..609fca2b 100644
--- a/sparse/finch_backend/__init__.py
+++ b/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",
+]
diff --git a/sparse/tests/test_backends.py b/sparse/tests/test_backends.py
index c27f3b27..0512b8a4 100644
--- a/sparse/tests/test_backends.py
+++ b/sparse/tests/test_backends.py
@@ -1,17 +1,29 @@
 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():
@@ -19,14 +31,24 @@ def test_finch_backend():
     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))

From d93a6cfff991ee0685016a59cf4b2cf67f772c43 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 27 Mar 2024 10:43:31 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 sparse/tests/test_backends.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/sparse/tests/test_backends.py b/sparse/tests/test_backends.py
index 0512b8a4..6d56e174 100644
--- a/sparse/tests/test_backends.py
+++ b/sparse/tests/test_backends.py
@@ -12,6 +12,7 @@ def test_backend_contex_manager(backend):
 
     if backend == sparse.BackendType.Finch:
         import finch
+
         def storage():
             return finch.Storage(finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0.0)))), order="C")
 
@@ -32,6 +33,7 @@ def test_finch_backend():
 
     with sparse.Backend(backend=sparse.BackendType.Finch):
         import finch
+
         finch_dense = finch.Tensor(np_eye)
 
         assert np.shares_memory(finch_dense.todense(), np_eye)