[SparseArrayInterface] Add tests for dot

Author: mtfishman

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

@@ -8,11 +8,14 @@ function ArrayLayouts.MemoryLayout(arraytype::Type{<:BlockSparseArrayLike})
   inner_layout = typeof(MemoryLayout(blocktype(arraytype)))
   return BlockLayout{outer_layout,inner_layout}()
 end
+
+# TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`.
 function ArrayLayouts.MemoryLayout(
   arraytype::Type{<:Adjoint{<:Any,<:AbstractBlockSparseVector}}
 )
   return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}()
 end
+# TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`.
 function ArrayLayouts.MemoryLayout(
   arraytype::Type{<:Transpose{<:Any,<:AbstractBlockSparseVector}}
 )

NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/arraylayouts.jl

@@ -1,9 +1,22 @@
-using ArrayLayouts: ArrayLayouts, Dot, MatMulMatAdd, MatMulVecAdd, MulAdd
+using ArrayLayouts: ArrayLayouts, Dot, DualLayout, MatMulMatAdd, MatMulVecAdd, MulAdd
+using LinearAlgebra: Adjoint, Transpose
+using ..TypeParameterAccessors: parenttype
 
 function ArrayLayouts.MemoryLayout(arraytype::Type{<:SparseArrayLike})
   return SparseLayout()
 end
 
+# TODO: Generalize to `SparseVectorLike`/`AnySparseVector`.
+function ArrayLayouts.MemoryLayout(arraytype::Type{<:Adjoint{<:Any,<:AbstractSparseVector}})
+  return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}()
+end
+# TODO: Generalize to `SparseVectorLike`/`AnySparseVector`.
+function ArrayLayouts.MemoryLayout(
+  arraytype::Type{<:Transpose{<:Any,<:AbstractSparseVector}}
+)
+  return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}()
+end
+
 function sparse_matmul!(m::MulAdd)
   α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C
   sparse_mul!(a_dest, a1, a2, α, β)

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

@@ -55,10 +55,11 @@ function sparse_mul!(
 end
 
 function sparse_dot(a1::AbstractArray, a2::AbstractArray)
-  # This assumes `a1` and `a2` have the same shape.
+  # This requires that `a1` and `a2` have the same shape.
   # TODO: Generalize (Base supports dot products of
   # arrays with the same length but different sizes).
-  @assert size(a1) == size(a2)
+  size(a1) == size(a2) ||
+    throw(DimensionMismatch("Sizes $(size(a1)) and $(size(a2)) don't match."))
   dot_dest = zero(Base.promote_op(dot, eltype(a1), eltype(a2)))
   # TODO: First check if the number of stored elements (`nstored`, to be renamed
   # `stored_length`) is smaller in `a1` or `a2` and use whicheven one is smallar

NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl

@@ -32,6 +32,10 @@ function LinearAlgebra.mul!(
   return a_dest
 end
 
+function LinearAlgebra.dot(a1::SparseArray, a2::SparseArray)
+  return SparseArrayInterface.sparse_dot(a1, a2)
+end
+
 # AbstractArray interface
 Base.size(a::SparseArray) = a.dims
 function Base.similar(a::SparseArray, elt::Type, dims::Tuple{Vararg{Int}})

NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl

@@ -1,5 +1,5 @@
 @eval module $(gensym())
-using LinearAlgebra: mul!, norm
+using LinearAlgebra: dot, mul!, norm
 using NDTensors.SparseArrayInterface: SparseArrayInterface
 include("SparseArrayInterfaceTestUtils/SparseArrayInterfaceTestUtils.jl")
 using .SparseArrayInterfaceTestUtils.AbstractSparseArrays: AbstractSparseArrays
@@ -299,6 +299,18 @@ using Test: @test, @testset
   @test a_dest isa SparseArray{elt}
   @test SparseArrayInterface.nstored(a_dest) == 2
 
+  # Dot product
+  a1 = SparseArray{elt}(4)
+  a1[1] = randn()
+  a1[3] = randn()
+  a2 = SparseArray{elt}(4)
+  a2[2] = randn()
+  a2[3] = randn()
+  a_dest = a1' * a2
+  @test a_dest isa elt
+  @test a_dest ≈ Array(a1)' * Array(a2)
+  @test a_dest ≈ dot(a1, a2)
+
   # In-place matrix multiplication
   a1 = SparseArray{elt}(2, 3)
   a1[1, 2] = 12

src/lib/ITensorsNamedDimsArraysExt/src/to_nameddimsarray.jl

@@ -34,17 +34,19 @@ function to_nameddimsarray(x::DiagTensor)
   return named(DiagonalArray(data(x), size(x)), name.(inds(x)))
 end
 
-using ..NDTensors: BlockSparseTensor
+using ..NDTensors.BlockSparseArrays.BlockArrays: BlockArrays, blockedrange
+using ..NDTensors: BlockSparseTensor, array, blockdim, datatype, nblocks, nzblocks
 using ..NDTensors.BlockSparseArrays: BlockSparseArray
+using ..NDTensors.TypeParameterAccessors: set_ndims
 # TODO: Delete once `BlockSparse` is removed.
 function to_nameddimsarray(x::BlockSparseTensor)
-  blockinds = map(i -> [blockdim(i, b) for b in 1:nblocks(i)], inds(x))
+  blockinds = map(i -> blockedrange([blockdim(i, b) for b in 1:nblocks(i)]), inds(x))
   blocktype = set_ndims(datatype(x), ndims(x))
   # TODO: Make a simpler constructor:
   # BlockSparseArray(blocktype, blockinds)
-  arraystorage = BlockSparseArray{eltype(x),ndims(x),blocktype}(blockinds)
+  arraystorage = BlockSparseArray{eltype(x),ndims(x),blocktype}(undef, blockinds)
   for b in nzblocks(x)
-    arraystorage[BlockArrays.Block(Tuple(b)...)] = x[b]
+    arraystorage[BlockArrays.Block(Int.(Tuple(b))...)] = array(x[b])
   end
   return named(arraystorage, name.(inds(x)))
 end