rename everything

Author: mcabbott

Project.toml

@@ -12,7 +12,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 ChainRulesCore = "1"
-Functors = "0.2.7"
+Functors = "0.2.8"
 julia = "1.6"
 
 [extras]

docs/src/api.md

@@ -46,6 +46,7 @@ Such restrictions are also obeyed by this function for flattening a model:
 
 ```@docs
 Optimisers.destructure
+Optimisers.Restructure
 ```
 
 ## Rule Definition

src/Optimisers.jl

@@ -4,8 +4,9 @@ using Functors: functor, fmap, isleaf
 using LinearAlgebra
 
 include("interface.jl")
+
 include("destructure.jl")
-export destructure
+export destructure, total, total2
 
 include("rules.jl")
 export Descent, ADAM, Momentum, Nesterov, RMSProp,

src/destructure.jl

@@ -19,20 +19,42 @@ julia> re([10,20,30])
 ```
 """
 function destructure(x)
-  flat, off, len = alpha(x)
-  flat, Restucture(x, off, len)
+  flat, off, len = _flatten(x)
+  flat, Restructure(x, off, len)
 end
 
-struct Restucture{T,S}
+"""
+    Restructure(Model, ..., length)
+
+This is what [`destructure`](@ref) returns, and `re(p)` will re-build the model with
+new parameters from vector `p`. If the model is callable, then `re(x, p)` .
+
+# Example
+```julia
+julia> using Flux, Optimisers
+
+julia> _, re = destructure(Dense([1 2; 3 4], [0, 0], sigmoid))
+([1, 3, 2, 4, 0, 0], Restructure(Dense, ..., 6))
+
+julia> m = re(-4:1)
+Dense(2, 2, σ)      # 6 parameters
+
+julia> m([0.2, 0.3]) ≈ re([0.2, 0.3], -4:1)
+true
+```
+"""
+struct Restructure{T,S}
   model::T
   offsets::S
   length::Int
 end
-(re::Restucture)(flat) = beta(re.model, re.offsets, flat; len = re.length)
-Base.show(io::IO, re::Restucture{T}) where T = print(io, "Restructure(", T.name.name, ", ..., ", re.length, ")")
+(re::Restructure)(flat::AbstractVector) = _rebuild(re.model, re.offsets, flat; len = re.length)
+(re::Restructure)(x, flat::AbstractVector) = re(flat)(x)
+Base.show(io::IO, re::Restructure{T}) where T = print(io, "Restructure(", T.name.name, ", ..., ", re.length, ")")
+Base.length(re::Restructure) = re.length
 
 # This flattens a model, and returns a web of offsets for later use:
-function alpha(x)
+function _flatten(x)
   isnumeric(x) && return vcat(vec(x)), 0, length(x)  # trivial case
   arrays = AbstractVector[]
   len = Ref(0)
@@ -45,14 +67,14 @@ function alpha(x)
   reduce(vcat, arrays), off, len[]
 end
 
-function ChainRulesCore.rrule(::typeof(alpha), x)
-  flat, off, len = alpha(x)
-  alpha_back((dflat, _)) = (NoT, beta(x, off, dflat; walk = _Tangent_biwalk, prune = NoT, len))
-  (flat, off, len), alpha_back
+function ChainRulesCore.rrule(::typeof(_flatten), x)
+  flat, off, len = _flatten(x)
+  _flatten_back((dflat, _)) = (NoT, _rebuild(x, off, dflat; walk = _Tangent_biwalk, prune = NoT, len))
+  (flat, off, len), _flatten_back
 end
 
 # This reconstructs either a model like x, or a gradient for it:
-function beta(x, off, flat::AbstractVector; len, walk = _trainable_biwalk, kw...)
+function _rebuild(x, off, flat::AbstractVector; len, walk = _trainable_biwalk, kw...)
   len == length(flat) || error("wrong length")
   fmap(x, off; exclude = isnumeric, walk, kw...) do y, o
     _getat(y, o, flat)
@@ -66,40 +88,41 @@ _getat(y::AbstractArray, o::Int, flat::AbstractVector) =
 function _trainable_biwalk(f, x, aux)
   ch, re = functor(typeof(x), x)
   au, _ = functor(typeof(x), aux)
-  trainmap(f, ch, _trainable(x), au) |> re
+  _trainmap(f, ch, _trainable(x), au) |> re
 end
 
-function trainmap(f, ch, tr, aux)
-  map(ch, tr, aux) do c, t, a
+function _trainmap(f, ch, tr, aux)
+  map(ch, tr, aux) do c, t, a  # isnothing(t) indicates non-trainable field, safe given isnumeric(c)??
     isnothing(t) ? c : f(t, a)
   end
 end
 
-function _Tangent_biwalk(f, x, aux)  # use with prune = true
+function _Tangent_biwalk(f, x, aux)  # use with prune = NoT
   ch, re = functor(typeof(x), x)
   au, _ = functor(typeof(x), aux)
-  y = trainmap(f, ch, _trainable(x), au)
+  y = _trainmap(f, ch, _trainable(x), au)
   y isa Tuple{} && return NoT
   Tangent{typeof(x), typeof(y)}(y)
 end
 
-function ChainRulesCore.rrule(::typeof(beta), x, off, flat; len)
+function ChainRulesCore.rrule(::typeof(_rebuild), x, off, flat; len)
   dflat = map!(zero, similar(flat, float(eltype(flat))), flat)
-  beta_back(dx) = (NoT, NoT, NoT, gamma!(x, dx, off, dflat))
-  beta(x, off, flat; len), beta_back
+  _rebuild_back(dx) = (NoT, NoT, NoT, _accumulate!(x, dx, off, dflat))
+  _rebuild(x, off, flat; len), _rebuild_back
 end
 
 # This is the gradient of model reconstruction, accumulating duplicates:
-function gamma!(x, dx, off, flat::AbstractVector)
+function _accumulate!(x, dx, off, flat::AbstractVector)
   x′, _ = functor(typeof(x), x)
   dx′, _ = functor(typeof(x), dx)
   off′, _ = functor(typeof(x), off)
-  foreach((xᵢ, dxᵢ, oᵢ) -> gamma!(xᵢ, dxᵢ, oᵢ, flat), x′, dx′, off′)
+  foreach((xᵢ, dxᵢ, oᵢ) -> _accumulate!(xᵢ, dxᵢ, oᵢ, flat), x′, dx′, off′)
   flat
 end
-function gamma!(x, dx, off::Integer, flat::AbstractVector)
-  @views flat[off .+ (1:length(x))] .+= dx  # must visit all tied nodes, hence no fmap.
+function _accumulate!(x, dx, off::Integer, flat::AbstractVector)
+  @views flat[off .+ (1:length(x))] .+= dx  # must visit all tied nodes
   flat
 end
-gamma!(x, dx::Zero, off, flat::AbstractVector) = nothing
-gamma!(x, dx::Zero, off::Integer, flat::AbstractVector) = nothing  # ambiguity
+_accumulate!(x, dx::Zero, off, flat::AbstractVector) = nothing
+_accumulate!(x, dx::Zero, off::Integer, flat::AbstractVector) = nothing  # ambiguity
+