google / jaxonnxruntimeLinks
A user-friendly tool chain that enables the seamless execution of ONNX models using JAX as the backend.
☆122Updated last month
Alternatives and similar repositories for jaxonnxruntime
Users that are interested in jaxonnxruntime are comparing it to the libraries listed below
Sorting:
- A stand-alone implementation of several NumPy dtype extensions used in machine learning.☆301Updated 2 weeks ago
- ☆190Updated 3 weeks ago
- jax-triton contains integrations between JAX and OpenAI Triton☆428Updated last week
- Tokamax: A GPU and TPU kernel library.☆94Updated this week
- JMP is a Mixed Precision library for JAX.☆208Updated 8 months ago
- Implementation of Flash Attention in Jax☆219Updated last year
- OpTree: Optimized PyTree Utilities☆196Updated last week
- If it quacks like a tensor...☆59Updated 11 months ago
- ☆115Updated last month
- JAX-Toolbox☆350Updated last week
- ☆52Updated last year
- ☆335Updated last month
- A functional training loops library for JAX☆88Updated last year
- A simple library for scaling up JAX programs☆144Updated 11 months ago
- Run PyTorch in JAX. 🤝☆304Updated last week
- JAX Synergistic Memory Inspector☆179Updated last year
- PyTorch centric eager mode debugger☆48Updated 10 months ago
- Named Tensors for Legible Deep Learning in JAX☆210Updated this week
- Minimal yet performant LLM examples in pure JAX☆186Updated last month
- Experiment of using Tangent to autodiff triton☆80Updated last year
- Einsum-like high-level array sharding API for JAX☆35Updated last year
- LoRA for arbitrary JAX models and functions☆141Updated last year
- A Python package of computer vision models for the Equinox ecosystem.☆109Updated last year
- ☆21Updated 7 months ago
- TorchFix - a linter for PyTorch-using code with autofix support☆149Updated 2 months ago
- Serialize JAX, Flax, Haiku, or Objax model params with 🤗`safetensors`☆47Updated last year
- ☆77Updated last week
- A Jax-based library for building transformers, includes implementations of GPT, Gemma, LlaMa, Mixtral, Whisper, SWin, ViT and more.☆295Updated last year
- Demo of the unit_scaling library, showing how a model can be easily adapted to train in FP8.☆45Updated last year
- ☆145Updated last week