diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 4b00a3de7..3ea9d9fa8 100644 --- a/src/array/map-reduce.jl +++ b/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) = diff --git a/test/array/mapreduce.jl b/test/array/mapreduce.jl index c3dac692d..e16263ce1 100644 --- a/test/array/mapreduce.jl +++ b/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