ChainRulesCore.jl icon indicating copy to clipboard operation
ChainRulesCore.jl copied to clipboard

Improve type stability of `ProjectTo(non-numeric array)`.

Open ToucheSir opened this issue 3 years ago • 4 comments

Use broadcasting instead of map for non-numeric array projectors. This improves type stability when elements hit ProjectTo(::Any).

I found this while trying to use https://github.com/JuliaDiff/ChainRules.jl/blob/v1.35.3/src/rulesets/Base/mapreduce.jl#L425 with a vector of callables. It's unfortunate that inference can't figure out the map call, but at least broadcasting works.

ToucheSir avatar Jun 21 '22 03:06 ToucheSir

Codecov Report

Merging #557 (543002d) into main (117910b) will not change coverage. The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #557   +/-   ##
=======================================
  Coverage   93.02%   93.02%           
=======================================
  Files          15       15           
  Lines         860      860           
=======================================
  Hits          800      800           
  Misses         60       60           
Impacted Files Coverage Δ
src/projection.jl 97.30% <100.00%> (ø)

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 117910b...543002d. Read the comment docs.

codecov-commenter avatar Jun 21 '22 03:06 codecov-commenter

Are there any relevant weird cases where these may differ?

An irrelevant one is this (because adjoint vectors have their own path):

julia> map(sqrt, [1,2,3]')
1×3 adjoint(::Vector{Float64}) with eltype Float64:
 1.0  1.41421  1.73205

julia> sqrt.([1,2,3]')
1×3 Matrix{Float64}:
 1.0  1.41421  1.73205

mcabbott avatar Jun 25 '22 21:06 mcabbott

I thought SparseArrays might, but it appears map and broadcast behave similarly for them. One concern I did find though while writing up https://github.com/FluxML/Zygote.jl/pull/1248#issuecomment-1166411374 is that this PR appears to add a few seconds of latency to the MWE there. Will have to do a proper investigation to confirm that.

ToucheSir avatar Jun 26 '22 15:06 ToucheSir

OK, that's interesting. Nothing about ProjectTo was closely looked at for latency, sadly.

mcabbott avatar Jun 27 '22 03:06 mcabbott