Merge pull request #39 from willow-ahrens/new_release

API: Add verbose mode and lazy indexing

pytest.ini

@@ -5,3 +5,4 @@ filterwarnings =
 testpaths =
     finch
 junit_family=xunit2
+xfail_strict=true

src/finch/compiled.py

@@ -28,7 +28,7 @@ def lazy(tensor: Tensor):
     return tensor
 
 
-def compute(tensor: Tensor):
+def compute(tensor: Tensor, *, verbose: bool = False):
     if not tensor.is_computed():
-        return Tensor(jl.Finch.compute(tensor._obj))
+        return Tensor(jl.Finch.compute(tensor._obj, verbose=verbose))
     return tensor

src/finch/julia.py

@@ -1,6 +1,6 @@
 import juliapkg
 
-_FINCH_VERSION = "0.6.23"
+_FINCH_VERSION = "0.6.24"
 _FINCH_HASH = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce"
 
 deps = juliapkg.deps.load_cur_deps()

src/finch/tensor.py

@@ -197,9 +197,15 @@ def _elemwise_op(self, op: str, other: Optional["Tensor"] = None) -> "Tensor":
     def __getitem__(self, key):
         if not isinstance(key, tuple):
             key = (key,)
-        key = _expand_ellipsis(key, self.shape)
-        key = _add_missing_dims(key, self.shape)
-        key = _add_plus_one(key, self.shape)
+
+        if None in key:
+            # lazy indexing mode
+            key = _process_lazy_indexing(key)
+        else:
+            # standard indexing mode
+            key = _expand_ellipsis(key, self.shape)
+            key = _add_missing_dims(key, self.shape)
+            key = _add_plus_one(key, self.shape)
 
         result = self._obj[key]
         if jl.isa(result, jl.Finch.SwizzleArray) or jl.isa(result, jl.Finch.LazyTensor):
@@ -877,7 +883,7 @@ def _add_plus_one(key: tuple, shape: tuple[int, ...]) -> tuple:
             idx = normalize_axis_tuple(idx, size)
             new_key += (jl.Vector([i + 1 for i in idx]),)
         elif idx is None:
-            raise IndexError("'None' in the index key isn't supported")
+            raise IndexError("`None` in the index is supported only in lazy indexing")
         else:
             new_key += (idx,)
     return new_key
@@ -911,7 +917,20 @@ def _expand_ellipsis(key: tuple, shape: tuple[int, ...]) -> tuple:
         key = new_key
     return key
 
+
 def _add_missing_dims(key: tuple, shape: tuple[int, ...]) -> tuple:
     for i in range(len(key), len(shape)):
         key = key + (jl.range(start=1, stop=shape[i]),)
     return key
+
+
+def _process_lazy_indexing(key: tuple) -> tuple:
+    new_key = ()
+    for idx in key:
+        if idx == slice(None):
+            new_key += (jl.Colon(),)
+        elif idx is None:
+            new_key += (jl.nothing,)
+        else:
+            raise ValueError(f"Invalid lazy index member: {idx}")
+    return new_key

tests/test_indexing.py

@@ -1,12 +1,17 @@
 import numpy as np
 from numpy.testing import assert_equal
 import pytest
+import juliacall as jc
 
 import finch
 
 
 @pytest.mark.parametrize(
-    "index", [..., 40, (32,), slice(None), slice(30, 60, 3), -10, slice(None, -10, -2)]
+    "index",
+    [
+        ..., 40, (32,), slice(None), slice(30, 60, 3), -10,
+        slice(None, -10, -2), (None, slice(None)),
+    ]
 )
 @pytest.mark.parametrize("order", ["C", "F"])
 def test_indexing_1d(arr1d, index, order):
@@ -23,7 +28,11 @@ def test_indexing_1d(arr1d, index, order):
 
 
 @pytest.mark.parametrize(
-    "index", [..., 0, (2,), (2, 3), slice(None), (..., slice(0, 4, 2)), (-1, slice(-1, None, -1))]
+    "index",
+    [
+        ..., 0, (2,), (2, 3), slice(None), (..., slice(0, 4, 2)),
+        (-1, slice(-1, None, -1)), (None, slice(None), slice(None)),
+    ]
 )
 @pytest.mark.parametrize("order", ["C", "F"])
 def test_indexing_2d(arr2d, index, order):
@@ -49,7 +58,11 @@ def test_indexing_2d(arr2d, index, order):
         (slice(None, None, -1), slice(None, None, -1), slice(None, None, -1)),
         (slice(None, -1, 1), slice(-1, None, -1), slice(4, 1, -1)),
         (-1, 0, 0), (0, -1, -2), ([1, 2], 0, slice(3, None, -1)),
-        (0, slice(1, 0, -1), 0),
+        (0, slice(1, 0, -1), 0), (slice(None), None, slice(None), slice(None)),
+        pytest.param(
+            (slice(None), slice(None), slice(None), None),
+            marks=pytest.mark.xfail(reason="https://github.com/willow-ahrens/Finch.jl/issues/528"),
+        ),
     ]
 )
 @pytest.mark.parametrize(
@@ -73,8 +86,14 @@ def test_indexing_3d(arr3d, index, levels_descr, order):
     assert_equal(actual, expected)
 
 
-def test_index_none(arr3d):
+def test_invalid_index_none(arr3d):
     arr_finch = finch.Tensor(arr3d)
 
-    with pytest.raises(IndexError, match="'None' in the index key isn't supported"):
+    with pytest.raises(ValueError, match="Invalid lazy index member: Ellipsis"):
         arr_finch[..., None]
+
+    with pytest.raises(
+        jc.JuliaError,
+        match="Cannot index a lazy tensor with more or fewer `:` dims than it had original dims.",
+    ):
+        arr_finch[None, :]

tests/test_ops.py

@@ -19,6 +19,7 @@ def test_eager(arr3d):
     assert_equal(result.todense(), np.multiply(arr3d, arr2d))
 
 
+@pytest.mark.xfail(reason="https://github.com/willow-ahrens/Finch.jl/issues/524")
 def test_lazy_mode(arr3d):
     A_finch = finch.Tensor(arr3d)
     B_finch = finch.Tensor(arr2d)