jax
There are 1378 repositories under jax topic.
keras-team/keras
Deep Learning for humans
jax-ml/jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
d2l-ai/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.
lukasmasuch/best-of-ml-python
🏆 A ranked list of awesome machine learning Python libraries. Updated weekly.
ivy-llc/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.
arogozhnikov/einops
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
google/trax
Trax — Deep Learning with Clear Code and Speed
probml/pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
google/flax
Flax is a neural network library for JAX that is designed for flexibility.
BoltzmannEntropy/interviews.ai
It is my belief that you, the postgraduate students and job-seekers for whom the book is primarily meant will benefit from reading it; however, it is my hope that even the most experienced researchers will find it fascinating as well.
sanchit-gandhi/whisper-jax
JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.
tensorflow/datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
google-research/scenic
Scenic: A Jax Library for Computer Vision Research and Beyond
lululxvi/deepxde
A library for scientific machine learning and physics-informed learning
iree-org/iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
alpa-projects/alpa
Training and serving large-scale neural networks with auto parallelization.
google-deepmind/dm-haiku
JAX-based neural network library
phlippe/uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
bethgelab/foolbox
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
explosion/thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
google/brax
Massively parallel rigidbody physics simulation on accelerator hardware.
grame-cncm/faust
Functional programming language for signal processing and sound synthesis
PennyLaneAI/pennylane
PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Built by researchers, for research.
NVIDIA/TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
google-deepmind/mctx
Monte Carlo tree search in JAX
pyro-ppl/numpyro
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
patrick-kidger/equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
young-geng/EasyLM
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
google/neural-tangents
Fast and Easy Infinite Neural Networks in Python
apple/axlearn
An Extensible Deep Learning Library
danijar/dreamerv3
Mastering Diverse Domains through World Models
n2cholas/awesome-jax
JAX - A curated list of resources https://github.com/google/jax
kubeflow/trainer
Distributed AI Model Training and Fine-Tuning on Kubernetes
AI-Hypercomputer/maxtext
A simple, performant and scalable Jax LLM!
google-deepmind/penzai
A JAX research toolkit for building, editing, and visualizing neural networks.