fixup

Author: mcabbott

docs/src/models/basics.md

@@ -216,7 +216,7 @@ m(5) # => 26
 Flux provides a set of helpers for custom layers, which you can enable by calling
 
 ```julia
-Flux.@functor Affine
+Flux.@layer Affine
 ```
 
 This enables a useful extra set of functionality for our `Affine` layer, such as [collecting its parameters](../training/optimisers.md) or [moving it to the GPU](../gpu.md).

src/Flux.jl

@@ -7,6 +7,7 @@ using MacroTools: @forward
 
 @reexport using NNlib
 using MLUtils
+const stack = MLUtils.stack  # now exported by Base
 import Optimisers: Optimisers, trainable, destructure  # before v0.13, Flux owned these functions
 
 using Zygote, ChainRulesCore

src/layers/basic.jl

@@ -338,7 +338,7 @@ struct SkipConnection{T,F}
   connection::F  #user can pass arbitrary connections here, such as (a,b) -> a + b
 end
 
-@layer SkipConnection  # should this be expand?
+@layer :expand SkipConnection
 
 function (skip::SkipConnection)(input)
   skip.connection(skip.layers(input), input)

src/layers/macro.jl

@@ -3,25 +3,50 @@
     @layer Dense
     @layer :expand Chain
     @layer BatchNorm trainable=(β,γ)
-    @layer Struct functor=(α,β) trainable=(β,)
+    @layer Struct children=(α,β) trainable=(β,)
     
 This macro replaces most uses of `@functor` in Flux 0.14. Its basic purpose is the same:
 When you define a new layer, this tells Flux to explore inside it
 to see the parameters it trains, and also to move them to the GPU, change precision, etc.
+Like `@functor`, this assumes your struct has the default constructor, to enable re-building.
 
 Some "keywords" allow control of the recursion:
 * If some fields look like parameters but should not be trained,
-  then `Optimisers.trainable` lets you specify fields to include, and ignore the rest.
-* We can likewise add restructions to `Functors.functor`, but not yet written.
-* In fact you can provide an arbitrary keyword with this syntax, and it will
-  overload this function alla `trainable`... that might be a terrible idea.
+  then `trainable` lets you specify fields to include, and ignore the rest.
+* You can likewise add restructions to Functors's `children` (although this is seldom a good idea).
 
 It also handles overloads of `show` for pretty printing.
 * By default, it adds methods to 3-arg `Base.show` to treat your layer much like `Dense` or `Conv`.
 * If your layer is a container, more like `Chain` or `Parallel`, then `:expand` makes `show` unfold its contents.
-* To disable all `show` overloads, maybe we want a `:ignore` option too.
+* To disable all `show` overloads, there is an `:ignore` option too.
 
 (You probably still want to define 2-arg `show(io::IO, x::Layer)`, the macro does not touch this.)
+
+Note that re-running the macro with different options does not overwrite all methods, you will need to restart.
+  
+# Example
+```jldoctest
+julia> struct Trio; a; b; c end
+
+julia> tri = Trio(Dense([1.1 2.2],), Dense([3.3;;], false), Dropout(0.4))
+Trio(Dense(1 => 1, tanh), Dense(1 => 1; bias=false), Dropout(0.4))
+
+julia> Flux.destructure(tri)  # parameters not visible to Flux
+(Bool[], Restructure(Trio, ..., 0))
+
+julia> Flux.@layer :expand Trio
+
+julia> Flux.destructure(tri)  # now gpu, train!, etc will see inside too
+([1.1, 2.2, 0.0, 3.3], Restructure(Trio, ..., 4))
+
+julia> tri
+Trio(
+  Dense(2 => 1),                        # 3 parameters
+  Dense(1 => 1; bias=false),            # 1 parameters
+  Dropout(0.4),
+)                   # Total: 3 arrays, 4 parameters, 224 bytes.
+```
+
 """
 macro layer(exs...)
   out = quote end
@@ -40,10 +65,10 @@ macro layer(exs...)
   end
 
   # This function exists only for depwarns when you use @functor directly
-  push!(out.args, :(Flux._check_new_macro(::$(esc(type))) = nothing)) # scope is weird ?? can't use $ on func name?
+  push!(out.args, :(Flux._check_new_macro(::$(esc(type))) = nothing))
 
-  i = findfirst(ex -> Meta.isexpr(ex, :(=)) && ex.args[1] == :functor, rest)
-  if isnothing(i)
+  i = findfirst(ex -> Meta.isexpr(ex, :(=)) && ex.args[1] == :children, rest)
+  if isnothing(i)  # then default like @functor Layer
     push!(out.args, _macro_functor(esc(type)))
   else
     push!(out.args, _macro_functor(esc(type), rest[i].args[2]))
@@ -52,54 +77,70 @@ macro layer(exs...)
     j == i && continue
     ex = rest[j]
     Meta.isexpr(ex, :(=)) || error("expected keyword = fields")
-    if ex.args[1] == :trainable
-      push!(out.args, _macro_trainable(type, trainable, ex.args[2]))  # pass the function "trainable" not the symbol
+
+    name = if ex.args[1] == :trainable
+      :(Optimisers.trainable)
     else
-      error()
-      # @warn "defining a method for $(ex.args[1]) in your scope"  # ??
-      # push!(out.args, _macro_trainable(type, esc(ex.args[1]), ex.args[2]))
+      @warn "trying to define a method for `$(ex.args[1])` in your scope... this is experimental" maxlog=1
+      esc(ex.args[1])
     end
+    push!(out.args, _macro_trainable(esc(type), name, ex.args[2]))
   end
 
   out
 end
 
-# Temporary depwarn function:
+# Temporary depwarn function, called within `params`, is also called by `show`.
 
 function _check_new_macro(x::T) where T
   Functors.isleaf(x) && return
-  @warn "you used @functor for this type, but should now use @layer" T maxlog=1 _id=hash(T)
+  @warn "This type should now use Flux.@layer instead of @functor" T maxlog=1 _id=hash(T)
 end
 _check_new_macro(::Tuple) = nothing  # defined by Functors.jl, not by users
 _check_new_macro(::NamedTuple) = nothing
-_check_new_macro(::Transpose) = nothing
-_check_new_macro(::Adjoint) = nothing
+_check_new_macro(::AbstractArray) = nothing
 _check_new_macro(::Ref) = nothing
 
 # @layer's code for Functors & Adapt
 # Unlike @functor, _default_functor doesn't need to eval anything
 
 function _macro_functor(type)
   quote
-    Functors.functor(::Type{T}, x) where {T<:$type} = _default_functor(T, x)
-    Adapt.adapt_structure(to, layer::$type) = fmap(adapt(to), layer)
+    Functors.functor(::Type{T}, x) where {T<:$type} = $_default_functor(T, x)
+    Adapt.adapt_structure(to, layer::$type) = $fmap($adapt(to), layer)
   end
 end
 
 function _macro_functor(type, fields)
-  error("the equivalent of @functor Layer (:x,) isn't written yet, sorry")
+  Meta.isexpr(fields, :tuple) || error("expected a tuple of field names")
+  symbols = Tuple(map(_noquotenode, fields.args))
+  quote
+    Functors.functor(::Type{T}, x) where {T<:$type} = $_custom_functor(T, x, Val($symbols))
+    Adapt.adapt_structure(to, layer::$type) = $fmap($adapt(to), layer)
+  end
 end
+_macro_functor(type, field::Union{Symbol,QuoteNode}) = _macro_functor(type, :(($field,)))  # lets you forget a comma
 
 function _default_functor(::Type{T}, x) where {T}
    if @generated
      F = fieldnames(T)
      args = map(sy -> :(getfield(x, $(QuoteNode(sy)))), F)
-     C = Base.typename(T).name  # constructor
+     C = Base.typename(T).wrapper  # constructor
      recon = VERSION > v"1.9-" ? :(Splat($C)) : :(Base.splat($C))
      :((NamedTuple{$F}(($(args...),)), $recon))
    else
      # Getting this parameterless type takes about 2μs, every time:
-     namedtuple(x), Base.splat(Base.typename(T).wrapper)
+     spl = VERSION > v"1.9-" ? Splat : Base.splat
+     namedtuple(x), spl(Base.typename(T).wrapper)
+   end
+end
+
+function _custom_functor(::Type{T}, x, ::Val{which}) where {T,which}
+   if false
+     # TODO write the @generated version
+   else
+     remake(nt) = Base.typename(T).wrapper(map(f -> f in which ? getfield(nt, f) : getfield(x, f), fieldnames(T))...)
+     NamedTuple{which}(map(s -> getfield(x, s), which)), remake
    end
 end
 
@@ -117,61 +158,12 @@ function _macro_trainable(type, fun, fields)
   quoted = map(QuoteNode, symbols)
   gets = [:(getfield(x, $f)) for f in quoted]
   quote
-    # $fun(x::$type) = NamedTuple{$names}(($(gets...),))
-    Flux.trainable(x::$type) = NamedTuple{$symbols}(($(gets...),))  # ?? scope is weird
+    $fun(x::$type) = NamedTuple{$symbols}(($(gets...),))
+    # Flux.trainable(x::$type) = NamedTuple{$symbols}(($(gets...),))  # ?? scope is weird
   end
 end
 _macro_trainable(type, fun, field::Union{Symbol,QuoteNode}) = _macro_trainable(type, fun, :(($field,)))  # lets you forget a comma
 
 _noquotenode(s::Symbol) = s
 _noquotenode(q::QuoteNode) = q.value  # lets you write trainable=(:x,:y) instead of (x,y)
 _noquotenode(ex) = error("expected a symbol, got $ex")
-
-
-
-
-
-
-# @big_show Chain
-# @big_show Parallel
-# @big_show SkipConnection
-# @big_show Recur
-# @big_show Maxout
-
-
-
-
-"""
-    @big_show MyContainer
-
-This macro lets you opt-in to Flux's fancy printing.
-
-When `model::MyContainer` is returned at the REPL it will be treated like `Chain`,
-and the printing routine will recursively unfold its children.
-This is triggered by adding a method to 3-arg `Base.show(io::IO, ::MIME"text/plain", l::MyContainer)`.
-
-Custom layers which do not contain other layers (more like `Dense` than like `Chain`)
-need not call this, and should simply define 2-arg `Base.show(io::IO, l::MyLayer)`.
-
-# Example
-```jldoctest
-julia> struct Trio{A,B,C}; a::A; b::B; c::C end
-
-julia> Flux.@functor Trio
-
-julia> Flux.@big_show Trio
-
-julia> tri = Trio(Dense(10=>5,tanh), Dense(5=>2), softmax)
-Trio(
-  Dense(10 => 5, tanh),                 # 55 parameters
-  Dense(5 => 2),                        # 12 parameters
-  NNlib.softmax,
-)                   # Total: 4 arrays, 67 parameters, 492 bytes.
-```
-
-Note that there is no automatic method for 2-arg `show`, and thus
-something like `(tri, tri)` will print all the type parameters. 
-
-However, `Chain(tri, tri)` will always use Flux's recursive printing,
-even without using this macro: `Chain` is the entry point.
-"""

src/layers/normalise.jl

@@ -178,6 +178,8 @@ end
 testmode!(m::AlphaDropout, mode=true) =
   (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
 
+Base.show(io::IO, d::AlphaDropout) = print(io, "AlphaDropout(", d.p, ")")
+
 """
     LayerNorm(size..., λ=identity; affine=true, ϵ=1fe-5)
 

src/layers/recurrent.jl

@@ -136,7 +136,6 @@ function (m::Recur)(x)
 end
 
 @layer :expand Recur trainable=(cell,)
-# trainable(a::Recur) = (; cell = a.cell)
 
 Base.show(io::IO, m::Recur) = print(io, "Recur(", m.cell, ")")
 

src/layers/show.jl

@@ -1,28 +1,30 @@
 @nospecialize  # just for this file, for startup time
 
-# This is called by @layer, on layers which should be treated like Chain
+# This is called by @layer :expand, on layers which should be treated like Chain, and returns an expression:
 function _macro_big_show(ex)
-    quote
-        # Entry point:
-        function Base.show(io::IO, m::MIME"text/plain", x::$ex)
-            if get(io, :typeinfo, nothing) === nothing  # e.g. top level in REPL
-              _big_show(io, x)
-            elseif !get(io, :compact, false)  # e.g. printed inside a Vector, but not a Matrix
-              _layer_show(io, x)
-            else
-              show(io, x)
-            end
-        end
-        
-        # Don't show Chain(Tuple(...)), always splat that:
-        _show_children(x::$ex) = _flat_children(x)
+  quote
+    # Entry point:
+    function Base.show(io::IO, m::MIME"text/plain", x::$ex)
+      if get(io, :typeinfo, nothing) === nothing  # e.g. top level in REPL
+        _big_show(io, x)
+      elseif !get(io, :compact, false)  # e.g. printed inside a Vector, but not a Matrix
+        _layer_show(io, x)
+      else
+        show(io, x)
+      end
     end
+    
+    # Don't show Chain(Tuple(...)), always splat that. And ignore Recur's non-trainable state:
+    Flux._show_children(x::$ex) = _flat_children(trainable(x))
+  end
 end
 
 function _big_show(io::IO, obj, indent::Int=0, name=nothing)
   pre, post = obj isa Chain{<:AbstractVector} ? ("([", "])") : ("(", ")")
   children = _show_children(obj)
   if all(_show_leaflike, children)
+    # This check may not be useful anymore: it tries to infer when to stop the recursion by looking for grandkids,
+    # but once all layers use @layer, they stop the recursion by defining a method for _big_show.
     _layer_show(io, obj, indent, name)
   else
     println(io, " "^indent, isnothing(name) ? "" : "$name = ", nameof(typeof(obj)), pre)
@@ -56,48 +58,32 @@ _show_leaflike(::Tuple{Vararg{<:AbstractArray}}) = true  # e.g. parameters of LS
 # _show_leaflike(::Scale) = true                           # appears inside LayerNorm
 _show_leaflike(::AbstractArray{<:Number}) = true         # e.g. transposed arrays
 
-_show_children(x) = trainable(x)  # except for layers which hide their Tuple:
-# _show_children(c::Chain) = c.layers
-# _show_children(m::Maxout) = m.layers
-# _show_children(p::Parallel) = (p.connection, p.layers...)
-# _show_children(f::PairwiseFusion) = (f.connection, f.layers...)
-
+_show_children(x) = trainable(x)
+# This used to have methods for Chain, Maxout, Parallel, PairwiseFusion. Now @layer instead
+# writes a method to use this function. It flattens the Tuple within Chain etc.
+# (The remaining special cases are for printing of layer names when a NamedTuple, above.)
 function _flat_children(x)
     alpha = map(f -> getfield(x, f), fieldnames(typeof(x)))
     beta = map(y -> y isa Union{Tuple, NamedTuple} ? y : (y,), alpha)
     gamma = ((beta...)...,)
 end
 
-# This is called by @layer, on layers which should be treated like Dense
+# This is called by @layer, on layers which should be treated like Dense, and returns an expression:
 function _macro_layer_show(ex)
-    quote
-        # Entry point:
-        function Base.show(io::IO, m::MIME"text/plain", x::$ex)
-            if !get(io, :compact, false)
-              _layer_show(io, x)
-            else
-              show(io, x)
-            end
-        end
-        
-        # Exit from _big_show recursion, do we need this and _show_leaflike?
-        _big_show(io::IO, obj::$ex, indent::Int=0, name=nothing) = _layer_show(io, obj, indent, name)
-        # Since this isn't a container, do not recurse into its children, if any:
-        _show_leaflike(::$ex) = true
+  quote
+    # Entry point:
+    function Base.show(io::IO, m::MIME"text/plain", x::$ex)
+      if !get(io, :compact, false)
+        _layer_show(io, x)
+      else
+        show(io, x)
+      end
     end
+
+    # Exit from _big_show recursion:
+    Flux._big_show(io::IO, obj::$ex, indent::Int=0, name=nothing) = _layer_show(io, obj, indent, name)
+  end
 end
-# for T in [
-#     :Conv, :ConvTranspose, :CrossCor, :Dense, :Scale, :Bilinear, :Embedding,
-#     :BatchNorm, :LayerNorm, :InstanceNorm, :GroupNorm,
-#   ]
-#   @eval function Base.show(io::IO, m::MIME"text/plain", x::$T)
-#     if !get(io, :compact, false)
-#       _layer_show(io, x)
-#     else
-#       show(io, x)
-#     end
-#   end
-# end
 
 function _layer_show(io::IO, layer, indent::Int=0, name=nothing)
   _str = isnothing(name) ? "" : "$name = "