From ffd074db5e2ffc00f0ef93a755f724040442f139 Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Wed, 7 Oct 2026 15:52:55 +0300 Subject: [PATCH 1/4] DArray: Keep var/std in the input precision with dims OnlineStats' `value(::Variance)` multiplies by a `Float64` `NaN` literal when it has no finite mean, so it is inferred as `Union{S,Float64}` for any other precision `S`. Convert the value back to `S`, so that a reduction over `dims` can carry a concrete eltype. Nothing visible changes yet: the result eltype is `Any` until the reduction infers its eltype through `f`. --- src/array/map-reduce.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 4b00a3de7..2c23d5935 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -190,7 +190,11 @@ function _onlinestats_mapreduce(f, x::DArray{T}, stat; dims=nothing) where 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) = From 422d3b51c10cc4f823007ebf69e94b4e2c410645 Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Wed, 7 Oct 2026 15:52:55 +0300 Subject: [PATCH 2/4] DArray: Infer the mapreduce result eltype through f `MapReduce` took its eltype from `op_outer(T, T)`, where `T` is the input eltype. But `op_outer` combines what `f` returns, not the input elements. When `f` changes the type, the result was labelled with the wrong eltype and `collect` converted to it. For a `Float64` matrix `A` and a `ComplexF64` matrix `C`: sum(x->x>0.5, A) # 13.0 instead of 13 sum(abs2, C) # ComplexF64 instead of Float64 sum(x->complex(x, 1.0), A) # InexactError When `op_outer` has no method for `T` at all, the eltype fell back to `Any`: `mean`/`var`/`std` with `dims` (`merge` of two numbers), and `maximum(abs, C; dims)` (`max` of two complex numbers). `_onlinestats_mapreduce` now takes the statistic as a `::Type{S}`. Julia 1.11 does not infer what the closure returns when it captures the type from an untyped argument. An `init` takes part in the inference, so `sum(f, A; init=0.0)` keeps the type of the `init` when that is the wider one. The existing tests compare with `isapprox`, which accepts a `Matrix{Any}` and a `13.0`, so the new ones also compare the types. --- src/array/map-reduce.jl | 15 +++++++++++---- test/array/mapreduce.jl | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 2c23d5935..fa193bdc8 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 = Base._return_type(f, Tuple{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 @@ -184,9 +191,9 @@ end _extrema_outer(x::Extrema, y::Extrema) = Extrema(min(x.min, y.min), max(x.max, y.max)) -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) diff --git a/test/array/mapreduce.jl b/test/array/mapreduce.jl index c3dac692d..070e9775b 100644 --- a/test/array/mapreduce.jl +++ b/test/array/mapreduce.jl @@ -51,3 +51,36 @@ 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) +end From 5c07e421497dd2fdcb6c81fd4913e947f14eedd8 Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Wed, 7 Oct 2026 15:52:55 +0300 Subject: [PATCH 3/4] DArray: Do not wrap a BitArray partial reduction The tree reduction broadcasts `op_outer` over two partial results. When those hold `Bool`s the broadcast returns a `BitArray`, which `to_array` did not take for an array, so it wrapped it in a 1-element `Array`. For a matrix split into several partitions, `all(f, A; dims=1)` returned a matrix of `BitMatrix`, and `any`, `prod` and `maximum` with a `Bool` valued `f` threw. --- src/array/map-reduce.jl | 2 +- test/array/mapreduce.jl | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index fa193bdc8..47af6f4d0 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -90,7 +90,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) diff --git a/test/array/mapreduce.jl b/test/array/mapreduce.jl index 070e9775b..9e3dfd5b2 100644 --- a/test/array/mapreduce.jl +++ b/test/array/mapreduce.jl @@ -84,3 +84,12 @@ end @testset "var" test_mapreduce_type(var, A) @testset "std" test_mapreduce_type(std, A) 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 From 2720bf3e38f3877b994248e88d0551c69e737d0c Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Wed, 7 Oct 2026 15:52:55 +0300 Subject: [PATCH 4/4] DArray: Give extrema with dims a concrete eltype `extrema` reduces `Extrema` values, which `_extrema_inner` builds from what `f` returns. `_extrema_outer` has no method for the bare values, so `extrema(A; dims=1)` had the eltype `Tuple{Any,Any}`. Tell `MapReduce` what the partial results hold. --- src/array/map-reduce.jl | 6 +++++- test/array/mapreduce.jl | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 47af6f4d0..3ea9d9fa8 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -55,7 +55,7 @@ struct MapReduce{T,N} <: ArrayOp{T,N} end function MapReduce(f, op_inner, op_outer, input::DArray{T,N}, dims, init) where {T,N} # `op_outer` combines the values `f` produces, not the input elements - T_f = Base._return_type(f, Tuple{T}) + 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}) @@ -70,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) @@ -190,6 +192,8 @@ 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}, ::Type{S}; dims=nothing) where {T,S} _f(x) = fit!(S(T), f(x)) diff --git a/test/array/mapreduce.jl b/test/array/mapreduce.jl index 9e3dfd5b2..e16263ce1 100644 --- a/test/array/mapreduce.jl +++ b/test/array/mapreduce.jl @@ -83,6 +83,8 @@ end @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`