diff --git a/Project.toml b/Project.toml index 5be41143..8bd54ba6 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.3.3" +version = "0.3.4" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/abstractblocksparsearray/views.jl b/src/abstractblocksparsearray/views.jl index ba3b85ec..f2451988 100644 --- a/src/abstractblocksparsearray/views.jl +++ b/src/abstractblocksparsearray/views.jl @@ -130,6 +130,21 @@ function Base.view( ) where {T,N} return viewblock(a, block) end + +# Fix ambiguity error with BlockArrays.jl for slices like +# `a = BlockSparseArray{Float64}(undef, [2, 2], [2, 2]); @view a[:, :]`. +function Base.view( + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{Base.Slice,BlockSlice{<:BlockRange{1}}},N}}, + }, + block::Block{N}, +) where {T,N} + return viewblock(a, block) +end + function Base.view( a::SubArray{ T, diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 1f1a8619..cdc4cbf6 100644 --- a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -361,14 +361,25 @@ function Base.replace_in_print_matrix( end # attempt to catch things that wrap GPU arrays -function Base.print_array(io::IO, X::AnyAbstractBlockSparseArray) - X_cpu = adapt(Array, X) - if typeof(X_cpu) === typeof(X) # prevent infinite recursion +function Base.print_array(io::IO, a::AnyAbstractBlockSparseArray) + a_cpu = adapt(Array, a) + if typeof(a_cpu) === typeof(a) # prevent infinite recursion # need to specify ndims to allow specialized code for vector/matrix @allowscalar @invoke Base.print_array( - io, X_cpu::AbstractArray{eltype(X_cpu),ndims(X_cpu)} + io, a_cpu::AbstractArray{eltype(a_cpu),ndims(a_cpu)} ) - else - Base.print_array(io, X_cpu) + return nothing end + Base.print_array(io, a_cpu) + return nothing +end + +using Adapt: Adapt, adapt +function Adapt.adapt_structure(to, a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray}) + # In the generic definition in Adapt.jl, `parentindices(a)` are also + # adapted, but is broken when the parent indices contained blocked unit + # ranges since `adapt` is broken on blocked unit ranges. + # TODO: Fix adapt for blocked unit ranges by making an AdaptExt for + # BlockArrays.jl. + return SubArray(adapt(to, parent(a)), parentindices(a)) end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 6f4200aa..843a1a70 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -13,6 +13,7 @@ using BlockArrays: blockcheckbounds, blockisequal, blocklengths, + blocklength, blocks, findblockindex using DerivableInterfaces: DerivableInterfaces, @interface, DefaultArrayInterface @@ -412,6 +413,7 @@ end to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block) to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks) +to_blocks_indices(I::Base.Slice{<:BlockedOneTo}) = Base.OneTo(blocklength(I.indices)) @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks( a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}} diff --git a/test/test_basics.jl b/test/test_basics.jl index 144128c3..dc55702d 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -34,7 +34,7 @@ using BlockSparseArrays: sparsemortar, view! using GPUArraysCore: @allowscalar -using JLArrays: JLArray +using JLArrays: JLArray, JLMatrix using LinearAlgebra: Adjoint, Transpose, dot, mul!, norm using SparseArraysBase: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, storedlength using TensorAlgebra: contract @@ -306,6 +306,27 @@ arrayts = (Array, JLArray) @test @views(at[Block(1, 2)]) isa Adjoint end end + @testset "adapt" begin + a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) + a_12 = randn(elt, 2, 2) + a[Block(1, 2)] = a_12 + a_jl = adapt(JLArray, a) + @test a_jl isa BlockSparseMatrix{elt,JLMatrix{elt}} + @test blocktype(a_jl) == JLMatrix{elt} + @test blockstoredlength(a_jl) == 1 + @test a_jl[Block(1, 2)] isa JLMatrix{elt} + @test adapt(Array, a_jl[Block(1, 2)]) == a_12 + + a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) + a_12 = randn(elt, 2, 2) + a[Block(1, 2)] = a_12 + a_jl = adapt(JLArray, @view(a[:, :])) + @test a_jl isa SubArray{elt,2,<:BlockSparseMatrix{elt,JLMatrix{elt}}} + @test blocktype(a_jl) == JLMatrix{elt} + @test blockstoredlength(a_jl) == 1 + @test a_jl[Block(1, 2)] isa JLMatrix{elt} + @test adapt(Array, a_jl[Block(1, 2)]) == a_12 + end @testset "Tensor algebra" begin a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) @views for b in [Block(1, 2), Block(2, 1)] @@ -1149,15 +1170,19 @@ arrayts = (Array, JLArray) # Not testing other element types since they change the # spacing so it isn't easy to make the test general. - a = BlockSparseMatrix{elt,arrayt{elt,2}}(undef, [2, 2], [2, 2]) - @allowscalar a[1, 2] = 12 - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n $(zero(eltype(a))) $(eltype(a)(12)) │ ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) │ ⋅ ⋅ \n ───────────┼──────────\n ⋅ ⋅ │ ⋅ ⋅ \n ⋅ ⋅ │ ⋅ ⋅ " + a′ = BlockSparseMatrix{elt,arrayt{elt,2}}(undef, [2, 2], [2, 2]) + @allowscalar a′[1, 2] = 12 + for a in (a′, @view(a′[:, :])) + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n $(zero(eltype(a))) $(eltype(a)(12)) │ ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) │ ⋅ ⋅ \n ───────────┼──────────\n ⋅ ⋅ │ ⋅ ⋅ \n ⋅ ⋅ │ ⋅ ⋅ " + end - a = BlockSparseArray{elt,3,arrayt{elt,3}}(undef, [2, 2], [2, 2], [2, 2]) - @allowscalar a[1, 2, 1] = 121 - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n[:, :, 1] =\n $(zero(eltype(a))) $(eltype(a)(121)) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 2] =\n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 3] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 4] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ " + a′ = BlockSparseArray{elt,3,arrayt{elt,3}}(undef, [2, 2], [2, 2], [2, 2]) + @allowscalar a′[1, 2, 1] = 121 + for a in (a′, @view(a′[:, :, :])) + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n[:, :, 1] =\n $(zero(eltype(a))) $(eltype(a)(121)) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 2] =\n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 3] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 4] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ " + end end end @testset "TypeParameterAccessors.position" begin