Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀
-
Updated
Jun 14, 2024 - Python
Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
This is the official repository for the paper "Flora: Low-Rank Adapters Are Secretly Gradient Compressors" in ICML 2024.
Variational Graph Autoencoder implemented using Jax & Jraph
Goal-conditioned reinforcement learning like 🔥
The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.
A reimplementation of Parallel DNN Training in JAX by Will Whitney using haiku and optax.
An implementation of MNIST classification using LeNet-300-100 in JAX (using Haiku and Optax).
This repository contains some of the code I wrote for the assignments in DSA4212 - Optimisation for Large-Scale Data-Driven Inference.
Stochastic Weight Averaging (SWA) transforms for Optax with JAX
dm-haiku implementation of hyperbolic neural networks
JAX/Flax implementation of finite-size scaling
JAX implementation of Classical and Quantum Algorithms for Orthogonal Neural Networks by (Kerenidis et al., 2021)
An implementation of adan optimizer for optax
Add a description, image, and links to the optax topic page so that developers can more easily learn about it.
To associate your repository with the optax topic, visit your repo's landing page and select "manage topics."