WIP: Implementation of Recurrent layers

[ptigwe]

Initial implementation of recurrent layers starting with GRUCell which ports the corresponding implementation from flax.nn.recurrent.GRUCell.

The API is still a WIP and not fixed yet, but open to discussion as development goes on.

At the moment, the GRUCell allows for initialization of the starting hidden state using either:

• initialize_carry: which is essentially operating in a similar fashion to that of flax
• init_carry: which requires the module to have been initialized but then only takes in a tuple representing the batch_dims as an argument.

[cgarciae]

Hey @ptigwe! Thanks for taking the initiative.

The first general comment I have is that this reimplements GRU core logic from scratch, we have avoided so far for the following reason: Flax is well supported and Treex could leverage Flax improvements as they come for free, reducing maintenance cost. If you want some inspiration take a look at BatchNorm.

[ptigwe]

Was in a bit of a rush and seem to have messed up my git history. Would redo this in a bit.