0.8.6

🚀 Features

• Weights and Biases Callback for Elegy
  • PR: #220

🐛 Fixes

• Docs typos
  • PR: #222
• Donate model's memory buffer to jit/pmap functions.
  • PR: #226
• Lazy load wandb
  • PR: #228

0.8.5

🚀 Features

• Handle distributed metrics with psum
  • PR: #218
• Improve CI
  • PR: #219

0.8.4

Changes

• @cgarciae: Fix import errors without dev dependencies (#214)

0.8.3

Changes

• @cgarciae: Fix import error when Haiku is not installed (#213)

0.8.2

Changes

• @cgarciae: Updates treex to 0.6.5 (#212)

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.

0.7.4

Changes

• Update optax, jax, jaxlib #177 (cgarciae)
• Fix readme #174 (cgarciae)

Note

Due to a mistake version 0.7.3 was skipped.

0.7.2

Merged pull requests:

• DataLoader Timeout #172 (alexander-g)
• Saved Model support: implements Model.saved_model that enables the user to serialize pred_step as a SavedModel #171 (cgarciae)

0.7.1

Merged pull requests:

• Make Transformers official #170 (cgarciae)
• More Docs: Expand documentation for the low-level API guides. #168 (cgarciae)
• Rich Summary: uses rich to style the summary console output #167 (cgarciae)

0.7.0

Features

• init now only called once internally and required to be called explicitly by the user under certain circumstances
• summary now uses jax.eval_shape under the hood so its super fast since it doesn't allocate memory or perform any computations on the device.

Merged pull requests:

• Fix notebook #166 (cgarciae)
• Single Initialization: Removes the current progressive initialization in favor of a single underlying call to init_step. #165 (cgarciae)