Reference preserving hooks

[cgarciae]

This PR implements Haiku's Module system in such a way that hooks preserve reference information. This fixes the following problems that are currently present:

1. Haiku's resctriction of not being able to instantiate modules outside of a transform, which lead to the awkward use of defer.
2. Since modules didn't carry their parameters and states with them (in fact you never got hold of a reference to any module) it made more difficult tasks like:

• Extracting a submodule an doing inference with it, for example, training a VAE and using the Decoder to generate new samples.
• Getting a pretrained model and adding more layers on top for transfer learning.
• Training different parts of a Model seperately (GANs).

While point 2 was solvable by modifying the parameters and states dictionaries according to each new architecture, there was no clear strategy on how to make this ergnonomic.

Reference preserving hooks use python objects to store references to submodules, parameters and states as is commonly done in Keras and Pytorch, therefore any change in code architecture leads to a change in parameter structure.

Submodule hooks are implemented by leveraging Metaclasses in such a way that all submodules created inside call are automatically registered inside the parent module.

[codecov-commenter]

Codecov Report

Merging #85 into master will decrease coverage by 3.04%.
The diff coverage is 67.06%.

@@            Coverage Diff             @@
##           master      #85      +/-   ##
==========================================
- Coverage   73.42%   70.38%   -3.05%     
==========================================
  Files          77       81       +4     
  Lines        2751     3761    +1010     
==========================================
+ Hits         2020     2647     +627     
- Misses        731     1114     +383     

Impacted Files	Coverage Δ
elegy/callbacks/early_stopping.py	77.04% <ø> (ø)
elegy/losses/categorical_crossentropy_test.py	100.00% <ø> (ø)
...egy/losses/sparse_categorical_crossentropy_test.py	100.00% <ø> (ø)
elegy/metrics/accuracy_test.py	100.00% <ø> (ø)
elegy/metrics/categorical_accuracy_test.py	100.00% <ø> (ø)
elegy/metrics/mean_absolute_error_test.py	100.00% <ø> (ø)
elegy/metrics/mean_squared_error_test.py	100.00% <ø> (ø)
elegy/metrics/sparse_categorical_accuracy_test.py	100.00% <ø> (ø)
elegy/nn/multi_head_attention.py	0.00% <0.00%> (ø)
elegy/nn/sequential_test.py	100.00% <ø> (ø)
... and 54 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 ae58afe...f90b4da. Read the comment docs.