JAX
216 items from n2cholas/awesome-jax ★2,167
-
Implementing NeRF in JAX by Soumik Rakshit and Saurav Maheshkar wandb.ai
A tutorial on 3D volumetric rendering of scenes represented by Neural Radiance Fields in JAX.
-
Writing a Training Loop in JAX + FLAX by Saurav Maheshkar and Soumik Rakshit wandb.ai
A tutorial on writing a simple end-to-end training and evaluation pipeline in JAX, Flax and Optax.
-
Deep Learning tutorials with JAX+Flax by Phillip Lippe uvadlc-notebooks.readthedocs.io
A series of notebooks explaining various deep learning concepts, from basics (e.g. intro to JAX/Flax, activiation functions) to recent advances (e.g., Vision Transformers, SimCLR), with translations to PyTorch.
-
Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020 slideslive.com
Tutorial created by Zico Kolter, David Duvenaud, and Matt Johnson with Colab notebooks avaliable in Deep Implicit Layers.
-
JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye… slideslive.com
JAX intro presentation in Program Transformations for Machine Learning workshop.
-
From PyTorch to JAX: towards neural net frameworks that purify stateful code by Sabrina J. Mielke sjmielke.com
Showcases how to go from a PyTorch-like style of coding to a more Functional-style of coding.
-
Evolving Neural Networks in JAX by Robert Tjarko Lange roberttlange.github.io
Explores how JAX can power the next generation of scalable neuroevolution algorithms.
-
Getting started with JAX (MLPs, CNNs & RNNs) by Robert Lange roberttlange.github.io
Neural network building blocks from scratch with the basic JAX operators.
-
Reddit reddit.com
-
Understanding Autodiff with JAX by Srihari Radhakrishna radx.in
Understand how autodiff works using JAX.
-
Elegy poets-ai.github.io
A High Level API for Deep Learning in JAX. Supports Flax, Haiku, and Optax.
-
White paper describing an early version of JAX, detailing how computation is traced and compiled.
-
Plugging Into JAX by Nick Doiron medium.com
Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge.
-
Solving y=mx+b with Jax on a TPU Pod slice - Mat Kelcey matpalm.com
A four part YouTube tutorial series with Colab notebooks that starts with Jax fundamentals and moves up to training with a data parallel approach on a v3-32 TPU Pod slice.
-
Out of distribution (OOD) detection by Mat Kelcey matpalm.com
Implements different methods for OOD detection.
-
Evolved channel selection by Mat Kelcey matpalm.com
Trains a classification model robust to different combinations of input channels at different resolutions, then uses a genetic algorithm to decide the best combination for a particular loss.
-
Ensemble networks by Mat Kelcey matpalm.com
Ensemble nets are a method of representing an ensemble of models as one single logical model.
-
Deterministic ADVI in JAX by Martin Ingram martiningram.github.io
Walk through of implementing automatic differentiation variational inference (ADVI) easily and cleanly with JAX.
-
Jax in Action manning.com
A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.
-
Exploring hyperparameter meta-loss landscapes with JAX by Luke Metz lukemetz.com
Demonstrates how to use JAX to perform inner-loss optimization with SGD and Momentum, outer-loss optimization with gradients, and outer-loss optimization using evolutionary strategies.
-
Simple PDE solver + Constrained Optimization with JAX by Philip Mocz levelup.gitconnected.com
A simple example of solving the advection-diffusion equations with JAX and using it in a constrained optimization problem to find initial conditions that yield desired result.
-
Learn JAX: From Linear Regression to Neural Networks by Rito Ghosh kaggle.com
A gentle introduction to JAX and using it to implement Linear and Logistic Regression, and Neural Network models and using them to solve real world problems.
-
How to add a progress bar to JAX scans and loops by Jeremie Coullon jeremiecoullon.com
Tutorial on how to add a progress bar to compiled loops in JAX using the host_callback module.
-
Writing an MCMC sampler in JAX by Jeremie Coullon jeremiecoullon.com
Tutorial on the different ways to write an MCMC sampler in JAX along with speed benchmarks.
-
Oryx github.com
Probabilistic programming language based on program transformations.
-
JAX, Flax & Transformers 🤗 github.com
3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics.
-
Reformer github.com
Implementation of the Reformer (efficient transformer) architecture.
-
JAX GitHub Discussions github.com
-
Flax NNX github.com
An evolution on Flax by the same team
-
VDVAE github.com
Adaptation of Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images, original code at openai/vdvae.
- next page of items loading…
