get-started-with-JAX

JAX tutorial series

A repository providing tutorials and resources to learn JAX, a popular alternative to PyTorch and TensorFlow for machine learning.

The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

GitHub

670 stars
9 watching
97 forks
Language: Jupyter Notebook
last commit: almost 3 years ago
Linked from 1 awesome list

deep-learningflaxhaikujaxjupyterlaxlearn-jaxmachine-learningnumpyoptaxpythontutorialxla

Backlinks from these awesome lists:

Related projects:

RepositoryDescriptionStars
jaxgaussianprocesses/gpjaxProvides a low-level interface to Gaussian process models in JAX for flexible extension and customisation467
ikostrikov/jaxrlProvides JAX implementations of various reinforcement learning algorithms with continuous action spaces.640
darshandeshpande/jax-modelsProvides a collection of deep learning models and utilities in JAX/Flax for research purposes.151
google-deepmind/chexA set of utilities for writing reliable JAX code797
matthias-wright/flaxmodelsProvides pre-trained deep learning models for the Jax/Flax ecosystem.240
dfm/extending-jaxThis project provides infrastructure to interface custom C++ and CUDA code with the JAX library for scientific computing379
hakky54/java-tutorialsA repository containing various Java tutorials on topics such as Elasticsearch, gRPC, security, and serialization.35
expectationmax/sklearn-jax-kernelsA set of composable kernels for scikit-learn implemented in JAX to accelerate computation and gradient calculation.42
jeremiecoullon/jax-tqdmLibrary that enhances JAX with support for dynamic progress bars in scans and loops100
google-deepmind/jraphA lightweight library for working with graph neural networks in jax.1,380
vict0rsch/deep_learningA collection of tutorials and resources on implementing deep learning models using Python libraries such as Keras and Lasagne.426
logicalclocks/hopsworks-tutorialsA collection of tutorials and notebooks providing hands-on experience with the Hopsworks Platform for machine learning development and data analysis.261
nvidia/jax-toolboxProvides optimized tools and infrastructure for JAX development on NVIDIA GPUs268
owickstrom/gi-gtk-declarativeA Haskell package providing a declarative programming framework for building GTK+ applications.288
kmheckel/spyxA JAX-based library for training and utilizing spiking neural networks104