DArray: Infer the result type of reductions through f

src/array/map-reduce.jl

@@ -54,7 +54,14 @@ struct MapReduce{T,N} <: ArrayOp{T,N}
     init
 end
 function MapReduce(f, op_inner, op_outer, input::DArray{T,N}, dims, init) where {T,N}
-    T_new = Base._return_type(op_outer, Tuple{T, T})
+    # `op_outer` combines the values `f` produces, not the input elements
+    T_f = _mapped_type(f, op_inner, T)
+    if !(init isa Base._InitialValue)
+        # ... starting from `init`, which may widen them
+        T_init = Base._return_type(op_outer, Tuple{typeof(init), T_f})
+        T_init === Union{} || (T_f = T_init)
+    end
+    T_new = Base._return_type(op_outer, Tuple{T_f, T_f})
     if T_new === Union{}
         T_new = Any
     end
@@ -63,6 +70,8 @@ function MapReduce(f, op_inner, op_outer, input::DArray{T,N}, dims, init) where
     return MapReduce{T_new,N_new}(f, op_inner, op_outer, input, dims, init)
 end
 
+_mapped_type(f, op_inner, ::Type{T}) where T = Base._return_type(f, Tuple{T})
+
 function stage(ctx::Context, r::MapReduce{T,N}) where {T,N}
     inp = stage(ctx, r.input)
 
@@ -83,7 +92,7 @@ function stage(ctx::Context, r::MapReduce{T,N}) where {T,N}
         A[1] = x
         return A
     end
-    to_array(x::Array, N) = x
+    to_array(x::AbstractArray, N) = x
     function treered_f(op, x, y, N)
         value = op.(x, y)
         return to_array(value, N)
@@ -183,14 +192,20 @@ function _extrema_inner(f, X; dims, init)
 end
 _extrema_outer(x::Extrema, y::Extrema) =
     Extrema(min(x.min, y.min), max(x.max, y.max))
+_mapped_type(f, ::typeof(_extrema_inner), ::Type{T}) where T =
+    Extrema{Base._return_type(f, Tuple{T})}
 
-function _onlinestats_mapreduce(f, x::DArray{T}, stat; dims=nothing) where T
-    _f(x) = fit!(stat(T), f(x))
-    init = stat(T)
+function _onlinestats_mapreduce(f, x::DArray{T}, ::Type{S}; dims=nothing) where {T,S}
+    _f(x) = fit!(S(T), f(x))
+    init = S(T)
     return _onlinestats_finish(mapreduce(_f, merge, x; dims, init), dims)
 end
 _onlinestats_finish(result, ::Union{Nothing,Colon}) = OnlineStats.value(result)
-_onlinestats_finish(result, ::Union{Int,Dims}) = map(OnlineStats.value, result)
+_onlinestats_finish(result, ::Union{Int,Dims}) = map(_onlinestats_value, result)
+_onlinestats_value(o) = OnlineStats.value(o)
+# `value(::Variance)` multiplies by a `Float64` `NaN` when it has no finite
+# mean, so it infers as `Union{S,Float64}` for any other precision
+_onlinestats_value(o::OnlineStats.Variance{T,S}) where {T,S} = convert(S, OnlineStats.value(o))
 
 mean(x::DArray; dims=nothing) = mean(identity, x; dims)
 mean(f::Function, x::DArray; dims=nothing) =

test/array/mapreduce.jl

@@ -51,3 +51,47 @@ end
 @testset "mean" test_mapreduce(mean, rand; zero_init=nothing, types=(Float32, Float64))
 @testset "var" test_mapreduce(var, rand; zero_init=nothing, types=(Float32, Float64))
 @testset "std" test_mapreduce(std, rand; zero_init=nothing, types=(Float32, Float64))
+
+# The result type follows `f`, not only the input eltype
+function test_mapreduce_type(f, A; cmp=isapprox)
+    X = distribute(A, Blocks(10, 10))
+    @testset "dims=$dims" for dims in (Colon(), 1, 2)
+        expected = f(A; dims)
+        if dims == Colon()
+            actual = f(X; dims)
+            @test typeof(actual) == typeof(expected)
+            @test cmp(actual, expected)
+        else
+            actual = f(X; dims)
+            @test eltype(actual) == eltype(expected)
+            @test cmp(collect(actual), expected)
+        end
+    end
+end
+@testset "result type follows f" begin
+    A = rand(20, 20)
+    C = rand(ComplexF64, 20, 20)
+    @testset "sum (Bool)" test_mapreduce_type((X; dims)->sum(x->x>0.5, X; dims), A)
+    @testset "sum (real of complex)" test_mapreduce_type((X; dims)->sum(abs2, X; dims), C)
+    @testset "sum (complex of real)" test_mapreduce_type((X; dims)->sum(x->complex(x, 1.0), X; dims), A)
+    @testset "prod (Float32)" test_mapreduce_type((X; dims)->prod(x->Float32(x)+1, X; dims), A)
+    @testset "mapreduce" test_mapreduce_type((X; dims)->mapreduce(abs2, +, X; dims), C)
+    @testset "maximum" test_mapreduce_type((X; dims)->maximum(abs, X; dims), C)
+    @testset "count" test_mapreduce_type((X; dims)->count(x->x>0.5, X; dims), A)
+    @testset "sum (Bool, Float64 init)" test_mapreduce_type((X; dims)->sum(x->x>0.5, X; dims, init=0.0), A)
+    @testset "sum (Int, Float64 init)" test_mapreduce_type((X; dims)->sum(X; dims, init=0.0), rand(1:9, 20, 20))
+    @testset "mean" test_mapreduce_type(mean, A)
+    @testset "var" test_mapreduce_type(var, A)
+    @testset "std" test_mapreduce_type(std, A)
+    @testset "extrema" test_mapreduce_type(extrema, A; cmp=Base.:(==))
+    @testset "extrema (real of complex)" test_mapreduce_type((X; dims)->extrema(abs, X; dims), C; cmp=Base.:(==))
+end
+
+# Reducing `Bool`s across partitions broadcasts into a `BitArray`
+@testset "Bool results" begin
+    A = rand(20, 20)
+    @testset "all" test_mapreduce_type((X; dims)->all(x->x>0.05, X; dims), A; cmp=Base.:(==))
+    @testset "any" test_mapreduce_type((X; dims)->any(x->x>0.95, X; dims), A; cmp=Base.:(==))
+    @testset "prod" test_mapreduce_type((X; dims)->prod(x->x>0.05, X; dims), A; cmp=Base.:(==))
+    @testset "maximum" test_mapreduce_type((X; dims)->maximum(x->x>0.95, X; dims), A; cmp=Base.:(==))
+end