Introduction to JAX

JAX, Flax & 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.

JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne

JAX intro presentation in Program Transformations for Machine Learning workshop

Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020

Tutorial created by Zico Kolter, David Duvenaud, and Matt Johnson with Colab notebooks avaliable in Deep Implicit Layers

NeurIPS 2020: JAX Ecosystem Meetup

JAX, its use at DeepMind, and discussion between engineers, scientists, and JAX core team