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 introductory tutorials or go straight for the documentation.
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.1.tar.gz
(22.4 kB
view hashes)
Built Distribution
flax_extra-0.0.1-py3-none-any.whl
(28.7 kB
view hashes)
Close
Hashes for flax_extra-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3008ddd54fa510db3cab5083ca24f5e76cbd2017ee639c859d9cadb582dfa9e7 |
|
MD5 | 7cff5e84904c6d0c0b551526bc4c26db |
|
BLAKE2b-256 | 01d81931e1efcff7461bae10689da5ee00dfcd60b28be90737a9b7227862bad0 |