Need to run an ONNX model while keeping JAX as the execution backend? JAX ONNX Runtime is a toolchain for seamlessly executing ONNX models with JAX. It fits the interoperability side of the JAX ecosystem: instead of rebuilding a model in JAX, you can use a model created in the Open Neural Network Exchange format. section 12.1.3 section 2.8
Use it when the model and execution environment come from different parts of the machine-learning stack:
- Start with an ONNX model and target JAX for execution. section 12.1.3
- Distinguish it from TF2JAX, which converts TensorFlow functions or graphs into JAX functions. section 12.1.3 section 2.8
- Treat it as an execution and interoperability tool, not as a demonstrated model-training workflow. section 12.1.3
