Mutable addition and substraction of sparse arrays

jump-dev · Apr 23, 2024 · 3bc025d · 3bc025d
1 parent e367b37
commit 3bc025d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 5 deletions.
diff --git a/src/implementations/LinearAlgebra.jl b/src/implementations/LinearAlgebra.jl
@@ -89,13 +89,13 @@ function _check_dims(A, B)
     return
 end
 
-function operate!(op::Union{typeof(+),typeof(-)}, A::Array, B::AbstractArray)
+function operate!(op::Union{typeof(+),typeof(-)}, A::AbstractArray, B::AbstractArray)
     _check_dims(A, B)
     return broadcast!(op, A, B)
 end
 
 function operate_to!(
-    output::Array,
+    output::AbstractArray,
     op::Union{typeof(+),typeof(-)},
     A::AbstractArray,
     B::AbstractArray,
@@ -116,7 +116,7 @@ end
 # We call `scaling_to_number` as `UniformScaling` do not support broadcasting
 function operate!(
     op::AddSubMul,
-    A::Array,
+    A::AbstractArray,
     B::AbstractArray,
     α::Vararg{Scaling,M},
 ) where {M}
@@ -126,7 +126,7 @@ end
 
 function operate!(
     op::AddSubMul,
-    A::Array,
+    A::AbstractArray,
     α::Scaling,
     B::AbstractArray,
     β::Vararg{Scaling,M},
@@ -137,7 +137,7 @@ end
 
 function operate!(
     op::AddSubMul,
-    A::Array,
+    A::AbstractArray,
     α1::Scaling,
     α2::Scaling,
     B::AbstractArray,

diff --git a/test/matmul.jl b/test/matmul.jl
@@ -119,13 +119,26 @@ end
             "Cannot sum or substract a matrix of axes `$(axes(B))` into matrix of axes `$(axes(A))`, expected axes `$(axes(B))`.",
         )
         @test_throws err MA.operate!(+, A, B)
+        A = spzeros(2)
+        B = spzeros(2, 1)
+        err = DimensionMismatch(
+            "Cannot sum or substract a matrix of axes `$(axes(B))` into matrix of axes `$(axes(A))`, expected axes `$(axes(B))`.",
+        )
+        @test_throws err MA.operate!(+, A, B)
         output = zeros(2)
         A = zeros(2, 1)
         B = zeros(2, 1)
         err = DimensionMismatch(
             "Cannot sum or substract matrices of axes `$(axes(A))` and `$(axes(B))` into a matrix of axes `$(axes(output))`, expected axes `$(axes(B))`.",
         )
         @test_throws err MA.operate_to!(output, +, A, B)
+        output = spzeros(2)
+        A = spzeros(2, 1)
+        B = spzeros(2, 1)
+        err = DimensionMismatch(
+            "Cannot sum or substract matrices of axes `$(axes(A))` and `$(axes(B))` into a matrix of axes `$(axes(output))`, expected axes `$(axes(B))`.",
+        )
+        @test_throws err MA.operate_to!(output, +, A, B)
     end
     @testset "unsupported_product" begin
         unsupported_product()
@@ -454,6 +467,21 @@ function test_array_sum(::Type{T}) where {T}
     return
 end
 
+function test_sparse_vector_sum(::Type{T}) where {T}
+    x = sparsevec([1, 3], T[5, 7])
+    y = copy(x)
+    z = copy(y)
+    alloc_test(() -> MA.operate!(+, y, z), 0)
+    alloc_test(() -> MA.operate!(-, y, z), 0)
+    alloc_test(() -> MA.add!!(y, z), 0)
+    alloc_test(() -> MA.sub!!(y, z), 0)
+    alloc_test(() -> MA.operate_to!(x, +, y, z), 0)
+    alloc_test(() -> MA.operate_to!(x, -, y, z), 0)
+    alloc_test(() -> MA.add_to!!(x, y, z), 0)
+    alloc_test(() -> MA.sub_to!!(x, y, z), 0)
+    return
+end
+
 @testset "Array sum" begin
     test_array_sum(Int)
 end