Stream helps developers build engaging apps that scale to millions with performant and flexible Chat, Feeds, Moderation, and Video APIs and SDKs powered by a global edge network and enterprise-grade infrastructure. Learn more →
Top 23 Python Jax Projects
-
transformers
🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.
python3.12 -m venv new_venv_312 source new_venv_312/bin/activate pip install --upgrade pip pip install https://github.com/huggingface/transformers/archive/main.zip torchaudio peft soundfile torchcodec ### and also pip install librosa
-
InfluxDB
InfluxDB – Built for High-Performance Time Series Workloads. InfluxDB 3 OSS is now GA. Transform, enrich, and act on time series data directly in the database. Automate critical tasks and eliminate the need to move data externally. Download now.
-
Keras 3 multi-backend
-
jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
-
diffusers
🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.
Diffusers Documentation
-
d2l-en
Interactive deep learning book with multi-framework code, math, and discussions. Adopted at 500 universities from 70 countries including Stanford, MIT, Harvard, and Cambridge.
-
Project mention: A ranked list of machine learning Python libraries. Updated weekly | news.ycombinator.com | 2025-01-31
-
-
Stream
Stream - Scalable APIs for Chat, Feeds, Moderation, & Video. Stream helps developers build engaging apps that scale to millions with performant and flexible Chat, Feeds, Moderation, and Video APIs and SDKs powered by a global edge network and enterprise-grade infrastructure.
-
wandb
The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.
Project mention: Best Practices for Ensuring AI Agent Performance and Reliability | dev.to | 2025-07-22Use tools like Weights & Biases, Labelbox, or Maxim’s data engine to version your datasets, track changes, and continuously add new edge cases and user feedback.
-
einops
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
-
-
-
Good read! I am developing PINNs at work and this certainly helped me recall important concepts. This post used deepxde library [2] to compose the PINN. Can anyone comment on how NVIDIA's modulus [2] compares to this? Modulus appears to be much more verbose and poorly documented.
[1]: https://github.com/lululxvi/deepxde
-
-
TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
-
foolbox
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
-
-
pennylane
PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Built by researchers, for research.
Project mention: Python library for quantum computing, quantum ML, and quantum chemistry | news.ycombinator.com | 2025-11-05 -
equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Project mention: An Introduction to Neural Ordinary Differential Equations [pdf] | news.ycombinator.com | 2025-01-11 -
-
numpyro
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
Project mention: Numpyro: Probabilistic programming with NumPy powered by Jax | news.ycombinator.com | 2024-11-16 -
EasyLM
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
-
-
-
SaaSHub
SaaSHub - Software Alternatives and Reviews. SaaSHub helps you find the best software and product alternatives
NOTE:
The open source projects on this list are ordered by number of github stars.
The number of mentions indicates repo mentiontions in the last 12 Months or
since we started tracking (Dec 2020).
Python Jax discussion
Python Jax related posts
-
"Why don't you use dependent types?"
-
Bringing NumPy's type-completeness score to nearly 90% – Pyrefly
-
Defeating Nondeterminism in LLM Inference
-
Axlearn: Extensible Deep Learning Library
-
Show HN: Localscope–Limit scope of Python functions for reproducible execution
-
AXLearn: Apple's Deep Learning library built on top of Jax
-
How to Install & Run VideoLLaMA3-7B Locally
-
A note from our sponsor - Stream
getstream.io | 15 Nov 2025
Index
What are some of the best open-source Jax projects in Python? This list will help you:
| # | Project | Stars |
|---|---|---|
| 1 | transformers | 152,508 |
| 2 | Keras | 63,551 |
| 3 | jax | 33,955 |
| 4 | diffusers | 31,604 |
| 5 | d2l-en | 26,601 |
| 6 | best-of-ml-python | 22,786 |
| 7 | ivy | 14,234 |
| 8 | wandb | 10,521 |
| 9 | einops | 9,271 |
| 10 | datasets | 4,499 |
| 11 | scenic | 3,707 |
| 12 | deepxde | 3,602 |
| 13 | dm-haiku | 3,117 |
| 14 | TransformerEngine | 2,912 |
| 15 | foolbox | 2,903 |
| 16 | thinc | 2,877 |
| 17 | pennylane | 2,874 |
| 18 | equinox | 2,665 |
| 19 | mctx | 2,558 |
| 20 | numpyro | 2,551 |
| 21 | EasyLM | 2,498 |
| 22 | dreamerv3 | 2,358 |
| 23 | axlearn | 2,288 |
Sponsored
InfluxDB – Built for High-Performance Time Series Workloads
InfluxDB 3 OSS is now GA. Transform, enrich, and act on time series data directly in the database. Automate critical tasks and eliminate the need to move data externally. Download now.
www.influxdata.com