-
-
Notifications
You must be signed in to change notification settings - Fork 24
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Type instability in Flux.setup
#162
Comments
|
Values from the cache are used when an object |
We could use |
function _setup(rule, x; cache)
if haskey(cache, x)
T1 = Base._return_type(init, Tuple{typeof(rule), typeof(x)})
T2 = Base._return_type(Leaf, Tuple{typeof(rule), T1})
return cache[x]::T2
end
if isnumeric(x)
ℓ = Leaf(rule, init(rule, x))
# as before... gives julia> @code_warntype test_setup(opt, s)
MethodInstance for test_setup(::Optimisers.Adam, ::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
from test_setup(opt, s) @ Main REPL[5]:1
Arguments
#self#::Core.Const(test_setup)
opt::Optimisers.Adam
s::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}
Locals
state::NamedTuple{(:layers,), <:Tuple{Tuple{NamedTuple, NamedTuple}}}
Body::NamedTuple{(:layers,), <:Tuple{Tuple{NamedTuple, NamedTuple}}}
1 ─ %1 = Flux.setup::Core.Const(Flux.Train.setup)
│ (state = (%1)(opt, s))
└── return state
julia> @code_warntype Optimisers.setup(opt, s)
MethodInstance for Optimisers.setup(::Optimisers.Adam, ::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
from setup(rule::AbstractRule, model) @ Optimisers ~/.julia/dev/Optimisers/src/interface.jl:29
Arguments
#self#::Core.Const(Optimisers.setup)
rule::Optimisers.Adam
model::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}
Locals
tree::NamedTuple{(:layers,), <:Tuple{Tuple{NamedTuple, NamedTuple}}}
cache::IdDict{Any, Any}
msg::String
kwargs::@NamedTuple{}
line::Int64
file::String
id::Symbol
logger::Union{Nothing, Base.CoreLogging.AbstractLogger}
_module::Module
group::Symbol
std_level::Base.CoreLogging.LogLevel
level::Base.CoreLogging.LogLevel
Body::NamedTuple{(:layers,), <:Tuple{Tuple{NamedTuple, NamedTuple}}}
1 ── (cache = Optimisers.IdDict())
│ %2 = (:cache,)::Core.Const((:cache,))
│ %3 = Core.apply_type(Core.NamedTuple, %2)::Core.Const(NamedTuple{(:cache,)})
... |
Looks like the inference path |
In the meantime, would it make sense to add a sentence like |
Yes, probably. Also to emphasise that the way to deal with this is a function barrier. You run setup exactly once & pass its result to something. If you are running it in a tight loop, you are probably doing it wrong. |
Output:
Julia version 1.9.3 and Flux version 0.14.6:
The text was updated successfully, but these errors were encountered: