[BlockSparseArrys] Zero dimensional block sparse array

Author: mtfishman

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -22,11 +22,16 @@ function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T,N}}) where {T,N
   return SparseArrayDOK{AbstractArray{T,N},N}
 end
 
-## # Specialized in order to fix ambiguity error with `BlockArrays`.
+# Specialized in order to fix ambiguity error with `BlockArrays`.
 function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N}
   return blocksparse_getindex(a, I...)
 end
 
+# Specialized in order to fix ambiguity error with `BlockArrays`.
+function Base.getindex(a::AbstractBlockSparseArray{<:Any,0})
+  return blocksparse_getindex(a)
+end
+
 ## # Fix ambiguity error with `BlockArrays`.
 ## function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Block{N}) where {N}
 ##   return ArrayLayouts.layout_getindex(a, I)
@@ -51,6 +56,12 @@ function Base.setindex!(
   return a
 end
 
+# Fix ambiguity error.
+function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value)
+  blocksparse_setindex!(a, value)
+  return a
+end
+
 function Base.setindex!(
   a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Block{1},N}
 ) where {N}

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -92,6 +92,10 @@ function Base.getindex(
 )
   return ArrayLayouts.layout_getindex(a, I...)
 end
+# Fixes ambiguity error.
+function Base.getindex(a::BlockSparseArrayLike{<:Any,0})
+  return ArrayLayouts.layout_getindex(a)
+end
 
 # TODO: Define `blocksparse_isassigned`.
 function Base.isassigned(
@@ -100,6 +104,11 @@ function Base.isassigned(
   return isassigned(blocks(a), Int.(index)...)
 end
 
+# Fix ambiguity error.
+function Base.isassigned(a::BlockSparseArrayLike{<:Any,0})
+  return isassigned(blocks(a))
+end
+
 function Base.isassigned(a::BlockSparseArrayLike{<:Any,N}, index::Block{N}) where {N}
   return isassigned(a, Tuple(index)...)
 end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blocksparsearray.jl

@@ -72,6 +72,10 @@ function BlockSparseArray{T,N}(axes::Tuple{Vararg{AbstractUnitRange,N}}) where {
   return BlockSparseArray{T,N,default_arraytype(T, axes)}(axes)
 end
 
+function BlockSparseArray{T,0}(axes::Tuple{}) where {T}
+  return BlockSparseArray{T,0,default_arraytype(T, axes)}(axes)
+end
+
 function BlockSparseArray{T,N}(dims::Tuple{Vararg{Vector{Int},N}}) where {T,N}
   return BlockSparseArray{T,N}(blockedrange.(dims))
 end
@@ -84,6 +88,10 @@ function BlockSparseArray{T}(axes::Tuple{Vararg{AbstractUnitRange}}) where {T}
   return BlockSparseArray{T,length(axes)}(axes)
 end
 
+function BlockSparseArray{T}(axes::Tuple{}) where {T}
+  return BlockSparseArray{T,length(axes)}(axes)
+end
+
 function BlockSparseArray{T}(dims::Vararg{Vector{Int}}) where {T}
   return BlockSparseArray{T}(dims)
 end
@@ -92,6 +100,10 @@ function BlockSparseArray{T}(axes::Vararg{AbstractUnitRange}) where {T}
   return BlockSparseArray{T}(axes)
 end
 
+function BlockSparseArray{T}() where {T}
+  return BlockSparseArray{T}(())
+end
+
 function BlockSparseArray{T,N,A}(
   ::UndefInitializer, dims::Tuple
 ) where {T,N,A<:AbstractArray{T,N}}

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -22,6 +22,13 @@ function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where
   return a[findblockindex.(axes(a), I)...]
 end
 
+# Fix ambiguity error.
+function blocksparse_getindex(a::AbstractArray{<:Any,0})
+  # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430
+  # is fixed.
+  return a[BlockIndex{0,Tuple{},Tuple{}}((), ())]
+end
+
 # a[1:2, 1:2]
 # TODO: This definition means that the result of slicing a block sparse array
 # with a non-blocked unit range is blocked. We may want to change that behavior,
@@ -77,6 +84,14 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N
   return a
 end
 
+# Fix ambiguity error.
+function blocksparse_setindex!(a::AbstractArray{<:Any,0}, value)
+  # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430
+  # is fixed.
+  a[BlockIndex{0,Tuple{},Tuple{}}((), ())] = value
+  return a
+end
+
 function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N}) where {N}
   i = Int.(Tuple(block(I)))
   a_b = blocks(a)[i...]
@@ -86,6 +101,15 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N
   return a
 end
 
+# Fix ambiguity error.
+function blocksparse_setindex!(a::AbstractArray{<:Any,0}, value, I::BlockIndex{0})
+  a_b = blocks(a)[]
+  a_b[] = value
+  # Set the block, required if it is structurally zero.
+  blocks(a)[] = a_b
+  return a
+end
+
 function blocksparse_fill!(a::AbstractArray, value)
   for b in BlockRange(a)
     # We can't use:

NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl

@@ -82,6 +82,11 @@ function sparse_getindex(a::AbstractArray, I::Vararg{Int})
   return sparse_getindex(a, CartesianIndex(I))
 end
 
+# Fix ambiguity error.
+function sparse_getindex(a::AbstractArray{<:Any,0})
+  return sparse_getindex(a, CartesianIndex())
+end
+
 # Linear indexing
 function sparse_getindex(a::AbstractArray, I::CartesianIndex{1})
   return sparse_getindex(a, CartesianIndices(a)[I])