writing type stable wrapper

[tpapp]

Flux.Tracker.forward is not type stable (cf FluxML/Flux.jl#167), so I am wondering what is the best strategy to make the wrapper below type stable:

import Flux

function value_and_gradient(f, x)
    y, back = Flux.Tracker.forward(f, x)
    yval = Flux.Tracker.data(y)
    yval, first(Flux.Tracker.data.(back(1)))
end

Currently eg

julia> @code_warntype value_and_gradient(x -> -sum(abs2, x), zeros(3))
Body::Tuple{Any,Any}
 2 1 ─ %1 = invoke Flux.Tracker.forward(_2::Function, _3::Array{Float64,1})::Tuple{Any,getfield(Flux.Tracker, Symbol("##9#11")){_1} where _1}
   │   %2 = (Base.getfield)(%1, 1)::Any                                                       │╻ indexed_iterate
   │   %3 = (Base.getfield)(%1, 2)::getfield(Flux.Tracker, Symbol("##9#11")){_1} where _1     │╻ indexed_iterate
 3 │   %4 = (Flux.Tracker.data)(%2)::Any                                                      │ 
 4 │   %5 = (%3)(1)::Any                                                                      │ 
   │   %6 = (Base.broadcasted)(Flux.Tracker.data, %5)::Base.Broadcast.Broadcasted{_1,Nothing,typeof(Flux.Tracker.data),_2} where _2 where _1
   │   %7 = (Base.materialize)(%6)::Any                                                       │ 
   │   %8 = (Main.first)(%7)::Any                                                             │ 
   │   %9 = (Core.tuple)(%4, %8)::Tuple{Any,Any}                                              │ 
   └──      return %9                                                                         │ 

[MikeInnes]

The best you can do here is to add a type assert. We can't do better than that with how Tracker is implemented.

[tpapp]

Thanks for the feedback. I am not sure this issue is still relevant (the package that uses this underwent a lot of changes), but I will look into it.

[tpapp]

I ended up adding a wrapper in https://github.com/tpapp/LogDensityProblems.jl/blob/d7a3e8e82f6090efca84668764e767e10a463c2f/src/AD_Flux.jl#L25

My understanding is that issue won't and cannot be fixed, so feel free to close it, or keep it open and add a wontfix label.