diff --git a/Project.toml b/Project.toml index b4dee05..d29259f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "BlockTensorKit" uuid = "5f87ffc2-9cf1-4a46-8172-465d160bd8cd" -version = "0.3.16" +version = "0.3.17" authors = ["Lukas Devos and contributors"] [deps] diff --git a/src/tensors/abstractblocktensor/abstractarray.jl b/src/tensors/abstractblocktensor/abstractarray.jl index 186720d..46ce188 100644 --- a/src/tensors/abstractblocktensor/abstractarray.jl +++ b/src/tensors/abstractblocktensor/abstractarray.jl @@ -100,20 +100,33 @@ end t::AbstractBlockTensorMap, indices::Vararg{Any, M} ) where {M} M == numind(t) || return _slice_getindex_single(t, indices...) - inds = Base.to_indices(t, indices) - inds isa NTuple{M, Int} && return parent(t)[inds...] + return _slice_getindex_full(t, Base.to_indices(t, indices)) +end + +# single block vs. sliced block tensor is chosen by dispatch, so the return type is inferrable +@propagate_inbounds _slice_getindex_full( + t::AbstractBlockTensorMap, inds::NTuple{M, Int} +) where {M} = parent(t)[inds...] +@propagate_inbounds function _slice_getindex_full(t::AbstractBlockTensorMap, inds::Tuple) tdst = similar(t, space(eachspace(t)[inds...])) length(tdst) == 0 && return tdst return _copyslice!(tdst, t, inds) end +# a mask may consume several dimensions, which `to_indices` cannot infer through the fillers below +@propagate_inbounds _slice_getindex_single( + t::AbstractBlockTensorMap, index::AbstractArray{Bool} +) = _slice_getindex_single(t, Base.to_index(t, index)) + # a single index is only supported when it selects a single nontrivial dimension @noinline function _slice_getindex_single( t::AbstractBlockTensorMap, indices::Vararg{Any, M} ) where {M} space(eachspace(t)[indices...]) # errors as before if unsupported d = something(findfirst(>(1), size(t)), 1) - return _slice_getindex(t, ntuple(i -> i == d ? only(indices) : 1, numind(t))...) + # `1:1` rather than `1` keeps `Int` out of the tuple, so the dispatch above stays static + inds = ntuple(i -> i == d ? only(indices) : (1:1), numind(t)) + return _slice_getindex_full(t, Base.to_indices(t, inds)) end # TODO: check if this fallback is fair diff --git a/test/abstracttensor/indexing.jl b/test/abstracttensor/indexing.jl index ef23edf..ceee0df 100644 --- a/test/abstracttensor/indexing.jl +++ b/test/abstracttensor/indexing.jl @@ -92,6 +92,45 @@ for (label, t1) in (("dense", rand(V ← one(V))), ("sparse", sprand(V ← one(V end end +# single-index slicing of a multi-index block tensor takes a different path than the case above +W = SumSpace(ℂ^2) +for (label, ts) in (("dense", rand(W ⊗ V ← W)), ("sparse", sprand(W ⊗ V ← W, 0.8))) + @testset "$label single-index slicing" begin + @test size(ts) == (1, 3, 1) + mid = 2:(length(ts) - 1) + # a single index must stay as inferrable as the equivalent full-rank slice + @test @inferred(ts[mid]) == ts[:, mid, :] + @test space(ts[mid]) == space(ts[:, mid, :]) + @test @inferred(ts[1]) isa TensorMap + @test @inferred(ts[:]) == ts + @test size(@inferred(ts[[1, 3]])) == (1, 2, 1) + @test ts[[1, 3]][1, 2, 1] == ts[1, 3, 1] + mask = [true, false, true] + @test @inferred(ts[mask]) == ts[:, mask, :] + end +end + +# a widened return type propagates into everything built from a slice, so keep every form concrete +infers_concretely(f, args...) = + isconcretetype(Core.Compiler.return_type(f, Tuple{map(typeof, args)...})) + +@testset "indexing inference ($label)" for (label, t) in ( + ("dense", rand(W ⊗ V ← W)), ("sparse", sprand(W ⊗ V ← W, 0.8)), + ("dense square", rand(V ⊗ V ← V)), ("sparse square", sprand(V ⊗ V ← V, 0.5)), + ("dense vector", rand(V ← one(V))), ("sparse vector", sprand(V ← one(V), 0.8)), + ) + for ind in (1, 2:3, 1:2:3, :, [1, 3], [true, false, true]) + @test infers_concretely(getindex, t, ind) + end + ndims(t) == 3 || continue + for inds in ( + (1, 2, 1), (:, :, :), (1, 2:3, 1), (1, [1, 3], 1), (1, [true, false, true], 1), + (1:1, 1:2:3, 1:1), ([1, 1, 3], :, :), + ) + @test infers_concretely(getindex, t, inds...) + end +end + # the parent array has its own slicing implementation @testset "parent array slicing" begin st = sprand(V ⊗ V ⊗ V, 0.5)