reduce_init undefined for Broadcasted

[chengchingwen]

Sum over a Broadcasted along some dimension result in MethodError for reduce_init.

We have mapreduce_dim defined for both Array and Broadcast which call reduce_init,

julia/base/reducedim.jl

Lines 336 to 337 in 58ffe7e

	_mapreduce_dim(f, op, ::_InitialValue, A::AbstractArrayOrBroadcasted, dims) =
	mapreducedim!(f, op, reducedim_init(f, op, A, dims), A)

but reduce_init only defined for Array.

julia/base/reducedim.jl

Line 108 in 58ffe7e

function reducedim_init(f, op::Union{typeof(+),typeof(add_sum)}, A::AbstractArray, region)

MWE:

bc = Base.broadcasted(+, randn(3,3), 2)
sum(bc) # this work correctly
sum(bc, dims=1) # result in MethodError

versioninfo:

julia> versioninfo()
Julia Version 1.6.1
Commit 6aaedecc44 (2021-04-23 05:59 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: Intel(R) Core(TM) i7-7800X CPU @ 3.50GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake-avx512)
Environment:
  JULIA_NUM_THREADS = 8

[mcabbott]

It seems easy to make this work, but maybe more care is needed to make it fast? With this branch: master...mcabbott:broadcast_init

julia> a2 = rand(100,100); b2 = rand(100,100);
julia> bc2i = Broadcast.instantiate(Broadcast.broadcasted(*, a2, b2));

julia> s1 = @btime sum($bc2i, dims=1);
  18.416 μs (1 allocation: 896 bytes)

julia> s2 = @btime sum($a2 .* $b2, dims=1);
  3.964 μs (3 allocations: 79.08 KiB)

I think scalar reductions were made to work in #31020. Reduction over matrices seems much slower than the example there for vectors (of the same length):

julia> @btime sum($bc2i);
  16.083 μs (0 allocations: 0 bytes)
  53.333 μs (0 allocations: 0 bytes)  # without Broadcast.instantiate

julia> @btime sum($a2 .* $b2);
  4.310 μs (2 allocations: 78.20 KiB)
  4.619 μs (0 allocations: 0 bytes)  # for LinearAlgebra.dot

julia> a = rand(10_000); b = rand(10_000);

julia> @btime sum(Broadcast.instantiate(Broadcast.broadcasted(*, $a, $b)));
  2.458 μs (0 allocations: 0 bytes)
  9.291 μs (0 allocations: 0 bytes)  # without Broadcast.instantiate

[chengchingwen]

@mcabbott It seems it is faster now for reducing with broadcasted. I get this with julia 1.8.2:

julia> s1 = @btime sum($bc2i, dims=1);
  2.767 μs (5 allocations: 1.03 KiB)

julia> s2 = @btime sum($a2 .* $b2, dims=1);
  5.888 μs (3 allocations: 79.05 KiB)

[chengchingwen]

It seems Broadcasted doesn't support linear indexing. The reduction sometimes fail because of the bound check.

[N5N3]

Base,has_fast_linear_indexing is basically broken for Broadcasted. (see #43736)

Use BroadcastedArray from LazyArrays.jl instead would fix it and resolve the original concern here.

[chengchingwen]

Interesting, didn't know that exist! I wonder why we can't make this work with Broadcasted? #43736 seems to be opened for a while.

btw I notice you also have an ExBroadcast.jl, do you have any plan about it?

[N5N3]

We can make this work, #43736 is a bug fix but it's still not reviewed.
Since this error is easy to meet so I'm asssuming this usage is not that popular. (The biggest block might be that it hasn't been documented?)

Even we get #43736 merged, I'm affaid that its usage is still too limited, as there's no "reliable" way to obtain the correct eltype of a Broadcasted before collecting it, so it might be hard to to make this efficient in all cases.
BroadcastedArray might be a better choice if you only work with inferable cases. (And the syntax should be more easy to use.)

As for ExBroadcast.jl, I don't think type piracy is a good choice here.