orbax-export

What it is and what it does

Orbax Export is a serialization library that bridges JAX and TensorFlow by converting JAX models into TensorFlow's SavedModel format. This enables JAX users to export trained models for deployment in TensorFlow-based production systems or to share models across frameworks. The package is part of the larger Orbax ecosystem and integrates with orbax-checkpoint for model persistence.

The package requires nine runtime dependencies centered on jax, jaxlib, jaxtyping, numpy, and serialization tools like protobuf and dataclasses-json. TensorFlow is not bundled by default to allow flexibility in version selection—users must either install the optional [all] extra or manage their own TensorFlow installation. The library targets Python 3.10 or later.

Use it for:

• Export a trained JAX model to SavedModel format for deployment in TensorFlow-based inference systems.
• Share JAX models with teams or collaborators who work primarily in TensorFlow environments.
• Convert research models from JAX to a standardized format for production deployment pipelines.
• Integrate JAX training workflows with existing TensorFlow-based model serving infrastructure.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Orbax Export serializes JAX models to TensorFlow SavedModel format, enabling JAX-trained models to be deployed and used in TensorFlow environments.

Yes, if you train models in JAX and need to deploy them in TensorFlow environments or share them across frameworks. The package is actively maintained, has no known vulnerabilities, and uses a permissive license. The main consideration is ensuring TensorFlow is installed separately and that your models fit the supported export patterns—check the documentation before committing to a production workflow.

Install

orbax-export on PyPI

pip install orbax-export

uv add orbax-export

poetry add orbax-export

Installing orbax-export

Before you install

Low install friction with a pure Python wheel. Active maintenance with recent commits and steady releases since first release. Requires Python 3.10 or later and nine runtime dependencies including jax, jaxlib, and orbax-checkpoint.

License in practice

Licensed under Apache Software License (permissive), allowing commercial and private use with minimal restrictions.

Quickstart

pip install orbax-export
import orbax.export
# Export a JAX model to SavedModel format
# (See documentation for specific export API usage)

Requires Python 3.10 or later. TensorFlow must be installed separately; use `pip install orbax-export[all]` for standard TensorFlow, or install your preferred TensorFlow version independently.

Verify before relying

• Specific export API surface and supported JAX model architectures beyond SavedModel capability
• Performance characteristics or limitations when exporting large or complex JAX models
• Compatibility matrix with specific JAX and TensorFlow versions

Package facts

Evidence: orbax_export-0.0.8-py3-none-any.whl

Keywords: JAX machine learning, serialization, export

Tags

More Artificial Intelligence packages