0.8.1

Elegy is now based on Treex 🎉

Changes

• Remove the module, nn, metrics, and losses from Elegy, instead Elegy reexports these modules from Treex.
• GeneralizedModule and friends are gone, to use Flax Modules use the elegy.nn.FlaxModule wrapper.
• Low level API is massively simplified:
  • States is removed, since Model is a pytree all parameters are tracked automatically thanks to Treex / Treeo.
  • All static state arguments (training, initializing) are removed, Modules can simply use self.training to pick their training state and self.initializing() to check whether they are initializing.
  • Signature for pred_step, test_step, and train_step now simply consists of inputs and labels, where labels is a dict that can contain additional keys like sample_weight or class_weight as required by the losses and metrics.
• Adds the DistributedStrategy class which currently has 3 instances
  • Eager: Runs model in a single device in eager mode (no jit)
  • JIT: Runs model in a single device with jit
  • DataParallel: Run the model in multiple devices using pmap.
• Adds methods to change the model's distributed strategy:
  • .distributed(strategy = DataParallel): changes the distributed strategy, DataParallel used by default.
  • .local(): changes the distributed strategy to JIT.
  • .eager(): changes the distributed strategy to Eager.
• Removes the .eager field in favor of the .eager() method.