keras-team / keras

Deep Learning for humans

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

unifyai / ivy

Convert Machine Learning Code Between Frameworks

wandb / wandb

The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.

HIPS / autograd

Efficiently computes derivatives of NumPy code.

tensorflow / datasets

TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...

danijar / dreamerv3

Mastering Diverse Domains through World Models

google-research / scenic

Scenic: A Jax Library for Computer Vision Research and Beyond

NVIDIA / 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.

PennyLaneAI / pennylane

PennyLane is an open-source quantum software platform for quantum computing, quantum machine learning, and quantum chemistry. Create meaningful quantum algorithms, from inspiration to implementation.

google-deepmind / dm-haiku

JAX-based neural network library

pyro-ppl / numpyro

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

google-deepmind / mctx

Monte Carlo tree search in JAX

EMI-Group / evox

Distributed GPU-Accelerated Framework for Evolutionary Computation. Comprehensive Library of Evolutionary Algorithms & Benchmark Problems.

AI-Hypercomputer / maxtext

A simple, performant, and scalable Jax LLM!

google-deepmind / optax

Optax is a gradient processing and optimization library for JAX.

deepmodeling / deepmd-kit

A deep learning package for many-body potential energy representation and molecular dynamics

magenta / magenta-realtime

Magenta RealTime 2: An Open-Weights Live Music Model

kubeflow / katib

Automated Machine Learning on Kubernetes

tensorly / tensorly

TensorLy: Tensor Learning in Python.

extropic-ai / thrml

Thermodynamic Hypergraphical Model Library in JAX

blackjax-devs / blackjax

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

Toni-SM / skrl

Modular Reinforcement Learning (RL) library (implemented in PyTorch, JAX, and NVIDIA Warp) with support for Gymnasium/Gym, NVIDIA Isaac Lab, MuJoCo Playground and other environments

google / jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.

probml / dynamax

A Python package for probabilistic state space modeling with JAX

keras-team / keras-hub

Pretrained model hub for Keras 3.

mosecorg / mosec

A high-performance ML model serving framework, offers dynamic batching and CPU/GPU pipelines to fully exploit your compute machine

instadeepai / jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

RobertTLange / evosax

Evolution Strategies in JAX 🦎

google / grain

Library for reading and processing ML training data.

deepmodeling / jax-fem

Differentiable Finite Element Method with JAX

google-research / jax3d

—

google-deepmind / torax

TORAX: Tokamak transport simulation in JAX

netket / netket

Machine learning algorithms for many-body quantum systems

ARahim3 / kaggle-tpu-lab

Frontier-class open models on a free Kaggle TPU v5e-8: GLM-5.3-Flash 320B MoE (~165 tok/s, our own JAX engine) and Qwen3.8-27B bf16 (~130 tok/s), 262k context, prefix caching. Works with Claude Code, Codex, opencode and pi.

QuantClimate / GPJax

Gaussian processes in JAX and Equinox.

araffin / sbx

SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms

patrick-kidger / lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

google-deepmind / simply

Minimal and scalable research codebase in JAX, designed for rapid iteration on frontier research in LLM and other autoregressive models.

zengxiao-he / tessera

From teacher to tiles — a from-scratch LLM distillation & serving engine: custom Triton/CUDA kernels, FSDP distillation, paged-KV continuous batching, speculative decoding, a Rust gateway, a JAX oracle, and interpretability tooling.

google / paxml

Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.

mpi4jax / mpi4jax

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python :zap:

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

iit-DLSLab / Quadruped-PyMPC

A model predictive controller for quadruped robots based on the single rigid body model and written in python. Gradient-based (acados) or Sampling-based (jax)

monatis / ggmlc

A multi-framework neural network compiler lowering PyTorch, JAX, Flax, and Keras models to portable, high-performance GGML execution

pmocz / noether

Automatically prove mathematical statements about your jax code in lean

ECLIPSE-AI4Science / gradsolve

Differentiable ensemble solvers for differential equations on GPUs, in JAX

jaxxchen003 / digital-human-video-production

Reusable skill for presenter-led digital human product videos

extropic-ai / sparse-transformers

Sparse Transformers in JAX