workaround for #35800, inference issue in mapreduce (#37105)

JuliaLang · Aug 18, 2020 · 903542b · 903542b
1 parent bf8a8e8
commit 903542b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/base/reduce.jl b/base/reduce.jl
@@ -956,7 +956,15 @@ julia> count([true, false, true, true])
 """
 count(itr) = count(identity, itr)
 
-count(f, itr) = mapreduce(_bool(f), add_sum, itr, init=0)
+count(f, itr) = _simple_count(f, itr)
+
+function _simple_count(pred, itr)
+    n = 0
+    for x in itr
+        n += pred(x)::Bool
+    end
+    return n
+end
 
 function count(::typeof(identity), x::Array{Bool})
     n = 0

diff --git a/base/reducedim.jl b/base/reducedim.jl
@@ -387,7 +387,10 @@ julia> count(<=(2), A, dims=2)
 ```
 """
 count(A::AbstractArrayOrBroadcasted; dims=:) = count(identity, A, dims=dims)
-count(f, A::AbstractArrayOrBroadcasted; dims=:) = mapreduce(_bool(f), add_sum, A, dims=dims, init=0)
+count(f, A::AbstractArrayOrBroadcasted; dims=:) = _count(f, A, dims)
+
+_count(f, A::AbstractArrayOrBroadcasted, dims::Colon) = _simple_count(f, A)
+_count(f, A::AbstractArrayOrBroadcasted, dims) = mapreduce(_bool(f), add_sum, A, dims=dims, init=0)
 
 """
     count!([f=identity,] r, A)

diff --git a/test/reducedim.jl b/test/reducedim.jl
@@ -4,6 +4,10 @@ using Random
 
 # main tests
 
+# issue #35800
+# tested very early since it can be state-dependent
+@test @inferred(mapreduce(x->count(!iszero,x), +, [rand(1)]; init = 0.)) == 1.0
+
 function safe_mapslices(op, A, region)
     newregion = intersect(region, 1:ndims(A))
     return isempty(newregion) ? A : mapslices(op, A, dims = newregion)