https://github.com/google/orbax

Orbax provides common utility libraries for JAX users.
https://github.com/google/orbax

checkpoint flax jax

Last synced: about 2 months ago
JSON representation

Orbax provides common utility libraries for JAX users.

• Host: GitHub
• URL: https://github.com/google/orbax
• Owner: google
• License: apache-2.0
• Created: 2022-05-06T17:36:22.000Z (over 2 years ago)
• Default Branch: main
• Last Pushed: 2024-04-13T08:11:14.000Z (5 months ago)
• Last Synced: 2024-04-13T17:12:30.236Z (5 months ago)
• Topics: checkpoint, flax, jax
• Language: Python
• Homepage: https://orbax.readthedocs.io/
• Size: 1.4 MB
• Stars: 221
• Watchers: 11
• Forks: 22
• Open Issues: 78
• Metadata Files:
  • Readme: README.md
  • Changelog: CHANGELOG.md
  • Contributing: CONTRIBUTING.md
  • License: LICENSE

Awesome Lists containing this project

README

        
# Orbax

[Orbax](https://orbax.readthedocs.io/en/latest/) is a

namespace providing common utility libraries for JAX users.

## [Documentation](https://orbax.readthedocs.io/en/latest/)

## Checkpointing

`pip install orbax-checkpoint` (latest PyPi release) OR

`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'` (from this repository, at HEAD)

`import orbax.checkpoint`

Orbax includes a checkpointing library oriented towards JAX users, supporting a

variety of different features required by different frameworks, including

asynchronous checkpointing, various types, and various storage formats.

We aim to provide a highly customizable and composable API which maximizes

flexibility for diverse use cases.

> [!NOTE]

> Please see [Announcements](https://orbax.readthedocs.io/en/latest/orbax_checkpoint_announcements.html) for important updates.

## Exporting

`pip install orbax-export` (latest PyPi release) OR

`pip install 'git+https://github.com/google/orbax/#subdirectory=export'` (from this repository, at HEAD)

`import orbax.export`

Orbax also includes a serialization library for JAX users, enabling the exporting of JAX models to the TensorFlow SavedModel format.

Note that `orbax-export` requires TensorFlow, but does not include it by default to allow for flexibility in version choice. If you wish to install with standard TensorFlow, please use `pip install orbax-export[all]`.

## Support

Contact [email protected] for help or with any questions about Orbax!

### History

Orbax was initially published as a catch-all package itself. In order to minimize dependency bloat for users, we have frozen that package at `orbax-0.1.6`, and will continue to release future changes under the domain-specific utilities detailed above (e.g. `orbax-checkpoint`). 

As we have preserved the orbax namespace, existing import statements can remain unchanged (e.g. `from orbax import checkpoint`).