Project description

JAX ONNX Runtime

JAX ONNX Runtime is a robust and user-friendly tool chain that enables the seamless execution of ONNX models using JAX as the backend.

More specifically, this tool chain has the abilities:

ONNX Model Conversion: Converts ONNX models into JAX format modules. Tested on popular large language models including GPT-2, BERT, and LLaMA.
Hardware Acceleration: Enable the jit mode of the converted JAX modules, which accelerates execution on GPU and/or TPU.
Compatibility with JAX ecosystem: E.g., export models by Orbax, and serve the saved models by Tensorflow Serving system.

We follow most of the interface definitions by onnx.backend here.
Please check a brief example on model conversion and forward calling in examples/imagenet/imagenet_main.py.

Thank you for taking the time to contribute! Please see the contribution guidelines.

This project is licensed under the Apache License.

These details have not been verified by PyPI

View statistics for this project via Libraries.io, or by using our public dataset on Google BigQuery

0.3.0

Aug 18, 2023

0.2.0

Jul 25, 2023

This version

0.1.0

Jun 20, 2023

0.0.2

May 19, 2023

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Uploaded Jun 20, 2023 Source

Uploaded Jun 20, 2023 Python 3

Hashes for jaxonnxruntime-0.1.0.tar.gz
Algorithm	Hash digest
SHA256	`385ec3e590115164280a04958fadb72027fafd5e7bb703e6954ffe859846d5f5`
MD5	`c4c37a5fc494c62e371bf21f63cf2781`
BLAKE2b-256	`8e7793e948c6320b6af4fe75c44ea6e74cfc4d0ec18ea62e11bf6996cbdd84f6`