Transformer A pure-functional implementation of a machine learning transformer model in Python/JAX A pure-functional implementation of a machine learning transformer model in Python/JAX 07 August 2022
JAX Training and Lightweighting Cookbook in JAX/FLAX Training and Lightweighting Cookbook in JAX/FLAX 31 July 2022
PyTorch The performance of Swin-Transformer v2 implemented in JAX and PyTorch The performance of Swin-Transformer v2 implemented in JAX and PyTorch 26 July 2022
JAX Pax - A framework to configure and run machine learning experiments on top of Jax Pax - A framework to configure and run machine learning experiments on top of Jax 28 June 2022
JAX CLOOB training (JAX) and inference (JAX and PyTorch) CLOOB training (JAX) and inference (JAX and PyTorch) 31 March 2022
JAX Official Pytorch and JAX implementation of Efficient-VDVAE: Less is more Official Pytorch and JAX implementation of Efficient-VDVAE: Less is more 31 March 2022
JAX A lossless neural compression framework built on top of JAX A lossless neural compression framework built on top of JAX 15 February 2022
JAX Little demo of how to instantiate a cluster of GPUs with JAX Little demo of how to instantiate a cluster of GPUs with JAX 22 January 2022
Neural Network Evolving neural network parameters in JAX Evolving neural network parameters in JAX 17 January 2022
JAX PEP 503 compliant repository index for Jax(lib) PEP 503 compliant repository index for Jax(lib) 12 January 2022
JAX Unofficial JAX implementation of Deep Learning models Unofficial JAX implementation of Deep Learning models 10 January 2022
JAX Advantage Actor Critic (A2C): jax + flax implementation Advantage Actor Critic (A2C): jax + flax implementation 31 December 2021
Neural Network A minimal TPU compatible Jax implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis A minimal TPU compatible Jax implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis 22 December 2021
JAX A demo of how to use JAX to create a simple gravity simulation A demo of how to use JAX to create a simple gravity simulation 03 December 2021
JAX A framework-agnostic Trainer interface for the Jax ecosystem A framework-agnostic Trainer interface for the Jax ecosystem 01 December 2021
JAX Implementation of different GAN in JAX/Haiku Implementation of different GAN in JAX/Haiku 21 November 2021
reinforcement learning Conservative Q Learning for Offline Reinforcement Reinforcement Learning in JAX Conservative Q Learning for Offline Reinforcement Reinforcement Learning in JAX 11 November 2021
JAX Extending JAX with custom C++ and CUDA code Extending JAX with custom C++ and CUDA code 09 November 2021
JAX V objective diffusion inference code for JAX V objective diffusion inference code for JAX 09 November 2021
Transformer Flaxformer: transformer architectures in JAX/Flax Flaxformer: transformer architectures in JAX/Flax 07 November 2021
Transformer Model parallel transformers in JAX and Haiku Model parallel transformers in JAX and Haiku 07 November 2021
JAX A JAX implementation of Broaden Your Views for Self-Supervised Video Learning, or BraVe for short A JAX implementation of Broaden Your Views for Self-Supervised Video Learning, or BraVe for short 15 October 2021
JAX Score-Based Generative Modeling through Stochastic Differential Equations Score-Based Generative Modeling through Stochastic Differential Equations 15 October 2021
JAX Turning SymPy expressions into JAX functions Turning SymPy expressions into JAX functions 28 September 2021
Optimization A Python package for solving the inverse problems that arise in scientific imaging applications SCICO is a Python package for solving the inverse problems that arise in scientific imaging applications 26 September 2021
PyTorch A prototype of JAX-like composable function transforms for PyTorch functorch is a prototype of JAX-like composable function transforms for PyTorch 26 September 2021
JAX A library for developing custom Federated Learning (FL) algorithms in JAX FedJAX is a library for developing custom Federated Learning (FL) algorithms in JAX. 08 September 2021
Tree A Pytree Module system for Deep Learning in JAX A Pytree-based Module system for Deep Learning in JAX 05 September 2021
JAX Solving reinforcement learning tasks which require language and vision JAX implementations of the following multimodal reinforcement learning approaches. 23 August 2021
JAX A Gaussian process (GP) library built in JAX (with objax) Newt is a Gaussian process (GP) library built in JAX (with objax), built and actively maintained by Will Wilkinson. 15 August 2021
OpenAI Plug-n-Play Reinforcement Learning in Python with OpenAI Gym and JAX Modular framework for Reinforcement Learning in python 10 August 2021
JAX A Python library for differentiable optimal control on accelerators A Python library for differentiable optimal control on accelerators. 07 August 2021
Computer Vision Scenic: A Jax Library for Computer Vision and Beyond Scenic is a codebase with a focus on research around attention-based models for computer vision. 07 August 2021
JAX Semantically Consistent Few-Shot View Synthesis Implementation This project attempted to implement the paper Putting NeRF on a Diet (DietNeRF) in JAX/Flax. 07 August 2021
Neural Network Callable PyTrees and filtered JIT/grad transformations neural networks in JAX Callable PyTrees and filtered JIT/grad transformations neural networks in JAX 05 August 2021
JAX A Mixed Precision library for JAX in python JMP is a Mixed Precision library for JAX. 04 August 2021