From 632405bf7fce8a572c94e2e78f18a6858a23b8bf Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 11:35:39 -0500 Subject: [PATCH 1/9] Rewrap blockview of wrapped array --- .../src/abstractblocksparsearray/views.jl | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl index 4742f7b781..15193c7b72 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl @@ -278,3 +278,17 @@ function BlockArrays.viewblock( ) where {T,N} return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) end + +# migrate wrapper layer for viewing `adjoint` and `transpose`. +for (f, F) in ((:adjoint, :Adjoint), (:transpose, :Transpose)) + @eval begin + function Base.view(A::$F{<:Any,<:AbstractBlockSparseVector}, b::Block{1}) + return $f(view(parent(A), b)) + end + + Base.view(A::$F{<:Any,<:AbstractBlockSparseMatrix}, b::Block{2}) = view(A, Tuple(b)...) + function Base.view(A::$F{<:Any,<:AbstractBlockSparseMatrix}, b1::Block{1}, b2::Block{1}) + return $f(view(parent(A), b2, b1)) + end + end +end \ No newline at end of file From 6adeb52919583ce64831f3274fc9c61f725abfbf Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 11:36:19 -0500 Subject: [PATCH 2/9] Excise `SparseTransposeBlocks` and `SparseAdjointBlocks` --- .../blocksparsearrayinterface.jl | 75 +------------------ .../wrappedabstractsparsearray.jl | 10 +++ 2 files changed, 12 insertions(+), 73 deletions(-) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 98694efe94..b6374671db 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -186,79 +186,8 @@ end reverse_index(index) = reverse(index) reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index))) -# Represents the array of arrays of a `Transpose` -# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Transpose`. -struct SparseTransposeBlocks{T,BlockType<:AbstractArray{T},Array<:Transpose{T}} <: - AbstractSparseMatrix{BlockType} - array::Array -end -function blocksparse_blocks(a::Transpose) - return SparseTransposeBlocks{eltype(a),blocktype(parent(a)),typeof(a)}(a) -end -function Base.size(a::SparseTransposeBlocks) - return reverse(size(blocks(parent(a.array)))) -end -function Base.getindex(a::SparseTransposeBlocks, index::Vararg{Int,2}) - return transpose(blocks(parent(a.array))[reverse(index)...]) -end -# TODO: This should be handled by generic `AbstractSparseArray` code. -function Base.getindex(a::SparseTransposeBlocks, index::CartesianIndex{2}) - return a[Tuple(index)...] -end -# TODO: Create a generic `parent_index` function to map an index -# a parent index. -function Base.isassigned(a::SparseTransposeBlocks, index::Vararg{Int,2}) - return isassigned(blocks(parent(a.array)), reverse(index)...) -end -function SparseArrayInterface.stored_indices(a::SparseTransposeBlocks) - return map(reverse_index, stored_indices(blocks(parent(a.array)))) -end -# TODO: Either make this the generic interface or define -# `SparseArrayInterface.sparse_storage`, which is used -# to defined this. -SparseArrayInterface.stored_length(a::SparseTransposeBlocks) = length(stored_indices(a)) -function SparseArrayInterface.sparse_storage(a::SparseTransposeBlocks) - return error("Not implemented") -end - -# Represents the array of arrays of a `Adjoint` -# wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Adjoint`. -struct SparseAdjointBlocks{T,BlockType<:AbstractArray{T},Array<:Adjoint{T}} <: - AbstractSparseMatrix{BlockType} - array::Array -end -function blocksparse_blocks(a::Adjoint) - return SparseAdjointBlocks{eltype(a),blocktype(parent(a)),typeof(a)}(a) -end -function Base.size(a::SparseAdjointBlocks) - return reverse(size(blocks(parent(a.array)))) -end -# TODO: Create a generic `parent_index` function to map an index -# a parent index. -function Base.getindex(a::SparseAdjointBlocks, index::Vararg{Int,2}) - return blocks(parent(a.array))[reverse(index)...]' -end -# TODO: Create a generic `parent_index` function to map an index -# a parent index. -# TODO: This should be handled by generic `AbstractSparseArray` code. -function Base.getindex(a::SparseAdjointBlocks, index::CartesianIndex{2}) - return a[Tuple(index)...] -end -# TODO: Create a generic `parent_index` function to map an index -# a parent index. -function Base.isassigned(a::SparseAdjointBlocks, index::Vararg{Int,2}) - return isassigned(blocks(parent(a.array)), reverse(index)...) -end -function SparseArrayInterface.stored_indices(a::SparseAdjointBlocks) - return map(reverse_index, stored_indices(blocks(parent(a.array)))) -end -# TODO: Either make this the generic interface or define -# `SparseArrayInterface.sparse_storage`, which is used -# to defined this. -SparseArrayInterface.stored_length(a::SparseAdjointBlocks) = length(stored_indices(a)) -function SparseArrayInterface.sparse_storage(a::SparseAdjointBlocks) - return error("Not implemented") -end +blocksparse_blocks(a::Transpose) = transpose(blocks(parent(a))) +blocksparse_blocks(a::Adjoint) = adjoint(blocks(parent(a))) # Represents the array of arrays of a `SubArray` # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`. diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index f9ec08acb7..1be9f6ca1c 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -7,3 +7,13 @@ const WrappedAbstractSparseArray{T,N,A} = WrappedArray{ const AnyAbstractSparseArray{T,N} = Union{ <:AbstractSparseArray{T,N},<:WrappedAbstractSparseArray{T,N} } + +function stored_indices(a::Adjoint) + return Iterators.map(I -> CartesianIndex(reverse(I.I)), stored_indices(parent(a))) +end +stored_length(a::Adjoint) = stored_length(parent(a)) + +function stored_indices(a::Transpose) + return Iterators.map(I -> CartesianIndex(reverse(I.I)), stored_indices(parent(a))) +end +stored_length(a::Transpose) = stored_length(parent(a)) \ No newline at end of file From 93655c3214441273e7fb517581f5bf973d19c1b0 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 12:05:04 -0500 Subject: [PATCH 3/9] Unbreak previously broken tests --- .../src/lib/BlockSparseArrays/test/test_basics.jl | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index e41f736783..4ee93acf68 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -70,12 +70,6 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype @test adjoint(a) isa Adjoint{elt,<:BlockSparseArray} @test_broken adjoint(a)[Block(1), :] isa Adjoint{elt,<:BlockSparseArray} # could also be directly a BlockSparseArray - - a = dev(BlockSparseArray{elt}([1], [1, 1])) - @allowscalar a[1, 2] = 1 - @test [a[Block(Tuple(it))] for it in eachindex(block_stored_indices(a))] isa Vector - ah = adjoint(a) - @test_broken [ah[Block(Tuple(it))] for it in eachindex(block_stored_indices(ah))] isa Vector end @testset "Constructors" begin # BlockSparseMatrix @@ -210,6 +204,15 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype ## @test b[Block()[]] == 2 end end + + # adjoint and transpose + a = dev(BlockSparseArray{elt}([1], [1, 1])) + @allowscalar a[1, 2] = 1 + @test [@views(a[it]) for it in block_stored_indices(a)] isa Vector + ah = adjoint(a) + @test [@views(ah[it]) for it in block_stored_indices(ah)] isa Vector + at = transpose(a) + @test [@views(at[it]) for it in block_stored_indices(at)] isa Vector end @testset "Tensor algebra" begin a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) From 3d7986dc1adc2098cb655f92b6bb128a75ad12b6 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 14:35:57 -0500 Subject: [PATCH 4/9] `Tuple(I::CartesianIndex)` instead of `I.I` Co-authored-by: Matt Fishman --- .../src/abstractsparsearray/wrappedabstractsparsearray.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index 1be9f6ca1c..5b6a0613ef 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -9,11 +9,11 @@ const AnyAbstractSparseArray{T,N} = Union{ } function stored_indices(a::Adjoint) - return Iterators.map(I -> CartesianIndex(reverse(I.I)), stored_indices(parent(a))) + return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end stored_length(a::Adjoint) = stored_length(parent(a)) function stored_indices(a::Transpose) - return Iterators.map(I -> CartesianIndex(reverse(I.I)), stored_indices(parent(a))) + return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end stored_length(a::Transpose) = stored_length(parent(a)) \ No newline at end of file From 35c9e979e5bb8c0f4fdb048b46259f48f1bff957 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 13 Nov 2024 14:37:31 -0500 Subject: [PATCH 5/9] Formatter --- .../lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl | 2 +- .../src/abstractsparsearray/wrappedabstractsparsearray.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl index 15193c7b72..7283b850d3 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl @@ -291,4 +291,4 @@ for (f, F) in ((:adjoint, :Adjoint), (:transpose, :Transpose)) return $f(view(parent(A), b2, b1)) end end -end \ No newline at end of file +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index 5b6a0613ef..5aae045a70 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -16,4 +16,4 @@ stored_length(a::Adjoint) = stored_length(parent(a)) function stored_indices(a::Transpose) return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end -stored_length(a::Transpose) = stored_length(parent(a)) \ No newline at end of file +stored_length(a::Transpose) = stored_length(parent(a)) From cd6f8ab8a6c8b828b34eb98c1de574211ac8df0a Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 14 Nov 2024 18:32:43 -0500 Subject: [PATCH 6/9] Add missing imports --- .../src/abstractsparsearray/wrappedabstractsparsearray.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index 5aae045a70..c59982a468 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -1,4 +1,5 @@ using Adapt: WrappedArray +using LinearAlgebra: Adjoint, Transpose const WrappedAbstractSparseArray{T,N,A} = WrappedArray{ T,N,<:AbstractSparseArray,<:AbstractSparseArray{T,N} From 3efdb893a0b1ca1d1492493a3cc083bfd2299084 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Fri, 15 Nov 2024 09:31:47 -0500 Subject: [PATCH 7/9] Add better tests --- .../lib/BlockSparseArrays/test/test_basics.jl | 59 +++++++++++++++---- 1 file changed, 49 insertions(+), 10 deletions(-) diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index 4ee93acf68..ddf0d46eb1 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -16,7 +16,7 @@ using BlockArrays: mortar using Compat: @compat using GPUArraysCore: @allowscalar -using LinearAlgebra: Adjoint, dot, mul!, norm +using LinearAlgebra: Adjoint, Transpose, dot, mul!, norm using NDTensors.BlockSparseArrays: @view!, BlockSparseArray, @@ -33,7 +33,7 @@ using NDTensors.GPUArraysCoreExtensions: cpu using NDTensors.SparseArrayInterface: stored_length using NDTensors.SparseArrayDOKs: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK using NDTensors.TensorAlgebra: contract -using Test: @test, @test_broken, @test_throws, @testset +using Test: @test, @test_broken, @test_throws, @testset, @inferred include("TestBlockSparseArraysUtils.jl") using NDTensors: NDTensors @@ -205,14 +205,53 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype end end - # adjoint and transpose - a = dev(BlockSparseArray{elt}([1], [1, 1])) - @allowscalar a[1, 2] = 1 - @test [@views(a[it]) for it in block_stored_indices(a)] isa Vector - ah = adjoint(a) - @test [@views(ah[it]) for it in block_stored_indices(ah)] isa Vector - at = transpose(a) - @test [@views(at[it]) for it in block_stored_indices(at)] isa Vector + @testset "Transpose" begin + a = dev(BlockSparseArray{elt}([2, 2], [3, 3, 1])) + a[Block(1, 1)] = dev(randn(elt, 2, 3)) + a[Block(2, 3)] = dev(randn(elt, 2, 1)) + + at = @inferred transpose(a) + @test at isa Transpose + @test size(at) == reverse(size(a)) + @test blocksize(at) == reverse(blocksize(a)) + @test stored_length(at) == stored_length(a) + @test block_stored_length(at) == block_stored_length(a) + for bind in block_stored_indices(a) + bindt = Block(reverse(Int.(Tuple(bind)))) + @test bindt in block_stored_indices(at) + end + + @test @views(at[Block(1, 1)]) == transpose(a[Block(1, 1)]) + @test @views(at[Block(1, 1)]) isa Transpose + @test @views(at[Block(3, 2)]) == transpose(a[Block(2, 3)]) + # TODO: BlockView == AbstractArray calls scalar code + @test @allowscalar @views(at[Block(1, 2)]) == transpose(a[Block(2, 1)]) + @test @views(at[Block(1, 2)]) isa Transpose + end + + @testset "Adjoint" begin + a = dev(BlockSparseArray{elt}([2, 2], [3, 3, 1])) + a[Block(1, 1)] = dev(randn(elt, 2, 3)) + a[Block(2, 3)] = dev(randn(elt, 2, 1)) + + at = @inferred adjoint(a) + @test at isa Adjoint + @test size(at) == reverse(size(a)) + @test blocksize(at) == reverse(blocksize(a)) + @test stored_length(at) == stored_length(a) + @test block_stored_length(at) == block_stored_length(a) + for bind in block_stored_indices(a) + bindt = Block(reverse(Int.(Tuple(bind)))) + @test bindt in block_stored_indices(at) + end + + @test @views(at[Block(1, 1)]) == adjoint(a[Block(1, 1)]) + @test @views(at[Block(1, 1)]) isa Adjoint + @test @views(at[Block(3, 2)]) == adjoint(a[Block(2, 3)]) + # TODO: BlockView == AbstractArray calls scalar code + @test @allowscalar @views(at[Block(1, 2)]) == adjoint(a[Block(2, 1)]) + @test @views(at[Block(1, 2)]) isa Adjoint + end end @testset "Tensor algebra" begin a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) From d1a5b53106c03387fc1db3b0a6614b44465e1c5c Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Fri, 15 Nov 2024 09:32:08 -0500 Subject: [PATCH 8/9] Add `sparse_storage` for `Adjoint` and `Transpose` --- .../src/abstractsparsearray/wrappedabstractsparsearray.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index c59982a468..ca53420809 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -13,8 +13,10 @@ function stored_indices(a::Adjoint) return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end stored_length(a::Adjoint) = stored_length(parent(a)) +sparse_storage(a::Adjoint) = Iterators.map(adjoint, sparse_storage(parent(a))) function stored_indices(a::Transpose) return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end stored_length(a::Transpose) = stored_length(parent(a)) +sparse_storage(a::Transpose) = Iterators.map(transpose, sparse_storage(parent(a))) \ No newline at end of file From 4c6a095c57df6a7888eb01a544d8b34f5635930d Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Fri, 15 Nov 2024 11:56:33 -0500 Subject: [PATCH 9/9] Formatter --- .../src/abstractsparsearray/wrappedabstractsparsearray.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl index ca53420809..a4ee4bebe3 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/wrappedabstractsparsearray.jl @@ -19,4 +19,4 @@ function stored_indices(a::Transpose) return Iterators.map(I -> CartesianIndex(reverse(Tuple(I))), stored_indices(parent(a))) end stored_length(a::Transpose) = stored_length(parent(a)) -sparse_storage(a::Transpose) = Iterators.map(transpose, sparse_storage(parent(a))) \ No newline at end of file +sparse_storage(a::Transpose) = Iterators.map(transpose, sparse_storage(parent(a)))