Merge pull request #55 from willow-ahrens/tweak-nonzero

Make `nonzero` return Tensors
finch-tensor · May 17, 2024 · d676a15 · d676a15
2 parents 4950921 + 2ad4378
commit d676a15
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "finch-tensor"
-version = "0.1.20"
+version = "0.1.21"
 description = ""
 authors = ["Willow Ahrens <willow.marie.ahrens@gmail.com>"]
 readme = "README.md"

diff --git a/src/finch/tensor.py b/src/finch/tensor.py
@@ -723,7 +723,7 @@ def nonzero(x: Tensor, /) -> tuple[np.ndarray, ...]:
     indices = jl.ffindnz(x._obj)[:-1]  # return only indices, skip values
     indices = tuple(np.asarray(i) - 1 for i in indices)
     sort_order = np.lexsort(indices[::-1])  # sort to row-major, C-style order
-    return tuple(i[sort_order] for i in indices)
+    return tuple(Tensor(i[sort_order]) for i in indices)
 
 
 def _reduce(x: Tensor, fn: Callable, axis, dtype=None):

diff --git a/tests/test_sparse.py b/tests/test_sparse.py
@@ -291,7 +291,7 @@ def test_nonzero(order, format_shape):
     actual = finch.nonzero(tns)
     expected = np.nonzero(arr)
     for actual_i, expected_i in zip(actual, expected):
-        assert_equal(actual_i, expected_i)
+        assert_equal(actual_i.todense(), expected_i)
 
 
 @pytest.mark.parametrize("dtype_name", ["int64", "float64", "complex128"])