Fix adapt, fix printing subarrays (#82)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.3.3"
+version = "0.3.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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,

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

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}}}

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