optax
There are 36 repositories under optax topic.
phlippe/uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
gordicaleksa/get-started-with-JAX
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.
erfanzar/EasyDeL
Accelerate, Optimize performance with streamlined training and serving options with JAX.
BorealisAI/flora-opt
This is the official repository for the paper "Flora: Low-Rank Adapters Are Secretly Gradient Compressors" in ICML 2024.
varun-ml/diffusion-models-tutorial
Experiment with diffusion models that you can run on your local jupyter instances
hamishs/JAX-RL
JAX implementations of various deep reinforcement learning algorithms.
hushon/JAX-ResNet-CIFAR10
Simple CIFAR10 ResNet example with JAX.
evanatyourservice/psgd_jax
Implementation of PSGD optimizer in JAX
yonesuke/jaxfss
JAX/Flax implementation of finite-size scaling
frankroeder/goal_conditioned_rl
Goal-conditioned reinforcement learning like 🔥
hr0nix/optax-adan
An implementation of adan optimizer for optax
JesseFarebro/flax-mup
Maximal Update Parametrization (μP) with Flax & Optax.
mzguntalan/h-former
H-Former is a VAE for generating in-between fonts (or combining fonts). Its encoder uses a Point net and transformer to compute a code vector of glyph. Its decoder is composed of multiple independent decoders which act on a code vector to reconstruct a point cloud representing a glpyh.
qdevpsi3/quantum-orthogonal-nn
JAX implementation of Classical and Quantum Algorithms for Orthogonal Neural Networks by (Kerenidis et al., 2021)
salfaris/vgae-jax
Variational Graph Autoencoder implemented using Jax & Jraph
NITHISHM2410/flax-pilot
A Simplistic trainer for Flax
ethanluoyc/td3_bc_jax
Direct port of TD3_BC to JAX using Haiku and optax.
evanatyourservice/sophia-schedulefree-jax
Sophia optimizer with ScheduleFree
Raffaelbdl/hyperbolic-nn-haiku
dm-haiku implementation of hyperbolic neural networks
NTT123/wavernn-16bit
The (unofficial) vanilla version of WaveRNN
ysngshn/ivon-optax
An Optax-based JAX implementation of the IVON optimizer for large-scale VI training of NNs (ICML'24 spotlight)
daniel-j-h/nedem
Neural implicit digital elevation model
salfaris/PriorVGAE
Oxford MSc thesis. PriorVAE with graph convolutional networks for learning locally-aware spatial prior distributions
activatedgeek/optax-swag
Stochastic Weight Averaging (SWA) transforms for Optax with JAX
bischtob/Opterax
An gradient-free optimization suite written in JAX. We conform to the optax interface and provide ensemble-based optimizers.
elttaes/VAE-MNIST-Haiku-Jax
Simple VAE example with Jax.
evanatyourservice/flat-sophia
sophia optimizer further projected towards flat areas of loss landscape
NTT123/haiku_trainer
A helper library for training dm-haiku models.
alexjackson1/hk-mnist
An implementation of MNIST classification using LeNet-300-100 in JAX (using Haiku and Optax).
alexjackson1/hk-parallel
A reimplementation of Parallel DNN Training in JAX by Will Whitney using haiku and optax.
BigBalloon8/Science_Bacc
SQA scottish Science Baccalaureate of Christopher Rae. Improving the preformance of distributed data parallelism of low wifi bandwidths
cevenkadir/ferminet_tum
A library which trains the Fermionic Neural Network to find the ground state wave functions of an atom or a molecule using neural network quantum states.
mzguntalan/red
A place for code to relax in happiness without the pressure of maintenance. Mostly related to Machine Learning and Deep Learning
xghan99/dsa4212-assignments
This repository contains some of the code I wrote for the assignments in DSA4212 - Optimisation for Large-Scale Data-Driven Inference.
ShawonAshraf/postagger-lstm-jax
LSTM POS Tagger implementation in Jax and Flax
yonesuke/gprax
Simple JAX implementation of Gaussian process regression