From 8b5c7ffc9ea83489419a1e0d8af1e2259980b75e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 08:37:35 -0400 Subject: [PATCH] Slice dense block tensors without allocating blocks Slicing a `BlockTensorMap` allocated every block of the destination through `similar`, only for `_copyslice!` to replace each of them by the block of the source it selects. Build the slice directly from the sliced parent array instead: the result shares the blocks of the source, as before, without allocating new ones. Co-Authored-By: Claude Opus 5.5 --- src/tensors/blocktensor.jl | 12 ++++++++++++ test/abstracttensor/indexing.jl | 12 ++++++++++++ 2 files changed, 24 insertions(+) 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