diff --git a/pytest.ini b/pytest.ini index 4dda7d2..a351985 100644 --- a/pytest.ini +++ b/pytest.ini @@ -5,3 +5,4 @@ filterwarnings = testpaths = finch junit_family=xunit2 +xfail_strict=true diff --git a/src/finch/compiled.py b/src/finch/compiled.py index bae0d20..e6617cb 100644 --- a/src/finch/compiled.py +++ b/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 diff --git a/src/finch/julia.py b/src/finch/julia.py index 201af3d..e844938 100644 --- a/src/finch/julia.py +++ b/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() diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 4ba0f7a..2aaf926 100644 --- a/src/finch/tensor.py +++ b/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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 6c95066..3d82c7c 100644 --- a/tests/test_indexing.py +++ b/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, :] diff --git a/tests/test_ops.py b/tests/test_ops.py index 560a3f9..a634209 100644 --- a/tests/test_ops.py +++ b/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)