JAX
216 items from n2cholas/awesome-jax ★2,167
-
TriMap github.com
Official JAX implementation of TriMap: Large-scale Dimensionality Reduction Using Triplets.
-
Symbolic Functionals github.com
Demonstration from Evolving symbolic density functionals.
-
Spin-weighted Spherical CNNs github.com
Adaptation of Spin-Weighted Spherical CNNs.
-
SNeRG github.com
Official implementation of Baking Neural Radiance Fields for Real-Time View Synthesis.
-
Distributed Shampoo github.com
Implementation of Second Order Optimization Made Practical.
-
Robust Loss github.com
Reference code for the paper A General and Adaptive Robust Loss Function.
-
RegNeRF github.com
Official implementation of RegNeRF: Regularizing Neural Radiance Fields for View Synthesis from Sparse Inputs.
-
Slot Attention github.com
Reference implementation for Differentiable Patch Selection for Image Recognition.
-
Protein LM github.com
Implements BERT and autoregressive models for proteins, as described in Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences and ProGen: Language Modeling for Protein Generation.
-
Persistent Evolution Strategies github.com
Code used for the paper Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies.
-
Performer github.com
Flax implementation of the Performer (linear transformer via FAVOR+) architecture.
-
MUSIQ github.com
Checkpoints and model inference code for the ICCV 2021 paper MUSIQ: Multi-scale Image Quality Transformer
-
Light Field Neural Rendering github.com
Official implementation of Light Field Neural Rendering.
-
Latent Programmer github.com
Code for the ICML 2021 paper Latent Programmer: Discrete Latent Codes for Program Synthesis.
-
JaxNeRF github.com
Implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis with multi-device GPU/TPU support.
-
JAX-DFT github.com
One-dimensional density functional theory (DFT) in JAX, with implementation of Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics.
-
IPA-GNN github.com
Official implementation of Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks.
-
Gumbel-max Causal Mechanisms github.com
Code for Learning Generalized Gumbel-max Causal Mechanisms, with extra code in GuyLor/gumbel_max_causal_gadgets_part2.
-
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 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.
- next page of items loading…
