Flux Extension

Project.toml

@@ -1,18 +1,25 @@
 name = "ParameterSchedulers"
 uuid = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"
 authors = ["Kyle Daruwalla"]
-version = "0.3.7"
+version = "0.3.8"
 
 [deps]
-Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 InfiniteArrays = "4858937d-0d70-526a-a4dd-2d5cb5dd786c"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
+
+[weakdeps]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+
+[extensions]
+ParameterSchedulersFluxExt = "Flux"
 
 [compat]
 Flux = "0.11.2, 0.12, 0.13, 0.14"
 InfiniteArrays = "0.10.4, 0.11, 0.12, 0.13"
 julia = "1.6"
 
 [extras]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [publish]
@@ -21,4 +28,4 @@ theme = "_flux-theme"
 title = "ParameterSchedulers.jl"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Flux"]

ext/ParameterSchedulersFluxExt.jl

@@ -0,0 +1,54 @@
+
+module ParameterSchedulersFluxExt
+
+using ParameterSchedulers, Flux
+
+
+mutable struct Scheduler{T, O, F} <: Flux.Optimise.AbstractOptimiser
+    state::IdDict{Any, Int}
+    schedule::T
+    optim::O
+    update_func::F
+
+    function Scheduler(state::IdDict{Any, Int},
+                       schedule::T,
+                       optim::O,
+                       update_func::F) where {T, O, F}
+        Base.depwarn("""`Scheduler` will transition to explicit Optimisers.jl style
+                        optimizers in the next release""", :Scheduler)
+
+        return new{T, O, F}(state, schedule, optim, update_func)
+    end
+end
+
+ParameterSchedulers.Scheduler(state, schedule, opt, update_func) =
+    Scheduler(state, schedule, opt, update_func)
+
+ParameterSchedulers.Scheduler(schedule, opt, update_func) =
+    Scheduler(IdDict{Any, Int}(), schedule, opt, update_func)
+
+Base.show(io::IO, s::Scheduler) =
+    print(io, "Scheduler(", s.schedule, ", ", s.optim, ")")
+
+function Flux.Optimise.apply!(opt::Scheduler, x, Δ)
+    # get iteration
+    t = get!(opt.state, x, 1)
+    opt.state[x] = t + 1
+
+    # set param
+    opt.update_func(opt.optim, opt.schedule(t))
+
+    # do normal apply
+    return Flux.Optimise.apply!(opt.optim, x, Δ)
+end
+
+for Opt in (Descent, Momentum, Nesterov, RMSProp,
+            Adam, RAdam, AdaMax, OAdam, AdaGrad,
+            AdaDelta, AMSGrad, NAdam, AdaBelief)
+    @eval begin
+        ParameterSchedulers.Scheduler(schedule, opt::$Opt; update_func = (o, s) -> (o.eta = s)) =
+            ParameterSchedulers.Scheduler(schedule, opt, update_func)
+    end
+end
+
+end # module

src/ParameterSchedulers.jl

@@ -1,7 +1,6 @@
 module ParameterSchedulers
 
 using Base.Iterators
-using Flux
 using InfiniteArrays: OneToInf
 
 include("interface.jl")
@@ -19,9 +18,8 @@ export Sequence, Loop, Interpolator, Shifted, ComposedSchedule
 
 include("utils.jl")
 
-# TODO
-# Remove this once Optimisers.jl has support
-# for schedules + optimizers
+# TODO: support Optimisers.jl as an ext
+
 """
     Scheduler{T, O, F}(schedule::AbstractSchedule, opt, update_func)
     Scheduler(schedule, opt; update_func = (o, s) -> (o.eta = s))
@@ -42,6 +40,10 @@ If `opt` does not have a field `eta`, then there is no default behavior
 - `update_func`: a mutating function of with inputs `(optim, param)`
                  that mutates `optim`'s fields based on the current `param` value
 
+!!! warning
+    `Scheduler` functionality depends on `Flux`. Load `Flux` with `using Flux`
+    before using `Scheduler`.
+
 # Examples
 ```julia
 # cosine annealing schedule for Descent
@@ -55,47 +57,11 @@ julia> opt = Scheduler(s, Momentum(); update_func = (o, s) -> o.rho = s)
 Scheduler(CosAnneal{Float64,Int64}(0.1, 0.8, 10), Momentum(0.01, 0.9, IdDict{Any,Any}()))
 ```
 """
-mutable struct Scheduler{T, O, F} <: Flux.Optimise.AbstractOptimiser
-    state::IdDict{Any, Int}
-    schedule::T
-    optim::O
-    update_func::F
-
-    function Scheduler(state::IdDict{Any, Int},
-                       schedule::T,
-                       optim::O,
-                       update_func::F) where {T, O, F}
-        Base.depwarn("""`Scheduler` will transition to explicit Optimisers.jl style
-                        optimizers in the next release""", :Scheduler)
-
-        return new{T, O, F}(state, schedule, optim, update_func)
-    end
-end
-Scheduler(schedule, opt, update_func) =
-    Scheduler(IdDict{Any, Int}(), schedule, opt, update_func)
-
-Base.show(io::IO, s::Scheduler) =
-    print(io, "Scheduler(", s.schedule, ", ", s.optim, ")")
+function Scheduler end
 
-function Flux.Optimise.apply!(opt::Scheduler, x, Δ)
-    # get iteration
-    t = get!(opt.state, x, 1)
-    opt.state[x] = t + 1
-
-    # set param
-    opt.update_func(opt.optim, opt.schedule(t))
-
-    # do normal apply
-    return Flux.Optimise.apply!(opt.optim, x, Δ)
+using PackageExtensionCompat
+function __init__()
+    @require_extensions
 end
 
-for Opt in (Descent, Momentum, Nesterov, RMSProp,
-            Adam, RAdam, AdaMax, OAdam, AdaGrad,
-            AdaDelta, AMSGrad, NAdam, AdaBelief)
-    @eval begin
-        Scheduler(schedule, opt::$Opt; update_func = (o, s) -> (o.eta = s)) =
-            Scheduler(schedule, opt, update_func)
-    end
 end
-
-end