Slowdown caused by broadcasting macro?

[mcabbott]

I expected g1 and g2 here to be identical:

using Zygote, BenchmarkTools
function g1(x)
    y = reshape(x,1,1,:)
    dropdims(sum(x .* x' .* y; dims=(2,3)), dims=(2,3))
end
function g2(x)
    y = reshape(x,1,1,:)
    dropdims(sum(@.(x * x' * y); dims=(2,3)), dims=(2,3))
end
@btime gradient(sum∘g1, $(rand(20))); #  33.147 μs (74 allocations: 266.53 KiB)
@btime gradient(sum∘g2, $(rand(20))); # 962.748 μs (56172 allocations: 2.39 MiB)

Kristoffer Carlsson points me to JuliaLang/julia#29120 as a possible source of this. It looks like I can also trigger it like this:

function g4(x)
    y = reshape(x,1,1,:)
    dropdims(sum(.*(x, x', y); dims=(2,3)), dims=(2,3))
end
@btime gradient(sum∘g4, $(rand(20))); # 960.827 μs (56172 allocations: 2.39 MiB)

[mbauman]

Yes, this is a dup of JuliaLang/julia#29120 and/or a potential Zygote improvement. Explicitly adding parens here fixes it:

sum(@.((x * x') * y)

I think a potential Zygote-level fix would be to add a 3-arg gradient override for broadcasted *.

Zygote.jl/src/lib/broadcast.jl

Lines 79 to 80 in 9ff9ef4

	@adjoint broadcasted(::typeof(*), x::Numeric, y::Numeric) = x.*y,
	z̄ -> (nothing, unbroadcast(x, z̄ .* conj.(y)), unbroadcast(y, z̄ .* conj.(x)))

[cossio]

How come adjoint is not being broadcasted in @.(x * x' * y)?

[cossio]

I think a potential Zygote-level fix would be to add a 3-arg gradient override for broadcasted *.

But would this cover products with more terms, @. A * B * C * D * E ....?
Perhaps we need an n-arg override,

@adjoint broadcasted(::typeof(*), x::Numeric...) = ...