Conversation
Codecov Report❌ Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
`pairs(A::AbstractVector)` uses linear indices, so `nonzero_pairs` yielded `Int` keys for the parent of a one-dimensional dense `BlockTensorMap`, inconsistent with `nonzero_keys`. Consumers comparing against `CartesianIndex` then silently matched nothing, which made `t[:]` and `t[1:2]` return uninitialized blocks. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Slicing `getindex` scanned the whole source index grid once per stored block,
making it `O(nnz * prod(length.(indices)))` and materializing the slice region
as an `Array{CartesianIndex}`. Invert the index maps per dimension instead, so
each stored block finds its destinations in `O(1)`:
n nnz before after speedup
16 16 0.0045 ms 0.0015 ms 3x
64 64 0.1452 ms 0.0036 ms 40x
256 256 8.9455 ms 0.0107 ms 836x
512 512 71.1836 ms 0.0188 ms 3786x
(banded sparse block tensor, sliced to every other channel; the dense
`BlockTensorMap` variant goes from 7.94 ms to 0.087 ms at n = 64.)
`SparseTensorArray` had a second copy of the same logic for slicing the parent
array; both now share `_copyslice!`.
Along the way this fixes four bugs in the old index handling:
- Repeated indices dropped blocks instead of duplicating them, and for a dense
destination the dropped slot was returned as uninitialized memory. They now
duplicate, as they do for `AbstractArray`.
- Single-index slicing of one-dimensional block tensors (`t[:]`, `t[1:2]`)
returned uninitialized blocks.
- `Integer` indices other than `Int` threw: `t[UInt(1), 1:2, 1]` a
`BoundsError`, `t[Int8(1), Int8(1), Int8(1)]` a `MethodError`.
- Logical masks threw `CanonicalIndexError` on the parent array path,
`parent(t)[[true, false, true], :, :]`.
The `Vararg{Strided.SliceIndex}` methods are kept: they disambiguate against
TensorKit's own `getindex`/`setindex!` for `AbstractTensorMap`, so removing
them turns `t[:, :, :]` into a `MethodError`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The three `SparseTensorArray` `copyto!` methods all walked the assigned region rather than the stored blocks, so `t[inds...] = v` and `cat` cost `O(prod(length.(inds)))` regardless of how little was stored — the mirror image of the `getindex` problem in #67. Reuse the inverted index maps instead: shape region nnz before after speedup 2d n=64 3_844 64 0.150 ms 0.015 ms 10x 2d n=256 64_516 256 2.267 ms 0.036 ms 63x 2d n=512 260_100 512 13.752 ms 0.051 ms 269x 4d n=24 234_256 24 9.459 ms 0.009 ms 1040x 4d n=48 4_477_456 48 140.997 ms 0.013 ms 10600x narrow 4 512 0.005 ms 0.004 ms 1x `copyto!(dst, view(parent(t), inds...))` improves likewise, 14.6 ms to 0.041 ms for a 260_100-block region. The narrow row matters: a small assignment into a large sparse array must not become proportional to the destination's stored blocks. `_deletemissing!` therefore sweeps whichever of the two is smaller, making it `O(nnz(src) + min(region, nnz(dst)))` — never worse than before. `copyto!(t::SparseTensorArray, ::SubArray)` needs no sweep at all: the destination spans exactly the viewed region, so everything not copied is dropped. Also fix the `nonzero_*` accessors for `SparseTensorArray`, which fell through to the `AbstractArray` fallbacks and reported every index rather than the stored ones. `nonzero_keys` and `nonzero_length` were always wrong this way; `nonzero_pairs` became wrong when it started going through `pairs(IndexCartesian(), A)`, which bypasses `Base.pairs(::SparseTensorArray)`. The visible effects were that slicing the parent array densified its result and that `copyto!` over regions stored explicitly-zero blocks. Verified semantics-preserving: the stored-key sets after slice assignment, `copyto!` from a view, region-to-region `copyto!` and `cat` are identical to the previous implementation once the accessors are fixed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #67.
1. Slicing
getindexis no longer quadraticgetindexscanned the whole source index grid once per stored block, making itO(nnz × prod(length.(indices)))and materializing the slice region as anArray{CartesianIndex}. The index maps are now inverted per dimension, so each stored blockfinds its destinations in
O(1).Banded sparse block tensor (
nnz = n, one block per row), sliced to every other channel —t[1:2:n, 1:1, 1:1, 1:2:n], minimum of 20 runs:Dense
BlockTensorMap, same slice:The "before" column reproduces the measurements in #67. The old column grows ~8× per doubling
of
n, the new one linearly.SparseTensorArrayhad a second copy of the same logic for slicing the parent array; both callsites now share
_copyslice!.Bugs fixed along the way
Normalizing indices with
Base.to_indicesbefore inverting them cures four separate problemsin the old index handling. All four were reproduced on
mainbefore the change.Repeated indices dropped blocks — as uninitialized memory for a dense destination.
findfirstreturns only the first match, so the second copy was never written, and a densedestination comes from
BlockTensorMap{TT}(undef, P):This is the one intended behaviour change: repeated indices now duplicate the selected tensors,
as they do for
AbstractArray. Notendims(t[:, 1:2, [1, 1]]) == 3already appears indocs/src/blocktensors.md, so the documented example was returning garbage.Single-index slicing of one-dimensional block tensors returned uninitialized blocks.
pairs(A::AbstractVector)uses linear indices, sononzero_pairsyieldedIntkeys for a 1-ddense parent while the scan compared against
CartesianIndex{1}— nothing ever matched. Fixedat the root in
nonzero_pairs(first commit), which also makes it consistent withnonzero_keys.Integerindices other thanIntthrew.ind isa Intmissed the otherBitIntegers:Logical masks threw on the parent-array path.
Base.LogicalIndex <: AbstractVector{Int}, soit reached a
findfirstthat needsgetindex, whichLogicalIndexdoes not define:2. Slice assignment no longer scales with the region
The mirror-image problem, on the
setindex!side. All threeSparseTensorArraycopyto!methods walked the assigned region rather than the stored blocks, so
t[inds...] = vandcatcost
O(prod(length.(inds)))no matter how little was stored. Same inverted index maps, reused.Banded sparse tensors,
t[inds...] = srcover the interior:1:2, :, :, 1:2into n = 512)copyto!(dst, view(parent(t), inds...))likewise: 14.582 ms → 0.041 ms at a 260 100-blockregion (357×).
The narrow row is the point of the design. A small assignment into a large sparse array is
already cheap, and must not become proportional to the destination's stored blocks. So
_deletemissing!sweeps whichever of the assigned region and the stored entries is smaller,making the whole thing
O(nnz(src) + min(region, nnz(dst)))— never worse than before, in anyshape.
copyto!(t::SparseTensorArray, ::SubArray)needs no sweep at all: the destination spansexactly the viewed region, so everything not copied is dropped.
A third bug: the
nonzero_*accessors onSparseTensorArraynonzero_keysandnonzero_lengthfell through to theAbstractArrayfallbacks(
eachindex(IndexCartesian(), A)andlength(A)) and so reported every index rather than thestored ones.
nonzero_pairswas accidentally correct viaBase.pairs(::SparseTensorArray)untilthe first commit here routed it through
pairs(IndexCartesian(), A), which bypasses thatspecialization.
Two visible consequences, both now covered by tests:
parent(st)[1:3, :, :]came back with everyblock stored;
copyto!stored explicitly-zero blocks, becauseRsrc[I] in nonzero_keys(src)was always true.Fixed with the four obvious specializations, mirroring
SparseBlockTensorMap's.Notes for review
Vararg{Strided.SliceIndex}twins are kept deliberately. They are not redundantspecializations: they disambiguate against TensorKit's own
getindex(::AbstractTensorMap, ::Vararg{SliceIndex})/setindex!(abstracttensor.jl:540and
:552), which slice theStridedViewof the trivial fusion tree.t[:, :, :]resolvesto the twin, so deleting either one turns it into a
MethodError. Only the bodies werecollapsed; the
setindex!dedup is pure code motion.t[:]andt[[1]]still throwArgumentError,t[:, :]and out-of-bounds indices still throwBoundsError(including under@inbounds,since
SumSpaceindexing bounds-checks unconditionally), and all existing@inferredtestsstill pass.
_invert_indexdeliberately carries no@inbounds/@propagate_inbounds, so an out-of-rangeindex can never become an unchecked write.
divremspecialization — it measured no faster and is the most error-prone arithmetic in thedesign.
copyto!rewrite is verified semantics-preserving. The stored-key sets after sliceassignment,
copyto!from a view, region-to-regioncopyto!, andcatare byte-identical tothe previous implementation once the accessors are fixed (compared over a seeded batch of
randomized cases, dumping sorted key tuples and diffing).
(
t[[1, 1], :] = v), the old code's result depended on iteration order, since it could writeand then delete the same parent entry. The new code deletes only entries the source has no
copy of, then writes — deterministic.
copyto!still does not clear destination entries where the source isstructurally zero. That is pre-existing behaviour and
catdepends on it, so it is leftalone.
test/abstracttensor/indexing.jlgrew from 62 to ~240 lines. The dense and sparse copies are nowone loop, and it covers repeated indices, logical masks,
BitVector, reversed and stepped ranges,empty slices, non-
Intintegers, 1-d block tensors, the parent-array path (including that it doesnot densify), the
nonzero_*accessors, randomized slice-assignment against an explicitkey-set/identity reference,
copyto!from a view, region-to-regioncopyto!with a steppedregion, and
cat.Two complexity guards, both checked to fail against the old code:
@allocated < 2 MBon a 260 100-block region; the oldRsrcaloneis ~8 MB), so it cannot be timing-flaky;
it is a wall-clock bound with a deliberately huge margin — a 4.5 M-block region where the old
code takes 141 ms and the new one 0.013 ms, asserted under 20 ms.
🤖 Generated with Claude Code