Skip to content

jax2onnx 🌟

CI PyPI version License

jax2onnx converts your JAX, Flax NNX, Flax Linen, Equinox functions directly into the ONNX format.

jax2onnx.svg

✨ Key Features

  • Simple API – Convert JAX callables using to_onnx(...).
  • Model structure preserved – With @onnx_function, submodules appear as named functions in the ONNX graph.
  • Dynamic input support – Use abstract dimensions like 'B' or pass scalars as runtime inputs.
  • Plugin-based extensibility – Add support for new primitives by writing small, local plugins.
  • onnx-ir native pipeline – Conversion, optimization, and post-processing all run on the typed onnx_ir toolkit.
  • Netron-friendly outputs – Generated graphs carry shape/type annotations and a clean hierarchy.

🚀 Getting Started

Check out the Getting Started guide to install and export your first model in minutes.

See Roadmap for planned features and Past Versions for the full release archive.

Happy converting! 🎉