Functional functional layers

[ToucheSir]

This is a continuation of my comments at FluxML/Flux.jl#1301 (comment).

Part of Flux's and design philosophy is that plain functions aren't scary and you can freely use them in a model without creating a custom forward pass or using some abomination like keras.layers.Lambda. For this reason however, most Flux layers are layers in the first place because they require more than can be provided by a stateless function. This usually means holding parameters or other state that is required during the forward/backwards pass. Note that I would also consider closures to fit this criteria as well given how they are represented in Julia.

[darsnack]

The fact that a) closures are structs and b) structs are callable makes me question how wide this gap really is. IMO defining a layer as a callable struct is writing a function. I would suggest (and I think you agree) that the discussion around structs/layers vs plain functions be separated from the discussion around explicit vs implicit state from the discussion around immutable vs mutable. Here's a taxonomy of sorts:

1. Structs vs Functions: this is the distinction between the following:

   f(W, x) = W * x
   
   struct F{T}
     W::T
   end
   (f::F)(x) = f.W * x

   It makes no requirements on whether the struct is mutable or how to go about changing the behavior of the struct on the same x.

2. Explicit vs Implicit State: this is related to mutability of parameters relevant to the layer's behavior. For example, if W above is something that changes over the course of the program, should it be stored in F and changed (i.e. implicit), or should it be passed into every call to a F (i.e. explicit)?

3. Immutable vs Mutable: this means when you need to change W, should it be done by making F mutable and adjusting the field, or should we make F immutable and require that a newly constructed F is returned?

My thoughts on the above:

1. Use the version that is syntactically clear. Generally, if the layer has many parameters (even if none are trainable), a struct will be cleaner. Also, if the parameters are not primitives, a struct will be cleaner. There is a blurred line, and I think this is a subjective distinction. Dogmatic adherence to one or the other is just less flexibility IMO.
2. If there is value is referring to the state outside of the layers regular execution, then make it explicit. Otherwise, a user probably prefers not to worry about the state. Explicit state does not always make the intent of the code more readable. It will all depend on how much ergonomic effort is required for a user to find the place that state is modified. In this sense, implicit state can be very helpful in consolidating that search.
3. This will largely depend on the cost of reinitializing F. It's also worth mentioning that as much as mutability can be a curse, it can also be the clearest way to express something in code. Lots of Julia code uses mutability, so we shouldn't shy away from it.

Also, "making functions first class" is not the same as "writing functional code." Dogmatic adherence to a functional programming paradigm (i.e. immutability always at the syntactic level, explicit state everywhere, minimal structures, lots of wrappers) will only hurt us. The beauty of Julia is that it is powerful enough to express multiple programming paradigms in the same codebase elegantly. Not everything is object-based vs functional, there is a lot in between. Let's just pick and choose.

[darsnack]

Another note on lifting things into the type parameters: I'd prefer to do this for things that don't change. If something is a field, and a user might change it over the course of the program, then don't lift it into the type. It really doesn't make the code more readable. Do it if it improve performance though.

Same goes for consolidating parameters into more generic "parameter structs" (e.g. ConvDims). Do it if it improves performance, but don't go crazy with it cause I think it makes the code more opaque.