ProjectTo(::Vector{Vector{Float64}}) type unstable

[willtebbutt]

MWE:

using ChainRulesCore: ProjectTo
x = [randn(2), randn(3)]
@code_warntype ProjectTo(x)

It appears to be the case that this happens because

@code_warntype map(ProjectTo, x)

isn't type-stable.

I'm not sure what a fix looks like here unfortunately.

[mcabbott]

I wonder if making ProjectTo look inside containers was a mistake. Maybe it should be content to deal with the outer array here, and assume that whatever rule made the inner arrays knew what it was doing (or called ProjectTo itself)?

That said I don't quite see where this instability is from, why doesn't map infer if the elements infer?

julia> @code_warntype map(ProjectTo, x)
MethodInstance for map(::Type{ProjectTo}, ::Vector{Vector{Float64}})
  from map(f, A::AbstractArray) in Base at abstractarray.jl:2933
Arguments
  #self#::Core.Const(map)
  f::Type{ProjectTo}
  A::Vector{Vector{Float64}}
Body::Union{Vector{ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}}, Vector{Any}}
...

julia> @code_warntype ProjectTo(x[1])
MethodInstance for ProjectTo(::Vector{Float64})
  from ProjectTo(x::AbstractArray{T}) where T<:Number in ChainRulesCore at /Users/me/.julia/packages/ChainRulesCore/ctmSK/src/projection.jl:200
Static Parameters
  T = Float64
Arguments
  #self#::Type{ProjectTo}
  x::Vector{Float64}
Body::ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}

[willtebbutt]

That said I don't quite see where this instability is from, why doesn't map infer if the elements infer?

I was wondering that. Usually map infers things nicely 🤷