Skip to content
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.

Sign up for GitHub

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

Controlling which optimizer rule you're adjusting. #202

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Controlling which optimizer rule you're adjusting. #202

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
346669f
Allowing control of which optimizer type you're updating
murrellb Dec 19, 2024
a655ca0
Merge pull request #2 from FluxML/master
murrellb Dec 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions src/adjust.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,3 +144,33 @@ function _adjust(r::T, nt::NamedTuple) where T <: AbstractRule
end
T(vals...) # relies on having the default constructor
end

###
#adjust with type control
###

adjust!(ℓ::Leaf, oT::Type, eta::Real) = (ℓ.rule = adjust(ℓ.rule, oT, eta); nothing)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can all the Types here be changed to Type{<:AbstractRule}? It would provide some more information for users of the interface and ensure we're not promising support for arbitrary types.

adjust!(ℓ::Leaf, oT::Type; kw...) = (ℓ.rule = adjust(ℓ.rule, oT; kw...); nothing)

adjust(ℓ::Leaf, oT::Type, eta::Real) = Leaf(adjust(ℓ.rule, oT, eta), ℓ.state, ℓ.frozen)
adjust(ℓ::Leaf, oT::Type; kw...) = Leaf(adjust(ℓ.rule, oT; kw...), ℓ.state, ℓ.frozen)

adjust!(tree, oT::Type, eta::Real) = foreach(st -> adjust!(st, oT, eta), tree)
adjust!(tree, oT::Type; kw...) = foreach(st -> adjust!(st, oT; kw...), tree)

adjust(r::AbstractRule, oT::Type, eta::Real) = ifelse(isa(r, oT), adjust(r, eta), r)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These overloads are allocating the output of adjust(r, ...) even if it doesn't end up being used. Unless ifelse brings a demonstrable type stability benefit over a plain conditional, I'd recommend using a ternary instead:

Suggested change
adjust(r::AbstractRule, oT::Type, eta::Real) = ifelse(isa(r, oT), adjust(r, eta), r)
adjust(r::AbstractRule, oT::Type{<:AbstractRule}, eta::Real) = isa(r, oT) ? adjust(r, eta) : r

Copy link
Member

@mcabbott mcabbott Dec 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would guess the branch is resolved when compiling? Possibly it will prefer this style:

Suggested change
adjust(r::AbstractRule, oT::Type, eta::Real) = ifelse(isa(r, oT), adjust(r, eta), r)
adjust(r::AbstractRule, ::Type{T}, eta::Real) where {T<:AbstractRule} = isa(r, T) ? adjust(r, eta) : r

I think you can also make dispatch do this, although whether it's a good idea I don't know... less obvious to read:

adjust(r::C, oT::Type{A}, eta::Real) where {C <: A <: AbstractRule}  where {C<:AbstractRule} = adjust(r, eta)
adjust(r:: AbstractRule, oT::Type{<:AbstractRule}, eta::Real) = r

Edit, less obscure:

Suggested change
adjust(r::AbstractRule, oT::Type, eta::Real) = ifelse(isa(r, oT), adjust(r, eta), r)
adjust(r::T, oT::Type{T}, eta::Real) where {T<:AbstractRule} = adjust(r, eta)
adjust(r:: AbstractRule, oT::Type{<: AbstractRule}, eta::Real) = r

adjust(r::AbstractRule, oT::Type; kw...) = ifelse(isa(r, oT), adjust(r; kw...), r)

adjust!(r::AbstractRule, oT::Type, eta::Real) = ifelse(isa(r, oT), adjust!(r, eta), r)
adjust!(r::AbstractRule, oT::Type; kw...) = ifelse(isa(r, oT), adjust!(r; kw...), r)

function adjust(tree, oT::Type, eta::Real)
t′ = fmap(copy, tree; exclude = maywrite)
adjust!(t′, oT, eta)
t′
end
function adjust(tree, oT::Type; kw...)
t′ = fmap(copy, tree; exclude = maywrite)
adjust!(t′, oT; kw...)
t′
end
3 changes: 3 additions & 0 deletions src/rules.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,6 +799,9 @@ end
adjust(ℓ::OptimiserChain, eta::Real) = OptimiserChain(map(opt -> adjust(opt, eta), ℓ.opts)...)
adjust(ℓ::OptimiserChain; kw...) = OptimiserChain(map(opt -> adjust(opt; kw...), ℓ.opts)...)

adjust(ℓ::OptimiserChain, oT::Type, eta::Real) = OptimiserChain(map(opt -> adjust(opt, oT, eta), ℓ.opts)...)
adjust(ℓ::OptimiserChain, oT::Type; kw...) = OptimiserChain(map(opt -> adjust(opt, oT; kw...), ℓ.opts)...)


"""
AccumGrad(n::Int)
Expand Down
Loading