jax
32 projects share this GitHub topic
Lines connect members that are measurably related to each other. Dot size reflects stars.
🧬 Members
keras
Deep Learning for humans
★ 64.2k
einops
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
★ 9.5k
pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
★ 7.1k
whisper-jax
JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.
★ 4.7k
datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
★ 4.6k
deepxde
A library for scientific machine learning and physics-informed learning
★ 4.3k
iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
★ 3.8k
scenic
Scenic: A Jax Library for Computer Vision Research and Beyond
★ 3.8k
dreamerv3
Mastering Diverse Domains through World Models
★ 3.6k
dm-haiku
JAX-based neural network library
★ 3.3k
brax
Massively parallel rigidbody physics simulation on accelerator hardware.
★ 3.2k
uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam…
★ 3.2k
foolbox
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
★ 3k
equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
★ 2.9k
mctx
Monte Carlo tree search in JAX
★ 2.6k
EasyLM
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning,…
★ 2.5k
axlearn
An Extensible Deep Learning Library
★ 2.4k
penzai
A JAX research toolkit for building, editing, and visualizing neural networks.
★ 1.9k
tensorly
TensorLy: Tensor Learning in Python.
★ 1.7k
magenta-realtime
Magenta RealTime 2: An Open-Weights Live Music Model
★ 1.7k
Treasure-of-Transformers
💁 Awesome Treasure of Transformers Models for Natural Language processing contains papers, videos, blogs,…
★ 1.2k
purejaxrl
Really Fast End-to-End Jax RL Implementations
★ 1.1k
Mava
🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX
★ 923
jumanji
🕹️ A diverse suite of scalable reinforcement learning environments in JAX
★ 851
fast-plate-ocr
Lightweight & fast OCR models for license plate text recognition.
★ 613
paxml
Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and…
★ 556
Quadruped-PyMPC
A model predictive controller for quadruped robots based on the single rigid body model and written in…
★ 491
dm_pix
PIX is an image processing library in JAX, for JAX.
★ 442
whisper-finetune
Fine-tune and evaluate Whisper models for Automatic Speech Recognition (ASR) on custom datasets or datasets…
★ 365
xland-minigrid
JAX-accelerated Meta-Reinforcement Learning Environments Inspired by XLand and MiniGrid 🏎️
★ 343
jaxlie
Rigid transforms + Lie groups for JAX
★ 336
gemma_penzai
A JAX Research Toolkit for Visualizing, Manipulating, and Understanding Gemma Models with Multi-modal Support…
★ 95
🔗 Related families
Measured from GitHub topics shared by both projects, weighted by how rare each topic is.