implement count using mapreduce

[stev47]

This creates the same calling interface for count as for e.g. sum,
namely allowing the dims keyword.
The implementation is also shorter than before without sacrificing
performance.
mapreduce with add_sum may even yield performance benefits through
chunking, though this was not observed in simple tests.

[stev47]

buildbot/tester_win32 fail is unrelated

[stev47]

added compat annotations and tests.
Question: Do we really want count(Any[]) == 0 as count(Any[1]) throws?

[stev47]

buildbot failures are unrelated

[StefanKarpinski]

Would it make sense for @tkf to review since he overhauled the mapreduce machinery recently?

[tkf]

Actually, I tweaked only mapfoldl so far. But this patch LGTM anyway.

[stev47]

this seems ready, freebsd test failure looks unrelated

[tkf]

It'd be nice to use A::AbstractArrayOrBroadcasted instead (with some tests).

[stev47]

thanks, I changed it. You already added a test for Broadcasted and count() in 2f90dde.
Do you still feel we need more?

[stev47]

test fail on linuxaarch64 is worrying but seems unrelated

[tkf]

LGTM. Re-running failed CI: https://build.julialang.org/#/builders/28/builds/10356

[Jutho]

I am not quite sure how this is possible, given the minimal code change of this PR, but git bisect points to this commit as the culprit for the following change in type inference:

using LinearAlgebra, Test
X = [randn(n,n) for n = 1:5];
f(iter, init) = mapreduce(sum, +, iter; init = init);
g(iter, init) = mapreduce(norm, +, iter; init = init);
@inferred f(X, 0.)
@inferred g(X, 0.)

f is always inferred before and after this commit, g fails to infer after this commit. However, type inference acts very strangely. In particular, it is important that you run this on a fresh Julia session. In the slightly more elaborate example:

using LinearAlgebra, Test
X = [randn(n,n) for n = 1:5]

f1(iter, init) = reduce(+, sum(b) for b in iter; init = init);
f2(iter, init) = mapreduce(sum, +, iter; init = init);
g1(iter, init) = reduce(+, norm(b) for b in iter; init = init);
g2(iter, init) = mapreduce(norm, +, iter; init = init);
h1(iter, p, init) = reduce(+, norm(b,p) for b in iter; init = init);
h2(iter, p, init) = mapreduce(b->norm(b,p), +, iter; init = init);

@inferred f1(X, 0.)
@inferred f2(X, 0.)
@inferred g1(X, 0.)
@inferred g2(X, 0.)
@inferred h1(X, 2, 0.)
@inferred h2(X, 2, 0.)

f1 and f2 are always inferred, but out of g1, g2, h1 and h2, the one that is ran first is not inferred, and then the others are.

So I guess the norm implementation also plays some role in this, but I am still trying to understand the actual reason for this happening.

[Jutho]

I expect the lines

elseif p == 0
        return typeof(float(norm(first(itr))))(count(!iszero, itr))

in norm in LinearAlgebra/src/generic.jl have something to do with this, though norm itself with arbitrary p is perfectly inferred and type stable.

[stev47]

It seems to work if you add type assertions:

        return typeof(float(norm(first(itr))))(count(!iszero, itr))::typeof(float(norm(first(itr))))

Also for me only g1 is failing, h1 works fine. Maybe the compiler is giving up on inference because stuff is getting too deeply nested with the mapreduce count being called within a mapreduce operation, I'm not sure though.

[Jutho]

Regarding failing and working, it is the one which is inferred firstly that fails. So if you change the order of the @inferred statements in my code snippet, such that h1 is ran/inferred first in a fresh Julia session, that one fails, but then g1 works fine. That's what I find extremely surprising/disturbing?

[tkf]

It seems to work if you add type assertions:

Maybe we should use convert? There is less guarantee for the type constructor to return the actual type.

That's what I find extremely surprising/disturbing?

Yeah, I agree. Maybe some kind of caching bug? I think it is worth opening a dedicated issue.