Add broadcast in the frule for *
#148
Conversation
src/rulesets/Base/base.jl
Outdated
| @@ -103,7 +103,7 @@ | |||
| # product rule requires special care for arguments where `mul` is non-commutative | |||
|
|
|||
| function frule(::typeof(*), x::Number, y::Number, _, Δx, Δy) | |||
| return x * y, Δx * y + x * Δy | |||
| return x * y, @. muladd(Δx, y, x * Δy) | |||
There was a problem hiding this comment.
I find @. surprisingly hard to read.
Took me 10 seconds to realize the muladd was getting broadcast also 😂
So this is the same as:
| return x * y, @. muladd(Δx, y, x * Δy) | |
| return x * y, muladd.(Δx, y, x .* Δy) # optimized version of `Δx .* y .+ x .* Δy |
For interest i did some benchmarking:
Writing out the broadcast seems consistently slightly faster.
Is pretty small though.
But I think for the Number case which this always is, its just going to hit the
muladd(x,y,z) = x*y+z definition.
julia> @btime $dx .* $y .+ $x .* $dy;
49.495 μs (2 allocations: 781.33 KiB)
julia> @btime muladd.($dx, $y, $x .* $dy);
49.468 μs (2 allocations: 781.33 KiB)
So I think might as well do the clearer:
| return x * y, @. muladd(Δx, y, x * Δy) | |
| return x * y, (Δx .* y .+ x .* Δy) |
There was a problem hiding this comment.
It is potentially more accurate on machines with FMA instructions since there are only two roundings, one in muladd/fma the other in *.
There was a problem hiding this comment.
🤷♂
They come out identical on my machine (which has fma) over 100_000 values in Δy and Δx
I think the accuracy only comes into play if y was a matrix.
Which it is not permitted to be here.
There was a problem hiding this comment.
Having muladd certainly won't hurt. Also, people could define their own number type, and overload muladd.
Ref JuliaDiff/ChainRulesCore.jl#93