The package provides extra flexibility to Flax using ideas originated at Trax
Project description
Flax extra
The package provides extra flexibility to Flax using ideas originated at Trax.
Extras include:
- Support of combinators (serial, parallel, branch, etc.) backed by Redex.
- A modular training loop working nicely with Optax optimizers and metrics.
- Pluggable logging to stdout.
- Parallel execution of training and evaluation tasks.
- Regular and the best metric checkpoints.
- Tensorboard integration.
- Additional linen modules.
Check out documentation, introductory tutorials, or examples that include:
- Perceiver IO
- Classification model pretrained on images from ImageNet.
- Masked-language model pretrained using a large text corpus obtained by combining English Wikipedia and C4.
- Autoencoder model pretrained on multimodal input (audio, video, and label) of the Kinetics-700-2020 dataset.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
flax-extra-0.0.2.tar.gz
(39.8 kB
view hashes)
Built Distribution
flax_extra-0.0.2-py3-none-any.whl
(61.6 kB
view hashes)
Close
Hashes for flax_extra-0.0.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | aa793d3364024cf38ba158212e6ef96ae4a76633030552ca463a0030b463ef19 |
|
MD5 | 9364a7fe1a0a05bc4d19c0b3c26279be |
|
BLAKE2b-256 | 76102f3b7195bd6b2c4084ad53a689a8425b78cf64cccde4f1df8dfa38221d17 |