Metadata-Version: 2.4
Name: jax2onnx
Version: 0.13.1
Summary: export JAX to ONNX
License-File: LICENSE
Author: enpasos
Author-email: matthias.unverzagt@enpasos.ai
Requires-Python: >=3.11
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Requires-Dist: dm-pix (>=0.4.4)
Requires-Dist: einops (>=0.8.1)
Requires-Dist: equinox (>=0.13.1)
Requires-Dist: flax (>=0.12.1)
Requires-Dist: huggingface-hub (>=0.26.0)
Requires-Dist: jax (>=0.7.2)
Requires-Dist: ml_dtypes (>=0.5.1)
Requires-Dist: netron (>=8.2.9)
Requires-Dist: onnx (>=1.20.1)
Requires-Dist: onnx-ir (>=0.2.1)
Requires-Dist: onnxruntime (>=1.24.1)
Requires-Dist: optax (>=0.2.4)
Requires-Dist: orbax-checkpoint (>=0.11.6,<0.11.37)
Requires-Dist: orbax-export (>=0.0.6)
Description-Content-Type: text/markdown

# jax2onnx 🌟

[![CI](https://github.com/enpasos/jax2onnx/actions/workflows/ci.yml/badge.svg)](https://github.com/enpasos/jax2onnx/actions/workflows/ci.yml)
[![PyPI version](https://img.shields.io/pypi/v/jax2onnx.svg)](https://pypi.org/project/jax2onnx/) 
[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
[![Documentation](https://img.shields.io/badge/docs-live-blue)](https://enpasos.github.io/jax2onnx/)

`jax2onnx` converts your [JAX](https://docs.jax.dev/),  [Flax NNX](https://flax.readthedocs.io/en/latest/), [Flax Linen](https://flax-linen.readthedocs.io/en/latest/), [Equinox](https://docs.kidger.site/equinox/) functions directly into the ONNX format.


![jax2onnx.svg](https://enpasos.github.io/jax2onnx/images/jax2onnx.svg)

## 📚 Documentation

**[Read the full documentation here](https://enpasos.github.io/jax2onnx/)**

## 🚀 Quick Install

```bash
pip install jax2onnx
```

## ⚡ Quick Usage

```python
from jax2onnx import to_onnx
from flax import nnx

model = MyFlaxModel(...)
to_onnx(model, [("B", 32)], return_mode="file", output_path="model.onnx")
```

## 🤝 Contributing

We warmly welcome contributions! Please check our [Developer Guide](https://enpasos.github.io/jax2onnx/developer_guide/plugin_system/) for plugin tutorials and architecture details.

## 📜 License

Apache License, Version 2.0. See [`LICENSE`](./LICENSE).


## 🌟 Special Thanks

A huge thank you to all [our contributors and the community](https://enpasos.github.io/jax2onnx/about/acknowledgements/) for their help and inspiration!

---

**Happy converting! 🎉**

