Flux.jl
Flux.jl copied to clipboard
cpu() type stability
This is coming from a discussion I started on discourse. It looks like there is a type stability issue when using the cpu()
function where the culprit seems to be Flux.fmap()
. I don't think this is having a major impact on the performance of my model, but I wanted to document the issue just so you were aware of it.
Example of the issue:
julia> using Flux #v0.12.9
julia> @code_warntype cpu(rand(10))
MethodInstance for Flux.cpu(::Vector{Float64})
from cpu(x) in Flux at ~/.julia/packages/Flux/qAdFM/src/functor.jl:146
Arguments
#self#::Core.Const(Flux.cpu)
x::Vector{Float64}
Locals
#136::Flux.var"#136#137"
Body::Any
1 ─ (#136 = %new(Flux.:(var"#136#137")))
│ %2 = #136::Core.Const(Flux.var"#136#137"())
│ %3 = Flux.fmap(%2, x)::Any
└── return %3
We could add an overload for [cg]pu
that skips the fmap
path for arrays, but until it's demonstrated to cause actual performance issues I'm not sure that'll be much of a priority. 99% of the time there should be no need to call those functions on a hot code path, and the the VAE model zoo example you linked should be updated to reflect that (PRs welcome).
Sounds fine to me. If I end up refactoring my VAE to not use cpu()
in the main loop, I'll do a pull request for the model zoo.