Fix implementation of fused_map_reduce (#171)

Author: odow

src/interface.jl

@@ -666,6 +666,15 @@ function promote_operation_fallback(
     return promote_operation(*, promote_operation(adjoint, A), B)
 end
 
+function promote_operation_fallback(
+    ::typeof(LinearAlgebra.dot),
+    ::Type{<:AbstractArray{A}},
+    ::Type{<:AbstractArray{B}},
+) where {A,B}
+    C = promote_operation(*, A, B)
+    return promote_operation(+, C, C)
+end
+
 function buffer_for(::typeof(add_dot), a::Type, b::Type, c::Type)
     return buffer_for(add_mul, a, promote_operation(adjoint, b), c)
 end

src/reduce.jl

@@ -36,9 +36,11 @@ function promote_map_reduce(op::Function, args::Vararg{Any,N}) where {N}
     )
 end
 
+_concrete_eltype(x) = isempty(x) ? eltype(x) : typeof(first(x))
+
 function fused_map_reduce(op::F, args::Vararg{Any,N}) where {F<:Function,N}
     _check_same_length(args...)
-    T = promote_map_reduce(op, eltype.(args)...)
+    T = promote_map_reduce(op, _concrete_eltype.(args)...)
     accumulator = neutral_element(reduce_op(op), T)
     buffer = buffer_for(op, T, eltype.(args)...)
     for I in zip(eachindex.(args)...)

test/dispatch.jl

@@ -29,4 +29,17 @@ end
     # On `DummyBigInt` allocates more on previous releases of Julia
     # as it's dynamically allocated
     dispatch_tests(DummyBigInt)
+
+    @testset "dot non-concrete vector" begin
+        x = [5.0, 6.0]
+        y = Vector{Union{Float64,String}}(x)
+        @test MA.operate(LinearAlgebra.dot, x, y) == LinearAlgebra.dot(x, y)
+        @test MA.operate(*, x', y) == x' * y
+    end
+
+    @testset "dot vector of vectors" begin
+        x = [5.0, 6.0]
+        z = [x, x]
+        @test MA.operate(LinearAlgebra.dot, z, z) == LinearAlgebra.dot(z, z)
+    end
 end