diff --git a/src/tensors/blocktensor.jl b/src/tensors/blocktensor.jl index 58146a9..d14cce2 100644 --- a/src/tensors/blocktensor.jl +++ b/src/tensors/blocktensor.jl @@ -99,6 +99,18 @@ function BlockTensorMap( return BlockTensorMap(data, codom ← dom) end +# Slicing +# ------- +# a dense slice holds the blocks of `t` itself, so its own are never allocated +function _slice_getindex_full(t::BlockTensorMap, inds::Tuple) + V = space(eachspace(t)[inds...]) + return typeof(t)(parent(t)[map(_keepdim, inds)...], V) +end +@propagate_inbounds _slice_getindex_full(t::BlockTensorMap, inds::NTuple{M, Int}) where {M} = + parent(t)[inds...] +_keepdim(i::Integer) = i:i +_keepdim(i) = i + # AbstractBlockTensorMap -> BlockTensorMap function BlockTensorMap(t::AbstractBlockTensorMap) t isa BlockTensorMap && return t # TODO: should this copy? diff --git a/test/abstracttensor/indexing.jl b/test/abstracttensor/indexing.jl index ceee0df..dae5cfc 100644 --- a/test/abstracttensor/indexing.jl +++ b/test/abstracttensor/indexing.jl @@ -236,6 +236,18 @@ end @test @allocated(tb[inds...]) < 2_000_000 end +# a dense slice shares the blocks of its parent, without allocating any of its own +@testset "dense slicing shares blocks" begin + D = 64 + t = rand(SumSpace(ℂ^16) ⊗ SumSpace(fill(ℂ^2, D)...) ← SumSpace(ℂ^16)) + inds = (1:1, 2:(D - 1), 1:1) + s = t[inds...] + @test space(s) == space(eachspace(t)[inds...]) + @test all(i -> s[1, i, 1] === t[1, i + 1, 1], 1:(D - 2)) + @test t[1, 3, :] isa BlockTensorMap && t[1, 3, :][1, 1, 1] === t[1, 3, 1] + @test @allocated(t[inds...]) < sum(b -> sizeof(b.data), parent(t)) ÷ 4 +end + # guard against assignment work proportional to the assigned region @testset "assignment scales with nnz" begin D = 48