0.4.0

[cgarciae]

Changes

• Optimizer now flattens its params to be agnostic to the static components of the pytree.
• Generic types containing TreeParts are no longer valid type annotation as types like Tuple[int, tx.State[int]] make it appear as if the first element of the tuple where static and the second dynamic, when in fact Treex would treat to whole field as dynamic. Now your only option is tx.State[Tuple[int, int]].
• Adds the tx.Hashable class to wrap non-hashable types like numpy or jax arrays when you want to use them in static fields of a TreeObject.
• Adds FlaxModule: can wrap any Flax Module into a Treex Module.
• tabulate now accepts a sample_input and will show the input and output columns.
• Refactors a lot of the functional API.
• Introduces .freeze(), .unfreeze() and .frozen similar to train/eval/training.
• Updates BatchNorm and Dropout to leverage .frozen.