diff --git a/src/implementations/LinearAlgebra.jl b/src/implementations/LinearAlgebra.jl index 00ff2b6..2a14178 100644 --- a/src/implementations/LinearAlgebra.jl +++ b/src/implementations/LinearAlgebra.jl @@ -94,6 +94,25 @@ function operate!(op::Union{typeof(+),typeof(-)}, A::Array, B::AbstractArray) return broadcast!(op, A, B) end +function operate_to!( + output::Array, + op::Union{typeof(+),typeof(-)}, + A::AbstractArray, + B::AbstractArray, +) + if axes(output) != promote_shape(A, B) + throw( + DimensionMismatch( + "Cannot sum or substract matrices of axes `$(axes(A))` and" * + " `$(axes(B))` into a matrix of axes `$(axes(output))`," * + " expected axes `$(promote_shape(A, B))`.", + ), + ) + end + # We don't have `MA.broadcast_to!` as it would be exactly `Base.broadcast!`. + return Base.broadcast!(op, output, A, B) +end + # We call `scaling_to_number` as `UniformScaling` do not support broadcasting function operate!( op::AddSubMul, diff --git a/test/matmul.jl b/test/matmul.jl index 029edcc..ea7a045 100644 --- a/test/matmul.jl +++ b/test/matmul.jl @@ -119,6 +119,13 @@ 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) + 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) end @testset "unsupported_product" begin unsupported_product() @@ -435,3 +442,18 @@ end @test B == [1 2] @test D == B * A end + +function test_array_sum(::Type{T}) where {T} + x = zeros(T, 2) + y = copy(x) + z = copy(y) + alloc_test(() -> MA.operate!(+, y, z), 0) + alloc_test(() -> MA.add!!(y, z), 0) + alloc_test(() -> MA.operate_to!(x, +, y, z), 0) + alloc_test(() -> MA.add_to!!(x, y, z), 0) + return +end + +@testset "Array sum" begin + test_array_sum(Int) +end