Compact layer factory

[darsnack]

Transferring over from FluxML/Flux.jl#2107.

@MilesCranmer proposed two implementations of a "layer factory." You can find examples of this in other frameworks: Flax's @compact, Haiku's transforms.

Ideally, we would be able to just use Julia's let block syntax if Functors.jl was opt-out instead of opt-in. This may happen in Functors.jl v0.5.

One option is making a factory struct which we can @functor. Like so:

struct LayerFactory{F<:Function,NT<:NamedTuple}
    forward::F
    layers::NT
end

LayerFactory(f; layers...) = LayerFactory(f, NamedTuple(layers))

function (f::LayerFactory)(args...)
    return f.forward(f.layers, args...)
end

@functor LayerFactory

The above code is just an illustration of the idea. A better implementation is

LayerFactory(f; layers...) = Base.Fix1(f, NamedTuple(layers))

which should work in Functors v0.4+ (since Base.Fix1 is @functored).

Another option is to use macros. Like so:

"""
    @Magic(forward::Function; construct...)

Creates a layer by specifying some code to construct the layer, run immediately,
and (usually as a `do` block) a function for the forward pass.
You may think of `construct` as keywords, or better as a `let` block creating local variables.
Their names may be used within the body of the `forward` function.

```
r = @Magic(w = rand(3)) do x
    w .* x
end
r([1,1,1])
r([10,10,10])  # same random numbers

d = @Magic(in=5, out=7, W=randn(out,in), b=zeros(out), act=relu) do x
    y = W * x
    act.(y .+ b)
end
d(ones(5, 10))  # 7×10 Matrix
```
"""
macro Magic(fex, kwexs...)
    # check input
    Meta.isexpr(fex, :(->)) || error("expects a do block")
    isempty(kwexs) && error("expects keyword arguments")
    all(ex -> Meta.isexpr(ex, :kw), kwexs) || error("expects only keyword argumens")
    
    # make strings
    layer = "@Magic"
    setup = join(map(ex -> string(ex.args[1], " = ", ex.args[2]), kwexs), ", ")
    input = join(fex.args[1].args, ", ")
    block = string(Base.remove_linenums!(fex).args[2])

    # edit expressions
    vars = map(ex -> ex.args[1], kwexs)
    assigns = map(ex -> Expr(:(=), ex.args...), kwexs)
    @gensym self
    pushfirst!(fex.args[1].args, self)
    addprefix!(fex, self, vars)

    # assemble
    quote
        let
            $(assigns...)
            $MagicLayer($fex, ($layer, $setup, $input, $block); $(vars...))
        end
    end |> esc
end

function addprefix!(ex::Expr, self, vars)
    for i in 1:length(ex.args)
        if ex.args[i] in vars
            ex.args[i] = :($self.$(ex.args[i]))
        else
            addprefix!(ex.args[i], self, vars)
        end
    end
end
addprefix!(not_ex, self, vars) = nothing

struct MagicLayer{F,NT<:NamedTuple}
    fun::F
    strings::NTuple{4,String}
    variables::NT
end
MagicLayer(f::Function, str::Tuple; kw...) = MagicLayer(f, str, NamedTuple(kw))
(m::MagicLayer)(x...) = m.fun(m.variables, x...)
MagicLayer(args...) = error("MagicLayer is meant to be constructed by the macro")
Flux.@functor MagicLayer

function Base.show(io::IO, m::MagicLayer)
    layer, setup, input, block = m.strings
    print(io, layer, "(", setup, ") do ", input)
    print(io, block[6:end])
end

[MilesCranmer]

Thanks!

I think the @Magic macro is my personal favorite, although I would prefer the string representation to be a bit easier to read. The current one is like so:

model = @Magic(mlp1=Chain(Dense(1, 5, relu), Dense(5, 1)), mlp2=Chain(Dense(1, 5, relu), Dense(5, 1))) do x
    mlp1(x) .+ mlp(2)
end

print(model) will print the macro (almost) exactly as you typed it. I think this should instead be something like this (borrowing Oceananigans.jl's style of printout)

MagicLayer
├── device: GPU()
├── attributes:
│   ├── mlp1:
│   |      Chain(
│   |        Dense(1 => 5, relu),                  # 10 parameters
│   |        Dense(5 => 1),                        # 6 parameters
│   |      )                   # Total: 4 arrays, 16 parameters, 320 bytes.
│   └── mlp2: 
│           Chain(
│             Dense(1 => 5, relu),                  # 10 parameters
│             Dense(5 => 1),                        # 6 parameters
│           )                   # Total: 4 arrays, 16 parameters, 320 bytes.
└── forward: (x) -> 
            mlp1(x) .+ mlp(2)

What do you think? @mcabbott?

You could also have it so that MagicLayer has a reserved name keyword - which will be printed instead of "MagicLayer".

[mcabbott]

Yea the printing is a rough first cut, sometime after midnight!

Flux has fancy recursive printing for models, and I think that ideally it would fit in with this. A magic layer used within a Chain and containing another Chain should ideally all unfold nicely, like Parallel does:

julia> Chain(Dense(2 => 10), Parallel(vcat; α = Dense(10 => 3, tanh), β = Chain(Dense(10 => 4), softmax)))
Chain(
  Dense(2 => 10),                       # 30 parameters
  Parallel(
    vcat,
    α = Dense(10 => 3, tanh),           # 33 parameters
    β = Chain(
      Dense(10 => 4),                   # 44 parameters
      NNlib.softmax,
    ),
  ),
)                   # Total: 6 arrays, 107 parameters, 908 bytes.

What this printing tries to preserve is that it's still code, pasting it back in creates the same model again. This doesn't really work if you define a function like Parallel((x,y) -> cat(x,y,dims=3), Dense(1=>1), Dense(1=>1)). But the ideal macro would (I think) print your function nicely, so that you can read it and know what you did.

[mcabbott]

MagicLayer has a reserved name keyword - which will be printed instead of "MagicLayer".

My version above stored a string layer = "@Magic". But IDK quite what to do with it.

[MilesCranmer]

Maybe something like

Chain(
  Dense(2 => 10),                       # 30 parameters
  @Magic(
      mlp1 = Chain(
            Dense(1, 5, relu),     # n parameters
            Dense(5, 1)
      ),
      mlp2 = Chain(
            Dense(1, 5, relu),     # n parameters
            Dense(5, 1)
      )
    ) do x
        mlp1(x) .+ mlp2(x)
    end
  )
)                   # Total: 6 arrays, 107 parameters, 908 bytes.

where the chains would be printed as they are normally in the recursive expansion (too lazy to manually write it out).

[mcabbott]

Yes exactly, surely this is possible. By overloading _big_show from https://github.com/FluxML/Flux.jl/blob/master/src/layers/show.jl for this type.

[mcabbott]

https://github.com/MasonProtter/LegibleLambdas.jl is also worth a look.

It might be possible to make the macro do much less here. If all components are independent (unlike my Dense example) then perhaps the whole struct is Fix2 with a special function. Maybe it can't be exactly that of LegibleLambdas.jl but similar.

[mcabbott]

BTW I think it would be great to collect examples of nontrivial uses of such factory layers, things which aren't trivially Parallel etc. Are they always basically data-flow between existing layers, or do they want raw arrays etc?

Transformers.jl has some macros etc. for constructing its models, which I don't understand well. Example from here:

julia> using Transformers

julia> N = 3; using Transformers: PositionEmbedding, Stack, @nntopo, Transformer, Dropout

julia> const encoder = todevice(Stack(
           @nntopo(e → pe:(e, pe) → x → x → $N),
           PositionEmbedding(512),
           .+,
           Dropout(0.1),
           [Transformer(512, 8, 64, 2048) for i = 1:N]...
       ))
Stack(PositionEmbedding(512), Base.Broadcast.BroadcastFunction(+), Dropout(0.1), Transformer(head=8, head_size=64, pwffn_size=2048, size=512), Transformer(head=8, head_size=64, pwffn_size=2048, size=512), Transformer(head=8, head_size=64, pwffn_size=2048, size=512), Transformer(head=8, head_size=64, pwffn_size=2048, size=512), Transformer(head=8, head_size=64, pwffn_size=2048, size=512))

julia> encoder.topo
NNTopo{"e → (pe:(e, pe) → (x → (x → 5)))"}
topo_func(model, e)
	pe = model[1](e)
	x = model[2](e, pe)
	x = model[3](x)
	x = model[4](x)
	x = model[5](x)
	x = model[6](x)
	x = model[7](x)
	x = model[8](x)
	x
end

[MilesCranmer]

I think in general, graph networks are unintuitive to represent in a single chain, and are easier to think about in terms of full control flow, because you are updating different parts of your latent space at different times.

(from https://arxiv.org/abs/1806.01261)

[MilesCranmer]

(Once you have a block, you could stack them up normally though)

[MilesCranmer]

But a lot of GNN research is focused on the internals of a single block.

[MilesCranmer]

I guess transformers can be thought of as a type of GNN (and vice versa) so maybe this example is similar to yours.

[chengchingwen]

That encoder.topo is just a simple DSL that would be converted into multiple lines of assignment and function call.