Add more Layers

[cgarciae]

The idea is to add more wrappers over Flax layers.

Guidelines

• Core computation should be done by Flax, unless there is a good reason to re-implement it.
• Argument names and types should be the same unless there is a good reason to deviate from Flax API.
• "params" and "batch_stats" are assigned to fields, take implementation of Linear and BatchNorm as references.
• Must take into account training and frozen state, checkout the implementation of BatchNorm and Dropout as references.
• Must add test for numerical equivalence, check out test_linear.py as a reference.

Layers List

• [ ] ConvTranspose
• [x] Embed - should we call it Embedding?
• [x] LayerNorm
• [x] GroupNorm
• [ ] MaxPooling - Module version
• [ ] AvgPooling - Module version
• [ ] LSTM - Follow Keras API, implement it using scan but add a flag for unrolling it.
• [x] GRU - Follow Keras API, implement it using scan but add a flag for unrolling it.

If you want to work on a layer, feel free to "claim" it via a comment here so others know you are working on it.

[ptigwe]

On the embedding layer, I think it'll be fine with Embed as opposed to Embedding primarily due to the Flax-like design, plus it seems to be the name used in especially given the first point you make on the guidelines. At the same time, I'm not opposed to Embedding as its used in most frameworks. To the best of my knowledge of the other JAX-based NN libraries, only Trax calls it Embedding, Haiku and Flax both call it Embed.

[cgarciae]

@ptigwe Agreed. I'd leave it up to the implementer to decide the name.