Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BlockSparseArrays] Permute and merge blocks #1514

Merged
merged 27 commits into from
Jul 1, 2024
Merged
Changes from 1 commit
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
e5c7fce
[BlockSparseArrays] Redesign block views again
mtfishman Jun 26, 2024
849051b
[NDTensors] Bump to v0.3.38
mtfishman Jun 26, 2024
4e7a280
[BlockSparseArrays] Merge blocks
mtfishman Jun 27, 2024
87b2ee1
Fix issues with GradedAxes
mtfishman Jun 27, 2024
31157d6
Merge branch 'BlockSparseArrays_redesign_blockviews_again' into Block…
mtfishman Jun 27, 2024
79bf2d4
Start to fix map
mtfishman Jun 27, 2024
1a22fd3
Initial version of block merging
mtfishman Jun 27, 2024
d322bf0
Merge branch 'main' into BlockSparseArrays_merge_blocks_2
mtfishman Jun 27, 2024
a3219fa
Fix some namespace issues
mtfishman Jun 27, 2024
e075825
Fix some tests
mtfishman Jun 28, 2024
e90dbec
Fix another test
mtfishman Jun 28, 2024
991fa07
Fix namespace issue
mtfishman Jun 28, 2024
ce44726
Output a more specific type when slicing with unit ranges
mtfishman Jun 28, 2024
17e274a
Fix some more tests
mtfishman Jun 30, 2024
29c5fbe
Fix some broken tests
mtfishman Jun 30, 2024
5bd85e3
Add tests
mtfishman Jun 30, 2024
60cda06
[NDTensors] Bump to v0.3.39
mtfishman Jun 30, 2024
710fd89
Start implementing block merging
mtfishman Jul 1, 2024
9dbac78
Slicing operation to permute and merge blocks
mtfishman Jul 1, 2024
b6ebf38
Preserve block labels during fusion
mtfishman Jul 1, 2024
88179e1
Fix some tests
mtfishman Jul 1, 2024
6d6b47f
Fix splitdims
mtfishman Jul 1, 2024
a1b0730
Fix for Julia 1.6
mtfishman Jul 1, 2024
7bfd24a
Fix typos
mtfishman Jul 1, 2024
da5abfa
Another fix for Julia 1.6
mtfishman Jul 1, 2024
a38ecdb
Another fix for Julia 1.6
mtfishman Jul 1, 2024
67bcd40
Another fix for Julia 1.6
mtfishman Jul 1, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix splitdims
mtfishman committed Jul 1, 2024
commit 6d6b47f225198cafdf7879b6098a399fe24f7199
Original file line number Diff line number Diff line change
@@ -66,7 +66,12 @@ function TensorAlgebra.splitdims(
return length(axis) ≤ length(axes(a, i))
end
blockperms = invblockperm.(blocksortperm.(axes_prod))
a_blockpermed = a[blockperms...]
# TODO: This is doing extra copies of the blocks,
# use `@view a[axes_prod...]` instead.
# That will require implementing some reindexing logic
# for this combination of slicing.
a_unblocked = a[axes_prod...]
a_blockpermed = a_unblocked[blockperms...]
return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...)
end

Original file line number Diff line number Diff line change
@@ -101,7 +101,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
@test a[1, 1, 1, 1] == m[1, 1]
@test a[2, 2, 2, 2] == m[4, 4]
@test blocksize(m) == (3, 3)
@test_broken a == splitdims(m, (d1, d2), (d1, d2))
@test a == splitdims(m, (d1, d2), (d1, d2))
end
@testset "dual axes" begin
r = gradedrange([U1(0) => 2, U1(1) => 2])
Original file line number Diff line number Diff line change
@@ -57,6 +57,10 @@ function reblock(
return @view parent(a)[map(I -> Vector(I.blocks), parentindices(a))...]
end

# TODO: Rewrite this so that it takes the blocking structure
# made by combining the blocking of the axes (i.e. the blocking that
# is used to determine `union_stored_blocked_cartesianindices(...)`).
# `reblock` is a partial solution to that, but a bit ad-hoc.
# TODO: Move to `blocksparsearrayinterface/map.jl`.
function SparseArrayInterface.sparse_map!(
::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}