Simplify API: Cleans-up some API details around Model and Module

[cgarciae]

Changes

• elegy.States now implements Mapping and __getattr__ so you can add arbitrary keys but still use states.{field} to access values. By not enforcing any particular set of fields States can be more flexibly used with the low-level API, however, the user must know what the are the standard names (rng, net_params, net_states, metrics_states, etc) and their behavior, this will be properly documented.
• Removed all names previously in States (rng, net_params, net_states, ...) from the Dependency Injection lists, now the users must request states and access the field of interest (e.g. states.rng), this is both simpler and more explicit.
• elegy.Module.init and elegy.Module.apply now return (y_pred, parameters, collections) where parameters = collections.pop("parameters", None). apply now expects parameters and collections as arguments.
• Added base_step to the low-level API, defines the initial set of elegy.States.
• Simplified the definition of PredStep, it is now only (y_pred, states).
• Removed the Uninitialized type.
• Introduces summary_step and moves summary to ModelCore.

[cgarciae]

@alexander-g added the cases above so its compatible with the previous output of get_default_parameters (a dict with a "parameters" key) and the new version (a tuple (params, collections)). Feel free to clean this up if you update the weights.

[codecov-io]

Codecov Report

Merging #158 (a828d02) into master (cb70cf5) will decrease coverage by 0.04%.
The diff coverage is 97.44%.

@@            Coverage Diff             @@
##           master     #158      +/-   ##
==========================================
- Coverage   84.90%   84.85%   -0.05%     
==========================================
  Files         131      131              
  Lines        6815     6840      +25     
==========================================
+ Hits         5786     5804      +18     
- Misses       1029     1036       +7     

Impacted Files	Coverage Δ
elegy/generalized_module/generalized_module.py	66.15% <0.00%> (ø)
elegy/model/model_base.py	83.47% <ø> (ø)
elegy/nn/moving_averages.py	85.71% <ø> (ø)
elegy/nets/resnet.py	93.39% <87.50%> (-0.67%)	⬇️
elegy/types.py	95.50% <94.44%> (-1.13%)	⬇️
elegy/model/model_core.py	87.62% <95.78%> (+2.71%)	⬆️
elegy/module.py	95.20% <97.61%> (-0.12%)	⬇️
elegy/generalized_module/elegy_module.py	91.17% <100.00%> (-4.75%)	⬇️
elegy/generalized_module/haiku_module.py	90.76% <100.00%> (-0.41%)	⬇️
elegy/generalized_module/linen_module.py	90.62% <100.00%> (-0.15%)	⬇️
... and 8 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 cb70cf5...a828d02. Read the comment docs.