JAX
216 items from n2cholas/awesome-jax ★2,167
-
GIFT github.com
Official implementation of Gradual Domain Adaptation in the Wild:When Intermediate Distributions are Absent.
-
GFSA github.com
Official implementation of Learning Graph Structure With A Finite-State Automaton Layer.
-
Flax Models github.com
Collection of models and methods implemented in Flax.
-
FNet github.com
Official implementation of FNet: Mixing Tokens with Fourier Transforms.
-
Dreamfields github.com
Official implementation of the ICLR 2022 paper Progressive Distillation for Fast Sampling of Diffusion Models.
-
D3PM github.com
Official implementation of Structured Denoising Diffusion Models in Discrete State-Spaces.
-
Combiner github.com
Official implementation of Combiner: Full Attention Transformer with Sparse Computation Cost.
-
BNN-HMC github.com
Implementation for the paper What Are Bayesian Neural Network Posteriors Really Like?.
-
ARDM github.com
Official implementation of Autoregressive Diffusion Models.
-
AQuaDem github.com
Official implementation of Continuous Control with Action Quantization from Demonstrations.
-
Accurate Quantized Training github.com
Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax.
-
Amortized Bayesian Optimization github.com
Code related to Amortized Bayesian Optimization over Discrete Spaces.
-
ALX github.com
Open-source library for distributed matrix factorization using Alternating Least Squares, more info in ALX: Large Scale Matrix Factorization on TPUs.
-
WikiGraphs github.com
Baseline code to reproduce results in WikiGraphs: A Wikipedia Text - Knowledge Graph Paired Datase.
-
OGB-LSC github.com
This repository contains DeepMind's entry to the PCQM4M-LSC (quantum chemistry) and MAG240M-LSC (academic graph) tracks of the OGB Large-Scale Challenge (OGB-LSC).
-
Normalizer-Free Networks github.com
Official Haiku implementation of NFNets.
-
MMV github.com
Code for the models in Self-Supervised MultiModal Versatile Networks.
-
Glassy Dynamics github.com
Open source implementation of the paper Unveiling the predictive power of static structure in glassy systems.
-
Gated Linear Networks github.com
GLNs are a family of backpropagation-free neural networks.
-
Bootstrap Your Own Latent github.com
Implementation for the paper Bootstrap your own latent: A new approach to self-supervised Learning.
-
Adversarial Robustness github.com
Reference code for Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples and Fixing Data Augmentation to Improve Adversarial Robustness.
-
jax-flows github.com
Normalizing flows in JAX.
-
JAX on Cloud TPUs | NeurIPS 2020 | Skye Wanderman-Milne and James Bradbury drive.google.com
Presentation of TPU host access with demo.
-
JaxLLM (Unofficial) Discord discord.com
-
Sampling Path Candidates with Machine Learning differt.eertmans.be
Official tutorial and implementation from the paper Towards Generative Ray Path Sampling for Faster Point-to-Point Ray Tracing.
-
Using JAX to accelerate our research by David Budden and Matteo Hessel deepmind.com
Describes the state of JAX and the JAX ecosystem at DeepMind.
-
Introduction to JAX by Kevin Murphy colab.research.google.com
Colab that introduces various aspects of the language and applies them to simple ML problems.
-
Sharpened Cosine Similarity in JAX by Raphael Pisoni colab.research.google.com
A JAX/Flax implementation of the Sharpened Cosine Similarity layer.
-
Achieving 4000x Speedups with PureJaxRL chrislu.page
A blog post on how JAX can massively speedup RL training through vectorisation.
-
Differentiable Path Tracing on the GPU/TPU by Eric Jang blog.evjang.com
Tutorial on implementing path tracing.
-
Normalizing Flows in 100 Lines of JAX by Eric Jang blog.evjang.com
Concise implementation of RealNVP.
-
Meta-Learning in 50 Lines of JAX by Eric Jang blog.evjang.com
Introduction to both JAX and Meta-Learning.
-
White paper describing the XLB library: benchmarks, validations, and more details about the library.
-
Uses JAX's JIT and VMAP to achieve faster differentially private than existing libraries.
-
Introduces JAX, M.D., a differentiable physics library which includes simulation environments, interaction potentials, neural networks, and more.
-
Two Player Auction Learning github.com
JAX implementation of the paper Auction learning as a two-player game.
