Support layers without parameters or state

[gdalle]

This PR proposes a solution to #83 by adding 2 booleans P and S to the layer types, specifying whether or not they need parameters / states as inputs.

That way, in addition to the generic syntax layer(x, ps, st):

• a layer without parameters (e.g. Dropout) supports layer(x, st)
• a layer without state (e.g. Dense) supports layer(x, ps)
• a layer without parameters or state (e.g. ReshapeLayer) supports layer(x)

[codecov]

Codecov Report

Merging #84 (37f4851) into main (f7ba291) will decrease coverage by 3.87%.
The diff coverage is 51.85%.

@@            Coverage Diff             @@
##             main      #84      +/-   ##
==========================================
- Coverage   71.18%   67.30%   -3.88%     
==========================================
  Files          14       14              
  Lines         989     1043      +54     
==========================================
- Hits          704      702       -2     
- Misses        285      341      +56     

Impacted Files	Coverage Δ
src/adapt.jl	34.37% <0.00%> (-28.13%)	⬇️
src/layers/dropout.jl	100.00% <ø> (ø)
src/layers/recurrent.jl	100.00% <ø> (ø)
src/core.jl	65.00% <42.10%> (-35.00%)	⬇️
src/layers/basic.jl	58.73% <100.00%> (-1.27%)	⬇️
src/layers/conv.jl	100.00% <100.00%> (ø)
src/layers/normalize.jl	87.50% <100.00%> (+36.21%)	⬆️
src/layers/display.jl	0.00% <0.00%> (-54.64%)	⬇️
... and 5 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f7ba291...37f4851. Read the comment docs.

[gdalle]

Sofar I haven't added specific tests for these new methods, and they are not documented either. This is just a first draft to agree on the syntax before moving forward

[gdalle]

The main issue I see is for constructs such as Parallel, BranchLayer, PairwiseFusion and Chain. Since those layers are built from potentially heterogeneous NamedTuples, I don't know how to statically determine whether they have parameters / states. One would need to apply a logical OR to the hasparams / hasstate of each element of the NamedTuple, which sounds like it has to be done dynamically since we don't even know its length.

[gdalle]

The main issue I see is for constructs such as Parallel, BranchLayer, PairwiseFusion and Chain.

I may have found a solution. A Chain (or any other similar construct) requires parameters / states if and only if at least one of its components does. If I'm right, the code below can decide that statically (just imagine B is either hasparams or hasstate):

julia> using Test

julia> struct MyLayer{B} end

julia> struct MyChain{B} end

julia> MyChain(t::Tuple{Vararg{MyLayer}}) = MyChainTuple{true}();

julia> MyChain(t::Tuple{Vararg{MyLayer{false}}}) = MyChainTuple{false}();

julia> MyChain(t...) = MyChain(t);

julia> l0, l1 = MyLayer{false}(), MyLayer{true}();

julia> c00 = @inferred MyChain(l0, l0)
MyChainTuple{false}()

julia> c01 = @inferred MyChain(l0, l1)
MyChainTuple{true}()

julia> c11 = @inferred MyChain(l1, l1)
MyChainTuple{true}()

[avik-pal]

Made a first pass over the PR.

[avik-pal]

I think there is a disconnect between these functions and the original goal. The idea of having these properties in the type, would be to simplify user code. So a stateless layer can define (::Layer)(x, ps). With this the user is still forced to define the original 3 argument api.

[gdalle]

After trying this out in the code, I don't necessarily agree.
For the beginners that never define custom layers, having access to the font end syntax layer(x, ps) or layer(x) is a nice feature. But for the advanced users, allowing them to define layers like this in the back end will cause major headaches.
Indeed, if we allow such layer definitions, then layer(x, ps) must only return one output y, whereas layer(x, st) must return two outputs y and new_st. In terms of global coherence and documentation, this is a nightmare. I'd much rather impose that all layers be defined as layer(x, ps, st), and then we add convenience overloads to use them in a simplified manner.

[avik-pal]

Also if we are to simplify user code, we will have to define functions like (l::AbstractExplicitLayer{false}(x, ps, st) = l(x, st), (l::AbstractExplicitLayer{hasparams, false}(x, ps, st) = l(x, ps), st, (l::AbstractExplicitLayer{false, false}(x, ps, st) = l(x), st

[avik-pal]

Also tagging @ChrisRackauckas, @ToucheSir for thoughts on this

[avik-pal]

One would need to apply a logical OR to the hasparams / hasstate of each element of the NamedTuple, which sounds like it has to be done dynamically since we don't even know its length.

That information is available during compilation, I don't think it will lead to a dynamic dispatch. Worst case if the compiler can't do that, we can just use a generated function. Also, I wouldn't care too much about it since it is never going to appear in a performance-critical section of the code.

[ToucheSir]

I don't see any scenario where you'd need a generated function either, as the aforementioned bool type logic happens at construction. Have you considered whether it makes sense to add type aliases for some of the more common combinations?

[ChrisRackauckas]

I'm not sure I like the high level call changes with this. There's something simple about always needing f(x,p,st). If you start being able to drop some in some cases, then a lot of docs will be like "oh yeah it's just f(x,p)" (most SciML examples), but then the user goes to put in a BatchNorm and suddenly their code throws an odd error which they don't understand. I think there's simplicity in having only one calling convention.

[gdalle]

I'm not sure I like the high level call changes with this.

@ChrisRackauckas I think that depending on your use of the library, those changes might make a lot of sense. For instance, the notion of layer state is not obvious to beginners, and when they don't need it (eg. for a simple CNN without dropout), carrying it along makes the code heavier and less readable.
As mentioned by @avik-pal in #83, it's mostly about quality of life, and to help people get started with Lux.jl. Of course, as soon as you get more familiar with the syntax, you will start adding a state everywhere. But to ease first contact with the code, leaving the state aside for simple networks is a clear improvement to me.

[gdalle]

That information is available during compilation, I don't think it will lead to a dynamic dispatch.

I tried the following approach:

struct Parallel{P, S, F, T <: NamedTuple} <:
       AbstractExplicitContainerLayer{(:layers,), P, S}
    connection::F
    layers::T
end

function Parallel(connection::F, layers...) where {F}
    names = ntuple(i -> Symbol("layer_$i"), length(layers))
    layers_nt = NamedTuple{names}(layers)
    P = maximum(hasparams, layers_nt)
    S = maximum(hasstate, layers_nt)
    T = typeof(layers_nt)
    return Parallel{P, S, F, T}(connection, layers_nt)
end

Apparently, the compiler infers P and T without issues. However, type inference for the NamedTuple layer field no longer works (as opposed to the previous constructor). This in turn made all the tests fail, probably due to the generated functions that are used to apply Parallel.

Any idea how to fix it?

[gdalle]

Also, I wouldn't care too much about it since it is never going to appear in a performance-critical section of the code.

I'm not sure I agree. There are cases where you want to create a neural network inside a hot loop, and faulty type inference at construction can mess up everything that follows.

[ChrisRackauckas]

As I mentioned, no example I have written (which now is like 30 tutorials or so that use Lux) would require state. So the SciML uses of the library would definitely fall into the category that you are trying to target. But even with that, I don't think it's a good idea because this is so pervasive that I think it will make the cases with state confusing. If no tutorials ever include the state, then how are you ever expected to know that it exists? What I foresee happening is that there will be some layers that users will get confused by, like BatchNorm will just throw "mysterious" errors, when in reality the mystery is that the tutorials are all focused on a special case.

Now this is somewhat hidden in some cases like NeuralPDE and DiffEqFlux since in those cases users are rarely doing the call themselves. But SciMLSensitivity will definitely have this as an extra documentation issue, and either the solution will be to change our tutorials to showcase state in what will then be a non-standard way (i.e. include state even when it's nothing), always use layers with state in tutorials as a matter of training for the most "robust" form, or just ignoring it as a Lux issue and answering a bunch of issues with "oh, that's a Lux thing. If you use xxx then you have st and ...". I hate doing the latter because it quickly leads to 20 emails to answer everyday with the same thing, and by now I've learned that styles which require answering people's questions are just not scalable.

[gdalle]

What I foresee happening is that there will be some layers that users will get confused by, like BatchNorm will just throw "mysterious" errors

I also think that this is a foreseeable problem, but one that can be fixed with explicit error messages on the Lux.jl side when the wrong methods are called. Something like

Your neural network contains a layer that relies on a state variable.
You probably tried to call it with `network(x, ps)`, but instead you need to use `network(x, ps, st)`.
See [...] in the documentation for more details.

Anyway, I only started working on this because @avik-pal agreed it would be a welcome feature. If he ends up agreeing it's not, then I'll simply close the PR

[gdalle]

Another justification for this feature would be easing the transition from Flux.jl to Lux.jl. Since you announced that SciML would be using Lux.jl, I've been convinced that this transition is a good thing. And to get people hooked to Lux.jl, the fewer visible differences with Flux.jl, the better. Then, when models get more complex and they're in too deep, we can drop the bomb with the additional state argument and they won't be able to escape 😅

[ChrisRackauckas]

The transition is pretty much done BTW, so check out

• https://sensitivity.sciml.ai/dev/
• https://diffeqflux.sciml.ai/dev/
• https://neuralpde.sciml.ai/dev/

which all bumped majors and removed FastChain from everything (well, DiffEqFlux is keeping the depwarns around for awhile to help the updates). 🎉

Almost everything is Lux now (SciML/NeuralPDE.jl#533 and SciML/NeuralPDE.jl#533 pretty much sealed the deal), except a notable issue with Lux on GPUs (due to ComponentArrays not broadcasting on GPUs correctly). There is an on-going thread about if we're happy with the new coding styles these updates gave, so I encourage you to chime in (https://julialang.slack.com/archives/CN04R7WKE/p1657113244013009).

Indeed, one of the things that was missing from FastChain was state, so when doing these updates we had to add state in. And right now it doesn't look nice, it's using globals, which Chris Elrod wasn't too keen on, and there are some @ignore_derivatives points which seem to be correct for every case I discussed with @avik-pal though I do wonder whether derivatives w.r.t. state are something that may be required for something (and would love to hear if there is a case). For some things like NeuralODE it would be good to disallow models that have some form of state that we know would be incorrect (for example, LSTMs in a NeuralODE are bad because the neural ODE doesn't step forward, and so the adaptivity means an LSTM-defined neural ODE does not have a well-defined solution if it rejects steps!), but disallowing all models with state goes too far (normalization layers for example should be fine in the theory that the state is fixed for each ODE solve).

However, even with all of those new issues arising because of state, I think we need to just fine what coding behaviors and habits do it well. For example, always make a callable struct to enclose the state in an ODE definition? We're still working it out, but I think it would be best to find what the nice styles are here and use it everywhere. Otherwise it might just be one "advanced" example that shows it, with other examples showing "worse" coding styles, or even worse, coding styles that have some unnecessary complexity when no state is there (but done anyways for consistency with the cases where it is?).

That said, I will be willing to give it a try since it would simplify docs and such, though I am worried that in all tutorials of these 3 libraries, we would just remove state to match this PR. And if we do that, 😅 state examples will seem like magic. A good error message may be enough though.

[gdalle]

The transition is pretty much done BTW

I know the transition is done for SciML (kudos on the massive change!), but there's lots of other people using Flux out there that we could still convince to switch 😉

That said, I will be willing to give it a try since it would simplify docs and such, though I am worried that in all tutorials of these 3 libraries, we would just remove state to match this PR.

I'm not necessarily advocating for the layer(x, ps) to become the default (recommended) syntax, as a matter of fact I'm not sure it should.
And as I stated in this comment, I think that for people who define new layers, sticking to the three-argument syntax makes more sense.
The change I'm proposing is mostly for people who use pre-defined stateless layers, to which we can give additional options.

[avik-pal]

I kind of agree with Chris here that having a consistent API, even if it is not beginner-friendly, might be the way to go. Flux is great in that you can pretty much throw anything its way, and it will work, but for Lux, that is not the case. It relies on a pretty strict convention on how you should be writing your julia/zygote code.

Originally I was thinking about how to make it easier for people writing custom layers, in which case managing states is quite cumbersome and, in most cases, can be auto-generated. See #46, which kind of already does that (my only concern with merging that PR even though its nearly complete, users will not check the generated code and slap it around places where it might not work correctly)

Also in my experience, it doesn't matter if you show users the "correct" way to do something (especially after introducing a partially correct easier version). They tend to gravitate towards the easier (and in this particular case) the wrong way to do things. Anecdotally, flax which also operates similarly to lux (actually the other way around), introduces a non-state managed API first. And even after two months of using it, I still can't remember how to handle states properly.

[gdalle]

Okay, so do we just close this and move on?

[gdalle]

I must say the discussion also swayed me towards keeping the state around

[avik-pal]

Okay, so do we just close this and move on?

I think yeah. Also, I should have tagged people in the other issue before you put all this effort in. My bad.